"""
    Authors: <NAME>, <NAME>
    Project: Graduation Thesis: GIAdog

    [TODO: DESCRIPTION]

"""
import os
import numpy as np
import tensorflow as tf

from gym import spaces
from tensorflow import keras
from ANN import ANN, Normalizer
from typing import Tuple, List, Dict
from __env__ import PRIVILIGED_DATA, HISTORIAL_DATA, NON_PRIVILIGED_DATA,  \
    STATE_SIZE, PRIVILEGED_STATE_SIZE


PRIVILIGED_SHAPE = sum(PRIVILIGED_DATA[d] for d in PRIVILIGED_DATA)
PROPIOCEPTIVE_SHAPE = sum(HISTORIAL_DATA[d] for d in HISTORIAL_DATA) +\
    sum(NON_PRIVILIGED_DATA[d] for d in NON_PRIVILIGED_DATA)


class TeacherValueANN(ANN):
    """ Teacher value ANN class. """
    
    def __init__(self, observation_space: spaces.Dict):
        """
            Initializes the teacher value ANN.

            Arguments:
            ----------
                observation_space: gym.spaces.Dict
                    The observation space of the environment
        """
        self.observation_space = observation_space

        inputs_x_t = keras.Input(
            shape=PRIVILEGED_STATE_SIZE, 
            name='priviliged_data'
        )

        inputs_o_t = keras.Input(
            shape=STATE_SIZE,
            name='propioceptive_data'
        )

        concat = keras.layers.Concatenate()([inputs_x_t, inputs_o_t])
        
        x = keras.layers.Dense(256, activation='tanh')(concat)
        x = keras.layers.Dense(128, activation='tanh')(x)
        x = keras.layers.Dense(64, activation='tanh')(x)
        x = keras.layers.Dense(16, activation='tanh')(x)
        outputs = keras.layers.Dense(1)(x)
        self.model = keras.Model(
            inputs= [inputs_x_t, inputs_o_t], 
            outputs=outputs, 
            name='critic_model'
        )
        self.model.compile(
            optimizer='nadam',
            loss='mse'
        )

        # Normalizers
        self.x_normalizer = Normalizer(PRIVILEGED_STATE_SIZE)
        self.o_normalizer = Normalizer(STATE_SIZE)

    def __format_states(
            self, 
            states: List[np.ndarray]
        ) -> Tuple[np.ndarray, np.ndarray]:
        """
            Get dictionary states and convert them to numpy arrays.
        """
        states = np.array(states)
        input_x_t = self.x_normalizer(states[:,:PRIVILEGED_STATE_SIZE])
        input_o_t = self.o_normalizer(states[:,PRIVILEGED_STATE_SIZE:])
        return input_x_t, input_o_t
    
    def __call__(self, states: List[np.ndarray]) -> tf.Tensor:
        """
            Computes the value of the states.

            Arguments:
            ----------
                states: List[numpy.array]
                    States to process

            Return:
            -------
                tensorf.Tensor
                    Value computed
        """
        # Get ANN input
        assert self.verify_states(states)
        input_x_t, input_o_t = self.__format_states(states)

        return self.model([input_x_t, input_o_t])

    def fit(self, states: List[np.ndarray], y: List[float], *args, **kwargs):
        """
            TODO
        """
        # Get ANN input
        assert self.verify_states(states)
        input_x_t, input_o_t = self.__format_states(states)

        self.model.fit(
            [input_x_t, input_o_t], 
            np.array(y, dtype=np.float32), 
            *args, 
            **kwargs
        )

    def save(self, path: str):
        os.makedirs(path, exist_ok=True)
        self.model.save_weights(f'{path}/model.h5')
        np.savez(
            f'{path}/x_normalizer.npz',
            x_normalizer_n = self.x_normalizer.n,
            x_normalizer_mean = self.x_normalizer.mean,
            x_normalizer_var = self.x_normalizer.var,
        )
        np.savez(
            f'{path}/o_normalizer.npz',
            o_normalizer_n = self.o_normalizer.n,
            o_normalizer_mean = self.o_normalizer.mean,
            o_normalizer_var = self.o_normalizer.var,
        )

    def load(self, path: str):
        self.model.load_weights(f'{path}/model.h5')

        x_data = np.load(f'{path}/x_normalizer.npz')
        self.x_normalizer.n = x_data['x_normalizer_n']
        self.x_normalizer.mean = x_data['x_normalizer_mean']
        self.x_normalizer.var = x_data['x_normalizer_var']

        o_data = np.load(f'{path}/o_normalizer.npz')
        self.o_normalizer.n = o_data['o_normalizer_n']
        self.o_normalizer.mean = o_data['o_normalizer_mean']
        self.o_normalizer.var = o_data['o_normalizer_var']

class StudentValueANN(ANN):
    pass
