from pathlib import Path
from collections import defaultdict
import logging
import typer
from typing import Annotated

import Bio
from Bio import Align
from Bio.PDB.MMCIFParser import MMCIFParser
from Bio.PDB.MMCIF2Dict import MMCIF2Dict
from Bio.SeqUtils import seq1


from pdbfixer import PDBFixer
from openmm import (
    app,
    unit,
    LangevinIntegrator,
    CustomExternalForce,
)


app = typer.Typer(help='MMCIF tools to Get Protein Sequence or Convert to PDB')


logger = logging.getLogger(__name__)


def convert_mmcif_to_pdb(infile: str, find_missing_residues: bool = False) -> None:
    fixer = PDBFixer(infile)

    prefix = Path(infile).stem
    parent = Path(infile).parent

    if find_missing_residues:
        fixer.findMissingResidues()
    else:
        fixer.missingResidues = {}  # bypass logic

    fixer.findMissingAtoms()
    fixer.addMissingAtoms()
    fixer.addMissingHydrogens()

    # solvate
    fixer.addSolvent(
        padding= 1.0 * unit.nanometer,
        ionicStrength= 0.15 * unit.molar
    )

    # TODO: what about ligand?

    # Create OpenMM system
    forcefield = app.ForceField(
        "amber14-all.xml",
        "amber14/tip3pfb.xml"
    )

    system = forcefield.createSystem(
        fixer.topology,
        nonbondedMethod=app.PME,
        nonbondedCutoff=1 * unit.nanometer,
        constraints=app.HBonds
    )

    # heavy-atom restraints
    restraint = CustomExternalForce("k*((x-x0)^2+(y-y0)^2+(z-z0)^2)")
    restraint.addGlobalParameter("k", 1000 * unit.kilojoule_per_mole/unit.nanometer**2)

    for p in ["x0", "y0", "z0"]:
        restraint.addPerParticleParameter(p)

    atoms = list(fixer.topology.atoms())
    for i, pos in enumerate(fixer.positions):
        if atoms[i].element.symbol != "H":
            restraint.addParticle(i, pos.value_in_unit(unit.nanometer))
    system.addForce(restraint)

    integrator = LangevinIntegrator(
        300 * unit.kelvin,
        1 / unit.picosecond,
        0.002 * unit.picoseconds
    )

    simulation = app.Simulation(
        fixer.topology,
        system,
        integrator
    )

    simulation.context.setPositions(fixer.positions)

    logger.info("energy minimization...")
    simulation.minimizeEnergy(
        tolerance = 0.1 * unit.kilojoule_per_mole / unit.nanometer,
        maxIterations= 2000
    )
    state = simulation.context.getState(getPositions=True)

    # Save minimized structure (without solvent)
    modeller = app.Modeller(simulation.topology, 
                            state.getPositions())

    # remove solvent/ions
    to_delete = []
    for res in modeller.topology.residues():
        if res.name in ["HOH", "WAT", "TIP3", "TIP3P", "NA", "CL", "K"]:
            to_delete.append(res)
    modeller.delete(to_delete)

    with open("minimized.pdb", "w") as f:
        app.PDBFile.writeFile(modeller.topology, modeller.positions, f)

    logger.info("energy minimization completed")

    # # Short relaxation MD
    # print("Short relaxation MD")
    # simulation.context.setVelocitiesToTemperature(300*unit.kelvin)
    # for i, k in enumerate([1000.0, 100.0, 10.0, 1.0]):
    #     print(f"NVT with k={k} kJ/mol/nm**2")
    #     restraint.setGlobalParameterDefaultValue(
    #         0, 
    #         k * unit.kilojoule_per_mole/unit.nanometer**2,
    #         )
    #     simulation.step(10000)  # ~20 ps
    #     state = simulation.context.getState(getPositions=True)
    #     with open(f"relaxed_{i}.pdb", "w") as f:
    #         app.PDBFile.writeFile(
    #             simulation.topology,
    #             state.getPositions(),
    #             f
    #         )
        

"""
    On chain ID:

    🔹 _atom_site.label_asym_id : PDB-assigned (standardized) chain ID

    Generated by the wwPDB during deposition/validation.
    Designed to be unique, consistent, and unambiguous.
    Ensures that every entity (polymer, ligand, ion, water) gets a clean chain identifier.
    Used by downstream software and many standardized pipelines.
    Stable even when authors use duplicate or non-standard chain IDs.
    This is the canonical chain ID within the mmCIF data model.

    🔹 _atom_site.auth_asym_id - Author-provided chain ID

    Comes exactly as the experimentalists wrote it in their original model or PDB file.
    May contain:
    Duplicate chain IDs
    Empty chain IDs
    Non-alphanumeric characters
    Chain IDs that differ from PDB conventions
    Useful for matching original publications, manuscripts, or deposited scripts.
    This is the raw chain ID as the authors intended.

    🔍 Why mmCIF stores both?
    PDB format (old .pdb files) allowed only a single-letter chain ID and had many legacy constraints. Authors often used quirky identifiers or had structural reasons for using duplicates.
    mmCIF removes the ambiguity by storing both:
    auth_asym_id → for historical and scientific reproducibility
    label_asym_id → for clean programmatic use
"""

