import os
import argparse, sys
import numpy as np
from Bio import PDB
import re
import numpy as np

# This function can be modified to produce bonds with the same colour and width if required.
# Simply delete the write commands that alter the color and dash width properties.
def write_pymol_script(contact_list, model):
    modelname = os.path.splitext(os.path.basename(model))[0]
    outname = modelname + "_contacts.pml"
    red_cutoff = 3.2
    
    with open(outname, "w") as f:
        f.write(f"# Contact visualization script for {modelname}\n")
        f.write(f"# Generated by contact.py\n\n")
        
        for idx, (c1, r1, n1, a1, c2, r2, n2, a2, dist) in enumerate(contact_list):
            sel1 = f"/{modelname}//{c1}/{r1}/{a1}"
            sel2 = f"/{modelname}//{c2}/{r2}/{a2}"
            obj_name = f"contact_{idx+1}"
            
            # Write a human-readable comment
            comment = (f"# {obj_name}: {c1}{r1} {n1}.{a1} ↔ {c2}{r2} {n2}.{a2} "
                       f"[{dist:.2f} Å]\n")
            f.write(comment)
            f.write(f"distance {obj_name}, {sel1}, {sel2}\n")
            # Color it red if under threshold
            if dist < red_cutoff:
                f.write(f"color red, {obj_name}\n")
            dash_width = max(1.0, min(4.0, 5.0 - dist))
            f.write(f"set dash_width, {dash_width:.2f}, {obj_name}\n")
            f.write("\n")
        
        # Group them all at the end
        f.write(f"group {modelname}_con, contact_*\n")
    
    print(f"PyMOL script written to {outname}")

def write_latex_table(contact_list, model):
    modelname = os.path.splitext(os.path.basename(model))[0]
    outname = modelname + ".tex"

    with open(outname, "w") as f:
        f.write("\\begin{tabular}{cccc|cccc|c}\n")
        f.write("Chain & Res & ResName & Atom & Chain & Res & ResName & Atom & Distance (\AA) \\\\\n")
        f.write("\\hline\n")
        for c1, r1, n1, a1, c2, r2, n2, a2, dist in contact_list:
            f.write(f"{c1} & {r1} & {n1} & {a1} & {c2} & {r2} & {n2} & {a2} & {dist:.2f} \\\\\n")
        f.write("\\end{tabular}\n")
    print(f"LaTeX table written to {outname}")

def parse_control_file(control_file):
    """ Parses a control file to extract atom types, molecule selections, and residue ranges. """
    selections1 = []
    selections2 = []
    distance = -1.0
    contact_atoms = 0
    elem = 0

    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:
            if line.strip().startswith("#"):
                continue

            parts = line.strip().split()
            if not parts:
                continue
            
            if parts[0] == "atom":
                if parts[1]== "*":
                    contact_atoms = "*"
                else:
                    contact_atoms = parts[1:]  # List of atom types to use
            elif parts[0] == "elem":
                if parts[1]== "*":
                    elem = "*"
                else:
                    elem = parts[1:]  # List of elements to use
            elif parts[0] == "grp2":
                selections2.append(parse_residue_range(parts[1:]))
            elif parts[0] == "grp1":
                selections1.append(parse_residue_range(parts[1:]))
            elif parts[0] == "dist":
                distance = float(parts[1])
            elif parts[0] == "end":
                break
    
    if (distance < 0.):
        raise ValueError("Error: distance not specified. Try 'dist 3.5'")
    
    if contact_atoms == 0:
        if elem == 0:
            raise ValueError("Error: No atom or element types specified in control file.")
        else: 
            contact_atoms = "*"
    if elem == 0:
        elem = "*"
        
    if (len(selections1) == 0):
        raise ValueError("Error: No atoms specified in group 1.")
    
    if (len(selections2) == 0):
        raise ValueError("Error: No atoms specified in group 2.")
    
    return contact_atoms, elem, distance, selections1, selections2

