import os
import shutil

import numpy as np
import pandas as pd
import deepchem as dc
import matplotlib.pyplot as plt
from tqdm import tqdm
from rdkit import Chem
from rdkit.Chem import Descriptors
from rdkit.ML.Descriptors.MoleculeDescriptors import MolecularDescriptorCalculator
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (
    confusion_matrix, accuracy_score, recall_score, 
    precision_score, f1_score, balanced_accuracy_score, 
    roc_auc_score
)


data_dir = "data"
histgram_dir = os.path.join(data_dir, "histgram")
low_std_image_dir = os.path.join(data_dir, "low_std")

def get_data() -> dict:
    """DeepChemからHIVデータセットを取得する
    """
    tasks, datasets, transformers = dc.molnet.load_hiv()
    train, val, test = datasets

    dataset = {}
    for data_name, data in zip(["train", "val", "test"], [train, val, test]):
        smiles = []
        labels = []
        for (x, y, weight, s) in data.iterbatches():
            smiles += s.flatten().tolist()
            labels += y.flatten().tolist()


        df = pd.DataFrame([])
        df["SMILES"] = smiles
        df["ACTIVITY"] = labels
        df["ACTIVITY"] = df["ACTIVITY"].astype(int)

        # データがない行を削除
        df.dropna(inplace=True)
        df = df[df["SMILES"]!=""]
        dataset[data_name] = df

    return dataset

def calc_descriptor(descriptor_name:str="descriptors", 
                    calc:bool=False) -> dict:
    """記述子を計算する

    Parameters
    ------
    descriptor_name : str
        記述子計算結果保存先/読み込み先のCSVファイルパス
    calc : bool
        記述子計算を行う. Falseなら保存している結果を読み込む.
    """
    if calc:
        data = get_data()

    dataset = {}
    for dataset_type in ["train", "val", "test"]:
        descriptor_file = f"{descriptor_name}_{dataset_type}.csv"

        if calc:
            df = data[dataset_type]

            # SMILESからmolオブジェクトを取得
            smiles = df["SMILES"]
            mols = [Chem.MolFromSmiles(s) for s in tqdm(smiles)]

            # 記述子計算
            descriptors = [d[0] for d in Descriptors._descList]
            descriptor_calculation = MolecularDescriptorCalculator(descriptors)
            df_descriptor = pd.DataFrame(
                    [descriptor_calculation.CalcDescriptors(mol) for mol in tqdm(mols)], 
                    columns=descriptors)
            df_descriptor["y"] = df["ACTIVITY"]
            
            # 記述子計算に長時間かかるので、CSV保存
            df_descriptor.to_csv(descriptor_file, index=None)
        else:
            df_descriptor = pd.read_csv(descriptor_file)
        
        dataset[dataset_type] = {}
        dataset[dataset_type]["y"] = df_descriptor["y"]
        dataset[dataset_type]["x"] = df_descriptor.drop("y", axis=1)

    return dataset

def fill(df:pd.DataFrame, values:float, plot:bool, 
         return_cols:bool=False, low_var_col:list=None) -> [pd.DataFrame, list]:
    """データフレームの値を補完する

    Parameters
    ------
    df : pd.DataFrame
        補完をするデータフレーム
    values : list
        データフレーム中の[nan, inf, -inf]に対して補完する値
    plot : bool
        記述子ごとのヒストグラムをプロットする
    return_cols : bool
        訓練データで使用するカラムを返す
    low_var_col : list
        分散の少ないカラム (訓練データのカラム)
    
    Returns
    ------
    pd.DataFrame : 各種値を補完したデータフレーム
    list : 訓練データのカラム名. 訓練データ以外ではNone
    """
    # 欠損値を平均値で補完
    df = df.fillna(values[0])

    # np.infを最大値で、-np.infを最小値で補完
    # drop_inf = df.replace([np.inf, -np.inf], np.nan)
    df = df.replace(np.inf, values[1])
    df = df.replace(-np.inf, values[2])

    if low_var_col is None:
        df = drop_low_var_col(df, thres_var=0, plot=plot)
    else:
        df = df[low_var_col]

    if return_cols:
        cols = list(df.columns)
    else:
        cols = None

    # 最大値が大きすぎるのでクリップ
    # ValueError: Input X contains infinity or a value too large for dtype('float32').
    df = df.clip(-1e30, 1e30)

    return df, cols

def main():
    plot = False
    calc = False
    random_state = 0

    # 記述子の計算
    dataset = calc_descriptor(calc=calc)
    x_train, y_train = dataset["train"]["x"], dataset["train"]["y"]
    x_val, y_val = dataset["val"]["x"], dataset["val"]["y"]
    x_test, y_test = dataset["test"]["x"], dataset["test"]["y"]

    # 記述子の値の調整
    drop_inf = x_train.replace([np.inf, -np.inf], np.nan)
    fill_values = [x_train.mean(), drop_inf.max(), drop_inf.min()]
    x_train, cols = fill(x_train, fill_values, plot, return_cols=True)
    x_val, _ = fill(x_val, fill_values, plot, low_var_col=cols)
    x_test, _ = fill(x_test, fill_values, plot, low_var_col=cols)

    # pd.DataFrame -> np.ndarrayに型変換
    x_train = x_train.to_numpy().astype(float)
    y_train = y_train.to_numpy().astype(int)
    x_val = x_val.to_numpy().astype(float)
    y_val = y_val.to_numpy().astype(int)
    x_test = x_test.to_numpy().astype(float)
    y_test = y_test.to_numpy().astype(int)

    # 学習
    clf = RandomForestClassifier(random_state=random_state)
    clf.fit(x_train, y_train)

    # validationデータで予測
    pred_val = clf.predict(x_val)

    # 評価
    acc = accuracy_score(y_val, pred_val)
    balanced_acc = balanced_accuracy_score(y_val, pred_val)
    precision = precision_score(y_val, pred_val)
    recall = recall_score(y_val, pred_val)
    f1 = f1_score(y_val, pred_val)
    roc = roc_auc_score(y_val, pred_val)
    print("accuracy: ", acc)
    print("balanced accuracy: ", balanced_acc)
    print("precision: ", precision)
    print("recall: ", recall)
    print("f value: ", f1)
    print("roc score: ", roc)