"""This module includes functions to format output displays of
one ConfMatrixModel instance,
in graphic and tabular form.

Some formatting parameters are defined by module global variables,
which can be changed by user.

*** Version History:

* Version 0.8:
2025-10-24, all results use pandas data, cleanup code and comments
2025-10-09, all .save methods take kwargs with figure_format and table_format requests

* Version 0.7:
2018-09-15, first functional version
2019-12-02, minor fix to suppress matplotlib warning in eps output
"""
import numpy as np
import matplotlib.pyplot as plt
from itertools import cycle  #, product
import logging

import pandas as pd

from ConfMatrixCalc.cm_file import Table

plt.rcParams.update({'figure.max_open_warning': 0})

logger = logging.getLogger(__name__)

# --------------------------- Default format dict:
FMT = {'colors': 'rbgk',
       'markers': 'oxs*_',
       'interaction_sep': '\u00D7',  # mult.sign. separating situation labels in result file names
       'max_plot_cases': 50,    # max cases in percentile plots
       'max_case_heads': 1,     # limit in plot legends
       'x_space': 0.2,          # clean space between x_tick categories
       'condition_sep': '_',    # separating attribute and its conditioning situation(s) in file names
       'diff_marker': '-diff',  # marker after diff-axis categories in file names
       'and_head': ('', ''),    # two-level heading for diff-table first column
       'and_label': 'and',      # logical and in diff-table first column
       }
# = module-global dict with default settings for formatting details
# that may be changed by user


def set_format_param(mpl_style=None, mpl_params=None, **kwargs):
    """Set / modify format parameters.
    Called from cm_displaybefore any displays are generated.
    :param mpl_style: (optional) matplotlib style sheet, or list of style sheets
    :param mpl_params: (optional) dict with matplotlib (k, v) rcParam settings
    :param kwargs: dict with any formatting variables to be stored in FMT
    :return: None
    """
    if mpl_style is not None:
        plt.style.use(mpl_style)
    if mpl_params is not None:
        plt.rcParams.update(mpl_params)
    other_fmt = dict()
    for (k, v) in kwargs.items():  # user settings
        k = k.lower()
        if k in FMT:
            FMT[k] = v
        else:
            other_fmt[k] = v
    if len(other_fmt) > 0:
        logger.warning(f'Parameters {other_fmt} unknown, not used.')


# ---------------------------------------- Result Classes
class ResultPlot:
    """Container for a single graph instance
    """
    def __init__(self, ax, name):
        """
        :param ax: matplotlib Axes instance containing the graph
        :param name: string identifying the plot; used as file name
        """
        self.ax = ax
        self.name = name

    def __repr__(self):  # ***** use a superclass ? *******
        return (f'ResultPlot(ax= {repr(self.ax)}, ' +
                f'name= {repr(self.name)})')

    @property
    def fig(self):
        return self.ax.figure

    def save(self, path,
             figure_format,
             **kwargs):
        """Save figure to given path
        :param path: Path to directory for saving self
        :param figure_format: figure-format string code -> file-name suffix
        :param kwargs (optional) any additional kwargs, *** NOT USED ***
        :return: None
        """
        # *** select subset of kwargs allowed by savefig() ?
        # NO, depends on Matplotlib backend!
        f = (path / self.name).with_suffix('.' + figure_format)
        try:
            self.fig.savefig(f)
        except Exception as e:  # any error, just warn and continue
            logger.warning(f'Could not save plot to {f}. Error: {e}')


class ResultTable(Table):
    """A pd.DataFrame table subclass, with a name and special save method
    """
    def __init__(self, df, name):
        """
        :param df: a Table(pd.DataFrame) instance
        :param name: file name for saving the table
        """
        super().__init__(df)
        self.name = name

    def save(self, path,
             table_format='txt',
             **kwargs):
        """Save table to file.
        :param path: Path to directory for saving self.
            suffix is determined by FMT['table_format'] anyway
        :param table_format: table-format string code -> file-name suffix
        :param kwargs: (optional) any additional arguments to pandas writer function
        :return: None
        """
        f = (path / self.name).with_suffix('.' + table_format)
        try:
            super().save(f, **kwargs)
        except Exception as e:  # any error, just warn and continue
            logger.warning(f'Could not save result table. Error: {e}')


