"""
plot histogram 2d
"""
# packages
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from scipy import stats
from typing import Optional
from enum import Enum

# qis
import qis.plots.utils as put


def plot_histplot2d(df: pd.DataFrame,
                    title: str = None,
                    a_min: float = None,
                    a_max: float = None,
                    xvar_format: str = '{:.1f}',
                    yvar_format: str = '{:.1f}',
                    add_corr_legend: bool = True,
                    legend_loc: Optional[str] = 'upper left',
                    color: str = 'navy',
                    fontsize: int = 10,
                    ax: plt.Subplot = None,
                    **kwargs
                    ) -> plt.Figure:

    if len(df.columns) != 2:
        raise ValueError(f"should be 2 columns")

    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=(8, 6))
    else:
        fig = None

    if a_min is not None or a_max is not None:
        df = np.clip(df, a_min=a_min, a_max=a_max)

    sns.histplot(data=df,
                 x=df.columns[0],
                 y=df.columns[1],
                 bins=100,
                 cbar=False,
                 stat='probability',
                 cbar_kws=dict(shrink=.75),
                 ax=ax)

    put.set_ax_ticks_format(ax=ax, xvar_format=xvar_format, yvar_format=yvar_format, **kwargs)

    if add_corr_legend:
        rho, pval = stats.spearmanr(df.to_numpy(), nan_policy='omit', axis=0)  # column is variable
        label = f"Rank corr={rho:0.2f}, p-val={pval:0.2f}"
        lines = [(label, {'color': color})]

        put.set_legend(ax=ax,
                       legend_loc=legend_loc,
                       fontsize=fontsize,
                       lines=lines,
                       **kwargs)

    put.align_xy_limits(ax=ax)

    if title is not None:
        ax.set_title(title, fontsize=fontsize, **kwargs)

    return fig


class LocalTests(Enum):
    TEST = 1


def run_local_test(local_test: LocalTests):
    """Run local tests for development and debugging purposes.

    These are integration tests that download real data and generate reports.
    Use for quick verification during development.
    """

    if local_test == LocalTests.TEST:
        np.random.seed(1)
        n_instruments = 1000
        exposures_nm = np.random.normal(0.0, 1.0, size=(n_instruments, 2))
        data = pd.DataFrame(data=exposures_nm, columns=[f"id{n+1}" for n in range(2)])

        fig, ax = plt.subplots(1, 1, figsize=(3.9, 3.4), tight_layout=True)
        global_kwargs = dict(fontsize=6, linewidth=0.5, weight='normal', first_color_fixed=True)
        plot_histplot2d(df=data, ax=ax, **global_kwargs)

    plt.show()


if __name__ == '__main__':

    run_local_test(local_test=LocalTests.TEST)