def test_chain_id(path: str | Path) -> None:
    mm = MMCIF2Dict(path)
    print("_atom_site.auth_asym_id" in mm)   # True/False
    print("_atom_site.label_asym_id" in mm)  # True/False


def populate_atom_xtra(structure, cif_path: str | Path) -> None:
    """
    Populate atom.xtra['label_asym_id'] and ['auth_asym_id'] using _atom_site columns.
    Call this immediately after parsing (before any filtering/copying).
    """
    mm = MMCIF2Dict(cif_path)

    label_asym = mm.get("_atom_site.label_asym_id")
    auth_asym  = mm.get("_atom_site.auth_asym_id")
    
    if not (label_asym and auth_asym):
        raise ValueError("CIF is missing _atom_site.label_asym_id or _atom_site.auth_asym_id.")

    # Optionally attach a few other useful columns if present
    extra_cols = {
        "id":               "_atom_site.id",
        "label_atom_id":    "_atom_site.label_atom_id",
        "auth_atom_id":     "_atom_site.auth_atom_id",
        "label_comp_id":    "_atom_site.label_comp_id",
        "auth_comp_id":     "_atom_site.auth_comp_id",
        "label_seq_id":     "_atom_site.label_seq_id",
        "auth_seq_id":      "_atom_site.auth_seq_id",
        "type_symbol":      "_atom_site.type_symbol",
    }
    extra_data = {k: mm[v] for k, v in extra_cols.items() if v in mm}

    # Attach by order
    n = len(label_asym)
    i = 0
    for atom in structure.get_atoms():
        if i >= n:
            break
        atom.xtra["label_asym_id"] = label_asym[i]
        atom.xtra["auth_asym_id"]  = auth_asym[i]
        for k, arr in extra_data.items():
            atom.xtra[k] = arr[i]
        i += 1

    if i != n:
        # If counts differ, the atom order was altered (e.g., you removed waters/altlocs before calling this).
        # In that case, call this function immediately after parsing, or use the robust matching method below.
        print(f"Warning: attached to {i} atoms, but CIF has {n} _atom_site rows.")
    


def get_nonpoly_chains_and_residues(path: str | Path, model_index: int = 0) -> dict:
    parser = MMCIFParser(QUIET=True)
    structure = parser.get_structure("struct", path)
    populate_atom_xtra(structure, path)

    nonpoly_residues = defaultdict(list)
    for mi, model in enumerate(structure):
        if mi != model_index:
            continue
        for chain in model:
            for residue in chain:
                hetflag = residue.id[0]
                # In Biopython:
                # Protein/nucleic acid residues → residue.id[0] == " " (a blank space)
                # Hetero residues (HETATM) → residue.id[0] != " " (ex. H_GDP, H_MG)
                # Examples: ligands, ions, cofactors, modified residues
                # Water molecules → residue.id[0] == "W"
                if hetflag.strip():  # HET or water (W)
                    # pick any atom from the residue
                    atom = next(residue.get_atoms())
                    label_asym = atom.xtra.get("label_asym_id", None)
                    auth_asym = atom.xtra.get("auth_asym_id", None)
                    # print(f"label= {label_asym} auth={auth_asym} chain={chain.id}")
                    chainId = label_asym or auth_asym or chain.id
                    nonpoly_residues[chainId].append({
                        "name": residue.get_resname(),    # three-letter code
                        "resseq": residue.id[1],
                        # "icode": residue.id[2],
                        # "hetflag": hetflag
                    })

    return dict(nonpoly_residues)



def get_residue_poly_sequences(path: str | Path, model_index: int = 0) -> dict:
    """Return sequences per chain from ATOM records as one-letter strings.

    Output Example:
        {'A': 'MTEYKLVVVGADGVGKSALTIQLIQNHFVDEYDPTIEDSYRKQVVIDGETCLLDILDTASAMRDQYMRTGEGFLCV
        FAINNTKSFEDIHHYREQIKRVKDSEDVPMVLVGNKCDLPSRTVDTKQAQDLARSYGIPFIETSAKTRQGVDDAFYTLVREIRKHKEK'}

    """
    parser = MMCIFParser(QUIET=True)
    structure = parser.get_structure("struct", path)
    model = list(structure)[model_index]

    seqs = {}
    for chain in model:
        # pick any atom from the chain
        atom = next(chain.get_atoms())
        auth_asym = atom.xtra.get("auth_asym_id", None)
        label_asym = atom.xtra.get("label_asym_id", None)
        chainId = label_asym or auth_asym or chain.id

        letters = []
        seen_resseq_ins = set()
        for res in chain:
            # skip non-standard groups (water, ligands)
            hetflag = res.id[0]
            if hetflag not in (" ", "W"):  # " " = standard residue, "W" = water
                continue
            # use resseq+icode to avoid duplicates from altlocs
            key = (res.id[1], res.id[2])
            if key in seen_resseq_ins:
                continue
            seen_resseq_ins.add(key)
            resname = res.get_resname()
            try:
                letters.append(seq1(resname, custom_map={"SEC":"C","MSE":"M"}))  # extend map as needed
            except KeyError:
                letters.append("X")
        seqs[chainId] = "".join(letters)

    return seqs



