import argparse
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from Bio.PDB import PDBParser
from rdkit import Chem
import pandas as pd
import random
from torch_scatter import scatter_mean
from openbabel import openbabel
openbabel.obErrorLog.StopLogging()  # suppress OpenBabel messages

import utils
from lightning_modules import LigandPocketDDPM
from constants import FLOAT_TYPE, INT_TYPE
from analysis.molecule_builder import build_molecule, process_molecule
from analysis.metrics import MoleculeProperties


def prepare_from_sdf_files(sdf_files, atom_encoder):

    ligand_coords = []
    atom_one_hot = []
    for file in sdf_files:
        rdmol = Chem.SDMolSupplier(str(file), sanitize=False)[0]
        ligand_coords.append(
            torch.from_numpy(rdmol.GetConformer().GetPositions()).float()
        )
        types = torch.tensor([atom_encoder[a.GetSymbol()] for a in rdmol.GetAtoms()])
        atom_one_hot.append(
            F.one_hot(types, num_classes=len(atom_encoder))
        )

    return torch.cat(ligand_coords, dim=0), torch.cat(atom_one_hot, dim=0)


def prepare_ligands_from_mols(mols, atom_encoder, device='cpu'):

    ligand_coords = []
    atom_one_hots = []
    masks = []
    sizes = []
    for i, mol in enumerate(mols):
        coord = torch.tensor(mol.GetConformer().GetPositions(), dtype=FLOAT_TYPE)
        types = torch.tensor([atom_encoder[a.GetSymbol()] for a in mol.GetAtoms()], dtype=INT_TYPE)
        one_hot = F.one_hot(types, num_classes=len(atom_encoder))
        mask = torch.ones(len(types), dtype=INT_TYPE) * i
        ligand_coords.append(coord)
        atom_one_hots.append(one_hot)
        masks.append(mask)
        sizes.append(len(types))

    ligand = {
        'x': torch.cat(ligand_coords, dim=0).to(device),
        'one_hot': torch.cat(atom_one_hots, dim=0).to(device),
        'size': torch.tensor(sizes, dtype=INT_TYPE).to(device),
        'mask': torch.cat(masks, dim=0).to(device),
    }

    return ligand


def prepare_ligand_from_pdb(biopython_atoms, atom_encoder):

    coord = torch.tensor(np.array([a.get_coord()
                                   for a in biopython_atoms]), dtype=FLOAT_TYPE)
    types = torch.tensor([atom_encoder[a.element.capitalize()]
                          for a in biopython_atoms])
    one_hot = F.one_hot(types, num_classes=len(atom_encoder))

    return coord, one_hot


def prepare_substructure(ref_ligand, fix_atoms, pdb_model):

    if fix_atoms[0].endswith(".sdf"):
        # ligand as sdf file
        coord, one_hot = prepare_from_sdf_files(fix_atoms, model.lig_type_encoder)

    else:
        # ligand contained in PDB; given in <chain>:<resi> format
        chain, resi = ref_ligand.split(':')
        ligand = utils.get_residue_with_resi(pdb_model[chain], int(resi))
        fixed_atoms = [a for a in ligand.get_atoms() if a.get_name() in set(fix_atoms)]
        coord, one_hot = prepare_ligand_from_pdb(fixed_atoms, model.lig_type_encoder)

    return coord, one_hot


