import os

import re

import tempfile

import numpy as np

import torch

from pathlib import Path

import argparse

import pandas as pd

from rdkit import Chem

from tqdm import tqdm

try:

    import utils

except ModuleNotFoundError as e:

    print(e)

def calculate_smina_score(pdb_file, sdf_file):

    # add '-o <name>_smina.sdf' if you want to see the output

    out = os.popen(f'smina.static -l {sdf_file} -r {pdb_file} '

                   f'--score_only').read()

    matches = re.findall(

        r"Affinity:[ ]+([+-]?[0-9]*[.]?[0-9]+)[ ]+\(kcal/mol\)", out)

    return [float(x) for x in matches]

def smina_score(rdmols, receptor_file):

    """

    Calculate smina score

    :param rdmols: List of RDKit molecules

    :param receptor_file: Receptor pdb/pdbqt file or list of receptor files

    :return: Smina score for each input molecule (list)

    """

    if isinstance(receptor_file, list):

        scores = []

        for mol, rec_file in zip(rdmols, receptor_file):

            with tempfile.NamedTemporaryFile(suffix='.sdf') as tmp:

                tmp_file = tmp.name

                utils.write_sdf_file(tmp_file, [mol])

                scores.extend(calculate_smina_score(rec_file, tmp_file))

    # Use same receptor file for all molecules

    else:

        with tempfile.NamedTemporaryFile(suffix='.sdf') as tmp:

            tmp_file = tmp.name

            utils.write_sdf_file(tmp_file, rdmols)

            scores = calculate_smina_score(receptor_file, tmp_file)

    return scores

def sdf_to_pdbqt(sdf_file, pdbqt_outfile, mol_id):

    os.popen(f'obabel {sdf_file} -O {pdbqt_outfile} '

             f'-f {mol_id + 1} -l {mol_id + 1}').read()

    return pdbqt_outfile

def calculate_qvina2_score(receptor_file, sdf_file, out_dir, size=20,

                           exhaustiveness=16, return_rdmol=False):

    receptor_file = Path(receptor_file)

    sdf_file = Path(sdf_file)

    if receptor_file.suffix == '.pdb':

        # prepare receptor, requires Python 2.7

        receptor_pdbqt_file = Path(out_dir, receptor_file.stem + '.pdbqt')

        os.popen(f'prepare_receptor4.py -r {receptor_file} -O {receptor_pdbqt_file}')

    else:

        receptor_pdbqt_file = receptor_file

    scores = []

    rdmols = []  # for if return rdmols

    suppl = Chem.SDMolSupplier(str(sdf_file), sanitize=False)

    for i, mol in enumerate(suppl):  # sdf file may contain several ligands

        ligand_name = f'{sdf_file.stem}_{i}'

        # prepare ligand

        ligand_pdbqt_file = Path(out_dir, ligand_name + '.pdbqt')

        out_sdf_file = Path(out_dir, ligand_name + '_out.sdf')

        if out_sdf_file.exists():

            with open(out_sdf_file, 'r') as f:

                scores.append(

                    min([float(x.split()[2]) for x in f.readlines()

                         if x.startswith(' VINA RESULT:')])

                )

        else:

            sdf_to_pdbqt(sdf_file, ligand_pdbqt_file, i)

            # center box at ligand's center of mass

            cx, cy, cz = mol.GetConformer().GetPositions().mean(0)

            # run QuickVina 2

            out = os.popen(

                f'qvina2.1 --receptor {receptor_pdbqt_file} '

                f'--ligand {ligand_pdbqt_file} '

                f'--center_x {cx:.4f} --center_y {cy:.4f} --center_z {cz:.4f} '

                f'--size_x {size} --size_y {size} --size_z {size} '

                f'--exhaustiveness {exhaustiveness}'

            ).read()

            # clean up

            ligand_pdbqt_file.unlink()

            if '-----+------------+----------+----------' not in out:

                scores.append(np.nan)

                continue

            out_split = out.splitlines()

            best_idx = out_split.index('-----+------------+----------+----------') + 1

            best_line = out_split[best_idx].split()

            assert best_line[0] == '1'

            scores.append(float(best_line[1]))

            out_pdbqt_file = Path(out_dir, ligand_name + '_out.pdbqt')

            if out_pdbqt_file.exists():

                os.popen(f'obabel {out_pdbqt_file} -O {out_sdf_file}').read()

                # clean up

                out_pdbqt_file.unlink()

        if return_rdmol:

            rdmol = Chem.SDMolSupplier(str(out_sdf_file))[0]

            rdmols.append(rdmol)

    if return_rdmol:

        return scores, rdmols

    else:

        return scores

if __name__ == '__main__':

    parser = argparse.ArgumentParser('QuickVina evaluation')

    parser.add_argument('--pdbqt_dir', type=Path,

                        help='Receptor files in pdbqt format')

    parser.add_argument('--sdf_dir', type=Path, default=None,

                        help='Ligand files in sdf format')

    parser.add_argument('--sdf_files', type=Path, nargs='+', default=None)

    parser.add_argument('--out_dir', type=Path)

    parser.add_argument('--write_csv', action='store_true')

    parser.add_argument('--write_dict', action='store_true')

    parser.add_argument('--dataset', type=str, default='moad')

    args = parser.parse_args()

    assert (args.sdf_dir is not None) ^ (args.sdf_files is not None)

    args.out_dir.mkdir(exist_ok=True)

    results = {'receptor': [], 'ligand': [], 'scores': []}

    results_dict = {}

    sdf_files = list(args.sdf_dir.glob('[!.]*.sdf')) \

        if args.sdf_dir is not None else args.sdf_files

    pbar = tqdm(sdf_files)

    for sdf_file in pbar:

        pbar.set_description(f'Processing {sdf_file.name}')

        if args.dataset == 'moad':

            """

            Ligand file names should be of the following form:

            <receptor-name>_<pocket-id>_<some-suffix>.sdf

            where <receptor-name> and <pocket-id> cannot contain any 

            underscores, e.g.: 1abc-bio1_pocket0_gen.sdf

            """

            ligand_name = sdf_file.stem

            receptor_name, pocket_id, *suffix = ligand_name.split('_')

            suffix = '_'.join(suffix)

            receptor_file = Path(args.pdbqt_dir, receptor_name + '.pdbqt')

        elif args.dataset == 'crossdocked':

            ligand_name = sdf_file.stem

            receptor_name = ligand_name[:-4]

            receptor_file = Path(args.pdbqt_dir, receptor_name + '.pdbqt')

        # try:

        scores, rdmols = calculate_qvina2_score(

            receptor_file, sdf_file, args.out_dir, return_rdmol=True)

        # except AttributeError as e:

        #     print(e)

        #     continue

        results['receptor'].append(str(receptor_file))

        results['ligand'].append(str(sdf_file))

        results['scores'].append(scores)

        if args.write_dict:

            results_dict[ligand_name] = {

                'receptor': str(receptor_file),

                'ligand': str(sdf_file),

                'scores': scores,

                'rmdols': rdmols

            }

    if args.write_csv:

        df = pd.DataFrame.from_dict(results)

        df.to_csv(Path(args.out_dir, 'qvina2_scores.csv'))

    if args.write_dict:

        torch.save(results_dict, Path(args.out_dir, 'qvina2_scores.pt'))