import os
import argparse, sys
import numpy as np
from Bio import PDB
import re
from scipy.spatial.transform import Rotation
import numpy as np

def compute_screw_axis_point(com_before, com_after, axis_direction, theta):
    """
    Compute a point p on the screw axis using Rodrigues's formula.

    Parameters:
        com_before (np.array): Centre of mass before transformation.
        com_after (np.array): Centre of mass after transformation.
        axis_direction (np.array): Unit vector of screw axis.
        rotation_matrix (np.array): 3x3 rotation matrix.

    Returns:
        np.array: A point on the screw axis.
    """
    # Compute rotation angle θ from the rotation matrix
    #cos_theta = (np.trace(rotation_matrix) - 1) / 2
    #theta = np.arccos(np.clip(cos_theta, -1.0, 1.0))  # Clip for numerical stability

    if np.isclose(theta, 0):  # No rotation case, just use midpoint
        return (com_before + com_after) / 2

    # Rodrigues' formula correction term
    com_shift = com_after - com_before
    correction = np.cross(com_shift, axis_direction) / np.tan(theta / 2)

    # Compute the true point p on the screw axis
    p = (com_before + com_after) / 2 + 0.5 * correction

    return p

def generate_pymol_arrows(p, com_before, com_after, axis_direction):
    """
    Generate PyMOL CGO arrow commands to visualize:
    1. The center of mass displacement.
    2. The perpendicular distances from COM to the screw axis.

    Parameters:
        com_before (np.array): Center of mass before rotation-translation.
        com_after (np.array): Center of mass after rotation-translation.
        axis_direction (np.array): The screw axis direction (unit vector).
        p: A point on the screw axis midway between initial and final positions.

    Returns:
        list: PyMOL CGO arrow commands as formatted strings.
    """
    # Step 1: Create empty list 
    arrows = []

    # Step 2: Find perpendicular projections onto the screw axis
    proj_com_before = p + axis_direction * np.dot(com_before - p, axis_direction)
    proj_com_after = p + axis_direction * np.dot(com_after - p, axis_direction)

    # Step 3: Draw an arrow from COM_before to the screw axis
    arrow_perp_before = f"cgo_arrow {[round(x, 2) for x in com_before]}, {[round(x, 2) for x in proj_com_before]}"
    arrow_params = ", 0.5, 0, -1, -1, cyan, COM1-to-axis"
    arrows.append(arrow_perp_before + arrow_params)

    # Step 4: Draw an arrow from the screw axis to the new COM_after
    arrow_perp_after = f"cgo_arrow {[round(x, 2) for x in proj_com_after]}, {[round(x, 2) for x in com_after]}"
    arrow_params = ", 0.5, 0, -1, -1, magenta, COM2-to-axis"
    arrows.append(arrow_perp_after + arrow_params)

    # Step 5: Draw the screw axis itself (extend in both directions)
    arrow_axis = f"cgo_arrow {[round(x, 2) for x in proj_com_before]}, {[round(x, 2) for x in proj_com_after]}"
    arrow_params = ", 0.5, 0, -1, -1, yellow, screw-axis"
    arrows.append(arrow_axis + arrow_params)

    # Step 6: Draw the COM movement
    arrow_COM = f"cgo_arrow {[round(x, 2) for x in com_before]},{[round(x, 2) for x in com_after]}"
    arrow_params = ", 0.5, 0, -1, -1, red, COM"
    arrows.append(arrow_COM + arrow_params)

    # Step 7: Draw COM midpoint to screw axis
    midpoint = (com_before + com_after)/2.0
    arrow_mid = f"cgo_arrow {[round(x, 2) for x in p]},{[round(x, 2) for x in midpoint]}"
    arrow_params = ", 0.5, 0, -1, -1, green, midpoint"
    arrows.append(arrow_mid + arrow_params)
    return arrows

