import numpy as np

import os

from datetime import datetime

from math import ceil, floor, log

import tensorflow as tf
import tensorflow.keras as keras
import keras as K


from data_loader import DataGenerator

import pandas as pd


def weighted_log_loss(y_true, y_pred):
    """
    Can be used as the loss function in model.compile()
    ---------------------------------------------------
    """

    class_weights = np.array([2., 1., 1., 1., 1., 1.])

    eps = K.backend.epsilon()

    y_pred = K.backend.clip(y_pred, eps, 1.0-eps)

    out = -(         y_true  * K.backend.log(      y_pred) * class_weights
            + (1.0 - y_true) * K.backend.log(1.0 - y_pred) * class_weights)

    return K.backend.mean(out, axis=-1)


def _normalized_weighted_average(arr, weights=None):
    """
    A simple K implementation that mimics that of
    numpy.average(), specifically for this competition
    """

    if weights is not None:
        scl = K.backend.sum(weights)
        weights = K.backend.expand_dims(weights, axis=1)
        return K.backend.sum(K.backend.dot(arr, weights), axis=1) / scl
    return K.backend.mean(arr, axis=1)


def weighted_loss(y_true, y_pred):
    """
    Will be used as the metric in model.compile()
    ---------------------------------------------

    Similar to the custom loss function 'weighted_log_loss()' above
    but with normalized weights, which should be very similar
    to the official competition metric:
        https://www.kaggle.com/kambarakun/lb-probe-weights-n-of-positives-scoring
    and hence:
        sklearn.metrics.log_loss with sample weights
    """

    class_weights = K.backend.variable([2., 1., 1., 1., 1., 1.])

    eps = K.backend.epsilon()

    y_pred = K.backend.clip(y_pred, eps, 1.0-eps)

    loss = -(        y_true  * K.backend.log(      y_pred)
            + (1.0 - y_true) * K.backend.log(1.0 - y_pred))

    loss_samples = _normalized_weighted_average(loss, class_weights)

    return K.backend.mean(loss_samples)

class PredictionCheckpoint(K.callbacks.Callback):

    def on_epoch_end(self, epoch, logs={}):
        """
        Save each epoch file in case of crash
        """
        print("Saving checkpoint")
        self.model.save("epoch{}.hdf5".format(epoch))

class MyDeepModel:

    def __init__(self, engine, input_dims, batch_size=5, num_epochs=4, learning_rate=1e-3,
                 decay_rate=1.0, decay_steps=1, weights="imagenet", verbose=1, train_image_dir="", model_filename=""):

        self.engine = engine
        self.input_dims = input_dims
        self.batch_size = batch_size
        self.num_epochs = num_epochs
        self.learning_rate = learning_rate
        self.decay_rate = decay_rate
        self.decay_steps = decay_steps
        self.weights = weights
        self.verbose = verbose
        self.model_filename = model_filename
        self.train_images_dir=train_image_dir
        self._build()

    def _build(self):


        engine = self.engine(include_top=False, weights=self.weights, input_shape=self.input_dims,
                             backend = K.backend, layers = K.layers,
                             models = K.models, utils = K.utils)

        x = K.layers.GlobalAveragePooling2D(name='avg_pool')(engine.output)
        x = K.layers.Dropout(0.3)(x)
#         x = keras.layers.Dense(keras.backend.int_shape(x)[1], activation="relu", name="dense_hidden_1")(x)
#         x = keras.layers.Dropout(0.1)(x)
        out = K.layers.Dense(6, activation="sigmoid", name='dense_output')(x)

        self.model = K.models.Model(inputs=engine.input, outputs=out)

        self.model.compile(loss="binary_crossentropy", optimizer=K.optimizers.Adam(), metrics=["categorical_accuracy", "accuracy", weighted_loss])

    def get_model_filename(self):

        return self.model_filename

    def fit_model(self, train_df, valid_df):

        # callbacks
        checkpointer = K.callbacks.ModelCheckpoint(filepath=self.model_filename, verbose=1, save_best_only=True)
        scheduler = K.callbacks.LearningRateScheduler(lambda epoch: self.learning_rate * pow(self.decay_rate, floor(epoch / self.decay_steps)))

        self.model.fit_generator(
            DataGenerator(
                train_df.index,
                train_df,
                self.batch_size,
                self.input_dims,
                self.train_images_dir
            ),
            epochs=self.num_epochs,
            verbose=self.verbose,
            validation_data=DataGenerator(
                valid_df.index,
                valid_df,
                self.batch_size,
                self.input_dims,
                self.train_images_dir
            ),
            use_multiprocessing=True,
            workers=4,
            callbacks=[PredictionCheckpoint(), scheduler, checkpointer]
        )

    def save(self, path):
        self.model.save(path)

    def load(self, path):
        self.model.load_weights(path)


def create_submission(model, data, test_df):

    print("Creating predictions on test dataset")

    pred = model.predict_generator(data, verbose=1)

    out_df = pd.DataFrame(pred, index=test_df.index, columns=test_df.columns)

    test_df = out_df.stack().reset_index()

    test_df.insert(loc=0, column='ID', value=test_df['Image'].astype(str) + "_" + test_df['Diagnosis'])

    test_df = test_df.drop(["Image", "Diagnosis"], axis=1)

    print("Saving submissions to submission.csv")
    test_df.to_csv('submission.csv', index=False)

    return test_df