import gc
import time
from pathlib import Path
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from importlib import import_module
from logging import getLogger
from clf_util import fast_auc
from classifier.base_classifier import BaseClassifier
from classifier.torch_util import create_tensor_dataloader, train_model, predict_with_model
from save_log import stop_watch, get_back_training, get_training_logger, train_one_round, get_version, send_message


class PyTorchClassifier(BaseClassifier):

    @stop_watch("PyTorchClassifier.train()")
    def train(self, feature_names):

        # Initialize parameters
        validity = None
        model_path = Path(__file__).absolute().parents[2] / "data" / "model" / str(get_version())
        Path.mkdir(model_path, exist_ok=True, parents=True)
        START_FOLD = 0
        if get_back_training():
            START_FOLD = len(list(model_path.glob('**/*.model')))
        END_FOLD = 5
        if train_one_round():
            START_FOLD = 0
            END_FOLD = 1
        if START_FOLD == END_FOLD:
            return None

        get_training_logger(get_version()).debug(
            "fold \t iteration \
            \t train BCELoss \t valid BCELoss \
            \t train AUC \t valid AUC")

        # Process for each fold
        for fold in range(START_FOLD, END_FOLD):
            # Measure start time of the classification of this fold
            start = time.time()
            getLogger(get_version()).info("\t >> {} folds start".format(fold))
            send_message("\t :fire: {} folds start".format(fold))
            valid = "valid{}".format(str(fold))

            # Generate train data
            getLogger(get_version()).info("\t \t Generating datasets...")
            send_message("\t \t Generating datasets...")

            trn_x = super().get_feature_df(feature_names, valid, "train")
            trn_x.set_index("MachineIdentifier", inplace=True)
            trn_y = trn_x["HasDetections"].values.astype(np.float32)

            val_x = super().get_feature_df(feature_names, valid, "validate")
            val_x.set_index("MachineIdentifier", inplace=True)
            val_y = val_x["HasDetections"].values.astype(np.float32)

            # Initialize variables for scoring
            if validity is None:
                validity = pd.DataFrame()
                validity["HasDetections"] = pd.concat(
                    [trn_x["HasDetections"].astype(np.float32),
                     val_x["HasDetections"].astype(np.float32)])
                validity["Predict"] = 0

            del trn_x["HasDetections"], val_x["HasDetections"]
            trn_x.fillna(0, inplace=True)
            val_x.fillna(0, inplace=True)
            normalized_trn_x = (trn_x - trn_x.min()) / (trn_x.max() - trn_x.min())
            normalized_trn_x.fillna(normalized_trn_x.mean(), inplace=True)
            normalized_val_x = (val_x - val_x.min()) / (val_x.max() - val_x.min())
            normalized_val_x.fillna(normalized_val_x.mean(), inplace=True)

            train_loader, valid_loader, dataset_sizes = create_tensor_dataloader(self.params, normalized_trn_x, trn_y, normalized_val_x, val_y)
            data_loaders = {"train": train_loader, "valid": valid_loader}

            getLogger(get_version()).info("\t \t Datasets were generated.")
            send_message("\t \t Datasets were generated.")

            # Define the Network
            num_epochs = self.params["num_epochs"]
            network_class = getattr(import_module("classifier.pytorch_network." + self.params["network"]), self.params["network"])
            network = network_class(len(trn_x.columns), self.params["network_params"])
            criterion = nn.BCELoss()    # Don't change the criterion function!
            optimizer_class = getattr(import_module('torch.optim'), self.params["optimizer"])
            optimizer = optimizer_class(network.parameters(), **self.params["optimizer_params"])
            scheduler_class = getattr(import_module('torch.optim.lr_scheduler'), self.params["scheduler"])
            scheduler = scheduler_class(optimizer, **self.params["scheduler_params"])
            model = train_model(fold, data_loaders, dataset_sizes, network,
                                criterion, optimizer, scheduler,
                                num_epochs=num_epochs,
                                early_stopping_rounds=self.params["early_stopping_rounds"],
                                verbose=self.params["verbose"])

            # Classify
            validity.loc[validity.index.isin(val_x.index), "Predict"] = predict_with_model(self.params, normalized_val_x, model)

            # Measure finish time of the classification of this fold
            elapsed_time = int(time.time() - start)
            minutes, sec = divmod(elapsed_time, 60)
            hour, minutes = divmod(minutes, 60)
            getLogger(get_version()).info(
                "\t >> {} folds finish: [elapsed_time] >> {:0>2}:{:0>2}:{:0>2}"
                .format(fold, hour, minutes, sec))
            send_message("\t :fire: {} folds finish: [elapsed_time] >> {:0>2}:{:0>2}:{:0>2}".format(fold, hour, minutes, sec))

            # Post-process this fold
            gc.collect()
            torch.save(model, str(model_path / "valid{}.model".format(fold)))

        # Output CV score
        validity = validity.reset_index()
        columns_order = ["MachineIdentifier", "HasDetections", "Predict"]
        validity = validity.sort_values("MachineIdentifier").reset_index(drop=True).loc[:, columns_order]
        cv_auc = (fast_auc(validity["HasDetections"], np.array(validity["Predict"])))
        getLogger(get_version()).info("\t >> CV Score (AUC):{}".format(cv_auc))
        send_message("\t :fire: CV Score (AUC):{}".format(cv_auc))

        # Post-process the training
        gc.collect()

        return validity

    @stop_watch("PyTorchClassifier.predict()")
    def predict(self, feature_names):
        """
        Input:
            feature_names: directionary of features' names
        Output:
            predict_df: Dataframe(["MachineIdentifier", "HasDetections")
        """
        model_directory_path = Path(__file__).absolute().parents[2] / "data" / "model" / str(get_version())
        preds = None
        FOLDS = 5
        predict_df = None
        for fold in range(FOLDS):
            model_path = model_directory_path / "valid{}.model".format(fold)
            model = torch.load(str(model_path))
            valid = "valid{}".format(fold)
            test_df = super().get_feature_df(feature_names, valid, "test")
            if predict_df is None:
                predict_df = test_df["MachineIdentifier"]
            test_df = test_df.set_index("MachineIdentifier")
            normalized_test_df = (test_df - test_df.min()) / (test_df.max() - test_df.min())

            if preds is None:
                preds = predict_with_model(self.params, normalized_test_df, model) / FOLDS
            else:
                preds += predict_with_model(self.params, normalized_test_df, model) / FOLDS

        predict_df = pd.DataFrame(predict_df)
        predict_df["HasDetections"] = preds
        return predict_df
