from deepchem.molnet.load_function import hiv_datasets

from deepchem.models import GraphConvModel
from deepchem.data import NumpyDataset
from sklearn.metrics import average_precision_score
import numpy as np

tasks, all_datasets, transformers = hiv_datasets.load_hiv(featurizer="GraphConv")
train, valid, test = [NumpyDataset.from_DiskDataset(x) for x in all_datasets]
model = GraphConvModel(1, mode="classification")
model.fit(train)

y_true = np.squeeze(valid.y)
y_pred = model.predict(valid)[:,0,1]
print("Average Precision Score:%s" % average_precision_score(y_true, y_pred))
sorted_results = sorted(zip(y_pred, y_true), reverse=True)
hit_rate_100 = sum(x[1] for x in sorted_results[:100]) / 100
print("Hit Rate Top 100: %s" % hit_rate_100)

"""
Explanation: Screening Zinc For HIV Inhibition
In this tutorial I will walk through how to efficiently screen a large compound library with DeepChem (ZINC).  Screening a large compound library using machine learning is a CPU bound pleasingly parrellel problem.  The actual code examples I will use assume the resources available are a single very big machine (like an AWS c5.18xlarge), but should be readily swappable for other systmes (like a super computing cluster).  At a high level what we will do is...

Create a Machine Learning Model Over Labeled Data
Transform ZINC into "Work-Units"
Create an inference script which runs predictions over a "Work-Unit"
Load "Work-Unit" into a "Work Queue"
Consume work units from "Work Queue"
Gather Results

1. Train Model On Labelled Data
We are just going to knock out a simple model here.  In a real world problem you will probably try several models and do a little hyper parameter searching.
End of explanation
"""


tasks, all_datasets, transformers = hiv_datasets.load_hiv(featurizer="GraphConv", split=None)

model = GraphConvModel(1, mode="classification", model_dir="/tmp/zinc/screen_model")
model.fit(all_datasets[0])
model.save()

"""
Explanation: Retrain Model Over Full Dataset For The Screen
End of explanation
"""


import os
work_units = os.listdir('/tmp/zinc/screen')
with open('/tmp/zinc/work_queue.sh', 'w') as fout:
    fout.write("#!/bin/bash\n")
    for work_unit in work_units:
        full_path = os.path.join('/tmp/zinc', work_unit)
        fout.write("python inference.py %s" % full_path)

"""
Explanation: 2. Create Work-Units

Download All of ZINC15.

Go to http://zinc15.docking.org/tranches/home and download all non-empty tranches in .smi format.
I found it easiest to download the wget script and then run the wget script.
For the rest of this tutorial I will assume zinc was downloaded to /tmp/zinc.
The way zinc downloads the data isn't great for inference.  We want "Work-Units" which a single CPU can execute that takes a resonable amount of time (10 minutes to an hour).  To accomplish this we are going to split the zinc data into files each with 500 thousand lines.
bash
mkdir /tmp/zinc/screen
find /tmp/zinc -name '*.smi' -exec cat {} \; | grep -iv "smiles" \
     | split -l 500000 /tmp/zinc/screen/segment
This bash command
1. Finds all smi files
2. prints to stdout the contents of the file
3. removes header lines
4. splits into multiple files in /tmp/zinc/screen that are 500k molecules long
3. Creat Inference Script
Now that we have work unit we need to construct a program which ingests a work unit and logs the result.  It is important that the logging mechanism is thread safe!
For this example we will get the work unit via a file-path, and log the result to a file.
An easy extensions to distribute over multiple computers would be to get the work unit via a url, and log the results to a distributed queue.
Here is what mine looks like
inference.py
```python
import sys
import deepchem as dc
import numpy as np
from rdkit import Chem
import pickle
import os
def create_dataset(fname, batch_size=50000):
    featurizer = dc.feat.ConvMolFeaturizer()
    fin = open(fname)
    mols, orig_lines = [], []
    for line in fin:
        line = line.strip().split()
        try:
            mol = Chem.MolFromSmiles(line[0])
            if mol is None:
                continue
            mols.append(mol)
            orig_lines.append(line)
        except:
            pass
        if len(mols) > 0 and len(mols) % batch_size == 0:
            features = featurizer.featurize(mols)
            y = np.ones(shape=(len(mols), 1))
            ds = dc.data.NumpyDataset(features, y)
            yield ds, orig_lines
            mols, orig_lines = [], []
    if len(mols) > 0:
        features = featurizer.featurize(mols)
        y = np.ones(shape=(len(mols), 1))
        ds = dc.data.NumpyDataset(features, y)
        yield ds, orig_lines
def evaluate(fname):
    fout_name = "%s_out.smi" % fname
    model = dc.models.TensorGraph.load_from_dir('screen_model')
    for ds, lines in create_dataset(fname):
        y_pred = np.squeeze(model.predict(ds), axis=1)
        with open(fout_name, 'a') as fout:
            for index, line in enumerate(lines):
                line.append(y_pred[index][1])
                line = [str(x) for x in line]
                line = "\t".join(line)
                fout.write("%s\n" % line)
if name == "main":
    evaluate(sys.argv[1])
```
4. Load "Work-Unit" into a "Work Queue"
We are going to use a flat file as our distribution mechanism.  It will be a bash script calling our inference script for every work unit.  If you are at an academic institution this would be queing your jobs in pbs/qsub/slurm.  An option for cloud computing would be rabbitmq or kafka.
End of explanation
"""