class DiffTable(ResultTable):
    """Special subclass suppressing index in save method
    """
    def save(self, path,
             **kwargs):
        """Save table to file.
        :param path: Path to directory for saving self.
            suffix is determined by FMT['table_format'] anyway
        :param kwargs: (optional) any additional arguments to pandas writer function
        :return: None
        """
        if 'index' not in kwargs:
            kwargs['index'] = False  # override Pandas default = True
        super().save(path, **kwargs)


# ----------------------------------------------- Plot Functions:

def fig_percentiles(df,
                    y_label='',
                    file_label='',
                    y_min=None,
                    y_max=None,
                    **kwargs):
    """Create a figure with percentile results from a pd.DataFrame instance,
    structured like the output of tab_percentile(...)
    :param df: pd.DataFrame instance with primary percentile data, with
        one row for each case category, as defined in df.index.values elements,
        one column for each percentile value.
    :param y_label: (optional) string for y-axis label
    :param file_label: (optional) string as first part of file name
    :param y_min: (optional) enforced lower limit of vertical axis
    :param y_max: (optional) enforced upper limit of vertical axis
    :param kwargs: (optional) dict with any additional keyword arguments for plot commands. *** NOT NEEDED ? ****
    :return: ResultPlot instance with plot axis with all results

    NOTE: plot will use df.index.level[0] categories as x-axis labels,
    and index.level[1:] as plot labels in the legend
    """
    # ******* simplify? by stack quantiles, unstack df.index.level[0] ? like fig_boxplot ???
    # ----------------------------------- set up plot design:
    if df is None:
        return None
    if df.index.nlevels == 1:
        x_label = df.index.name
        x_tick_labels = list(df.index.values)
        case_head = ()
        case_list = [()]
    elif df.index.nlevels == 2:
        x_label = df.index.names[0]
        x_tick_labels = list(dict.fromkeys([c[0] for c in df.index.values]))
        case_head = df.index.names[1]
        case_list = list(dict.fromkeys([c[1] for c in df.index.values]))
    else:
        x_label = df.index.names[0]
        x_tick_labels = list(dict.fromkeys([c[0] for c in df.index.values]))
        case_head = tuple(df.index.names[1:])
        case_list = list(dict.fromkeys([c[1:] for c in df.index.values]))
        # in order as appearing in df
    if df.shape[0] > FMT['max_plot_cases']:
        logger.warning(f'Too many {x_label}, {case_head} cases to plot. '
                       + f'Set max_plot_cases >= {df.shape[0]} if really needed.')
        return None
    n_cases = len(case_list)
    dx = (1. - FMT['x_space']) / n_cases
    # = x step between range plots for separate cases
    x_offset = {c: (i - (n_cases - 1) / 2) * dx
                for (i, c) in enumerate(case_list)}
    case_color = {c: col for (c, col) in zip(case_list, cycle(FMT['colors']))}
    case_marker = {c: mark for (c, mark) in zip(case_list, cycle(FMT['markers']))}
    # ------------------------------------------------------------------------
    def plot_one(ax, row_index, y):
        """
        :param ax: axis for plot
        :param row_index: index label(s) for ONE row in df  *****  not used ? ******
        :param y: 1D array with corresponding y = values in ONE row
            y[p] = p-th percentile value
        :return: None
        """
        if df.index.nlevels == 1:
            x_tick = row_ind
            x_case = ()
        elif df.index.nlevels == 2:
            x_tick = row_ind[0]
            x_case = row_ind[1]
        else:
            x_tick = row_ind[0]
            x_case = row_ind[1:]
        x = x_tick_labels.index(x_tick) + x_offset[x_case]
        x = x * np.ones_like(y)
        c = case_color[x_case]
        m = case_marker[x_case]
        if len(y) == 1:  # only single marker
            line = ax.plot(x, y,
                           linestyle='', color=c,
                           marker=m, markeredgecolor=c, markerfacecolor='w',
                           **kwargs)  # ******** ?
        elif len(y) == 2:  # vertical range, no markers
            line = ax.plot(x, y,
                           # linestyle='solid',  # *** use Matplotlib rcParams ?
                           color=c,
                           marker=m, markeredgecolor=c, markerfacecolor='w',
                           **kwargs)
        else:  # vertical range, and markers for intermediate percentiles
            y = sorted(list(y))
            ax.plot([x[0], x[-1]], [y[0], y[-1]],
                    # linestyle='solid',
                    color=c,
                    **kwargs)
            line = ax.plot(x[1:-1], y[1:-1],
                           # linestyle='solid',
                           color=c,  # to get line+marker into label
                           marker=m, markeredgecolor=c, markerfacecolor='w',
                           **kwargs)
        if x_tick == x_tick_labels[0]:
            if type(case_head) is tuple and len(case_head) > FMT['max_case_heads']:
                line[0].set_label(str(x_case))
            else:
                line[0].set_label(str(case_head) + '=' + str(x_case))

    # ------------------------------------------------------------------
    fig, ax = plt.subplots()
    for (row_ind, y) in df.iterrows():
        plot_one(ax, row_ind, y)

    (x_min, x_max) = ax.get_xlim()
    x_min = min(x_min, -FMT['x_space'])
    x_max = max(x_max, len(x_tick_labels) - 1 + FMT['x_space'])
    ax.set_xlim(x_min, x_max)
    # if cat_limits is not None:
    #     _plot_response_intervals(ax, cat_limits)
    ax.set_xticks(np.arange(len(x_tick_labels)))
    xticks = [str(c) for c in x_tick_labels]
    ax.set_xticklabels(xticks,
                       **_x_tick_style(xticks))
    (y0, y1) = ax.get_ylim()
    if y_min is not None:
        y0 = y_min
    if y_max is not None:
        y1 = y_max
    ax.set_ylim(y0, y1)
    ax.set_ylabel(y_label)
    ax.set_xlabel(x_label)
    if len(case_list) > 1:
        ax.legend()  # (loc='best') ***************
    if len(file_label) > 0:
        file_label += FMT['condition_sep']
    f_name = file_label + FMT['interaction_sep'].join(df.index.names)
    # fig.set_tight_layout(tight=True)  # **** let rcParams control this  ****
    return ResultPlot(ax, name=f_name)