def least_squares_fit(fixed_coords, moving_coords):
    """
    Compute the optimal rotation and translation matrices that align `moving_coords` onto `fixed_coords`

    Parameters:
        fixed_coords (numpy.ndarray): Fixed reference Nx3 array.
        moving_coords (numpy.ndarray): Moving Nx3 array to be aligned.

    Returns:
        R (numpy.ndarray): 3x3 Rotation matrix.
        t (numpy.ndarray): 1x3 Translation vector.
        aligned_coords (numpy.ndarray): Transformed coordinates of moving set.
    """
    assert fixed_coords.shape == moving_coords.shape, "Input matrices must have the same shape (Nx3)."
    assert fixed_coords.shape[0] > 3, "More than 3 atoms required for overlay."

    # Compute centroids
    centroid_fixed = np.mean(fixed_coords, axis=0)
    centroid_moving = np.mean(moving_coords, axis=0)

    # Center the points
    fixed_centered = fixed_coords - centroid_fixed
    moving_centered = moving_coords - centroid_moving

    # Compute covariance matrix
    H = moving_centered.T @ fixed_centered  # (3xN) * (Nx3) = 3x3 matrix

    # Compute SVD
    U, S, Vt = np.linalg.svd(H)

    # Compute rotation matrix
    R = Vt.T @ U.T

    # Ensure a proper right-handed coordinate system (det(R) = 1)
    if np.linalg.det(R) < 0:
        Vt[2, :] *= -1
        R = Vt.T @ U.T

    # Compute translation
    t = centroid_fixed - R @ centroid_moving

    # Apply transformation
    aligned_coords = (R @ moving_coords.T).T + t
    centroid_RT = np.mean(aligned_coords, axis=0)

    # Compute RMSD
    diff = aligned_coords - fixed_coords  # Difference between aligned and target
    rmsd = np.sqrt(np.mean(np.sum(diff**2, axis=1)))  # RMSD formula

    return R, t, rmsd, centroid_moving, centroid_RT

def extract_coordinates(atoms):
    """
    Convert a list of Bio.PDB atom objects into an Nx3 NumPy coordinate array.

    Parameters:
        atoms (list): List of Bio.PDB atom objects.

    Returns:
        numpy.ndarray: Nx3 NumPy array of atomic coordinates.
    """
    veclist = [x.get_vector() for x in atoms]
    return np.array([[v[0], v[1], v[2]] for v in veclist])

def parse_control_file(control_file):
    """ Parses a control file to extract atom types, molecule selections, and residue ranges. """
    fix_selections = []
    mov_selections = []
    align_atoms = []
    refine = True

    if not os.path.exists(control_file):
        raise FileNotFoundError(f"Error: Control file '{control_file}' not found.")

    with open(control_file, "r") as file:
        for line in file:
            parts = line.strip().split()
            if not parts:
                continue
            
            if parts[0] == "atom":
                align_atoms = parts[1:]  # List of atom types to use
            elif parts[0] == "mol2":
                fix_selections.append(parse_residue_range(parts[1:]))
            elif parts[0] == "mol1":
                mov_selections.append(parse_residue_range(parts[1:]))
            elif parts[0] == "refi":
                refine = parts[1].upper() == "Y"
            elif parts[0] == "end":
                break
    
    if not align_atoms:
        raise ValueError("Error: No atom types specified in control file.")

    return align_atoms, fix_selections, mov_selections, refine

def parse_residue_range(parts):
    """ Parses residue range information from control file (e.g., 'a30 a159'). """
    if len(parts) != 2:
        raise ValueError("Error: Residue range must have exactly two values (e.g., 'a30 a159').")

    try:
        start_res = int(re.sub(r"[a-zA-Z]", "", parts[0]))  # Extract numbers from e.g., 'a30'
        end_res = int(re.sub(r"[a-zA-Z]", "", parts[1]))  # Extract numbers from e.g., 'a159'
    except ValueError:
        raise ValueError(f"Error: Invalid residue range values: {parts[0]} {parts[1]}")

    chain_id = parts[0][0].upper()  # Extract chain letter (e.g., 'A' from 'a30')

    return {'chain': chain_id, 'residues': (start_res, end_res)}

