from math import ceil
import os
from pathlib import Path
from tempfile import TemporaryDirectory
import numpy as np
import tensorflow as tf

from rnpp.utils.etc.data_handling import split_train_test_and_preprocess
from rnpp.utils.etc.data_handling import preprocess_data
from rnpp.utils.etc.utils import Bunch

# from rnpp.npp.recurrent_point.service_time_models import RASBetaStaticCritic as RasModel
from rnpp.npp.recurrent_point.service_time_models import RecurrentAdversarialServiceTimeBeta as RasModel

from prolothar_queue_mining.model.job import Job
from prolothar_queue_mining.inference.sklearn.job_to_vector_transformer_utils import create_job_to_vector_transformer_for_neural_model

class ServiceTimeModel:

    def __init__(
            self, categorical_attribute_names: list[str] = None,
            numerical_attribute_names: list[str] = None,
            log_directory: str|None = None):
        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.__job_to_vector_transformer = None
        if log_directory is None:
            #sometimes, tensorflow seems to block deletion
            self.__temporary_directory = TemporaryDirectory(prefix='service-', ignore_cleanup_errors=True)
        else:
            self.__temporary_directory = log_directory
        self.__path_to_model = None
        self.__training_params = None
        self.__max_service_time = None

    def train(
            self, observed_arrivals: list[tuple[Job, int]],
            observed_departures: list[tuple[Job, int]],
            train_hidden_states: np.ndarray, test_hidden_states: np.ndarray,
            parameters: dict):
        self.__job_to_vector_transformer = create_job_to_vector_transformer_for_neural_model(
            [job for job,_ in observed_arrivals],
            self.__categorical_attribute_names,
            self.__numerical_attribute_names)
        params = Bunch(parameters)
        self.__training_params = params
        params.input_dim = 1 + len(self.__job_to_vector_transformer.get_feature_names_of_vector_components())
        params.covariates_bool = params.input_dim > 1
        if isinstance(self.__temporary_directory, str):
            params.store_dir = os.path.join(self.__temporary_directory, 'save')
            params.log_dir = os.path.join(self.__temporary_directory, 'log')
        else:
            params.store_dir = os.path.join(self.__temporary_directory.name, 'save')
            params.log_dir = os.path.join(self.__temporary_directory.name, 'log')

        departure_time_per_job = dict(observed_departures)
        service_times = np.array([
            [departure_time_per_job.get(job, arrival_time) - arrival_time]
            for job,arrival_time in observed_arrivals
        ])
        self.__max_service_time = service_times.max()
        service_times = service_times / self.__max_service_time
        covariates = np.array([
            self.__job_to_vector_transformer.transform(job)
            for job,_ in observed_arrivals
        ])
        data = np.concatenate((
            covariates,
            service_times), axis=1)
        data = np.array_split(data, ceil(data.shape[0] / params.num_sequences))

        train_set, test_set = split_train_test_and_preprocess(
            data, params.batch_size, params.num_steps, params.test_size)
        params.batch_size_test = test_set[0].shape[0]

        tf.compat.v1.reset_default_graph()
        with tf.compat.v1.Graph().as_default():
            model = RasModel(params.__dict__)
            with tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(log_device_placement=False)) as sess:
                model.train(train_set, test_set, train_hidden_states, test_hidden_states, sess)

        path_to_model = os.path.join(params.store_dir, os.listdir(params.store_dir)[0])
        path_to_model = sorted(Path(path_to_model).iterdir(), key=os.path.getctime)[-1]
        self.__path_to_model = os.path.join(path_to_model, 'best-model')

    def sample(self, hidden_states, jobs: list[Job]) -> np.ndarray:
        tf.compat.v1.reset_default_graph()
        with tf.compat.v1.Graph().as_default():
            with tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(log_device_placement=False)) as sess:
                covariates = np.array([
                    self.__job_to_vector_transformer.transform(job)
                    for job in jobs
                ])

                data = np.concatenate((
                    covariates, np.zeros((covariates.shape[0], 1))), axis=1)
                data = np.array_split(data, 1)
                data = preprocess_data(data, len(data), self.__training_params.num_steps)
                data = [t[0] for t in data]

                model = RasModel.restore(
                    sess, self.__path_to_model, batch_size_test=data[0].shape[0])

                y = model.sample(data, hidden_states, sess)
                y = y * self.__max_service_time / y.max()
                return y