def fig_boxplot(df,
                x_cases,
                y_label='',
                file_label='',
                # y_min=None,  # *************
                # y_max=None,
                **kwargs):
    """Boxplot for one individual performance measure
    :param df: a pandas.Series instance,
        with data grouped by categories in a MultiIndex.
    :param x_cases: sequence with ONE or MORE df index name(s) for data to be plotted
        x_cases[0] -> x_label
        df.index categories of x_cases[0] -> x_ticklabels;
        df.index levels x_cases[1:] -> sub-case factors
    :param y_label: string for y-axis label
    :param file_label: string for y-axis label
    :param kwargs: (optional) dict with any additional keyword arguments for boxplot function.
    :return: a FigureRef instance
    """
    # ***** try violin plot instead ? *********
    def plot_one_case(ax, df,
                      x_offset=0.,
                      widths=0.5,
                      case=(),
                      case_head=(),
                      color=FMT['colors'][0],
                      marker=FMT['markers'][0]):
        """
        :param ax: plt.Axes instance
        :param df: DataFrame with columns to be box-plotted
        :param x_offset: offset from integer xticks values
        :param widths: default box widths
        :param case: tuple with case-axis categorie(s) for this set of boxplots
        :param case_head: tuple with case-head(s); len(case_head) == len(case)
        :param color: (optional) color
        :param marker: (optional) marker
        :return: None
        """
        x_pos = np.arange(len(df.columns.values)) + x_offset
        ax.boxplot(df.to_numpy(),
                   label=_plot_label(case_head, case),
                   positions=x_pos,
                   widths=widths,
                   manage_ticks=False,  # done later
                   boxprops=dict(linestyle='-', color=color),
                   flierprops=dict(marker=marker,
                                   markerfacecolor='w',
                                   markeredgecolor=color),
                   whiskerprops=dict(color=color),
                   medianprops=dict(color=color, marker=marker,
                                    markeredgecolor=color,
                                    markerfacecolor='w'),
                   capprops=dict(color=color),
                   **kwargs)
    # ---------------------------------------------------------------
    if len(x_cases) == 0:  # should not happen, handled by cm_display before call
        x_label = '-'  # fake index name already defined in cm_display
        case_head = ()
        df = df.to_frame(name='-')  # fake category name defined in cm_display
    else:
        x_label = x_cases[0]
        case_head = x_cases[1:]
        df = df.unstack(level=x_label)
    x_ticklabels = df.columns.values
    x_ticks = np.arange(len(x_ticklabels))
    (fig, ax) = plt.subplots()
    if len(case_head) == 0:  # df has plain Index
        plot_one_case(ax, df)
    else:   # df has MultiIndex with several levels
        case_levels = [l for l in df.index.levels if l.name in case_head]
        n_cases = np.prod([len(ind) for ind in case_levels], dtype=int)
        if n_cases > FMT['max_plot_cases']:
            logger.warning(f'Too many {x_label}, {case_head} cases to plot. '
                           + f'Set max_plot_cases >= {n_cases} if really needed.')
            return None
        dx = (1. - FMT['x_space']) / n_cases
        if len(case_head) == 1:
            group_by = case_head[0] # avoid future warning
        else:
            group_by = case_head  # keep as tuple
        for ((i, (case, df_c)), color, marker) in zip(enumerate(df.groupby(level=group_by)),
                                                     cycle(FMT['colors']),
                                                     cycle(FMT['markers'])):
            x_offset = dx * (i - (n_cases - 1) / 2)
            plot_one_case(ax, df_c, x_offset,
                          widths=dx * 0.7,
                          case_head=case_head,
                          case=case,
                          color=color,
                          marker=marker)
    ax.set_xlabel(x_label)
    ax.set_xticks(x_ticks)
    ax.set_xticklabels(x_ticklabels,
                       **_x_tick_style(x_ticklabels))
    ax.set_ylabel(y_label)
    if len(case_head) > 0:
        ax.legend()
    if len(file_label) > 0:
        file_label += FMT['condition_sep']
    f_name = file_label + FMT['interaction_sep'].join((x_label,) + case_head)
    return ResultPlot(ax, name=f_name)


