import os
import random
import numpy as np
import warnings
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from torch.utils.data import Dataset, DataLoader
import gc
import streamlit as st

warnings.filterwarnings("ignore")

 
st.title('ReactionT5 task forward')
st.markdown('''
##### At this space, you can predict the products of reactions from their inputs.
##### The code expects input_data as a string or CSV file that contains an "input" column. 
##### The format of the string or contents of the column should be "REACTANT:{reactants}REAGENT:{reagents}".
##### If there is no reagent, fill the blank with a space. For multiple compounds, concatenate them with ".".
##### The output contains SMILES of predicted products and the sum of log-likelihood for each prediction, ordered by their log-likelihood (0th is the most probable product).
''')

display_text = 'input the reaction smiles (e.g. REACTANT:COC(=O)C1=CCCN(C)C1.O.[Al+3].[H-].[Li+].[Na+].[OH-]REAGENT:C1CCOC1)'

st.download_button(
                label="Download demo_input.csv",
                data=pd.read_csv('demo_input.csv').to_csv(index=False),
                file_name='demo_input.csv',
                mime='text/csv',
            )

class CFG():
    num_beams = st.number_input(label='num beams', min_value=1, max_value=10, value=5, step=1)
    num_return_sequences = num_beams
    uploaded_file = st.file_uploader("Choose a CSV file")
    input_data = st.text_area(display_text)
    model_name_or_path = 'sagawa/ReactionT5v2-forward'
    input_column = 'input'
    input_max_length = 400
    model = 't5'
    seed = 42
    batch_size=1

def seed_everything(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True



def prepare_input(cfg, text):
    inputs = tokenizer(
        text,
        return_tensors="pt",
        max_length=cfg.input_max_length,
        padding="max_length",
        truncation=True,
    )
    dic = {"input_ids": [], "attention_mask": []}
    for k, v in inputs.items():
        dic[k].append(torch.tensor(v[0], dtype=torch.long))
    return dic


class ProductDataset(Dataset):
    def __init__(self, cfg, df):
        self.cfg = cfg
        self.inputs = df[cfg.input_column].values

    def __len__(self):
        return len(self.inputs)

    def __getitem__(self, idx):
        return prepare_input(self.cfg, self.inputs[idx])


def predict_single_input(input_compound):
    inp = tokenizer(input_compound, return_tensors="pt").to(device)
    with torch.no_grad():
        output = model.generate(
            **inp,
            num_beams=CFG.num_beams,
            num_return_sequences=CFG.num_return_sequences,
            return_dict_in_generate=True,
            output_scores=True,
        )
    return output


def decode_output(output):
    sequences = [
        tokenizer.decode(seq, skip_special_tokens=True).replace(" ", "").rstrip(".")
        for seq in output["sequences"]
    ]
    if CFG.num_beams > 1:
        scores = output["sequences_scores"].tolist()
        return sequences, scores
    return sequences, None


def save_single_prediction(input_compound, output, scores):
    output_data = [input_compound] + output + (scores if scores else [])
    columns = (
        ["input"]
        + [f"{i}th" for i in range(CFG.num_beams)]
        + ([f"{i}th score" for i in range(CFG.num_beams)] if scores else [])
    )
    output_df = pd.DataFrame([output_data], columns=columns)
    return output_df


def save_multiple_predictions(input_data, sequences, scores):
    output_list = [
        [input_data.loc[i // CFG.num_return_sequences, CFG.input_column]]
        + sequences[i : i + CFG.num_return_sequences]
        + scores[i : i + CFG.num_return_sequences]
        for i in range(0, len(sequences), CFG.num_return_sequences)
    ]
    columns = (
        ["input"]
        + [f"{i}th" for i in range(CFG.num_return_sequences)]
        + ([f"{i}th score" for i in range(CFG.num_return_sequences)] if scores else [])
    )
    output_df = pd.DataFrame(output_list, columns=columns)
    return output_df


if st.button('predict'):
    with st.spinner('Now processing. If num beams=5, this process takes about 15 seconds per reaction.'):

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        seed_everything(seed=CFG.seed)
        
        tokenizer = AutoTokenizer.from_pretrained(CFG.model_name_or_path, return_tensors="pt")
        model = AutoModelForSeq2SeqLM.from_pretrained(CFG.model_name_or_path).to(device)
        model.eval()
        
        if CFG.uploaded_file is None:
            input_compound = CFG.input_data
            output = predict_single_input(input_compound)
            sequences, scores = decode_output(output)
            output_df = save_single_prediction(input_compound, sequences, scores)
        else:
            input_data = pd.read_csv(CFG.uploaded_file)
            dataset = ProductDataset(CFG, input_data)
            dataloader = DataLoader(
                dataset,
                batch_size=CFG.batch_size,
                shuffle=False,
                num_workers=4,
                pin_memory=True,
                drop_last=False,
            )
        
            all_sequences, all_scores = [], []
            for inputs in dataloader:
                inputs = {k: v[0].to(device) for k, v in inputs.items()}
                with torch.no_grad():
                    output = model.generate(
                        **inputs,
                        num_beams=CFG.num_beams,
                        num_return_sequences=CFG.num_return_sequences,
                        return_dict_in_generate=True,
                        output_scores=True,
                    )
                sequences, scores = decode_output(output)
                all_sequences.extend(sequences)
                if scores:
                    all_scores.extend(scores)
                del output
                torch.cuda.empty_cache()
                gc.collect()
        
            output_df = save_multiple_predictions(input_data, all_sequences, all_scores)
        
        @st.cache
        def convert_df(df):
            return df.to_csv(index=False)
        
        csv = convert_df(output_df)
        
        st.download_button(
            label="Download data as CSV",
            data=csv,
            file_name='output.csv',
            mime='text/csv',
        )