import random
from pathlib import Path
from ipdb import launch_ipdb_on_exception, set_trace
import numpy as np
from sklearn.metrics import roc_auc_score, accuracy_score
from keras.callbacks import EarlyStopping, ModelCheckpoint
from keras.callbacks import ReduceLROnPlateau, LearningRateScheduler
from keras.callbacks import CSVLogger, Callback
from keras.utils import plot_model


class AUCLoggingCallback(Callback):
    """Custom logger callback for logging and roc-auc."""

    def __init__(self, logger, valid_data, out_size):
        super().__init__()
        self.logger = logger
        self.valid_X = valid_data[0]
        self.valid_y = valid_data[1]
        self.out_size = out_size

    def on_epoch_end(self, epoch, logs={}):
        pred = self.model.predict(self.valid_X)
        if self.out_size is None:
            acc = logs['acc']
            loss = logs['loss']
            val_loss = logs['val_loss']
            val_acc = logs['val_acc']

            auc = roc_auc_score(self.valid_y, pred)

            mes = f"Epoch {epoch:>2}: loss {loss:.4f}, acc {acc:.4f}, "
            mes += f"val_loss {val_loss:.4f}, val_acc {val_acc:.4f}, "
            mes += f"val_auc {auc:.4f}"
            self.logger.info(mes)
        else:
            aucs = list()
            for i in range(self.out_size):
                auc = roc_auc_score(self.valid_y[i], pred[i])
                self.logger.info(f"val_auc {i}: {auc:.4f}")
                aucs.append(auc)
            auc = max(aucs)
        logs['val_auc'] = auc

    def on_train_end(self, logs={}):
        pred = self.model.predict(self.valid_X)
        if self.out_size is None:
            test_acc = accuracy_score(self.valid_y, pred > 0.5)
            test_auc = roc_auc_score(self.valid_y, pred)
            mes = f"Test: acc {test_acc:.4f}, auc {test_auc:.4f}"
        else:
            test_auc = roc_auc_score(self.valid_y[9], pred[9])
            mes = f"Test: auc {test_auc:.4f}"
        self.logger.info(mes)