# ----------------------------------------- table display functions:

def tab_percentiles_pd(q_ds, file_label=''):
    """Re-format a pd.DataFrame with percentile results.
    :param q_ds: a pd.Series instance with quantiles as last index axis,
        and cases in other index axes.
    :param file_label: (optional) string as first part of file name
    :return: a ResultTable(pd.DataFrame) instance,
        with one column for each percentile,
        and one row for each combination product of cases considered.
    """
    tab_perc = q_ds.unstack()
    case_head = tab_perc.index.names
    tab_perc.columns = [f'{q:.1%}' for q in tab_perc.columns]
    if len(file_label) > 0:
        file_label += FMT['condition_sep']
    f_name = file_label + FMT['interaction_sep'].join(case_head)
    return ResultTable(tab_perc, name=f_name)


def tab_credible_diff(diff,
                      diff_head,
                      cred_head,
                      case_head=(),
                      y_label='',
                      file_label='',
                      # and_label=FMT['and_label'],  # not arg use FMT explicitly? ******'
                      # and_head=FMT['and_head']
                      ):
    """Create table with credible differences among results -- pandas version
    :param diff: list of tuples (((i,j), c0,...), p),
        defining jointly credible differences, indicating that
        prob{ quality of diff category i > quality of category j, given case category c
        AND all previous pairs } == p
        i, j are either a string label or a tuple of such labels
        c0,... is one or more case labels
        ***** better?: c is either a string label or a tuple of such labels, like i, j ******
    :param diff_head: tuple of keys for heading of diff_labels column in table
        len(diff_head) == len(i) == len(j) if tuple, or len(diff_head) == 1
    :param cred_head: string for header of Credibility column
    :param case_head: (optional) tuple of case keys, one for each case-dimension table column
        len(case_head) == len(c0,...)
    :param y_label: (optional) string with label of tabulated attribute
    :param file_label: (optional) string for first part of file name
    :return: ResultTable object with header lines + one line for each credible difference
    """
    if len(diff) == 0:
        return None
    y_head_i = y_label + ' >'
    y_head_j = y_label
    # --------------------- table columns as dicts:
    col = { FMT['and_head']:  [' '] + [ FMT['and_label']] * (len(diff) - 1)}  # first column with only AND flags
    diff_i = [(ijc[0][0] if type(ijc[0][0]) is tuple else (ijc[0][0],)) for (ijc, p) in diff]
    diff_j = [(ijc[0][1] if type(ijc[0][1]) is tuple else (ijc[0][1],)) for (ijc, p) in diff]
    diff_c = [ijc[1:] for (ijc, p) in diff]
    diff_p = [p for (ijc, p) in diff]
    # --------- column(s) for higher results:
    col |= {(y_head_i, d_head_k): [d[k] for d in diff_i]
            for (k, d_head_k) in enumerate(diff_head)}
    # --------- column(s) for lower results:
    col |= {(y_head_j, d_head_k): [d[k] for d in diff_j]
            for (k, d_head_k) in enumerate(diff_head)}  # cols for lesser results
    # --------- column(s) for optional case labels:
    if len(case_head) > 0:
        # diff_c = [case_labels[d[0][2]]
        #           for d in diff]
        col |= {('', c_head_k): [d[k] for d in diff_c]
                for (k, c_head_k) in enumerate(case_head)}
    # --------- credibility column:
    col |= {(' ', cred_head): diff_p}
    df = pd.DataFrame(col)  # do all this in samppy.credibility_pd ? ***********
    # each column name is a tuple with two elements -> MultiIndex with two levels
    # df = df.reindex(columns=pd.MultiIndex.from_tuples(df.columns))  # ****** Needed ?
    if len(file_label) > 0:
        file_label += FMT['condition_sep']
    f_name = file_label + FMT['interaction_sep'].join(diff_head) + FMT['diff_marker']
    if len(case_head) > 0:
        f_name += FMT['condition_sep'] + FMT['interaction_sep'].join(case_head)
    return DiffTable(df, name=f_name)  # *** can use col as arg ? *******


