# Copyright 2023 Cristian Grosu
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from math import sqrt
import tensorflow as tf
import os
import numpy as np
import json
import matplotlib.pyplot as plt

from ..Logger import Logger, LoggerInterface

TRAINING_HISTORY_FILE_NAME = "history.json"
SEED = 42

class AbstractNN:
    """ 
    Defines an abstract class for the neural networks used by the Decision Boundary Mapper.

    Public methods:
        load \n
        save \n
        show_predictions \n

    Methods to be implemented by the classes that inherit from this class.
        encode: Encodes the input data.
    """

    def __init__(self,
                 folder_path: str,
                 nn_name: str,
                 logger: LoggerInterface | None = None):
        """ 
        Initializes an NN object.

        Args:
            folder_path (str): the folder path where the model will be saved/loaded.
            nn_name (str): Name of the neural network that uses the classifier.
            logger (LoggerInterface, optional): Defaults to console logging.

        Raises:
            Exception: If the classifier is not provided.
        """
        if logger is not None:
            self.console = logger
        else:
            self.console = Logger(name=nn_name)

        self.save_folder_path = folder_path
        self.nn_name = nn_name
        self.neural_network = None

        try:
            self.load()
        except Exception as e:
            self.console.warn("Model not found, {}".format(e))
            self.console.warn("The model will be built and trained from scratch.")

    def load(self):
        """
            Loads an auto encoder from the specified folder path. With the .h5 extension.
            Args:
                folder_path (str): The path to the folder where the model is saved.
        """
        try:
            self.neural_network = tf.keras.models.load_model(os.path.join(self.save_folder_path, self.nn_name), compile=False)
            self.console.log("NN loaded successfully")
        except Exception as e:
            self.console.log(
                f"NN not found. Please check the path folder {self.save_folder_path} and make sure the model is saved there")
            raise e

    def save(self, history=None):
        """
            Saves the model to the specified folder path. With the .tf extension.
        """
        folder_path = self.save_folder_path
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)

        path = os.path.join(folder_path, self.neural_network.name)  # type: ignore
        self.neural_network.save(path, save_format="tf")            # type: ignore
        self.console.log(f"Model saved to {folder_path}")

        if history is None:
            return

        with open(os.path.join(folder_path, TRAINING_HISTORY_FILE_NAME), "w") as f:
            f.write(json.dumps(history.history))
            self.console.log(f"History saved to {folder_path}")

    def show_predictions(self, dataNd: np.ndarray, labels: np.ndarray, data2d: np.ndarray | None = None, n_samples: int = 16):
        """ Shows the predictions of the model. By taking first n_samples data points and comparing them to the actual labels.

        Args:
            dataNd (np.ndarray): N-dimensional data points.
            data2d (np.ndarray, optional): 2-dimensional data points. Defaults to None.
            labels (np.ndarray): Actual labels of the data points
            n_samples (int, optional): Number of samples to show. Defaults to 16.
        """
        if data2d is None:
            data2d = self.encode(dataNd)

        decoded, predicted = self.decode(data2d)
        predicted_labels = np.argmax(predicted, axis=1)

        plt.figure(figsize=(20, 10))
        m = int(sqrt(n_samples))
        n = m * m

        for i in range(1, 2*n, 2):
            plt.subplot(m, 2*m, i)
            plt.imshow(dataNd[i], cmap='gray')
            plt.title(f"Actual: {labels[i]}", color='green' if predicted_labels[i]== labels[i] else 'red', fontsize=12)
            plt.axis('off')
            plt.subplot(m, 2*m, i+1)
            plt.imshow(decoded[i], cmap='gray')
            plt.title(f"Predicted: {predicted_labels[i]}", color='green' if predicted_labels[i] == labels[i] else 'red', fontsize=12)
            plt.axis('off')

        plt.show()

    def decode(self, data: np.ndarray, verbose: int = 0) -> np.ndarray:
        """ Decodes the data points.

        Args:
            data (np.ndarray): The data points to decode.
            verbose (int, optional): Verbosity mode. Defaults to 0.

        Returns:
            Xnd, labels_predictions: The decoded data points and the predictions of the classifier.
        """
        return self.neural_network.predict(data, verbose=verbose)  # type: ignore

    def encode(self, data: np.ndarray, verbose: int = 0) -> np.ndarray:
        """ Encodes the data points.

        Args:
            data (np.ndarray): The data points to encode.
            verbose (int, optional): Verbosity mode. Defaults to 0.

        Returns:
            np.ndarray: The encoded 2D data points.
        """
        raise NotImplementedError("The encode method is not implemented.")
