import json
import os

import numpy as np
from matplotlib import pyplot as plt
from scipy.stats import norm, pareto, lognorm, expon, gamma, ks_2samp


def take_arrivals_from_folder(model_dir, num_realizations=10):
    arrivals_covariates = []
    realizations_dir = [a for a in os.listdir(model_dir)]
    for realization_dir in realizations_dir[:num_realizations]:
        arrivals_files = [a for a in os.listdir(model_dir + realization_dir) if "arrivals" in a]
        ARRIVALS_REALIZATION = []
        for arrival_file in arrivals_files:
            ARRIVALS_REALIZATION.extend(np.loadtxt(model_dir + realization_dir + "/" + arrival_file))
        arrivals_covariates.append(np.asarray(ARRIVALS_REALIZATION))
    arrivals_covariates = np.asarray(arrivals_covariates)
    return arrivals_covariates


def service_from_covariates(ARRIVALS_COVARIATES_LIST, arrivals_time_index=1):
    service_times = []
    proportions = []
    for arrivals_covariates in ARRIVALS_COVARIATES_LIST:
        positive_index = np.where(arrivals_covariates[:, -1] > 0)[0]
        services = arrivals_covariates[positive_index, -1] - arrivals_covariates[positive_index, arrivals_time_index]
        service_times.append(services)
        uncensor_proportion = (len(positive_index) / float(arrivals_covariates.shape[0]))
        proportions.append(uncensor_proportion)
    return service_times, proportions


def all_metadata_from_folders(model_dir):
    realizations_dir = [a for a in os.listdir(model_dir)]
    all_metadata = [json.loads(open(model_dir + mdir + "/metadata.json", "r").read()) for mdir in realizations_dir]
    all_num_of_arrivals = [int(metadata["total_arrivals"]) for metadata in all_metadata]
    return all_metadata, all_num_of_arrivals


def loglog_distribution_plots(service, service_simulated=None,
                              plot_dist=["pareto", "norm", "expon", "gamma", "lognorm"], x_lim=None,
                              y_lim=None, ax=None, title=None):
    if x_lim is None:
        xMin = min(service)
        xMax = max(service)
    else:
        xMin, xMax = x_lim
    dX = (xMax - xMin) * 0.001
    support = np.arange(xMin, xMax, dX)

    if "pareto" in plot_dist:
        dist_param = pareto.fit(service)
        myDist = pareto(*dist_param)
        Y_pareto = myDist.pdf(support)
        print "Pareto: " + str(ks_2samp(service, Y_pareto)[0])

    if "norm" in plot_dist:
        dist_param = norm.fit(service)
        myDist = norm(*dist_param)
        Y_norm = myDist.pdf(support)
        print "Norm: " + str(ks_2samp(service, Y_norm)[0])

    if "expon" in plot_dist:
        dist_param = expon.fit(service)
        myDist = expon(*dist_param)
        Y_exp = myDist.pdf(support)
        print "Exp: " + str(ks_2samp(service, Y_exp)[0])

    if "lognorm" in plot_dist:
        dist_param = lognorm.fit(service)
        myDist = lognorm(*dist_param)
        Y_lognorm = myDist.pdf(support)
        print "LogNorm: " + str(ks_2samp(service, Y_lognorm)[0])

    if "gamma" in plot_dist:
        dist_param = gamma.fit(service)
        myDist = gamma(*dist_param)
        Y_gamma = myDist.pdf(support)
        print "Gamma: " + str(ks_2samp(service, Y_gamma)[0])

    if "lognorm" in plot_dist:
        if ax is None:
            plt.plot(support, Y_lognorm, label="ln", linewidth=6.0)
        else:
            ax.plot(support, Y_lognorm, linewidth=6.0)
    if "gamma" in plot_dist:
        if ax is None:
            plt.plot(support, Y_gamma, label="gamma", linewidth=6.0)
        else:
            ax.plot(support, Y_gamma, linewidth=6.0)

    if "norm" in plot_dist:
        if ax is None:
            plt.plot(support, Y_norm, linewidth=6.0)
        else:
            ax.plot(support, Y_norm, linewidth=6.0)
    if "pareto" in plot_dist:
        if ax is None:
            plt.plot(support, Y_pareto, label="pareto", linewidth=6.0)
        else:
            ax.plot(support, Y_pareto, linewidth=6.0)
    if "expon" in plot_dist:
        if ax is None:
            plt.plot(support, Y_exp, linewidth=6.0)
        else:
            ax.plot(support, Y_exp, linewidth=6.0)

    if x_lim is not None:
        if ax is None:
            plt.xlim(x_lim)
        else:
            ax.set_xlim(x_lim)
    if y_lim is not None:
        if ax is None:
            plt.ylim(y_lim)
        else:
            ax.set_ylim(y_lim)

    if ax is None:
        plt.yscale("log")
        plt.xscale("log")
    else:
        ax.set_yscale("log")
        ax.set_xscale("log")

    if ax is None:
        plt.ylabel("Frequency")
        plt.xlabel("Service Time (s)")
    else:
        ax.set_ylabel("Frequency")
        ax.set_xlabel("Service Time (s)")
    count, bins = np.histogram(service, bins=100, normed=True)
    if ax is None:
        plt.plot(bins[1:], count, "o", label="empirical", markersize=10)
    else:
        ax.plot(bins[1:], count, "o", label="empirical", markersize=10)
    if service_simulated is not None:
        count, _ = np.histogram(service_simulated, bins=bins, normed=True)
        if ax is None:
            plt.plot(bins[1:], count, "D", label="model", markersize=10)
        else:
            ax.plot(bins[1:], count, "D", label="model", markersize=10)
    if ax is None:
        plt.grid(True)
        plt.legend(bbox_to_anchor=(0, 1.02, 1, 0.2), loc="lower left",
                   mode="expand", borderaxespad=0, ncol=2)
    else:
        ax.grid(True)
        ax.legend(bbox_to_anchor=(0, 1.02, 1, 0.2), loc="lower left",
                  mode="expand", borderaxespad=0, ncol=2)
    if title is not None:
        if ax is None:
            plt.title(title)
        else:
            ax.set_title(title)
    return ax


