# 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 src.decision_boundary_mapper import Logger
from src.decision_boundary_mapper.Logger import LoggerInterface


class NNinterface:
    """ 
    Defines the interface for the neural networks used by the Decision Boundary Mapper.
    
    Methods to be implemented by the classes that inherit from this class.
    fit: Fits the model to the specified data.
    encode: Encodes the input data. (optional)
    """
    
    def __init__(self, 
                 folder_path: str,
                 classifier = None,
                 nn_name: str = "invNN", 
                 logger: LoggerInterface = None):
        """ Initializes the NN interface.

        Args:
            folder_path (str): the folder path where the model will be saved/loaded.
            classifier (tf.keras.model): The classifier model. Defaults to None.
            nn_name (str, optional): Name of the neural network that uses the classifier. Defaults to "invNN".
            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)
        
        if classifier is None:
            self.console.error("Please provide a classifier.")
            raise Exception("Please provide a classifier.")
        
        self.save_folder_path = folder_path
        self.classifier = classifier
        self.nn_name = nn_name
        self.neural_network = None
        
        try:
            self.load()
        except Exception as e:
            self.console.error("Error loading the model: {}".format(e))
            self.console.warn("The model will be built and trained from scratch.")
            
    def fit(self):
        """ Fits the model to the specified data."""
        raise NotImplementedError("The fit method is not implemented.")
    
    
    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):
        """
            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)
            
        self.neural_network.save(os.path.join(folder_path, self.neural_network.name), save_format="tf")
        self.console.log(f"Model saved to {folder_path}")
        
        with open(os.path.join(folder_path, "history.json"), "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, 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):
        """ 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)
    
    def encode(self, data: np.ndarray, verbose: int = 0):
        """ 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.
        """
        return NotImplementedError("The encode method is not implemented.")