from argparse import ArgumentParser
from fnmatch import fnmatch
from math import sqrt
import os
from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
import numpy as np

from .results import get_result_dataframe, get_icsbep_dataframe, abbreviated_name


def plot(files, labels=None, plot_type='keff', match=None, show_mean=True,
         show_shaded=True, show_uncertainties=True):
    """For all benchmark cases, produce a plot comparing the k-effective mean
    from the calculation to the experimental value along with uncertainties.

    Parameters
    ----------
    files : iterable of str
        Name of a results file produced by the benchmarking script.
    labels: iterable of str
        Labels for each dataset to use in legend
    plot_type : {'keff', 'diff'}
        Type of plot to produce. A 'keff' plot shows the ratio of the
        calculation k-effective mean to the experimental value (C/E). A 'diff'
        plot shows the difference between C/E values for different
        calculations. Default is 'keff'.
    match : str
        Pattern to match benchmark names to
    show_mean : bool
        Whether to show bar/line indicating mean value
    show_shaded : bool
        Whether to show shaded region indicating uncertainty of mean
    show_uncertainties : bool
        Whether to show uncertainties for individual cases

    Returns
    -------
    matplotlib.axes.Axes
        A matplotlib.axes.Axes object

    """
    if labels is None:
        labels = [Path(f).name for f in files]

    # Read data from spreadsheets
    dataframes = {}
    for csvfile, label in zip(files, labels):
        dataframes[label] = get_result_dataframe(csvfile).dropna()

    # Get model keff and uncertainty from ICSBEP
    icsbep = get_icsbep_dataframe()

    # Determine common benchmarks
    base = labels[0]
    index = dataframes[base].index
    for df in dataframes.values():
        index = index.intersection(df.index)

    # Applying matching as needed
    if match is not None:
        cond = index.map(lambda x: fnmatch(x, match))
        index = index[cond]

    # Setup x values (integers) and corresponding tick labels
    n = index.size
    x = np.arange(1, n + 1)
    xticklabels = index.map(abbreviated_name)

    fig, ax = plt.subplots(figsize=(17, 6))

    if plot_type == 'diff':
        # Check that two results files are specified
        if len(files) < 2:
            raise ValueError('Must provide two or more files to create a "diff" plot')

        kwargs = {'mec': 'black', 'mew': 0.15, 'fmt': 'o'}

        keff0 = dataframes[base]['keff'].loc[index]
        stdev0 = 1.96*dataframes[base]['stdev'].loc[index]
        for i, label in enumerate(labels[1:]):
            df = dataframes[label]
            keff_i = df['keff'].loc[index]
            stdev_i = 1.96*df['stdev'].loc[index]

            diff = (keff_i - keff0) * 1e5
            err = np.sqrt(stdev_i**2 + stdev0**2) * 1e5
            kwargs['label'] = labels[i + 1] + ' - ' + labels[0]
            if show_uncertainties:
                ax.errorbar(x, diff, yerr=err, color=f'C{i}', **kwargs)
            else:
                ax.plot(x, diff, color=f'C{i}', **kwargs)

            # Plot mean difference
            if show_mean:
                mu = diff.mean()
                if show_shaded:
                    sigma = diff.std() / sqrt(n)
                    verts = [(0, mu - sigma), (0, mu + sigma), (n+1, mu + sigma), (n+1, mu - sigma)]
                    poly = Polygon(verts, facecolor=f'C{i}', alpha=0.5)
                    ax.add_patch(poly)
                else:
                    ax.plot([-1, n], [mu, mu], '-', color=f'C{i}', lw=1.5)

        # Define y-axis label
        ylabel = r'$\Delta k_\mathrm{eff}$ [pcm]'

    else:
        for i, (label, df) in enumerate(dataframes.items()):
            # Calculate keff C/E and its standard deviation
            coe = (df['keff'] / icsbep['keff']).loc[index]
            stdev = 1.96 * df['stdev'].loc[index]

            # Plot keff C/E
            kwargs = {'color': f'C{i}', 'mec': 'black', 'mew': 0.15, 'label': label}
            if show_uncertainties:
                ax.errorbar(x, coe, yerr=stdev, fmt='o', **kwargs)
            else:
                ax.plot(x, coe, 'o', **kwargs)

            # Plot mean C/E
            if show_mean:
                mu = coe.mean()
                sigma = coe.std() / sqrt(n)
                if show_shaded:
                    verts = [(0, mu - sigma), (0, mu + sigma), (n+1, mu + sigma), (n+1, mu - sigma)]
                    poly = Polygon(verts, facecolor=f'C{i}', alpha=0.5)
                    ax.add_patch(poly)
                else:
                    ax.plot([-1, n], [mu, mu], '-', color=f'C{i}', lw=1.5)

        # Show shaded region of benchmark model uncertainties
        unc = icsbep['stdev'].loc[index]
        vert = np.block([[x, x[::-1]], [1 + unc, 1 - unc[::-1]]]).T
        poly = Polygon(vert, facecolor='gray', edgecolor=None, alpha=0.2)
        ax.add_patch(poly)

        # Define axes labels and title
        ylabel = r'$k_\mathrm{eff}$ C/E'

    # Configure plot
    ax.set_axisbelow(True)
    ax.set_xlim((0, n+1))
    ax.set_xticks(x)
    ax.set_xticklabels(xticklabels, rotation='vertical')
    ax.tick_params(axis='x', which='major', labelsize=10)
    ax.tick_params(axis='y', which='major', labelsize=14)
    ax.set_xlabel('Benchmark case', fontsize=18)
    ax.set_ylabel(ylabel, fontsize=18)
    ax.grid(True, which='both', color='lightgray', ls='-', alpha=0.7)
    ax.legend(numpoints=1)
    return ax


def main():
    """Produce plot of benchmark results"""
    parser = ArgumentParser()
    parser.add_argument('files', nargs='+', help='Result CSV files')
    parser.add_argument('--labels', help='Comma-separated list of dataset labels')
    parser.add_argument('--plot-type', choices=['keff', 'diff'], default='keff')
    parser.add_argument('--match', help='Pattern to match benchmark names to')
    parser.add_argument('--show-mean', action='store_true', help='Show line/bar indicating mean')
    parser.add_argument('--no-show-mean', dest='show_mean', action='store_false',
                        help='Do not show line/bar indicating mean')
    parser.add_argument('--show-uncertainties', action='store_true',
                        help='Show uncertainty bars on individual cases')
    parser.add_argument('--no-show-uncertainties', dest='show_uncertainties', action='store_false',
                        help='Do not show uncertainty bars on individual cases')
    parser.add_argument('--show-shaded', action='store_true',
                        help='Show shaded region indicating uncertainty of mean C/E')
    parser.add_argument('--no-show-shaded', dest='show_shaded', action='store_false',
                        help='Do not show shaded region indicating uncertainty of mean C/E')
    parser.add_argument('-o', '--output', help='Filename to save to')
    parser.set_defaults(show_uncertainties=True, show_shaded=True, show_mean=True)
    args = parser.parse_args()

    if args.labels is not None:
        args.labels = args.labels.split(',')

    ax = plot(
        args.files,
        labels=args.labels,
        plot_type=args.plot_type,
        match=args.match,
        show_mean=args.show_mean,
        show_shaded=args.show_shaded,
        show_uncertainties=args.show_uncertainties,
    )
    if args.output is not None:
        plt.savefig(args.output, bbox_inches='tight', transparent=True)
    else:
        plt.show()
