# -*- coding: utf-8 -*-
"""Untitled94.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/15_dmQuUdLkUUISrOBktAAP7HDWebFJmB
"""

import pandas as pd
from tqdm import tqdm
import numpy as np
from nona.nona import nona
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import Ridge
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from pytorch_tabnet.tab_model import TabNetClassifier
from pytorch_tabnet.pretraining import TabNetPretrainer
from imblearn.over_sampling import SMOTE
import torch
from sklearn.metrics import precision_score, recall_score, f1_score


def process_features(file_path='./MIMIC/icu/chartevents.csv', output_path='df_31.csv', item_ids={
    225309: "ART BP Systolic",
    220045: "HR",
    220210: "RR",
    223762: "Temperature C"}):
    """ Collecting features of the dataset

    Keyword arguments:
    file_path -- path to the file
    output_path -- path to the output file
    item_ids -- list of item ids

    """
    item_ids_set = set(map(str, item_ids.keys()))

    result = {}

    with open(file_path) as f:
        headers = f.readline().replace('\n', '').split(',')
        i = 0
        for line in tqdm(f):
            values = line.replace('\n', '').split(',')
            subject_id = values[0]
            item_id = values[6]
            valuenum = values[8]
            if item_id in item_ids_set:
                if subject_id not in result:
                    result[subject_id] = {}
                result[subject_id][item_id] = valuenum
            i += 1

    table = pd.DataFrame.from_dict(result, orient='index')
    table['subject_id'] = table.index

    table.to_csv(output_path, index=False)
    return table


def add_diagnosis_column(drgcodes_path, merged_data_path, output_path):
    """Add diagnosis column to dataset.

    Keyword arguments:
    drgcodes_path -- path to the drgcodes file
    merged_data_path -- path to the file with features
    output_path -- path to the output file

    """
    drgcodes = pd.read_csv(drgcodes_path)
    merged_data = pd.read_csv(merged_data_path)
    drgcodes['drg_code'] = pd.to_numeric(drgcodes['drg_code'], errors='coerce')
    merged_data['diagnosis'] = 0
    target_subjects = drgcodes.loc[drgcodes['drg_code'].isin([870, 871, 872]), 'subject_id']
    merged_data.loc[merged_data['subject_id'].isin(target_subjects), 'diagnosis'] = 1
    merged_data.to_csv(output_path, index=False)


def impute_data(input_path, output_path, features):
    """Fill in the blanks.

    Keyword arguments:
    input_path -- path to the input file
    output_path -- path to the output file
    features -- list of features

    """
    df = pd.read_csv(input_path)
    X = df[features].copy()
    nona(
        data=X,
        algreg=make_pipeline(StandardScaler(with_mean=False), Ridge(alpha=0.1)),
        algclass=RandomForestClassifier(max_depth=2, random_state=0)
    )
    # df[features] = X
    df.loc[:, features] = X.values
    df.to_csv(output_path, index=False)


def prepare_and_save_data(input_path, test_size, random_state, features, target, resampled_output_path,
                          test_output_path):
    """We divide the data into train and test. Getting rid of class imbalance in train.

    Keyword arguments:
    input_path -- path to the input file
    test_size -- size of the test set
    random_state -- random state
    features -- list of features
    target -- target column

    """
    df = pd.read_csv(input_path)
    train_data, test_data = train_test_split(df, test_size=test_size, random_state=random_state)
    X_train = train_data[features]
    y_train = train_data[target]
    smote = SMOTE(random_state=random_state)
    X_resampled, y_resampled = smote.fit_resample(X_train, y_train)
    df_resampled = pd.concat(
        [pd.DataFrame(X_resampled, columns=X_train.columns), pd.DataFrame(y_resampled, columns=[target])], axis=1)
    df_resampled.to_csv(resampled_output_path, index=False)
    test_data.to_csv(test_output_path, index=False)
    return df_resampled, test_data


