import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import numpy as np

from munch import Munch
from .kaux import mergemunch

def plot_df(df, plot_order, plot_t, fb_states, fig=None, options=None):
    """
    options: Munch dictionary with keys:
        show_beginning: 0: don't show beginning of trial trace
            1: mark beginning by a symbol
            > 1: show first number of samples with a heavy line
        show_end: 0: don't show end of trial trace
            1: mark end by a symbol
            > 1: show last number of samples with a heavy line
        show_thresholds: T: will use thresholds in fb_states to plot horizontal lines
        show_median: T: will show median lines for groups
    """
    default_options = Munch({'show_beginning': 1, 'show_end': 2, 'show_thresholds': True, 'show_median':True})
    if options is None:
        options = default_options
    else:
        options = Munch(mergemunch(default_options, options))

    if fig is None:
        fig = plt.figure(34, figsize=(16, 4))
        fig.clf()

    n_rows = np.max([x.r for x in plot_order.p]) + 1
    n_c = np.max([x.c for x in plot_order.p]) + 1
    n_cols = len(plot_order.s) * n_c

    axs = fig.subplots(n_rows, n_cols, False, True, squeeze=False)

    for i, s in enumerate(plot_order.s):
        good_sample_idx = f'{s.dcol}_good'
        xtr_n = plot_t[s.dcol].twin
        xtr_n = xtr_n - xtr_n[0]
        xtr_n = xtr_n / xtr_n[-1]

        i_c = i
        i_r = 0

        if s.dcol == 'stimulus_start_samples':
            stim_end_t = np.mean((df.stimulus_stop - df.stimulus_start) / 3e4)

            # rect = patches.Rectangle((0.0,0.0), stim_end_t, 1, linewidth=1, edgecolor='none', facecolor=(.01,.01,.01,.14))
            # axs[i_c].add_patch(rect)

            for (st, vals) in fb_states.items():
                axs[i_r, i_c].plot([0, stim_end_t], [vals[0], vals[0]], lw=4, ls='-', c=plot_order.state[st].c,
                                   alpha=.99, clip_on=False)
                axs[i_r, i_c].text(0, vals[0] + .02, plot_order.state[st].d, c=plot_order.state[st].c, alpha=1)

        for jj, pp in enumerate(plot_order.p):
            pdd = df[pp.sel(df)]
            if len(pdd) == 0:
                continue
            i_c = i + pp.c * len(plot_order.s)
            i_r = pp.r

            axs[i_r, i_c].spines['top'].set_visible(False)
            axs[i_r, i_c].spines['right'].set_visible(False)
            # if s.dcol == 'stimulus_start_samples':
            #     stim_end_t = np.mean((df.stimulus_stop - df.stimulus_start) / 3e4)
            #
            #     axs[i_r, i_c].plot([0, stim_end_t], [fb_states[pp.t][0], fb_states[pp.t][0]], lw=4, ls='-',
            #                        c=plot_order.state[pp.t].c, alpha=.99, clip_on=False)
            #     axs[i_r, i_c].text(0, fb_states[pp.t][0] + .02, plot_order.state[pp.t].d, c=plot_order.state[pp.t].c,
            #                        alpha=1)

            for ix, pd_row in pdd.iterrows():
                if pd_row[s.dcol] is None:
                    continue
                y = pd_row[s.dcol][pd_row[good_sample_idx]]
                xtr = plot_t[s.dcol].twin[pd_row[good_sample_idx]]
                if s.norm:
                    xtr = xtr - xtr[0]
                    xtr = xtr / xtr[-1]
                axs[i_r, i_c].plot(xtr, y, c=pp.col, alpha=.3, lw=.8)
                if options.show_beginning == 1:
                    axs[i_r, i_c].plot(xtr[0], y[0], c=pp.col, alpha=.3, lw=.8, marker='o', markersize=2)
                elif options.show_beginning > 1:
                    axs[i_r, i_c].plot(xtr[:(options.show_beginning - 1)], y[:(options.show_beginning - 1)], c=pp.col, alpha=.8, lw=1.5)
                if options.show_end == 1:
                    axs[i_r, i_c].plot(xtr[-1], y[-1], c=pp.col, alpha=.3, lw=.8, marker='o', markersize=2)
                elif options.show_end > 1:
                    axs[i_r, i_c].plot(xtr[(-options.show_end ):], y[(-options.show_end ):], c=pp.col, alpha=.8, lw=1.5)
                # axs[i_r, i_c].plot(xtr_n, pd_row['response_start_samples_norm'], c=plot_order.cols[i_col], alpha=.3, lw=1)
            if options.show_median and not pdd[s.dcol].empty:
                if s.norm:
                    all_traces = np.vstack(pdd[f'{s.dcol}_norm'].to_numpy())
                    mean_tr = np.median(all_traces, 0)
                    axs[i_r, i_c].plot(xtr_n, mean_tr, c=pp.col, alpha=.6, lw=4, ls=pp.ls)
                else:
                    all_traces = np.vstack(pdd[s.dcol].to_numpy())
                    mean_tr = np.nanmedian(all_traces, 0)
                    axs[i_r, i_c].plot(plot_t[s.dcol].twin, mean_tr, c=pp.col, alpha=.6, lw=4, ls=pp.ls)
            axs[i_r, i_c].set_title(f'{s.title} {pp.desc} n={len(pdd)}')
            if options.show_thresholds:
                thr = fb_states.up
                axs[i_r, i_c].plot(plot_t[s.dcol].twin[[0,-1]], [thr[1], thr[1]], linestyle='-', c="#333333", lw=0.5, alpha=.4)
                thr = fb_states.down
                axs[i_r, i_c].plot(plot_t[s.dcol].twin[[0,-1]], [thr[2], thr[2]], linestyle='-', c="#333333", lw=0.5, alpha=.4)
            axs[i_r, i_c].set_xlim(plot_t[s.dcol].t)
    return fig

