import os
from math import ceil
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.rpp_models import RnnHawkes

from prolothar_queue_mining.model.job import Job

class RecurrentPointProcessArrivalModel:

    def __init__(self, log_directory: str|None = None):
        #sometimes, tensorflow seems to block deletion
        if log_directory is None:
            self.__temporary_directory = TemporaryDirectory(prefix='arrival-', ignore_cleanup_errors=True)
        else:
            self.__temporary_directory = log_directory
        self.__path_to_model = None
        self.__max_arrival_time = None
        self.__params = None

    def train(self, observed_arrivals: list[tuple[Job, int]], parameters: dict):
        params = Bunch(parameters)
        self.__params = params
        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')
        observed_arrivals = [(None,0)] + observed_arrivals
        arrival_times = np.array([[arrival_time] for _,arrival_time in observed_arrivals])
        self.__max_arrival_time = arrival_times.max()
        arrival_times = arrival_times / self.__max_arrival_time
        inter_arrival_times = np.diff(arrival_times, axis=0)
        data = np.concatenate((
            arrival_times[1:,:],
            inter_arrival_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)

        with tf.compat.v1.Graph().as_default():
            model = RnnHawkes(params.__dict__)
            with tf.compat.v1.Session() as sess:
                model.train(train_set, test_set, 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')
        with tf.compat.v1.Graph().as_default():
            with tf.compat.v1.Session() as sess:
                model = RnnHawkes.restore(sess, self.__path_to_model)
                train_hidden_states, test_hidden_states = model.extract_hidden_states(
                    train_set, test_set, sess)

        return train_hidden_states, test_hidden_states

    def compute_hidden_states(self, arrivals: list[tuple[Job, int]]) -> np.ndarray:
        arrivals = [(None,0)] + arrivals
        arrival_times = np.array([[arrival_time] for _,arrival_time in arrivals])
        arrival_times = arrival_times / self.__max_arrival_time
        inter_arrival_times = np.diff(arrival_times, axis=0)
        data = np.concatenate((
            arrival_times[1:,:],
            inter_arrival_times), axis=1)
        data = np.array_split(data, 1)
        data = preprocess_data(data, len(data), self.__params.num_steps)
        data = [t[0] for t in data]

        tf.compat.v1.reset_default_graph()
        with tf.compat.v1.Graph().as_default():
            with tf.compat.v1.Session() as sess:
                model = RnnHawkes.restore(sess, self.__path_to_model)
                return model.extract_hidden_states_test(data, sess)