def get_atoms(structure, chain_id, residue_range, atom_types):
    """ Extracts specified atom types from a given chain and residue range. """
    model = structure[0]  # Assume first model
    atoms = []
    
    for res_id in range(residue_range[0], residue_range[1] + 1):
        try:
            res = model[chain_id][res_id]
            if res.is_disordered():
                print(f"Warning: Disordered residue. Chain {chain_id} Residue Number {res_id}")
            for atom_name in atom_types:
                if atom_name in res:
                    atoms.append(res[atom_name])
                    # Returning an atom object, not the coords as a vector
                else:
                    print(f"Warning: Atom {atom_name} missing from residue {res_id} of chain {chain_id}.")
        except KeyError:
            print(f"Warning: Residue {res_id} in chain {chain_id} not found.")

    if not atoms:
        raise ValueError(f"Error: No valid atoms found in chain {chain_id}, range {residue_range}, using atoms {atom_types}.")

    return atoms

def compute_screw_axis(rotation_matrix, translation_vector):
    """Computes the screw axis direction, rotation angle (degrees), and screw translation magnitude."""
    # Compute cos(theta) 
    cos_theta = (np.trace(rotation_matrix) - 1) / 2  

    # Return zero values when rotation is very small
    if (1 - cos_theta) < 1e-6:
        return np.array([0, 0, 0]), 0, 0, 0 # Again, ensure scalar screw_translation_mag
    
    # Convert rotation matrix to a Rotation object
    rotation = Rotation.from_matrix(rotation_matrix)

    # Extract rotation vector (axis * angle)
    axis_angle = rotation.as_rotvec()
    
    # Compute rotation angle theta in radians
    theta = np.linalg.norm(axis_angle)  
    
    # Again, if rotation is negligible, return zero values
    if theta < 1e-6:
        return np.array([0, 0, 0]), 0, 0, 0  # Ensure screw_translation_mag is scalar
    
    # Normalize rotation axis (unit vector)
    axis_direction = axis_angle / theta  

    # Compute the translation magnitude along the screw axis 
    screw_translation = np.dot(translation_vector, axis_direction) * axis_direction
    screw_translation_mag = np.linalg.norm(screw_translation)
    
    return axis_direction, theta, screw_translation_mag, screw_translation  