def plot_df_combined(df, plot_order, plot_t, fb_states, fig=None, options=None):
    """
    options: Munch dictionary with keys:
        show_beginning: 0: don't show beginning of trial trace
            1: mark beginning by a symbol
            > 1: show first number of samples with a heavy line
        show_end: 0: don't show end of trial trace
            1: mark end by a symbol
            > 1: show last number of samples with a heavy line
        show_thresholds: T: will use thresholds in fb_states to plot horizontal lines
        show_median: T: will show median lines for groups
    """
    default_options = Munch({'show_beginning': 1, 'show_end': 2, 'show_thresholds': True, 'show_median':True})
    if options is None:
        options = default_options
    else:
        options = Munch(mergemunch(default_options, options))

    if fig is None:
        fig = plt.figure(35, figsize=(16, 4))
        fig.clf()

    n_cols = len(plot_order.s)

    axs = fig.subplots(1, n_cols, False, True, squeeze=False)

    for i, s in enumerate(plot_order.s):
        good_sample_idx = f'{s.dcol}_good'
        xtr_n = plot_t[s.dcol].twin
        xtr_n = xtr_n - xtr_n[0]
        xtr_n = xtr_n / xtr_n[-1]
        i_c = i
        i_r = 0
        axs[i_r, i_c].set_title(f'{s.title}')
        
        if options.show_thresholds:
            thr = fb_states.up
            axs[i_r, i_c].plot(plot_t[s.dcol].twin[[0,-1]], [thr[1], thr[1]], linestyle='-', c="#333333", lw=0.5, alpha=.4)
            thr = fb_states.down
            axs[i_r, i_c].plot(plot_t[s.dcol].twin[[0,-1]], [thr[2], thr[2]], linestyle='-', c="#333333", lw=0.5, alpha=.4)
        axs[i_r, i_c].set_xlim(plot_t[s.dcol].t)
        axs[i_r, i_c].spines['top'].set_visible(False)
        axs[i_r, i_c].spines['right'].set_visible(False)

        for jj, pp in enumerate(plot_order.p):
            pdd = df[pp.sel(df)]

            if s.dcol == 'stimulus_start_samples':
                stim_end_t = np.mean((df.stimulus_stop - df.stimulus_start) / 3e4)

                axs[i_r, i_c].plot([0, stim_end_t], [fb_states[pp.t][0], fb_states[pp.t][0]], lw=4, ls='-',
                                   c=plot_order.state[pp.t].c, alpha=.99, clip_on=False)
                axs[i_r, i_c].text(0, fb_states[pp.t][0] + .02, plot_order.state[pp.t].d, c=plot_order.state[pp.t].c,
                                   alpha=1)

            my_label = f"{pp.desc} (n = {len(pdd)})"
            for ix, pd_row in pdd.iterrows():
                y = pd_row[s.dcol][pd_row[good_sample_idx]]
                xtr = plot_t[s.dcol].twin[pd_row[good_sample_idx]]
                if s.norm:
                    xtr = xtr - xtr[0]
                    xtr = xtr / xtr[-1]
                axs[i_r, i_c].plot(xtr, y, c=pp.col, alpha=.3, lw=.8, ls=pp.ls, label=my_label)
                my_label = None
                if options.show_beginning == 1:
                    axs[i_r, i_c].plot(xtr[0], y[0], c=pp.col, alpha=.3, lw=.8, marker='o', markersize=2)
                elif options.show_beginning > 1:
                    axs[i_r, i_c].plot(xtr[:(options.show_beginning - 1)], y[:(options.show_beginning - 1)], c=pp.col, alpha=.8, lw=1.5)
                if options.show_end == 1:
                    axs[i_r, i_c].plot(xtr[-1], y[-1], c=pp.col, alpha=.3, lw=.8, marker='o', markersize=2)
                elif options.show_end > 1:
                    axs[i_r, i_c].plot(xtr[(-options.show_end ):], y[(-options.show_end ):], c=pp.col, alpha=.8, lw=1.5)
                # axs[i_r, i_c].plot(xtr_n, pd_row['response_start_samples_norm'], c=plot_order.cols[i_col], alpha=.3, lw=1)
            if options.show_median and not pdd[s.dcol].empty:
                if s.norm:
                    all_traces = np.vstack(pdd[f'{s.dcol}_norm'].to_numpy())
                    mean_tr = np.median(all_traces, 0)
                    axs[i_r, i_c].plot(xtr_n, mean_tr, c=pp.col, alpha=.6, lw=4, ls=pp.ls)
                else:
                    all_traces = np.vstack(pdd[s.dcol].to_numpy())
                    mean_tr = np.nanmedian(all_traces, 0)
                    axs[i_r, i_c].plot(plot_t[s.dcol].twin, mean_tr, c=pp.col, alpha=.6, lw=4, ls=pp.ls, label=pp.desc)
        axs[i_r, i_c].legend()
    return fig

