""" This file contains new functions that rewrite from each class """
from Bio.PDB.SASA import ShrakeRupley
import numpy as np
from Bio.PDB import Model, Chain, Residue, Atom
from scipy.spatial.distance import cdist
from concurrent.futures import ThreadPoolExecutor, as_completed
from env_info import num_workers
from LRFPCA import LRFEstimatorPCA
amino_acids = {
        "ALA",
        "ARG",
        "ASN",
        "ASP",
        "CYS",
        "GLN",
        "GLU",
        "GLY",
        "HIS",
        "HSE",
        "ILE",
        "LEU",
        "LYS",
        "MET",
        "PHE",
        "PRO",
        "SER",
        "THR",
        "TRP",
        "TYR",
        "VAL",
    }

def sort_residues(chain):
      # Sort residues from a complex
    protein_residues = []
    lig_residues = []

    for res in chain.get_residues():
        # Check if ligand is present and select only the first one
        if res.resname not in amino_acids and not lig_residues:
            lig_residues.append(res)
        else:
            protein_residues.append(res)
    return protein_residues, lig_residues

def get_residues(structure, model_id, chain_id):
    # Get residues from one subunit, chain, and ligand in the complex

    # Get specified model or chain if selected
    if model_id is not None and chain_id is not None:
        model = structure[model_id]
        chain = model[chain_id]

        # Sort residues
        protein_residues, lig_residues = sort_residues(chain)

    else:
        # Iterate through model/chain to retrieve FIRST one with a bound ligand
        found_ligand = False
        for model in structure.get_models():
            if found_ligand:
                break

            for chain in model.get_chains():
                protein_residues, lig_residues = sort_residues(chain)

                # Check if this chain has both protein AND ligand
                if lig_residues:
                    found_ligand = True
                    break

        if not found_ligand:
            raise ValueError("No ligand found in any model/chain of the structure")
        else:
            print(f"Found ligand {lig_residues[0].id} in model: {model.id}, chain: {chain.id}")
    return protein_residues, lig_residues

def get_coords_matrix(coords: list):
    return np.vstack(coords)

def get_atoms_from_res(residue):
    # Get non-hydrogen atoms and their corresponding coordinates from a residue
    coords = []
    atoms = []
    for atom in residue.get_atoms():
        if atom.element != 'H':
            # Assign atom to the highest occupancy conformer if multiple are present
            if atom.is_disordered():
                atom = max(atom.child_dict.values(), key=lambda x: x.get_occupancy() or 0)
            atoms.append(atom)
            coords.append(atom.get_vector().get_array())
    return coords, atoms

def get_atoms(residues):
    # Get non-hydrogen atoms and coordinates from residue(s)
    if isinstance(residues, list):
        all_coords = []
        all_atoms = []
        for residue in residues:
            coords, atoms = get_atoms_from_res(residue)
            all_coords.append(coords)
            all_atoms.extend(atoms)
    else:
        all_coords, all_atoms = get_atoms_from_res(residues)
    all_coords_matrix = get_coords_matrix(all_coords)
    
    return all_coords_matrix, all_atoms

def atoms_to_model(SASA_atoms, chain_id: str = 'A'):
    # convert 
    res_dict = {}
    for atom in SASA_atoms:
        resname = atom.get_parent().get_resname()
        seg_id = atom.get_parent().get_id()[0]
        res_id = atom.get_parent().get_id()

        # Ensure duplicate residues are not created
        res_key = (res_id, resname, seg_id)
        if res_key not in res_dict.keys():
            # Create new reisdue object
            current_residue = Residue.Residue(res_id, resname, seg_id)
            res_dict[res_key] = current_residue

        else:
            # Retrieve existing residue
            current_residue = res_dict[res_key]

        # Add atom to current residue
        # if atom.get_id() not in [a for a in current_residue.get_atoms()]:
        current_residue.add(atom)
    res_list = list(res_dict.values())
    # convert res to model
    model = Model.Model(id=0)
    chain = Chain.Chain(chain_id)
    seen_res_id = []
    for res in res_list:
        if res.get_id()[1] not in seen_res_id:
            seen_res_id.append(res.get_id()[1])
            chain.add(res)
    model.add(chain)
    return model

def get_protein_atoms(structure):
      atoms = []
      coords = []
      for atom in structure[0].get_atoms():
        atoms.append(atom)
        coords.append(atom.get_coord())
      coords_matrix = np.vstack(coords)
      return atoms, coords_matrix

