import os
from ast import literal_eval
from typing import Iterable

from prolothar_queue_mining.model.distribution import Distribution
from prolothar_queue_mining.model.distribution import PoissonDistribution
from prolothar_queue_mining.model.distribution import GammaDistribution
from prolothar_queue_mining.model.distribution import ExponentialDistribution
from prolothar_queue_mining.model.distribution import DiscreteDegenerateDistribution
from prolothar_queue_mining.model.distribution import GeometricDistribution
from prolothar_queue_mining.model.distribution import NegativeBinomialDistribution
from prolothar_queue_mining.model.population import Population
from prolothar_queue_mining.model.population import InfinitePopulation
from prolothar_queue_mining.model.job import Job

def parse_distribution(distribution: str, seed: int=None) -> Distribution:
    if distribution.startswith('D'):
        return DiscreteDegenerateDistribution(int(distribution[1:]))
    elif distribution.startswith('P'):
        return PoissonDistribution(
            float(distribution[1:]) - 1, shift=1, seed=seed)
    elif distribution.startswith('E'):
        return ExponentialDistribution(
            float(distribution[1:]), seed=seed)
    elif distribution.startswith('G-'):
        _, shape, rate = distribution.split('-')
        return GammaDistribution(
            float(shape), float(rate), seed=seed)
    elif distribution.startswith('Geo'):
        return GeometricDistribution(
            float(distribution[3:]), seed=seed
        )
    elif distribution.startswith('NB'):
        _, nr_of_successes, success_probability = distribution.split('-')
        return NegativeBinomialDistribution(
            float(nr_of_successes), float(success_probability), seed=seed
        )
    else:
        raise NotImplementedError(distribution)

def parse_population(
        nr_of_categorical_features: int,
        nr_of_categories: int,
        nr_of_numerical_features: int,
        seed: int=None) -> Population:
    return InfinitePopulation(
        categorical_feature_names=[
            f'cat_feature_{i}' for i in range(nr_of_categorical_features)],
        numerical_feature_names=[
            f'num_feature_{i}' for i in range(nr_of_numerical_features)],
        nr_of_categories=nr_of_categories,
        seed=seed
    )

def create_job_from_string(s: str) -> Job:
    #Job(435890_5896_1-2, {'BLECHID': '435890_5896_1', 'SPARTE': 101, 'QUALITAETSGRUPPE': 3, 'QUALITAET_LEGIERUNG': 'Y', 'LAENGE_IN_MM': 1269, 'Start': Timestamp('2019-01-11 22:27:26'), 'Duration': Timedelta('0 days 06:07:26')})
    job_id = s[4:s.index(', ')]
    feature_str = s[s.index('{'):s.index('}')+1]
    return Job(job_id, literal_eval(feature_str))

def load_dataset_with_waiting_times(dataset_path: str):
    job_id_dict = {}
    with open(os.path.join(dataset_path, 'jobs.txt')) as f:
        for line in f.readlines():
            if line:
                job = create_job_from_string(line)
                job_id_dict[job.job_id] = job

    arrivals = []
    with open(os.path.join(dataset_path, 'arrivals.txt')) as f:
        for line in f.readlines():
            if line:
                arrival_time, job_id = line.strip().split(';')
                arrival_time = int(arrival_time)
                arrivals.append((job_id_dict[job_id], arrival_time))

    departures = []
    with open(os.path.join(dataset_path, 'departures.txt')) as f:
        for line in f.readlines():
            if line:
                departure_time, job_id = line.strip().split(';')
                departure_time = int(departure_time)
                departures.append((job_id_dict[job_id], departure_time))

    waiting_times = {}
    if os.path.exists(os.path.join(dataset_path, 'waiting_times.txt')):
        with open(os.path.join(dataset_path, 'waiting_times.txt')) as f:
            for line in f.readlines():
                if line:
                    job_id, waiting_time = line.strip().split(';')
                    waiting_time = int(waiting_time)
                    waiting_times[job_id_dict[job_id]] = waiting_time

    return arrivals, departures, waiting_times

def load_dataset(dataset_path: str):
    arrivals, departures, _ = load_dataset_with_waiting_times(dataset_path)
    return arrivals, departures

def save_dataset(
        dataset_path: str, observed_arrivals: list[Job, int],
        observed_departures: list[Job, int], all_jobs: Iterable[Job],
        waiting_time_per_job: dict[Job, int] = None):
    with open(os.path.join(dataset_path, 'arrivals.txt'), 'w', encoding='utf-8') as f:
        i = 0
        for job, arrival_time in observed_arrivals:
            f.write(f'{arrival_time};{job.job_id}\n')
    with open(os.path.join(dataset_path, 'departures.txt'), 'w', encoding='utf-8') as f:
        for job, departure_time in observed_departures:
            f.write(f'{departure_time};{job.job_id}\n')
    with open(os.path.join(dataset_path, 'jobs.txt'), 'w', encoding='utf-8') as f:
        for job in all_jobs:
            f.write(str(job))
            f.write('\n')
    if waiting_time_per_job:
        with open(os.path.join(dataset_path, 'waiting_times.txt'), 'w', encoding='utf-8') as f:
            for job, waiting_time in waiting_time_per_job.items():
                f.write(f'{job.job_id};{waiting_time}\n')