def loglog_qq_plots(service, service_rnn=None, service_neural=None, x_lim=None, y_lim=None, qq_lim=1.,
                    neural_label="NS-G", plot_ylabels=True):
    if x_lim is None:
        xMin = min(service)
        xMax = max(service)
    else:
        xMin, xMax = x_lim
    dX = (xMax - xMin) * 0.001
    fig, ax = plt.subplots(2, 1)
    if x_lim is not None:
        ax[0].set_xlim(x_lim)
    if y_lim is not None:
        ax[0].set_ylim(y_lim)

    ax[0].set_yscale("log")
    ax[0].set_xscale("log")
    if plot_ylabels:
        ax[0].set_ylabel("Frequency")
    ax[0].set_xlabel("Service Time (s)")
    count_emp, bins = np.histogram(service, bins=2000, normed=True)
    ax[0].plot(bins[1:], count_emp, "o", label="Data", markersize=10, color="royalblue", )[0]
    count_rnn, _ = np.histogram(service_rnn, bins=bins, normed=True)
    ax[0].plot(bins[1:], count_rnn, "D", label="RAS", markersize=10, color="orange")[0]

    l = np.linspace(0, 1., 500)

    ax[1].scatter(np.quantile(service, l),
                  np.quantile(service_rnn, l), c='orange', marker="D", s=100, label="RAS")

    if service_neural is not None:
        ax[1].scatter(np.quantile(service, l),
                      np.quantile(service_neural, l), c='forestgreen', marker="*", s=100, label=neural_label)

    ll = np.linspace(0, qq_lim, 400)

    ax[1].plot(ll, ll, ls=(0, (5, 10)), alpha=0.4)
    ax[1].set_yscale("log")
    ax[1].set_xscale("log")
    ax[1].set_xlim(0, qq_lim)
    ax[1].set_ylim(0, qq_lim)
    if plot_ylabels:
        ax[1].set_ylabel("Model")
    ax[1].set_xlabel("Data")

    ax[0].legend()
    ax[1].legend()