def get_entity_poly_sequences(path):
    three_to_one = {
        # minimal common mapping; extend as needed
        "ALA":"A","ARG":"R","ASN":"N","ASP":"D","CYS":"C","GLN":"Q","GLU":"E","GLY":"G",
        "HIS":"H","ILE":"I","LEU":"L","LYS":"K","MET":"M","PHE":"F","PRO":"P","SER":"S",
        "THR":"T","TRP":"W","TYR":"Y","VAL":"V","SEC":"U","PYL":"O","MSE":"M"
    }

    mm = MMCIF2Dict(path)

    # Rows of entity_poly_seq are parallel lists:
    entity_ids = mm["_entity_poly_seq.entity_id"]
    mon_ids    = mm["_entity_poly_seq.mon_id"]
    num_list   = mm["_entity_poly_seq.num"]  # residue index within the polymer

    # Group by entity_id and order by .num
    per_entity = {}
    for e, mon, num in zip(entity_ids, mon_ids, num_list):
        per_entity.setdefault(e, []).append((int(num), mon))

    for e in per_entity:
        per_entity[e].sort(key=lambda x: x[0])

    # Convert to one-letter
    per_entity_seq = {e: "".join(three_to_one.get(mon, "X") for _, mon in rows)
                      for e, rows in per_entity.items()}

    # Map entity → chain(s) using _entity_poly.pdbx_strand_id (can be comma-separated)
    strand_id_key = "_entity_poly.pdbx_strand_id"
    entity_id_key = "_entity_poly.entity_id"
    chains_by_entity = {}
    if strand_id_key in mm and entity_id_key in mm:
        for e, strands in zip(mm[entity_id_key], mm[strand_id_key]):
            # strands can be like "A,B" or "A"
            chains = [s.strip() for s in str(strands).split(",")]
            chains_by_entity[e] = chains

    # Expand entity sequences to per-chain sequences
    per_chain_seq = {}
    for e, seq in per_entity_seq.items():
        chains = chains_by_entity.get(e, [e])  # fall back to entity id
        for ch in chains:
            per_chain_seq[ch] = seq

    return per_chain_seq



def get_info(path: str | Path) -> None:
    print("Biopython version:", Bio.__version__)
    poly = get_residue_poly_sequences(path)
    non_poly = get_nonpoly_chains_and_residues(path)
    for k, v in poly.items():
        print(f'id: {k}\nsequence: {v}')

    for k, arr in non_poly.items():
        print(f'\nid: {k}')
        for w in arr:
            for kk, vv in w.items():
                print(f'  {kk}: {vv}', end=" ")
        print()



@app.command()
def cif2seq(filename: Annotated[str, typer.Argument(help="Input .cif filename.")]):
    """Get missing residue(s) and sequence from a MMCIF(.cif) file.

    - Full sequence
    - Coordinate sequence with missing residues (`-`)

    Args:
        cif_path (str): .cif file path.
    """
    coor_seq = get_residue_poly_sequences(filename)
    auth_seq = get_entity_poly_sequences(filename)

    aligner = Align.PairwiseAligner()
    aligner.mode = 'local'
    aligner.match_score = 1.0
    aligner.open_gap_score = -1
    aligner.extend_gap_score = 0
    
    # Perform global alignment (simple, without scoring)
    alignments = aligner.align(auth_seq['1'], coor_seq['A'])
    alignment = alignments[0]
    print(alignment[0])
    print(alignment[1])
    print(f'aligner algorithm= {aligner.algorithm}')
    print(f'  length={alignment.length}')
    print(f'  aligned={alignment.aligned}')

    ligand = get_nonpoly_chains_and_residues(filename)
    print(ligand)



@app.command()
def cif2pdb(filename: Annotated[str, typer.Argument(help="Input .cif filename.")]):
    """Convert to PDB file.

    OpenMM energy minimization with heavy atom restraints

    Args:
        cif_path (str): .cif file path.
    """
    print(f'converting {filename} to .pdb')
    convert_mmcif_to_pdb(filename)



if __name__ == '__main__':
    app()