from dataclasses import dataclass

from prolothar_queue_mining.model.job import Job

from prolothar_queue_mining.inference.queue.ras.arrival_model import RecurrentPointProcessArrivalModel
from prolothar_queue_mining.inference.queue.ras.service_time_model import ServiceTimeModel
from prolothar_queue_mining.prediction.departure_time.ras_departure_time_predictor import RasDepartureTimePredictor

@dataclass
class _LearningParameters:
    batch_size: int = 32
    num_epochs: int = 100
    hidden_units: int = 32

@dataclass
class ArrivalLearningParameters(_LearningParameters):
    learning_rate: float = 0.00001

@dataclass
class ServiceTimeLearningParameters(_LearningParameters):
    censor_penalty: float = 50
    decay_beta_censor: float = 0.0
    n_mc_samples: int = 20
    batch_normalization: bool = False
    lambda_grad_penalty: float = 10
    generator_dimensions: int = 50
    discriminator_dimensions: int = 100
    learning_rate_gan: float = 0.0001
    nr_of_layers_generator: int = 3
    nr_of_layers_disciminator: int = 3
    dropout: bool = False
    dropout_probability: float = 0.8
    n_discriminator_updates: int = 10
    n_generator_updates: int = 1
    noisy_layers: bool = True
    clip_dis: float = 0.01

class RecurrentAdversarialServiceTimes:

    def __init__(
            self, random_seed: int = None,
            categorical_attribute_names: list[str] = None,
            numerical_attribute_names: list[str] = None,
            test_size: float = 0.1,
            num_steps: int = 100,
            num_sequences: int = 100,
            arrival_learning_parameters: ArrivalLearningParameters = ArrivalLearningParameters(),
            service_time_learning_parameters: ServiceTimeLearningParameters = ServiceTimeLearningParameters(),
            log_directory_arrivals: str = None,
            log_directory_service_times: str = None):
        self.__random_seed = random_seed
        if categorical_attribute_names is not None:
            self.__categorical_attribute_names = categorical_attribute_names
        else:
            self.__categorical_attribute_names = []
        if numerical_attribute_names is not None:
            self.__numerical_attribute_names = numerical_attribute_names
        else:
            self.__numerical_attribute_names = []
        self.__test_size = test_size
        self.__num_steps = num_steps
        self.__arrival_learning_parameters = arrival_learning_parameters
        self.__service_time_learning_parameters = service_time_learning_parameters
        self.__log_directory_arrivals = log_directory_arrivals
        self.__log_directory_service_times = log_directory_service_times
        self.__num_sequences = num_sequences

    def train(
            self, observed_arrivals: list[tuple[Job, int]],
            observed_departures: list[tuple[Job, int]]) -> RasDepartureTimePredictor:
        arrival_model = RecurrentPointProcessArrivalModel(log_directory=self.__log_directory_arrivals)
        train_hidden_states, test_hidden_states = arrival_model.train(
            observed_arrivals, self.__create_arrival_params_dict())

        service_time_model = ServiceTimeModel(
            categorical_attribute_names=self.__categorical_attribute_names,
            numerical_attribute_names=self.__numerical_attribute_names,
            log_directory=self.__log_directory_service_times
        )
        service_time_model.train(
            observed_arrivals, observed_departures,
            train_hidden_states, test_hidden_states,
            self.__create_service_times_params_dict())

        return RasDepartureTimePredictor(arrival_model, service_time_model)

    def __create_common_params_dict(self) -> dict:
        return {
            'num_runs': 1,
            'num_sequences': self.__num_sequences,
            'num_steps': self.__num_steps,
            'name': 'ras',
            'test_size': self.__test_size,
            'numpy_seed': self.__random_seed,
            'tf_seed': self.__random_seed,
            'cell': 'GRU',
            'initializer': 'normal',
            'init_value': 0.0001,
            'num_quadrature_nodes': 30,
            'quadrature_window': 10,
            'inverse_sampling_size': 1000,
            'log': 'INFO',
            'logging_interval': 10
        }

    def __create_arrival_params_dict(self) -> dict:
        return self.__create_common_params_dict() | {
            'batch_size': self.__arrival_learning_parameters.batch_size,
            'num_epochs': self.__arrival_learning_parameters.num_epochs,
            'input_dim': 2,
            'hidden_units': self.__arrival_learning_parameters.hidden_units,
            'M': 0,
            'K': 0,
            'learning_rate': self.__arrival_learning_parameters.learning_rate,
        }

    def __create_service_times_params_dict(self) -> dict:
        return self.__create_common_params_dict() | {
            'events_dim': 1,
            'batch_size': self.__service_time_learning_parameters.batch_size,
            'num_epochs': self.__service_time_learning_parameters.num_epochs,
            'hidden_units': self.__service_time_learning_parameters.hidden_units,
            'censor_penalty': self.__service_time_learning_parameters.censor_penalty,
            'decay_beta_censor': self.__service_time_learning_parameters.decay_beta_censor,
            'n_mc_samples': self.__service_time_learning_parameters.n_mc_samples,
            'batch_normalization': self.__service_time_learning_parameters.batch_normalization,
            'lambda_grad_penalty': self.__service_time_learning_parameters.lambda_grad_penalty,
            'gen_dim': self.__service_time_learning_parameters.generator_dimensions,
            'dis_dim': self.__service_time_learning_parameters.discriminator_dimensions,
            'learning_rate_gan': self.__service_time_learning_parameters.learning_rate_gan,
            'n_layers_gen': self.__service_time_learning_parameters.nr_of_layers_generator,
            'n_layers_dis': self.__service_time_learning_parameters.nr_of_layers_disciminator,
            'dropout': self.__service_time_learning_parameters.dropout,
            'dropout_prob': self.__service_time_learning_parameters.dropout_probability,
            'n_discriminator_updates': self.__service_time_learning_parameters.n_discriminator_updates,
            'n_generator_updates': self.__service_time_learning_parameters.n_generator_updates,
            'noisy_layers': self.__service_time_learning_parameters.noisy_layers,
            'clip_dis': self.__service_time_learning_parameters.clip_dis,
            'random_mixing_labels': False,
            'instance_noise': False,
            'arrivals_hidden_state_bool': True,
            'use_hidden': True
        }