def plot_df_avg(df, plot_order, plot_t, fb_states, fig=None, show_stimulus_start=False):
    if fig is None:
        fig = plt.figure(35, figsize=(16, 4))
        fig.clf()

    n_rows = 1
    # n_c = len(plot_order.c)
    n_cols = len(plot_order.s)
    axs = fig.subplots(n_rows, n_cols, False, True, squeeze=False)
    i_r = 0

    for i, s in enumerate(plot_order.s):
        good_sample_idx = f'{s.dcol}_good'
        xtr_n = plot_t[s.dcol].twin
        xtr_n = xtr_n - xtr_n[0]
        xtr_n = xtr_n / xtr_n[-1]
        i_c = i
        i_r = 0
        axs[i_r, i_c].spines['top'].set_visible(False)
        axs[i_r, i_c].spines['right'].set_visible(False)
        if s.dcol == 'stimulus_start_samples':
            stim_end_t = np.mean((df.stimulus_stop - df.stimulus_start) / 3e4)

            # rect = patches.Rectangle((0.0,0.0), stim_end_t, 1, linewidth=1, edgecolor='none', facecolor=(.01,.01,.01,.14))
            # axs[i_c].add_patch(rect)

            for (st, vals) in fb_states.items():
                axs[i_r, i_c].plot([0, stim_end_t], [vals[0], vals[0]], lw=4, ls='-', c=plot_order.state[st].c,
                                   alpha=.99, clip_on=False)
                axs[i_r, i_c].text(0, vals[0] + .02, plot_order.state[st].d, c=plot_order.state[st].c, alpha=1)

        for jj, pp in enumerate(plot_order.p):
            pdd = df[pp.sel(df)]

            i_r = 0

            if not pdd[s.dcol].empty:
                if s.norm:
                    all_traces = np.vstack(pdd[f'{s.dcol}_norm'].to_numpy())
                    avg_x = xtr_n
                else:
                    all_traces = np.vstack(pdd[s.dcol].to_numpy())
                    avg_x = plot_t[s.dcol].twin

                mean_tr = np.nanmedian(all_traces, 0)
                perct_tr = np.nanpercentile(all_traces, [25, 75], axis=0)

                axs[i_r, i_c].plot(avg_x, mean_tr, c=pp.col, alpha=1, lw=pp.lw, clip_on=False, ls=pp.ls)
                if pp.show_var:
                    axs[i_r, i_c].plot(avg_x, perct_tr[0, :], c=pp.col, alpha=.1, lw=pp.lw, clip_on=False, ls=pp.ls)
                    axs[i_r, i_c].plot(avg_x, perct_tr[1, :], c=pp.col, alpha=.1, lw=pp.lw, clip_on=False, ls=pp.ls)
                    axs[i_r, i_c].fill_between(avg_x, perct_tr[0, :], perct_tr[1, :],
                                               alpha=0.20, facecolor=pp.col, edgecolor='none', ls=pp.ls, lw=1,
                                               antialiased=True, clip_on=False)

                if show_stimulus_start and s.dcol == 'stimulus_stop_samples':
                    stim_start_offset = (pdd.stimulus_start - pdd.stimulus_stop) / 3e4
                    for x in stim_start_offset:
                        axs[i_r, i_c].axvline(x, c=[0, 0, 0], alpha=.1)
            thr = fb_states.up
            # axs[i_r, i_c].plot(plot_t[s.dcol].twin[[0,-1]], [thr[1], thr[1]], linestyle='--', c="#333333")
            thr = fb_states.down

            # axs[i_r, i_c].plot(plot_t[s.dcol].twin[[0,-1]], [thr[2], thr[2]], linestyle='--', c="#333333")
        # axs[i_c].set_ylim([0,1])
        if s.norm:
            axs[i_r, i_c].set_xlim([0, 1])
            axs[i_r, i_c].set_xlabel('t (normalized)')

        else:
            axs[i_r, i_c].set_xlim(plot_t[s.dcol].t)
            axs[i_r, i_c].set_xlabel('t [s]')

        if i_c == 0:
            axs[i_r, i_c].set_ylabel('normalized neural activity')
        axs[i_r, i_c].set_title(f'{s.title}')

    custom_lines = [Line2D([0], [0], color=(.3, .3, .3), lw=2, ls='-'),
                    Line2D([0], [0], color=(.3, .3, .3), lw=2, ls=':')]
    axs[i_r, n_cols - 1].legend(custom_lines, ['Correct Trials', 'Error Trials'])

    fig.show()
    return fig