# -------------------------------------- private help functions
def _x_tick_style(labels):
    """Select xtick properties to avoid tick-label clutter
    :param labels: list of tick label strings
    :return: dict with keyword arguments for set_xticklabels
    """
    maxL = max(len(l) for l in labels)
    rotate_x_label = maxL * len(labels) > 75  # ad hoc criterion
    if rotate_x_label:
        style = dict(rotation=15, horizontalalignment='right')
    else:
        style = dict(rotation='horizontal', horizontalalignment='center')
    return style


def _plot_label(group_by, case):
    if len(group_by) == 0:
        return None
    elif len(group_by) == 1:
        return group_by[0] + '=' + str(case)
    elif len(group_by) > FMT['max_case_heads']:
        return str(case)
    else:
        return str(group_by) + '=' + str(case)


# ----------------------------------------------------------- TEST:
if __name__ == '__main__':
    from scipy.stats import norm

    print('*** Testing fig_boxplot')
    n_samples = 100
    mu = [[0., 0.5, 1.],
          [0., -1., 0.5],
          [0., -1., -2.]]
    n_mu = 3
    len_x = 3
    x_labels = [f'x{i}' for i in range(len_x)]
    case_labels = [c for c in 'ABCDEF'[:n_mu]]
    x = np.array([norm.rvs(loc=m, size=(n_samples, len_x))
                  for m in mu])
    x2 = x.reshape((-1, len_x))
    ind = pd.MultiIndex.from_product([case_labels, range(n_samples)],
                                     names=['Case', '_sample'])
    x = pd.DataFrame(x2, index=ind, columns=x_labels)
    x.columns.set_names('x_element', inplace=True)
    x = x.stack()
    boxp = fig_boxplot(x,
                       x_cases=('x_element', 'Case'),
                       y_label='test',
                       file_label='test',
                       )

    boxp = fig_boxplot(x,
                       x_cases=('x_element',),
                       y_label='test',
                       file_label='test',
                       )
    plt.show()
