from abc import abstractmethod
from enum import Enum

import numpy as np

from neural_net.epoch import Epoch
from neural_net.transform_layer import Layer


class ModelData:
    def __init__(self, training_inputs, training_targets, test_inputs, test_targets):
        self.is_loaded = False
        self.training_inputs = training_inputs
        self.training_labels = training_targets
        self.test_inputs = test_inputs
        self.test_labels = test_targets


# class TrainingSession:
#     def __init__(self, training_data: ModelData, learning_rate: float, nr_epochs: int, batch_size: int = 1000):
#         self.training_data = training_data
#         self.learning_rate = learning_rate
#         self.nr_epochs = nr_epochs
#         self.batch_size = batch_size
#         self.epochs: [Epoch] = []
#         for i in range(self.nr_epochs):
#             self.epochs.append(
#                 Epoch(i, self.training_data.training_inputs, self.training_data.training_labels, self.batch_size))
#
#     def get_total_training_duration(self):
#         duration = 0.0
#         for epoch in self.epochs:
#             duration += epoch.duration
#         return duration


class NeuralNet:
    def __init__(self, layers: [Layer]):
        self.layers = layers
        self.last_loss = None
        self.last_accuracy = None

    def forward(self, inputs):
        outputs = inputs
        for layer in self.layers:
            outputs = layer.forward(outputs)
        return outputs

    def reset(self):
        for layer in self.layers:
            layer.reset()

    def backward(self, dL_dout, epoch):
        layer_dl_gradients = []
        layer_dl_bias = []
        layer_weights = []
        layer_biases = []

        for idx, layer in reversed(list(enumerate(self.layers))):
            dL_dout, dl_gradients, dl_biases, weights, biases = layer.backward(dL_dout, epoch.learning_rate)

            if dl_gradients is not None:
                layer_dl_gradients.append(dl_gradients)
            if dl_biases is not None:
                layer_dl_bias.append(dl_biases)
            if weights is not None:
                layer_weights.append(weights)
            if biases is not None:
                layer_biases.append(biases)

        return layer_dl_gradients, layer_dl_bias, layer_weights, layer_biases

    # def train(self, training_run: TrainingRun):
    #     self.training_runs.append(training_run)
    #
    #     for epoch in training_run.epochs:
    #         epoch.start()
    #
    #         for batch in epoch.batches:
    #             batch.predictions = self.forward(batch.inputs)
    #             dL_dout = self.loss_derivative(batch.predictions, batch.labels)
    #
    #             layer_dl_gradients, layer_dl_biases, layer_weights, layer_biases = self.backward(dL_dout, training_run.learning_rate, epoch)
    #             epoch.layer_dl_gradients.append(layer_dl_gradients)
    #             epoch.layer_dl_biases.append(layer_dl_biases)
    #
    #         epoch.finish()
    #         epoch.loss = self.loss(epoch.all_predictions(), epoch.all_labels())
    #
    #         if training_run.epoch_callback is not None:
    #             training_run.epoch_callback(training_run, epoch)
    #
    #     self.recalculate_loss(training_run.training_data.test_inputs, training_run.training_data.test_labels)
    #     self.recalculate_loss(training_run.training_data.test_inputs, training_run.training_data.test_labels)

    def get_all_weights(self):
        all_weights = []
        for layer in self.layers:
            if hasattr(layer, 'weights'):
                all_weights.append(layer.weights)
        return all_weights

    def recalculate_accuracy(self, inputs, labels):
        raw_outputs = self.forward(inputs)
        predictions = raw_outputs.argmax(axis=1)
        num_correct_predictions = 0
        for idx, prediction in enumerate(predictions):
            if prediction == labels[idx]:
                num_correct_predictions += 1
        self.last_accuracy = num_correct_predictions / len(predictions)
        return self.last_accuracy

    def recalculate_loss(self, inputs, labels):
        raw_outputs = self.forward(inputs)
        self.last_loss = self.loss(np.array(raw_outputs), np.array(labels))
        return self.last_loss

    @abstractmethod
    def loss(self, outputs: np.array, labels: np.array):
        pass

    @abstractmethod
    def loss_derivative(self, outputs: np.array, labels: np.array):
        pass

    def predict(self, inputs):
        return self.forward(inputs)