import torch
import numpy as np
import matplotlib.pyplot as plt
import time
from EMRI_DET.utilities import norm_inputs, unnorm, get_script_path
import seaborn as sns
sns.set_theme()


def run_on_dataset(model, test_data, n_batches=1, device=None, y_transform_fn=None, runtime=False):
    """
    Get the re-processed output of the supplied model on a set of supplied test data.

    Args:
        model (LinearModel): Model to test on `test_data`
        test_data (2-tuple/list): Tuple or list of the 'features' and their corresponding 'labels'.
        n_batches (int, optional): Number of batches to process the input data in. Defaults to 1 (the entire dataset).
        device (string, optional): Device to attach the input model to, if it is not attached already.
        y_transform_fn (function, optional): If the labels/ydata have been pre-processed with a function (e.g. log),
                                             this function must be supplied to undo this if comparison is to be made
                                             in the unaltered function space.
        runtime (bool, optional): If True, return timing statistics for the model on the provided dataset.
                                  Defaults to False.

    Returns:
        tuple containing:
            output_data (double): root-mean-square-error between the truth and the model output for this dataset.
            total_time (double): total time taken by the model to process the input dataset.
                                 Only if `runtime` is set to True.
            per_point (double): mean time taken by the model per sample in the input dataset.
                                 Only if `runtime` is set to True.
    """
    if device is not None:
        model = model.to(device)
    model.eval()

    xdata, ydata = test_data

    xmeanstd = np.load(get_script_path() + f'/../models/{model.name}_xdata_mean_std.npy')
    ymeanstd = np.load(get_script_path() + f'/../models/{model.name}_ydata_mean_std.npy')

    test_input = torch.Tensor(xdata)
    normed_input = norm_inputs(test_input, ref_mean=xmeanstd[0], ref_std=xmeanstd[1]).float().to(device)

    if runtime:
        st = time.perf_counter()
    with torch.no_grad():
        out = []
        for i in range(n_batches):
            output = model(normed_input[i * ydata.size // n_batches: (i + 1) * ydata.size // n_batches])
            out.append(output.detach().cpu().numpy())

    if runtime:
        et = time.perf_counter()
        total_time = et - st
        per_point = (et - st) / ydata.size

    output = np.concatenate(out)
    out_unnorm = unnorm(output, ref_mean=ymeanstd[0], ref_std=ymeanstd[1]).flatten()

    if y_transform_fn is not None:
        out_unnorm = y_transform_fn(out_unnorm)

    outputs = (out_unnorm,)

    if runtime:
        outputs += (total_time, per_point,)

    return outputs


def compute_rmse(comparison_sets):
    """
    Get the root-mean-square-error difference between two sets of data.
    Args:
        comparison_sets (2-tuple): Tuple of the list of true values, followed by the list of predicted values.
    Returns:
        rmse (double): root-mean-square-error between the truth and the model output for this dataset.
    """
    truth, pred = comparison_sets
    rmse = np.sqrt(np.mean((truth - pred) ** 2))
    return rmse


def test_threshold_accuracy(comparison_sets, threshold):
    truth, pred = comparison_sets

    out_classified = np.zeros(shape=pred.size)
    out_classified[pred.flatten() >= threshold] = 1

    truth_classified = np.zeros(shape=truth.size)
    truth_classified[truth >= threshold] = 1

    return 1 - np.mean(np.abs(out_classified - truth_classified))


def plot_histograms(comparison_sets, model_name, xlabel, title=None, title_kwargs={}, xlabel_kwargs={}, log=True,
                    fig_kwargs={}, plot_kwargs={}, save_kwargs={}, legend_kwargs={}):
    truth, pred = comparison_sets
    if log:
        truth = np.log10(truth)
        pred = np.log10(pred)

    fig, ax = plt.subplots(**fig_kwargs)

    try:
        bins = plot_kwargs.pop('bins')
    except KeyError:
        bins = 'auto'

    ax.hist(truth, bins=bins, label='Truth', **plot_kwargs)
    ax.hist(pred, bins=bins, label='Prediction', **plot_kwargs)

    ax.set_xlabel(xlabel, **xlabel_kwargs)

    if title is not None:
        ax.set_title(title, **title_kwargs)

    ax.legend(**legend_kwargs)
    plt.savefig(get_script_path() + f'/../models/{model_name}_hists.png', **save_kwargs)
    plt.close()


def plot_difference_histogram(comparison_sets, model_name, xlabel, title=None, title_kwargs={}, xlabel_kwargs={},
                              log=True, one_sided=False, ratio=False, fig_kwargs={}, plot_kwargs={}, save_kwargs={}):
    truth, pred = comparison_sets
    if log:
        truth = np.log10(truth)
        pred = np.log10(pred)

    fig, ax = plt.subplots(**fig_kwargs)

    try:
        bins = plot_kwargs.pop('bins')
    except KeyError:
        bins = 'auto'

    if ratio:
        out = pred / truth
        if log:
            out = np.log10(out)
            if one_sided:
                out = abs(out)
    else:
        out = pred - truth
        if log:
            out = np.log10(abs(out))

    ax.hist(out, bins=bins, **plot_kwargs)

    ax.set_xlabel(xlabel, **xlabel_kwargs)

    if title is not None:
        ax.set_title(title, **title_kwargs)

    if ratio:
        diff = 'rel'
    else:
        diff = 'abs'
    if log:
        logname = 'log'
    else:
        logname = ''

    plt.savefig(get_script_path() + f'/../models/{model_name}_hist_{logname}{diff}diff.png', **save_kwargs)
    plt.close()


def grid_heatmap_corner(dataframe, truth_column, pred_column, log=True, ratio=False, save=True, save_fp='heatmap_corner.png', savefig_kwargs={}):
    """
    Only for grids (in long form as a pandas dataframe, not Nd array)! Don't use for uniform data.
    Args:
        dataframe:
        truth_column:
        pred_column:
        log:
        ratio:

    Returns:

    """
    subframe = dataframe.drop([truth_column,pred_column],axis=1)
    # dimensions = np.array([len(np.unique(col)) for col in subframe])  # dimensions of the grid
    param_mins = [np.min(subframe[col].to_numpy()) for col in subframe]

    param_maxes = [np.max(subframe[col].to_numpy()) for col in subframe]
    titles = list(subframe)  # column headings
    nparams = len(titles)
    # plotmaps = np.zeros(shape=tuple(dimensions))
    plotmaps = []
    for i in range(nparams):
        col1 = subframe[titles[i]].to_numpy()
        gridvalues1 = np.unique(col1)
        for j in range(i+1,nparams):
            col2 = subframe[titles[j]].to_numpy()
            gridvalues2 = np.unique(col2)

            heatmap_here = np.zeros((gridvalues1.size,gridvalues2.size))
            for k,gv1 in enumerate(gridvalues1):
                for l,gv2 in enumerate(gridvalues2):
                    inds_to_keep = np.where(np.logical_and(col1 == gv1,col2 == gv2))
                    truths = dataframe[truth_column].to_numpy()[inds_to_keep]
                    preds = dataframe[pred_column].to_numpy()[inds_to_keep]

                    if ratio:
                        temp = preds/truths
                        if log:
                            temp = np.log10(temp)
                    else:
                        temp = preds - truths
                        if log:
                            temp = np.log10(abs(temp))

                    heatmap_here[k,l] = np.mean(temp)
            plotmaps.append(heatmap_here)

    vmin = np.min([np.min(hmap) for hmap in plotmaps])
    vmax = np.max([np.max(hmap) for hmap in plotmaps])

    fig, ax = plt.subplots(ncols=5, nrows=5, figsize=(8 + int(2*nparams),6 + int(2*nparams)))
    plt.subplots_adjust(wspace=0.05, hspace=0.05)
    num = 0
    for i in range(nparams):
        for j in range(i+1,nparams):
            temp_im = ax[j,i].imshow(plotmaps[num].T, origin='lower', aspect='auto',
                                     extent=[param_mins[i],param_maxes[i],param_mins[j],param_maxes[j]],
                                     vmin=vmin, vmax=vmax)
            ax[i,j].axis('off')
            ax[j,i].grid(False)
            if i != 0:
                ax[j,i].yaxis.set_visible(False)
            if j != 4:
                ax[j,i].xaxis.set_visible(False)

            num += 1

    for i in range(nparams):
        ax[-1, i].set_xlabel(titles[i])

    for i in range(1, nparams):
        ax[i, 0].set_ylabel(titles[i])


    singles = []
    for i in range(nparams):
        col1 = subframe[titles[i]].to_numpy()
        gridvalues1 = np.unique(col1)

        this_line = np.zeros(len(gridvalues1))
        for k,gv1 in enumerate(gridvalues1):
            inds_to_keep = np.where(col1 == gv1)
            truths = dataframe[truth_column].to_numpy()[inds_to_keep]
            preds = dataframe[pred_column].to_numpy()[inds_to_keep]

            if ratio:
                temp = preds/truths
                if log:
                    temp = np.log10(temp)
            else:
                temp = preds - truths
                if log:
                    temp = np.log10(abs(temp))
            this_line[k] = np.mean(temp)

        singles.append(this_line)

    for i in range(nparams):
        xvals = np.unique(subframe[titles[i]].to_numpy())
        ax[i, i].plot(xvals, singles[i], linestyle='dashed',marker='x',c='k')
        ax[i,i].set_title(titles[i])
        ax[i,i].yaxis.tick_right()
        ax[i,i].grid('on')
        ax[i, i].tick_params(axis="both", which='both', direction='in')

        if i != nparams-1:
            ax[i,i].tick_params(labelbottom=False)

    cbarlabel = ''
    if log:
        cbarlabel += 'log10'
    if ratio:
        cbarlabel+='(ratio)'
    else:
        cbarlabel+='|diff|'
    fig.colorbar(temp_im, ax=ax.ravel().tolist(), location='left', label=cbarlabel, pad=0.08)

    if save:
        fig.savefig(save_fp, **savefig_kwargs)
    else:
        plt.show()