def get_SASA_biomolecule(pdb_name, structure, distance: float = 8, model_id=None, chain_id=None):
    # Get atoms on solvent accessible surface area (SASA) of a biomolecule object
    
    # extract atoms and coordinates from the whole structure
    whole_atoms, whole_coords_matrix = get_protein_atoms(structure)

    protein_residues, lig_residues = get_residues(structure, model_id, chain_id)
    if len(lig_residues) != 1:
        raise Exception(f"{pdb_name} must be bound to one small molecule ligand.")
    
    lig_coords, lig_atoms = get_atoms(lig_residues)
    lig_coords_matrix = get_coords_matrix(lig_coords)

    site_coords = []
    site_atoms = []
    site_residues = []
    seen_atom_ids = set() # Keep track of added atoms for debugging purposes
    # all_protein_atoms = [] # get site indexes from the whole structure
    for protein_residue in protein_residues:
        # print(f"Processing residue: {protein_residue.get_resname()} {protein_residue.get_id()}")

        res_coords_matrix, res_atoms = get_atoms([protein_residue])
        # all_protein_atoms.extend(res_atoms) # get site indexes from the whole structure
        
        # Calculate all pairwise distances using cdist
        # dist_matrix shape: (m, n) where m = current residue atoms, n = ligand atoms
        dist_matrix = cdist(res_coords_matrix, lig_coords_matrix, metric='euclidean')

        # Find atoms within distance of ANY ligand atom
        lig_distance_mask = np.any(dist_matrix <= distance, axis=1)

        if np.any(lig_distance_mask):
            # Get indices where mask is True (e.g., distance is within threshold)
            idx = np.where(lig_distance_mask)[0]
            residue_added = False # Keep track of residues added

            for i in idx:
                atom = res_atoms[i]
                atom_id = (atom.get_id(), atom.get_parent().get_id(), atom.get_parent().get_resname())

                if atom_id not in seen_atom_ids:
                    site_atoms.append(atom)
                    seen_atom_ids.add(atom_id)
                    site_coords.append(res_coords_matrix[i])

                    # Add residue once
                    if not residue_added:
                        site_residues.append(protein_residue)
                        residue_added = True

    site_coords_matrix = get_coords_matrix(site_coords)
    sr = ShrakeRupley()
    sr.compute(structure, level='A')
    SASA_atoms = []
    SASA_coords = []
    SASA_indexes_from_whole = []
    SASA_indexes_from_site = []
    for i, atom in enumerate(structure.get_atoms()):
        if atom.sasa > 0:
            coord = atom.get_vector().get_array()
            matches = np.where(np.all(site_coords == coord, axis=1))[0]
            if matches.size > 0:
                idx_in_biomol = int(matches[0])
                SASA_atoms.append(atom)
                SASA_coords.append(coord)
                SASA_indexes_from_whole.append(i)
                SASA_indexes_from_site.append(idx_in_biomol)
    
    SASA_coords_matrix = get_coords_matrix(SASA_coords)
    return site_coords_matrix, SASA_atoms, SASA_coords_matrix, SASA_indexes_from_whole, SASA_indexes_from_site

def get_coordinate_lengths(coords_1, coords_2):
    # Compare cooridnate lengths of two biomolecules for use in later function(s)
    # Note, that the biomolecule with fewer atoms should always be 'biomolecule A'
    downsampling = True
    switch = True
    if len(coords_1) < len(coords_2):
        coords_A = coords_1
        coords_B = coords_2
        switch = False
    elif len(coords_1) > len(coords_2):
        coords_A = coords_2
        coords_B = coords_1
    else:
        coords_A = coords_1
        coords_B = coords_2
        downsampling = False
        switch = False
    return coords_A, coords_B, downsampling, switch

def Kabsch_Umeyama(coords_A, coords_B):
    # Get optimal rotation and translation matrices using the Kabsch Umeyama algorithm
    # Source(s): https://web.stanford.edu/class/cs273/refs/umeyama.pdf, https://zp l.fi/aligning-point-patterns-with-kabsch-umeyama-algorithm/

    # Center coordinate matrices
    centroid_A = np.mean(coords_A, axis=0)
    centroid_B = np.mean(coords_B, axis=0)
    Am = coords_A - centroid_A
    Bm = coords_B - centroid_B

    # Compute SVD of covariance matrix
    m, n = coords_A.shape
    H = Am.T @ Bm
    U, D, vt = np.linalg.svd(H)

    # Recover reflections
    det = np.linalg.det(U) * np.linalg.det(vt.T)
    if det >= 0:
        S = np.diag([1] * (n - 1) + [1])
    else:
        S = np.diag([1] * (n - 1) + [-1])

    # Recover optimal rotation and translation matrices
    R = U @ S @ vt  # rotation matrix
    t = centroid_A - R @ centroid_B  # translation vector
    return R, t

def align_coordinates(R, t, coords_B):
    # Align biomolecule coordinates after optimal rotation and translation are computed
    return np.array([t + R @ b for b in coords_B])
    
def get_alignment_attributes(coords_A, coords_B):
    # Get defined attributes associated with transforming biomolecules
    R, t = Kabsch_Umeyama(coords_A, coords_B)
    aligned_coords = align_coordinates(R, t, coords_B)
    RMSD = np.sqrt(np.mean((coords_A - aligned_coords) ** 2))
        
    return aligned_coords, RMSD, R, t