def perform_lsq_fit(fixed_pdb, moving_pdb, control_file, output_pdb="RT.pdb"):
    """ Performs least-squares fitting based on control file specifications. """
    parser = PDB.PDBParser(QUIET=True)

    # Validate PDB file existence
    for pdb_file in [fixed_pdb, moving_pdb]:
        if not os.path.exists(pdb_file):
            raise FileNotFoundError(f"Error: PDB file '{pdb_file}' not found.")

    # Parse the control file
    align_atoms, fix_selections, mov_selections, refine = parse_control_file(control_file)
    
    print("Aligned atom types: ", end='')
    for i in align_atoms:
        print(f"{i} ", end='')
    print()
    print("Fixed selections")
    for i in fix_selections:
        print(i["chain"], i["residues"])

    print("Moving selections")
    for i in mov_selections:
        print(i["chain"], i["residues"])

    # Load structures
    fixed_structure = parser.get_structure("fixed", fixed_pdb)
    moving_structure = parser.get_structure("moving", moving_pdb)

    # Get atoms for alignment
    fixed_atoms = []
    moving_atoms = []
    for i in fix_selections:
        temp = get_atoms(fixed_structure, i['chain'], i['residues'], align_atoms)
        for j in temp:
            fixed_atoms.append(j)
    
    for i in mov_selections:
        temp = get_atoms(moving_structure, i['chain'], i['residues'], align_atoms)
        for j in temp:
            moving_atoms.append(j)

    if len(fixed_atoms) != len(moving_atoms):
        raise ValueError("Error: Mismatch in the number of selected atoms between fixed and moving structures.")
    else:
        print(f"Comparing {len(fixed_atoms)} atom positions.")

    if not(refine):
        sum = 0.0
        n = len(fixed_atoms)
        for i in range(n):
            distance = fixed_atoms[i] - moving_atoms[i]
            sum += distance * distance
        sum = np.sqrt(sum/n)
        print(f"RMSD between the models in place is {sum:.3f}")
        return 0

    fixed_array = extract_coordinates(fixed_atoms)
    moving_array = extract_coordinates(moving_atoms)

    # Perform LSQ fit
    # Extract transformation
    # rotation_matrix, translation_vector = super_imposer.rotran
    rotation_matrix, translation_vector, rmsdev, COMb, COMa = least_squares_fit(fixed_array, moving_array)
    print("\n=== Rotation Matrix ===")
    print(np.array_str(rotation_matrix, precision=6, suppress_small=True))
    print("\n=== Translation Vector ===")
    print(np.array_str(translation_vector, precision=3, suppress_small=True))
    print(f"\nRMSD = {rmsdev:.4f}")

    # Compute screw axis
    screw_axis, theta, screw_translation_mag, screw_translation = compute_screw_axis(rotation_matrix, translation_vector)
    kappa =  theta * 180.0 / np.pi
    print("\n=== Kappa ===")
    print(f"Kappa: {round(kappa,2)} degrees.")
    if (screw_translation_mag > 0.2):
        print("\n=== Screw Axis Direction ===")
        print(np.array_str(screw_axis, precision=6, suppress_small=True))
        print("\n=== Screw Translation Magnitude ===")
        print(f"{screw_translation_mag:.3f} Angstroms")
        print("\n=== Screw Translation Vector ===")
        print(np.array_str(screw_translation, precision=3, suppress_small=True))
        if kappa > 5:
            p = compute_screw_axis_point(COMb, COMa, screw_axis, theta)
            CGO = generate_pymol_arrows(p, COMb, COMa, screw_axis)
            for a in CGO:
                print(a)
    else:
        print("Screw translation under 0.2 Angstroms.")

    # Save aligned moving structure using BioPython 
    super_imposer = PDB.Superimposer()
    super_imposer.set_atoms(fixed_atoms, moving_atoms)
    super_imposer.apply(moving_structure.get_atoms())
    io = PDB.PDBIO()
    io.set_structure(moving_structure)
    if output_pdb is None:
        output_pdb = "RT.pdb"
    io.save(output_pdb)
    print(f"Aligned structure saved as {output_pdb}")

# Example usage:
# python3 lsqfit.py -f fixed.pdb -m moving.pdb -o RTout.pdb -c control.txt
# fixed_pdb_file = "fixed.pdb"
# moving_pdb_file = "moving.pdb"
# output_pdb   = "RTout.pdb"
# control_file = "control.txt"

def main():
    parser = argparse.ArgumentParser(
    description="Calculates LSQ fit of residues in given PDB structures.")
    parser.add_argument("-f", metavar="PDBFile1", dest="PDBFile1", help="fixed PDB structure")
    parser.add_argument("-m", metavar="PDBFile2", dest="PDBFile2", help="moving PDB structure")
    parser.add_argument("-o", metavar="PDBFile3", dest="PDBFile3", help="output PDB structure")
    parser.add_argument("-c", metavar="control_file", dest="control_file", 
                        help="Residue ranges and atom types for comparison.")

    # parse CLI
    if len(sys.argv) == 1:
        parser.print_help(sys.stderr)
        sys.exit(1)

    args = parser.parse_args()
    fixed_pdb_file = args.PDBFile1
    moving_pdb_file = args.PDBFile2
    control_file = args.control_file
    output_pdb = args.PDBFile3

    try:
        perform_lsq_fit(fixed_pdb_file, moving_pdb_file, control_file, output_pdb)
    except Exception as e:
        print(f"Terminating program: {e}")

if __name__ =='__main__':
    main()