from rdkit import Chem
from rdkit.Chem.Draw import IPythonConsole
from IPython.display import SVG
from rdkit.Chem.Draw import rdMolDraw2D 
best_mols = [Chem.MolFromSmiles(x.strip().split()[0]) for x in open('/tmp/zinc/screen/top_100k.smi').readlines()[:100]]
best_scores = [x.strip().split()[2] for x in open('/tmp/zinc/screen/top_100k.smi').readlines()[:100]]

print(best_scores[0])
best_mols[0]

print(best_scores[0])
best_mols[1]

print(best_scores[0])
best_mols[2]

print(best_scores[0])
best_mols[3]

"""
Explanation: 5. Consume work units from "distribution mechanism"
We will consume work units from our work queue using a very simple Process Pool.  It takes lines from our "Work Queue" and runs them, running as many processes in parrallel as we have cpus.  If  you are using a supercomputing cluster system like pbs/qsub/slurm it will take care of this for you.  The key is to use one CPU per work unit to get highest throughput.  We accomplish that here using the linux utility "taskset".
Using an c5.18xlarge on aws this will finish overnight. 
process_pool.py
```python
import multiprocessing
import sys
from multiprocessing.pool import Pool
import delegator
def run_command(args):
  q, command = args
  cpu_id = q.get()
  try:
    command = "taskset -c %s %s" % (cpu_id, command)
    print("running %s" % command)
    c = delegator.run(command)
    print(c.err)
    print(c.out)
  except Exception as e:
    print(e)
  q.put(cpu_id)
def main(n_processors, command_file):
  commands = [x.strip() for x in open(command_file).readlines()]
  commands = list(filter(lambda x: not x.startswith("#"), commands))
  q = multiprocessing.Manager().Queue()
  for i in range(n_processors):
    q.put(i)
  argslist = [(q, x) for x in commands]
  pool = Pool(processes=n_processors)
  pool.map(run_command, argslist)
if name == "main":
  processors = multiprocessing.cpu_count()
  main(processors, sys.argv[1])
```
```bash


python process_pool.py /tmp/zinc/work_queue.sh
```


6. Gather Results
Since we logged our results to *_out.smi we now need to gather all of them up and sort them by our predictions.  The resulting file wile be > 40GB.  To analyze the data further you can use dask, or put the data in a rdkit postgres cartridge.
Here I show how to join the and sort the data to get the "best" results.
```bash
find /tmp/zinc -name '*_out.smi' -exec cat {} \; > /tmp/zinc/screen/results.smi
sort -rg -k 3,3 /tmp/zinc/screen/results.smi > /tmp/zinc/screen/sorted_results.smi
Put the top 100k scoring molecules in their own file
head -n 50000 /tmp/zinc/screen/sorted_results.smi > /tmp/zinc/screen/top_100k.smi
```
/tmp/zinc/screen/top_100k.smi is now a small enough file to investigate using standard tools like pandas.
End of explanation
"""


