from io import BytesIO
import os
from pathlib import Path
import sys

import gradio as gr
import pandas as pd
from rdkit import Chem
from rdkit.Chem import RDConfig, Descriptors, Lipinski, Crippen

sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score'))
import sascorer


root = Path.cwd()
task_list = [f.stem for f in root.parent.joinpath("configs/task").iterdir() if f.suffix == ".yaml"]
preset_list = [f.stem for f in root.parent.joinpath("configs/preset").iterdir() if f.suffix == ".yaml"]
predictor_list = [f.stem for f in root.parent.joinpath("configs/model/predictor").iterdir() if f.suffix == ".yaml"]
drug_encoder_list = [f.stem for f in root.parent.joinpath("configs/model/predictor/drug_encoder").iterdir() if
                     f.suffix == ".yaml"]
drug_featurizer_list = [f.stem for f in root.parent.joinpath("configs/data/drug_featurizer").iterdir() if
                        f.suffix == ".yaml"]
protein_encoder_list = [f.stem for f in root.parent.joinpath("configs/model/predictor/protein_encoder").iterdir() if
                        f.suffix == ".yaml"]
protein_featurizer_list = [f.stem for f in root.parent.joinpath("configs/data/protein_featurizer").iterdir() if
                           f.suffix == ".yaml"]
classifier_list = [f.stem for f in root.parent.joinpath("configs/model/predictor/decoder").iterdir() if
                   f.suffix == ".yaml"]


def load_csv(file):
    return pd.read_csv(BytesIO(file))


def drug_target_predict(
        predict_data,
        task,
        preset,
):
    try:
        return load_csv(predict_data)
    except Exception as e:
        raise gr.Error(str(e))


# Define a function that takes a CSV output and a list of analytical utility functions as inputs
def analyze_csv(df_predict, score_list, filter_list):
    df_report = df_predict.copy()
    try:
        # Loop through the list of functions and apply them to the dataframe
        for filter_name in filter_list:
            gr.Info(f'Applying {filter_name}...')

        for score_name in score_list:
            gr.Info(f'Calculating {score_name}...')
            # Apply the function to the dataframe and assign the result to a new column
            df_report[score_name] = df_report.apply(score_map[score_name], axis=1)
        # Return the dataframe as a table
        return df_report
    except Exception as e:
        raise gr.Error(str(e))


def sa_score(row):
    return sascorer.calculateScore(Chem.MolFromSmiles(row['X1']))


def mw(row):
    return Chem.Descriptors.MolWt(Chem.MolFromSmiles(row['X1']))


def hbd(row):
    return Lipinski.NumHDonors(Chem.MolFromSmiles(row['X1']))


def hba(row):
    return Lipinski.NumHAcceptors(Chem.MolFromSmiles(row['X1']))


def logp(row):
    return Crippen.MolLogP(Chem.MolFromSmiles(row['X1']))


score_map = {
    'SAscore': sa_score,
    'RAscore': None,  # https://github.com/reymond-group/RAscore
    'SCScore': None,  # https://pubs.acs.org/doi/10.1021/acs.jcim.7b00622
    'LogP': logp,  # https://www.rdkit.org/docs/source/rdkit.Chem.Crippen.html
    'MW': mw,  # https://www.rdkit.org/docs/source/rdkit.Chem.Descriptors.html
    'HBD': hbd,  # https://www.rdkit.org/docs/source/rdkit.Chem.Lipinski.html
    'HBA': hba,  # https://www.rdkit.org/docs/source/rdkit.Chem.Lipinski.html
    'TopoPSA': None,  # http://mordred-descriptor.github.io/documentation/master/api/mordred.TopoPSA.html
}

filter_map = {
    'PAINS filter': None,
    "Lipinski's rule of five": None,  # https://gist.github.com/strets123/fdc4db6d450b66345f46
    'ADMET filter': None,
    'TCL filter': None
}

# cbg = gr.CheckboxGroup(choices=list(df.columns), label="Select columns")
#
# df_report = gr.Dataframe(type="pandas", interactive=False)

df_predict = gr.Dataframe(type="pandas", interactive=False, height=500, visible=False)

with gr.Blocks(theme=gr.themes.Soft(spacing_size="sm", text_size='lg'), title='DeepScreen') as demo:
    with gr.Tab(label='Inference') as inference:
        gr.Markdown('''
        # <center>DeepScreen Inference Service</center>
    
        DeepScreen for predicting drug-target interaction. 
        ''')
        # Upload a prediction dataset, choose a prediction task
        # (`binary` for interaction, and `regression` for affinity), and select a model preset to submit a job.
        with gr.Row():
            predict_data = gr.File(file_count="single", label='Prediction dataset file', type='binary')

        # predict_data.change(fn=load_csv, inputs=predict_data, outputs=df_predict)
        with gr.Row():
            task = gr.Dropdown(task_list, label='Task')
            preset = gr.Dropdown(preset_list + [None], label='Preset')

        with gr.Row():
            gr.ClearButton()
            predict_btn = gr.Button("Predict", variant="primary")

        predict_btn.click(fn=drug_target_predict, inputs=[predict_data, task, preset], outputs=df_predict)

    with gr.Tab(label='Report') as report:
        gr.Markdown('''
            # <center>DeepScreen Virtual Screening Report</center>
        
            Analytic report for virtual screening predictions. 
            ''')
        with gr.Row():
            score_menu = gr.CheckboxGroup(score_map.keys(), label='Scores')
            filter_menu = gr.CheckboxGroup(filter_map.keys(), label='Filters')

        df_report = gr.Dataframe(type="pandas", interactive=False, visible=True, height=500)

        with gr.Row():
            gr.ClearButton()
            analyze_btn = gr.Button("Analyze", variant="primary")