def diversify_ligands(model, pocket, mols, timesteps,
                    sanitize=False,
                    largest_frag=False,
                    relax_iter=0):
    """
    Diversify ligands for a specified pocket.
    
    Parameters:
        model: The model instance used for diversification.
        pocket: The pocket information including coordinates and types.
        mols: List of RDKit molecule objects to be diversified.
        timesteps: Number of denoising steps to apply during diversification.
        sanitize: If True, performs molecule sanitization post-generation (default: False).
        largest_frag: If True, only the largest fragment of the generated molecule is returned (default: False).
        relax_iter: Number of iterations for force field relaxation of the generated molecules (default: 0).
    
    Returns:
        A list of diversified RDKit molecule objects.
    """

    ligand = prepare_ligands_from_mols(mols, model.lig_type_encoder, device=model.device)

    pocket_mask = pocket['mask']
    lig_mask = ligand['mask']

    # Pocket's center of mass
    pocket_com_before = scatter_mean(pocket['x'], pocket['mask'], dim=0)

    out_lig, out_pocket, _, _ = model.ddpm.diversify(ligand, pocket, noising_steps=timesteps)

    # Move generated molecule back to the original pocket position
    pocket_com_after = scatter_mean(out_pocket[:, :model.x_dims], pocket_mask, dim=0)

    out_pocket[:, :model.x_dims] += \
        (pocket_com_before - pocket_com_after)[pocket_mask]
    out_lig[:, :model.x_dims] += \
        (pocket_com_before - pocket_com_after)[lig_mask]

    # Build mol objects
    x = out_lig[:, :model.x_dims].detach().cpu()
    atom_type = out_lig[:, model.x_dims:].argmax(1).detach().cpu()

    molecules = []
    for mol_pc in zip(utils.batch_to_list(x, lig_mask),
                      utils.batch_to_list(atom_type, lig_mask)):

        mol = build_molecule(*mol_pc, model.dataset_info, add_coords=True)
        mol = process_molecule(mol,
                               add_hydrogens=False,
                               sanitize=sanitize,
                               relax_iter=relax_iter,
                               largest_frag=largest_frag)
        if mol is not None:
            molecules.append(mol)

    return molecules


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--checkpoint', type=Path, default='checkpoints/crossdocked_fullatom_cond.ckpt')
    parser.add_argument('--pdbfile', type=str, default='example/5ndu.pdb')
    parser.add_argument('--ref_ligand', type=str, default='example/5ndu_linked_mols.sdf')
    parser.add_argument('--objective', type=str, default='sa', choices={'qed', 'sa'})
    parser.add_argument('--timesteps', type=int, default=100)
    parser.add_argument('--population_size', type=int, default=100)
    parser.add_argument('--evolution_steps', type=int, default=10)
    parser.add_argument('--top_k', type=int, default=7)
    parser.add_argument('--outfile', type=Path, default='output.sdf')
    parser.add_argument('--relax', action='store_true')


    args = parser.parse_args()

    pdb_id = Path(args.pdbfile).stem

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    population_size = args.population_size
    evolution_steps = args.evolution_steps
    top_k = args.top_k

    # Load model
    model = LigandPocketDDPM.load_from_checkpoint(
        args.checkpoint, map_location=device)
    model = model.to(device)

    # Prepare ligand + pocket
    # Load PDB
    pdb_model = PDBParser(QUIET=True).get_structure('', args.pdbfile)[0]
    # Define pocket based on reference ligand
    residues = utils.get_pocket_from_ligand(pdb_model, args.ref_ligand)
    pocket = model.prepare_pocket(residues, repeats=population_size)


    if args.objective == 'qed':
        objective_function = MoleculeProperties().calculate_qed
    elif args.objective == 'sa':
        objective_function = MoleculeProperties().calculate_sa
    else:
        ### IMPLEMENT YOUR OWN OBJECTIVE
        ### FUNCTIONS HERE 
        raise ValueError(f"Objective function {args.objective} not recognized.")

    ref_mol = Chem.SDMolSupplier(args.ref_ligand)[0]

    # Store molecules in history dataframe 
    buffer = pd.DataFrame(columns=['generation', 'score', 'fate' 'mol', 'smiles'])

    # Population initialization
    buffer = buffer.append({'generation': 0,
                            'score': objective_function(ref_mol),
                            'fate': 'initial', 'mol': ref_mol,
                            'smiles': Chem.MolToSmiles(ref_mol)}, ignore_index=True)

    for generation_idx in range(evolution_steps):

        if generation_idx == 0:
            molecules = buffer['mol'].tolist() * population_size
        else:
            # Select top k molecules from previous generation
            previous_gen = buffer[buffer['generation'] == generation_idx]
            top_k_molecules = previous_gen.nlargest(top_k, 'score')['mol'].tolist()
            molecules = top_k_molecules * (population_size // top_k)

            # Update the fate of selected top k molecules in the buffer
            buffer.loc[buffer['generation'] == generation_idx, 'fate'] = 'survived'

            # Ensure the right number of molecules
            if len(molecules) < population_size:
                molecules += [random.choice(molecules) for _ in range(population_size - len(molecules))]


        # Diversify molecules
        assert len(molecules) == population_size, f"Wrong number of molecules: {len(molecules)} when it should be {population_size}"
        print(f"Generation {generation_idx}, mean score: {np.mean([objective_function(mol) for mol in molecules])}")
        molecules = diversify_ligands(model,
                                    pocket,
                                    molecules,
                                timesteps=args.timesteps,
                                sanitize=True,
                                relax_iter=(200 if args.relax else 0))
        
        
        # Evaluate and save molecules
        for mol in molecules:
            buffer = buffer.append({'generation': generation_idx + 1,
            'score': objective_function(mol),
            'fate': 'purged',
            'mol': mol,
            'smiles': Chem.MolToSmiles(mol)}, ignore_index=True)


    # Make SDF files
    utils.write_sdf_file(args.outfile, molecules)
    # Save buffer
    buffer.drop(columns=['mol'])
    buffer.to_csv(args.outfile.with_suffix('.csv'))