def get_downsampled_attributes(coords_A, coords_B):
    # Retreive alignment attributes associated with minimum RMSD
    # In this case, coords_B should be a larger matrix than coords_A
    RMSDs = []
    Rs = []
    ts = []
    slice_size = coords_A.shape[0]

    # sliding windows
    for i in range(coords_B.shape[0] - slice_size + 1):
        # Downsample larger coordinate matrix
        coords_BB = coords_B[i : i + slice_size]
        _, RMSD, R, t = get_alignment_attributes(
            coords_A, coords_BB
        )
        RMSDs.append(RMSD)
        Rs.append(R)
        ts.append(t)

    min_idx = np.argmin(RMSDs)
    RMSD, R, t = RMSDs[min_idx], Rs[min_idx], ts[min_idx]

    # Apply transformation(s) to undownsampled coordinates
    aligned_coords = align_coordinates(R, t, coords_B)
    return aligned_coords, RMSD, R, t
    
def get_downsampled_attributes_parallel(coords_A, coords_B):
    # Retreive alignment attributes associated with minimum RMSD
    # In this case, coords_B should be a larger matrix than coords_A
    slice_size = coords_A.shape[0]
    n_slices = coords_B.shape[0] - slice_size + 1

    if n_slices <= 0:
        # Fallback to direct alignment if coords_B is not larger than coords_A
        aligned_coords, RMSD, R, t = get_alignment_attributes(coords_A, coords_B)
        return aligned_coords, RMSD, R, t

    RMSDs = np.empty(n_slices)
    Rs = [None] * n_slices
    ts = [None] * n_slices

    def _compute(i):
        coords_BB = coords_B[i:i + slice_size]
        _, RMSD, R, t = get_alignment_attributes(coords_A, coords_BB)
        return i, RMSD, R, t

    with ThreadPoolExecutor(max_workers=num_workers) as ex:
        futures = [ex.submit(_compute, i) for i in range(n_slices)]
        for fut in as_completed(futures):
            i, RMSD, R, t = fut.result()
            RMSDs[i] = RMSD
            Rs[i] = R
            ts[i] = t

    min_idx = int(np.argmin(RMSDs))
    RMSD = float(RMSDs[min_idx])
    R = Rs[min_idx]
    t = ts[min_idx]

    # Apply transformation(s) to undownsampled coordinates
    aligned_coords = align_coordinates(R, t, coords_B)
    return aligned_coords, RMSD, R, t

def compute_lrf(pdb_1_binding_coords, pdb_2_binding_coords, pdb_set1_single_SASA_coord, pdb_set2_single_SASA_coord, pdb_set1_single_SASA_indexes_from_site, pdb_set2_single_SASA_indexes_from_site):
    if len(pdb_1_binding_coords) <= len(pdb_2_binding_coords):
        aligned_coords_2_site, RMSD, R, t = get_downsampled_attributes_parallel(pdb_1_binding_coords, pdb_2_binding_coords)
        pdb_2_binding_coords = aligned_coords_2_site
        aligned_coords = align_coordinates(R, t, pdb_set2_single_SASA_coord)
        pdb_set2_single_SASA_coord = aligned_coords
        rotate = 2
    else:
        aligned_coords_1_site, RMSD, R, t = get_downsampled_attributes_parallel(pdb_2_binding_coords, pdb_1_binding_coords)
        pdb_1_binding_coords = aligned_coords_1_site
        aligned_coords = align_coordinates(R, t, pdb_set1_single_SASA_coord)
        pdb_set1_single_SASA_coord = aligned_coords
        rotate = 1
    
    radius = get_dynamic_radius(pdb_set1_single_SASA_coord, pdb_set2_single_SASA_coord)
    max_radius = radius
    lrf_1, lrf_2 = get_lrf(pdb_1_binding_coords, pdb_set1_single_SASA_indexes_from_site, pdb_2_binding_coords, pdb_set2_single_SASA_indexes_from_site, 10, max_radius, True)
    return lrf_1, lrf_2, max_radius, rotate, R, t

def get_lrf(site_A_coords, indices_A, aligned_site_B_coords, indices_B, k, radius, lrf_radius=True):
    k1 = max(k, int(0.9 * len(site_A_coords)))
    k2 = max(k, int(0.9 * len(aligned_site_B_coords)))
    k1 = min(k1, len(site_A_coords) - 1)
    k2 = min(k2, len(aligned_site_B_coords) - 1)
    lrf_A = LRFEstimatorPCA.get_lrf(site_A_coords, k=k1, indices=indices_A, radius=radius, use_ball_neighbors_within_radius=lrf_radius)
    lrf_B = LRFEstimatorPCA.get_lrf(aligned_site_B_coords, k=k2, indices=indices_B, radius=radius, use_ball_neighbors_within_radius=lrf_radius)
    return lrf_A, lrf_B

def get_dynamic_radius(coords_1, coords_2):
    D_site_1 = cdist(coords_1, coords_1, 'euclidean')
    D_site_2 = cdist(coords_2, coords_2, 'euclidean')
    mean_radius = max(np.mean(D_site_1), np.mean(D_site_2))
    return mean_radius*3