def parse_residue_range(parts):
    """Parses 'a30 a159' or 'a30-159' formats into chain and residue range."""
    if len(parts) == 1 and '-' in parts[0]:
        parts = re.split('[-]', parts[0])
        chain = parts[0][0].upper()
        start_res = int(parts[0][1:])
        end_res = int(parts[1])
    elif len(parts) == 2:
        start_res = int(re.sub(r"[a-zA-Z]", "", parts[0]))
        end_res = int(re.sub(r"[a-zA-Z]", "", parts[1]))
        chain = parts[0][0].upper()
    else:
        raise ValueError("Error: Invalid residue range format.")
    return {'chain': chain, 'residues': (start_res, end_res)}

def get_atoms(structure, chain_id, residue_range, atom_types, elem):
    """ 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}")
            if (atom_types=="*"):
                newreslist = res.get_atoms()
                for i in newreslist:
                    atoms.append(i)
            else:
                for atom_name in atom_types:
                    if atom_name in res:
                        atoms.append(res[atom_name])
                #else:
                #    print(f"Warning: Atom {atom_name} missing in 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}.")
    if elem=="*":
        return atoms
    else:
        atoms2 = []
        for i in atoms:
            j = i.element
            if (j in elem):
                atoms2.append(i)
        return atoms2

def find_contacts(PDBin, control_file):
    """ Performs least-squares fitting based on control file specifications. """
    parser = PDB.PDBParser(QUIET=True)

    # Validate PDB file existence
    if not os.path.exists(PDBin):
        raise FileNotFoundError(f"Error: PDB file '{PDBin}' not found.")

    # Parse the control file
    contact_atoms, elem, distance, sel1, sel2 = parse_control_file(control_file)
    
    if contact_atoms != 0:
        print("Contact atom types: ", end='')
        for i in contact_atoms:
            print(f"{i} ", end='')
        print(" ")
    if elem != 0:
        print("Contact element types: ", end='')
        for i in elem:
            print(f"{i} ", end='')
        print(" ")
    print("selection 1")
    for i in sel1:
        print(i["chain"], i["residues"])

    print("selection 2")
    for i in sel2:
        print(i["chain"], i["residues"])

    # Load structure
    structure = parser.get_structure("PDBfile", PDBin)

    # Get atoms for alignment
    atoms1 = []
    atoms2 = []
    for i in sel1:
        temp = get_atoms(structure, i['chain'], i['residues'], contact_atoms, elem)
        for j in temp:
            atoms1.append(j)
    
    for i in sel2:
        temp = get_atoms(structure, i['chain'], i['residues'], contact_atoms, elem)
        for j in temp:
            atoms2.append(j)

    # Perform contact search
    contact_list = []
    for i in atoms1:
        for j in atoms2:
            sep = (i-j)
            if (sep < distance):
                name1 = i.get_parent().get_resname()
                str1 = i.get_full_id()[2] + str(i.get_full_id()[3][1])
                name2 = j.get_parent().get_resname()
                str2 = j.get_full_id()[2] + str(j.get_full_id()[3][1])
                print(f"{str1.rjust(6)}:{name1} {i.fullname} ---", end='')
                print(f"{str2.rjust(6)}:{name2} {j.fullname}  {sep:.3f}")
                res1 = i.get_parent()
                res2 = j.get_parent()
                chain1 = i.get_full_id()[2]
                resi1 = i.get_full_id()[3][1]
                atom1 = i.get_name()
                name1 = res1.get_resname()
                chain2 = j.get_full_id()[2]
                resi2 = j.get_full_id()[3][1]
                atom2 = j.get_name()
                name2 = res2.get_resname()
                contact_list.append((chain1, resi1, name1, atom1,chain2, resi2, name2, atom2,round(sep, 2)))
    write_latex_table(contact_list, PDBin)
    write_pymol_script(contact_list, PDBin)

def main():
    parser = argparse.ArgumentParser(
    description="Calculates contacts between atom selections in a given PDB structure.")
    parser.add_argument("-f", metavar="PDBFile1", dest="PDBFile1", help="Input 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()
    PDBin = args.PDBFile1
    control_file = args.control_file

    try:
        find_contacts(PDBin, control_file)
    except Exception as e:
        print(f"Terminating program: {e}")

if __name__ =='__main__':
    main()