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
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


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_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 inpaint_ligand(model, pdb_file, n_samples, ligand, fix_atoms,
                   add_n_nodes=None, center='ligand', sanitize=False,
                   largest_frag=False, relax_iter=0, timesteps=None,
                   resamplings=1, save_traj=False):
    """
    Generate ligands given a pocket
    Args:
        model: Lightning model
        pdb_file: PDB filename
        n_samples: number of samples
        ligand: reference ligand given in <chain>:<resi> format if the ligand is
                contained in the PDB file, or path to an SDF file that
                contains the ligand; used to define the pocket
        fix_atoms: ligand atoms that should be fixed, e.g. "C1 N6 C5 C12"
        center: 'ligand' or 'pocket'
        add_n_nodes: number of ligand nodes to add, sampled randomly if 'None'
        sanitize: whether to sanitize molecules or not
        largest_frag: only return the largest fragment
        relax_iter: number of force field optimization steps
        timesteps: number of denoising steps, use training value if None
        resamplings: number of resampling iterations
        save_traj: save intermediate states to visualize a denoising trajectory
    Returns:
        list of molecules
    """
    if save_traj and n_samples > 1:
        raise NotImplementedError("Can only visualize trajectory with "
                                  "n_samples=1.")
    frames = timesteps if save_traj else 1
    sanitize = False if save_traj else sanitize
    relax_iter = 0 if save_traj else relax_iter
    largest_frag = False if save_traj else largest_frag

    # Load PDB
    pdb_model = PDBParser(QUIET=True).get_structure('', pdb_file)[0]

    # Define pocket based on reference ligand
    residues = utils.get_pocket_from_ligand(pdb_model, ligand)
    pocket = model.prepare_pocket(residues, repeats=n_samples)

    # Get fixed ligand substructure
    x_fixed, one_hot_fixed = prepare_substructure(ligand, fix_atoms, pdb_model)
    n_fixed = len(x_fixed)

    if add_n_nodes is None:
        num_nodes_lig = model.ddpm.size_distribution.sample_conditional(
            n1=None, n2=pocket['size'])
        num_nodes_lig = torch.clamp(num_nodes_lig, min=n_fixed)
    else:
        num_nodes_lig = torch.ones(n_samples, dtype=int) * n_fixed + add_n_nodes

    ligand_mask = utils.num_nodes_to_batch_mask(
        len(num_nodes_lig), num_nodes_lig, model.device)

    ligand = {
        'x': torch.zeros((len(ligand_mask), model.x_dims),
                         device=model.device, dtype=FLOAT_TYPE),
        'one_hot': torch.zeros((len(ligand_mask), model.atom_nf),
                               device=model.device, dtype=FLOAT_TYPE),
        'size': num_nodes_lig,
        'mask': ligand_mask
    }

    # fill in fixed atoms
    lig_fixed = torch.zeros_like(ligand_mask)
    for i in range(n_samples):
        sele = (ligand_mask == i)

        x_new = ligand['x'][sele]
        x_new[:n_fixed] = x_fixed
        ligand['x'][sele] = x_new

        h_new = ligand['one_hot'][sele]
        h_new[:n_fixed] = one_hot_fixed
        ligand['one_hot'][sele] = h_new

        fixed_new = lig_fixed[sele]
        fixed_new[:n_fixed] = 1
        lig_fixed[sele] = fixed_new

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

    # Run sampling
    xh_lig, xh_pocket, lig_mask, pocket_mask = model.ddpm.inpaint(
        ligand, pocket, lig_fixed, center=center,
        resamplings=resamplings, timesteps=timesteps, return_frames=frames)

    # Treat intermediate states as molecules for downstream processing
    if save_traj:
        xh_lig = utils.reverse_tensor(xh_lig)
        xh_pocket = utils.reverse_tensor(xh_pocket)

        lig_mask = torch.arange(xh_lig.size(0), device=model.device
                                ).repeat_interleave(len(lig_mask))
        pocket_mask = torch.arange(xh_pocket.size(0), device=model.device
                                   ).repeat_interleave(len(pocket_mask))

        xh_lig = xh_lig.view(-1, xh_lig.size(2))
        xh_pocket = xh_pocket.view(-1, xh_pocket.size(2))

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

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

    # Build mol objects
    x = xh_lig[:, :model.x_dims].detach().cpu()
    atom_type = xh_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)
    parser.add_argument('--pdbfile', type=str)
    parser.add_argument('--ref_ligand', type=str, default=None)
    parser.add_argument('--fix_atoms', type=str, nargs='+', default=None)
    parser.add_argument('--center', type=str, default='ligand', choices={'ligand', 'pocket'})
    parser.add_argument('--outfile', type=Path)
    parser.add_argument('--n_samples', type=int, default=20)
    parser.add_argument('--add_n_nodes', type=int, default=None)
    parser.add_argument('--relax', action='store_true')
    parser.add_argument('--sanitize', action='store_true')
    parser.add_argument('--resamplings', type=int, default=20)
    parser.add_argument('--timesteps', type=int, default=50)
    parser.add_argument('--save_traj', action='store_true')
    args = parser.parse_args()

    pdb_id = Path(args.pdbfile).stem

    device = 'cuda' if torch.cuda.is_available() else 'cpu'

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

    molecules = inpaint_ligand(model, args.pdbfile, args.n_samples,
                               args.ref_ligand, args.fix_atoms,
                               args.add_n_nodes, center=args.center,
                               sanitize=args.sanitize,
                               largest_frag=False,
                               relax_iter=(200 if args.relax else 0),
                               timesteps=args.timesteps,
                               resamplings=args.resamplings,
                               save_traj=args.save_traj)

    # Make SDF files
    utils.write_sdf_file(args.outfile, molecules)