# Copyright (C) Insilico Medicine, Inc - All Rights Reserved
# Unauthorized copying of this file, via any medium is strictly prohibited
# Proprietary and confidential
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from darwin.molecules import FingerPrintsTransformer, FingerPrints2DTransformer, FingerPrints3DTransformer, \
    MCFTransformer, SimilarityTransformerGPU, DrugLikenessFilterTransformer, StructuralSimilarityTransformerGPU, \
    TopologicalSimilarityTransformerGPU, MCFFilterTransformer
from darwin.database.molecules.db_interface import DBInterface
import pandas
from darwin.training.flow import Transform
import numpy
from darwin.training import PickleDataset, Dataset, pipe
from darwin.summary import PipeScorer

# from my_callback import *   # Due to the bug with pipe unpickling

db = DBInterface(db_name='db_gentlemen', username='gentleman', password='pudge23', port=27018)
X, _ = db.datasets('bbb').get_x(mol_repr='smiles')
data = pandas.concat([X])["smiles"].values
X = X["smiles"].values.squeeze()
mcf = db.mcf.get_dataframe()

y = X.copy()
data1 = Dataset(X, y)

X, _ = db.datasets('bbb').get_x(mol_repr='smiles')
X = X["smiles"].values.squeeze()
y = X.copy()

data2 = Dataset(X, y)

n_jobs = 10
maximum = 0.8
minimum = 0.75


@pipe
class mcf_filter:
    data = data1
    model = MCFFilterTransformer()
    model_mcf = mcf
    model_n_jobs = n_jobs
    model_on_the_fly = True
    model_return_dataframe = False
    context = {
        'train': {'top': data1},
        'io': Transform(top=data1, bottom='druglikeness_filter'),
        'io_output': Transform(top=data1)
    }


@pipe
class druglikeness_filter:
    model = DrugLikenessFilterTransformer()
    model_n_jobs = n_jobs
    model_on_the_fly = True
    model_return_dataframe = False
    context = {
        'train': {'top': 'mcf_filter__io'},
        'io': Transform(top='mcf_filter__io', bottom='fingerprints'),
        'io_output': Transform(top='mcf_filter__io')
    }


@pipe
class fingerprints:
    model = FingerPrintsTransformer()
    model_n_jobs = n_jobs
    model_return_dataframe = False
    model_fingerprint_type_to_ndarray = True
    context = {
        'train': {'top': 'druglikeness_filter__io'},
        'io': Transform(top='druglikeness_filter__io', bottom='structural_similarity'),
        'io_output': Transform(top='druglikeness_filter__io')
    }


@pipe
class fingerprints_external:
    model = FingerPrintsTransformer()
    model_n_jobs = n_jobs
    model_return_dataframe = False
    model_fingerprint_type_to_ndarray = True
    context = {
        'train': {'top': data2},
        'io': Transform(top=data2, bottom='structural_similarity'),
        'io_output': Transform(top=data2)
    }


@pipe
class structural_similarity:
    model = StructuralSimilarityTransformerGPU()
    model_maximum = maximum
    model_return_dataframe = False
    context = {
        'train': {'top_1': 'fingerprints__io', 'top_2': 'fingerprints_external__io'},
        'io': Transform(top_1='fingerprints__io', top_2='fingerprints_external__io',
                        bottom='fingerprints2d'),
        'io_output': Transform(top_1='fingerprints__io', top_2='fingerprints_external__io')
    }


@pipe
class fingerprints2d:
    model = FingerPrints2DTransformer()
    model_n_jobs = n_jobs
    model_return_dataframe = False
    model_fingerprint_type_to_ndarray = True
    context = {
        'train': {'top': 'structural_similarity__io'},
        'io': Transform(top='structural_similarity__io', bottom='topological_similarity'),
        'io_output': Transform(top='structural_similarity__io')
    }


@pipe
class fingerprints2d_external:
    model = FingerPrints2DTransformer()
    model_n_jobs = n_jobs
    model_return_dataframe = False
    model_fingerprint_type_to_ndarray = True
    data = data1
    context = {
        'train': {'top': data2},
        'io': Transform(top=data2, bottom='topological_similarity'),
        'io_output': Transform(top=data2)
    }


@pipe
class topological_similarity:
    model = TopologicalSimilarityTransformerGPU()
    model_minimum = minimum
    model_return_dataframe = False
    context = {
        'train': {'top_1': 'fingerprints2d__io', 'top_2': 'fingerprints2d_external__io'},
        'io_output': Transform(top_1='fingerprints2d__io', top_2='fingerprints2d_external__io',
                               bottom='topological_similarity'),
        'io_scorer': Transform(top_1='fingerprints2d__io', top_2='fingerprints2d_external__io',
                               bottom='scorer')
    }


@pipe
class scorer:
    model = PipeScorer()
    model_dataset_name = 'new_pred_test'
    context = {
        'train': {
            'top': 'topological_similarity__io_scorer'
        }
    }


# result = topological_similarity.make('io_output')
result = scorer.make('train')