import csv
import io
import logging
import os
import sys

try:
  import freewilson as fw
except:
  path = os.path.abspath(os.curdir)
  sys.path.insert(0, os.path.join(path, ".."))
  import freewilson as fw

import importlib
fw = importlib.reload(fw)

from rdkit import Chem, rdBase
from rdkit.Chem import rdFMCS, rdMolAlign

PATH = os.path.join(os.path.dirname(fw.__file__), 'data')
assert os.path.exists(PATH), PATH


def test_chembl():
  logging.getLogger().setLevel(logging.INFO)
  smilesfile = os.path.join(PATH, "CHEMBL2321810.smi")
  scaffoldfile = os.path.join(PATH, "CHEMBL2321810_scaffold.mol")
  csvfile = os.path.join(PATH, "CHEMBL2321810_act.csv")
  assert os.path.exists(smilesfile)
  mols = []
  for line in open(smilesfile):
    smiles, name = line.strip().split()
    m = Chem.MolFromSmiles(smiles)
    m.SetProp("_Name", name)
    mols.append(m)

  scaffold = Chem.MolFromMolBlock(open(scaffoldfile).read())
  data = {k: float(v) for k, v in list(csv.reader(open(csvfile)))[1:]}

  scores = [data[m.GetProp("_Name")] for m in mols]
  assert mols and len(mols) == len(scores)

  with rdBase.BlockLogs():
    free = fw.FWDecompose(scaffold, mols, scores)
  # let's make sure the r squared is decent
  assert free.r2 > 0.8, str(free.r2)

  # assert we get something
  preds = list(fw.FWBuild(free))
  assert len(preds)

  # check to see that the prediction filters work
  preds2 = list(fw.FWBuild(free, pred_filter=lambda x: x > 8))
  assert len(preds2)
  assert len([p for p in preds if p.prediction > 8]) == len(list(preds2))

  # check to see that the R groups are output in order, i.e. R10 after R3
  s = io.StringIO()
  fw.predictions_to_csv(s, free, preds2)
  assert s.getvalue()

  s2 = io.StringIO(s.getvalue())
  for i, row in enumerate(csv.reader(s2)):
    if i == 0:
      assert row == ['smiles', 'prediction', 'Core_smiles', 'R1_smiles', 'R3_smiles', 'R10_smiles']
  assert i > 0


def test_multicore():
  # test that we can add rgroups for later cores and not throw an exception
  scaffolds = [Chem.MolFromSmiles("c1ccccc1[*].NC=O"), Chem.MolFromSmiles("C1CCCCC1")]
  mols = [
    Chem.MolFromSmiles(x) for x in [
      'c1ccccc1CC2CNC2C(=O)N', 'Cc1ccccc1CC2CNC2C(=O)N', 'Cc1ccccc1CC2CNCC(=O)NC2',
      'C3c1ccccc1CC2CNC2C(=O)N3', 'C1CCCCC1F', 'ClC1CCCCC1F'
    ]
  ]
  decomp = fw.FWDecompose(scaffolds, mols, [1, 2, 3, 4, 5, 6])
  s = io.StringIO()
  fw.predictions_to_csv(s, decomp, fw.FWBuild(decomp))

def test_fep_benchmark_3D():
  benchmark = os.path.join(PATH, "cmet_ligands.sdf")
  mols = list(Chem.SDMolSupplier(benchmark, removeHs=False))
  smis = {Chem.MolToSmiles(m):m for m in mols if m}
  assert smis
  match = rdFMCS.FindMCS(mols)

  scores = [float(m.GetProp("r_exp_dg")) for m in mols]
  with rdBase.BlockLogs():
    free = fw.FWDecompose(match.queryMol, mols, scores)
  preds2 = list(fw.FWBuild(free, keep_training_set=True))
  assert preds2

  count = 0
  found = 0
  # we use an rms of 0.7 because we are assembling rgroups from different
  #  ligands together which may not exactly align
  for p in preds2:
    if not p.is_training: continue
    
    smi = Chem.MolToSmiles(p.mol)

    count += 1
    if smi not in smis: continue
    rms = rdMolAlign.GetBestRMS(smis[smi], p.mol)
    assert rms < 0.7