def resample_test_val_data(input_path, test_size, random_state, features, target, test_output_path, val_output_path):
    """We divide the data into test and validation. Getting rid of class imbalance in test and in validation.

    Keyword arguments:
    input_path -- path to the input file
    test_size -- size of the test set
    random_state -- random state
    features -- list of features
    target -- target column

    """
    df = pd.read_csv(input_path)
    data = df.drop(['subject_id'], axis=1)
    test_data, val_data = train_test_split(data, test_size=test_size, random_state=random_state)
    smote1 = SMOTE(random_state=42)
    X1 = test_data[features]
    y1 = test_data[target]
    X_resampled1, y_resampled1 = smote1.fit_resample(X1, y1)
    df_resampled1 = pd.concat(
        [pd.DataFrame(X_resampled1, columns=features), pd.DataFrame(y_resampled1, columns=[target])], axis=1)
    df_resampled1.to_csv(test_output_path, index=False)
    smote2 = SMOTE(random_state=42)
    X2 = val_data[features]
    y2 = val_data[target]
    X_resampled2, y_resampled2 = smote2.fit_resample(X2, y2)
    df_resampled2 = pd.concat(
        [pd.DataFrame(X_resampled2, columns=features), pd.DataFrame(y_resampled2, columns=[target])], axis=1)
    df_resampled2.to_csv(val_output_path, index=False)


def train_tabnet_model(train_path, val_path, feature_importances_path, model_save_path, optimizer_params,
                       scheduler_params, pretraining_lr=0.05, training_lr=0.05, mask_type='sparsemax',
                       pretraining_ratio=1.0, max_epochs=200, patience=50):
    """Model training, preserving the importance of features and saving the model.

    Keyword arguments:
    train_path -- path to the train data
    val_path -- path to the validation data
    feature_importances_path -- path to the file with feature importances
    model_save_path -- path to the model
    optimizer_params -- optimizer parameters
    scheduler_params -- scheduler parameters
    pretraining_lr -- learning rate for pretraining (default 0.05)
    training_lr -- learning rate for training (default 0.05)
    mask_type -- mask type (default 'sparsemax')
    pretraining_ratio -- ratio of pretraining data (default 1.0)
    max_epochs -- maximum number of epochs (default 200)
    patience -- patience (default 50)

    """
    train_data = pd.read_csv(train_path)
    val_data = pd.read_csv(val_path)
    X_train = train_data.drop(['diagnosis'], axis=1)
    y_train = train_data['diagnosis']
    X_val = val_data.drop(['diagnosis'], axis=1)
    y_val = val_data['diagnosis']
    # unsupervised_model = TabNetPretrainer(
    #     optimizer_fn=torch.optim.Adam,
    #     optimizer_params=dict(lr=pretraining_lr),
    #     mask_type=mask_type
    # )
    # unsupervised_model.fit(
    #     X_train=X_train.values,
    #     eval_set=[X_val.values],
    #     pretraining_ratio=pretraining_ratio,
    # )
    clf = TabNetClassifier(
        optimizer_fn=torch.optim.AdamW,
        optimizer_params=dict(lr=training_lr),
        scheduler_params=scheduler_params,
        scheduler_fn=torch.optim.lr_scheduler.MultiStepLR,
        mask_type=mask_type
    )
    clf.fit(
        X_train=X_train.values, y_train=y_train.values,
        eval_set=[(X_val.values, y_val.values)],
        eval_metric=['auc'],
        max_epochs=max_epochs,
        patience=patience
        # from_unsupervised=unsupervised_model
    )
    with open(feature_importances_path, 'w') as f:
        f.write(f'{clf.feature_importances_}')
    clf.save_model(model_save_path)


def evaluate_tabnet_model(model_path, test_data_path, metrics_output_path):
    """Evaluate the model.

    Keyword arguments:
    model_path -- path to the model
    test_data_path -- path to the test data
    metrics_output_path -- path to the output file

    """
    loaded_clf = TabNetClassifier()
    loaded_clf.load_model(model_path)
    test_data = pd.read_csv(test_data_path)
    X_test = test_data.drop(['diagnosis'], axis=1)
    y_test = test_data['diagnosis']
    result = loaded_clf.predict(X_test.values)
    accuracy = (result == y_test.values).mean()
    precision = precision_score(y_test.values, result)
    recall = recall_score(y_test.values, result)
    f1 = f1_score(y_test.values, result)
    with open(metrics_output_path, 'w') as file:
        file.write(f'Accuracy: {accuracy}\n')
        file.write(f'Precision: {precision}\n')
        file.write(f'Recall: {recall}\n')
        file.write(f'F1-score: {f1}\n')
        file.close()