class NNManager():
    """setting for Neural Net learning
    """

    def __init__(self,
                 dir_path,
                 model,
                 batch_size,
                 train_cols,
                 logger,
                 out_size=None,
                 onehot_size=None,
                 pos_neg_ratio=1,
                 proc_per_gen=None):

        self.dir_path = dir_path
        self.model = model
        self.batch_size = batch_size
        self.train_cols = train_cols

        if proc_per_gen:
            self.proc_per_gen = proc_per_gen
        else:
            self.proc_per_gen = lambda x: x

        self.logger = logger
        self.pos_neg_ratio = pos_neg_ratio
        self.out_size = out_size
        self.onehot_size = onehot_size

        Path(dir_path).mkdir(exist_ok=True, parents=True)

        self.dump_path = dir_path / 'weights.hdf5'
        self.csv_log_path = dir_path / 'epochs.csv'

        plot_model(self.model, to_file=dir_path / 'arch.png')
        with open(dir_path / 'entity_nn.txt', 'w') as fp:
            self.model.summary(print_fn=lambda x: fp.write(x + '\n'))

        self.callbacks = [
            # EarlyStopping(
            #     monitor='val_loss',
            #     patience=10,
            #     verbose=1,
            #     min_delta=0.001,
            #     mode='min'),
            # ReduceLROnPlateau(
            #     monitor='val_loss',
            #     factor=0.1,
            #     patience=5,
            #     verbose=1,
            #     min_delta=0.001,
            #     mode='min'),
            LearningRateScheduler(
                lambda e: 0.001 * (0.1**(e // 20)), verbose=1),
            ModelCheckpoint(
                monitor='val_loss',
                filepath=str(self.dump_path),
                save_best_only=True,
                save_weights_only=True,
                mode='min'),
            CSVLogger(self.csv_log_path)
        ]

    def dup_y(self, y):
        y_list = list()
        for _ in range(self.out_size):
            y_list.append(y)
        return y_list

    def balanced_data_gen(self, X_train, y_train, sample_size=None):
        """balanced data generator for training
        """

        pos_idx = list(X_train[y_train == 1].index.values)
        neg_idx = list(X_train[y_train == 0].index.values)
        neg_size = int(sample_size * self.pos_neg_ratio)
        print(len(pos_idx), len(neg_idx), neg_size)

        while True:
            pos_sample = random.sample(pos_idx, sample_size)
            neg_sample = random.sample(neg_idx, neg_size)
            sample_idx = np.concatenate([pos_sample, neg_sample])
            base_df = X_train.loc[sample_idx].reset_index(drop=True)

            base_y = y_train.loc[sample_idx].reset_index(drop=True)
            base_id = np.random.permutation(len(base_df))

            for start in range(0, len(base_df), self.batch_size):
                end = min(start + self.batch_size, len(base_df))
                batch_df_id = base_id[start:end]
                batch_df = base_df.iloc[batch_df_id]
                x_batch = batch_df[self.train_cols]

                if self.proc_per_gen:
                    with launch_ipdb_on_exception():
                        x_batch = self.proc_per_gen(x_batch)

                y_batch = base_y.iloc[batch_df_id]
                if self.out_size:
                    y_batch = self.dup_y(y_batch)
                if self.onehot_size:
                    y_batch = np.identity(self.onehot_size)[y_batch.values]

                yield x_batch, y_batch

    def normal_data_gen(self, X_train, y_train=None, test=False):
        """normal data generator for training
        """
        while True:
            if not test:
                base_id = np.random.permutation(range(len(X_train)))
            else:
                base_id = list(range(len(X_train)))
            for start in range(0, len(X_train), self.batch_size):
                end = min(start + self.batch_size, len(X_train))
                idx = base_id[start:end]
                batch_df = X_train.iloc[idx]
                x_batch = batch_df[self.train_cols]

                if self.proc_per_gen:
                    with launch_ipdb_on_exception():
                        x_batch = self.proc_per_gen(x_batch)

                if y_train is not None:
                    y_batch = y_train.iloc[idx]
                    if self.out_size:
                        y_batch = self.dup_y(y_batch)
                    if self.onehot_size:
                        y_batch = np.identity(self.onehot_size)[y_batch.values]
                    yield x_batch, y_batch
                else:
                    yield x_batch

    def calc_steps(self, df_size):
        return int(np.ceil(df_size / self.batch_size))

    def learn(self,
              X_train,
              y_train,
              X_valid,
              y_valid,
              valid_sampling=None,
              epochs=20):
        """main func fot NN training
        """

        # generator for training
        pos_size = len(y_train[y_train == 1])
        print(y_train.shape)
        train_gen = self.normal_data_gen(X_train, y_train)
        # pos_size * (1 + pos_neg_ratio)
        data_size = int(pos_size * (1 + self.pos_neg_ratio))
        step_per_epoch = self.calc_steps(data_size)

        # validation data
        X_valid = X_valid
        if self.out_size:
            y_valid = self.dup_y(y_valid)
        if self.onehot_size:
            y_valid = np.identity(self.onehot_size)[y_valid.values]
        valid_sample = (self.proc_per_gen(X_valid), y_valid)

        # to calculate auc score on each end of epoch
        # this part is very slow processing!!!
        # auc_callback = AUCLoggingCallback(self.logger, valid_sample,
        #                                   self.out_size)
        # self.callbacks = [auc_callback] + self.callbacks
        
        import ipdb
        with ipdb.launch_ipdb_on_exception():
            history = self.model.fit_generator(
                generator=train_gen,
                steps_per_epoch=step_per_epoch,
                epochs=epochs,
                verbose=1,
                callbacks=self.callbacks,
                validation_data=valid_sample)
        return history

    def predict_generator(self, test_df):
        self.model.load_weights(self.dump_path)
        step_per_epoch = self.calc_steps(len(test_df))
        test_gen = self.normal_data_gen(test_df, test=True)
        pred = self.model.predict_generator(test_gen,
                                            steps=step_per_epoch,
                                            verbose=1)
        return pred
