"""Visualization utilities for NGBoost and gradient boosting models.

This module provides visualization functions for NGBoost models including
distribution plots and learning curves for gradient boosting models.
"""

import sys
import warnings
from decimal import Decimal
from math import ceil
from typing import Optional

import matplotlib.pyplot as plt
import numpy as np
from sklearn.utils import check_array

from yikit.helpers import is_installed
from yikit.visualize._utils import (
    COLORS,
    set_font,
    with_custom_matplotlib_settings,
)

if sys.version_info >= (3, 10):
    from types import NoneType
else:
    NoneType = type(None)  # type: ignore[assignment,misc]

if is_installed("tqdm"):
    from tqdm.auto import tqdm
else:
    from yikit.helpers import dummy_tqdm as tqdm

if is_installed("lightgbm"):
    from lightgbm import (  # type: ignore[reportMissingImports]
        LGBMClassifier,
        LGBMRegressor,
    )
else:
    LGBMClassifier = NoneType  # type: ignore[misc,assignment]
    LGBMRegressor = NoneType  # type: ignore[misc,assignment]


if is_installed("ngboost"):
    from ngboost import (  # type: ignore[reportMissingImports]
        NGBClassifier,
        NGBRegressor,
    )
    from ngboost.distns import Distn  # type: ignore[reportMissingImports]
else:
    NGBClassifier = NoneType  # type: ignore[misc,assignment]
    NGBRegressor = NoneType  # type: ignore[misc,assignment]
    Distn = NoneType  # type: ignore[misc,assignment]


@with_custom_matplotlib_settings()
def get_dist_figure(
    y_dist,
    y_true=None,
    keep_y_range=True,
    return_axis=False,
    verbose=True,
    titles=[],
    fontfamily: Optional[str] = None,
):
    """get distribution figure.

    Parameters
    ----------
    y_dist : ngboost.distns.Distn object
        generated by ngb.pred_dist(X)
    y_true : 1d array, optional
        [description], by default None
    keep_y_range : bool, optional
        keep y range equal or not, by default True
    return_axis : bool, optional
        return axis or not, by default False
    verbose : bool, optional
        drwaing progress bar or not, by default True
    titles : 1d list, optional
        titles whose lengh == n_samples, by default []
    fontfamily : str, optional
        fontfamily, by default None

    Returns
    -------
    matplotlib.pyplot.Figure | (matplotlib.pyplot.Figure, matplotlib.pyplot.Axis)
    """
    # check
    if not isinstance(y_dist, Distn):
        raise TypeError("`y_dist` is not {0} object".format(Distn.__name__))
    y_pred = y_dist.mean()
    y_pred = check_array(y_pred, ensure_2d=False)
    n_samples = len(y_pred)
    if len(titles) not in (0, n_samples):
        raise ValueError(
            "`titles`'s lengh must be 0 or n_samples({}).".format(n_samples)
        )

    offset = np.ptp(y_pred) * 0.05
    y_range = np.linspace(
        min(y_pred) - offset, max(y_pred) + offset, 200
    ).reshape((-1, 1))
    dist_values = y_dist.pdf(y_range).transpose()

    n_rows = ceil(Decimal(n_samples).sqrt())
    n_cols = n_samples // n_rows + int(n_samples % n_rows > 0)

    set_font(fontfamily=fontfamily)

    fig, axes = plt.subplots(
        n_rows,
        n_cols,
        facecolor="white",
        dpi=72,
        figsize=(6.4 * n_cols, 4.8 * n_rows),
    )

    if verbose:
        pbar = tqdm(
            total=n_cols * n_rows + 1, desc="Drawing distribution figure"
        )
    # 一つずつ分布を書いていく．
    for idx in range(n_samples):
        ax = axes[idx // n_cols][idx % n_cols]
        ax.plot(y_range, dist_values[idx], c="#022C5E")

        prob_max_temp = max(
            dist_values[idx]
        )  # このidxにおける確率密度の最大値
        if y_true is not None:
            y_true = check_array(y_true, ensure_2d=False)
            ax.vlines(
                y_true[idx], 0, prob_max_temp, "#ab4e15", label="ground truth"
            )
        ax.vlines(y_pred[idx], 0, prob_max_temp, "#dc143c", label="pred")
        ax.set_ylabel("Probability density")
        ax.legend(loc="best", facecolor="#f0f0f0", edgecolor="None")
        if titles:
            ax.set_title("{0}".format(titles[idx]))
        else:
            ax.set_title("idx: {0}".format(idx))

        ax.set_xlim(y_range[0], y_range[-1])
        if keep_y_range:
            ax.set_ylim(None, np.max(dist_values) * 1.05)

        if verbose:
            pbar.update(1)
    # 余ったところを消す．
    for idx in range(n_samples, n_rows * n_cols):
        ax = axes[idx // n_cols][idx % n_cols]
        ax.axis("off")
        if verbose:
            pbar.update(1)

    fig.tight_layout()

    if verbose:
        pbar.update(1)
        pbar.set_description(desc="Completed")
        pbar.close()
    if return_axis:
        return fig, axes
    else:
        return fig


def is_correct_dist(y_pred, y_dist):
    b = np.allclose(y_dist.mean(), y_pred)
    if not b:
        warnings.warn("`y_dist.mean()` and `y_pred` is not close.")
    return b


def get_learning_curve_gb(
    estimator, fontfamily: Optional[str] = None, return_axis: bool = False
):
    """Plot learning curve for gradient boosting models.

    This function visualizes the learning curve (evaluation scores over iterations)
    for NGBoost or LightGBM models. It plots the evaluation scores from the
    evals_result attribute of the estimator.

    Parameters
    ----------
    estimator : NGBRegressor, NGBClassifier, LGBMRegressor, or LGBMClassifier
        A fitted gradient boosting model with an evals_result or evals_result_
        attribute containing evaluation history.
    fontfamily : str or None, optional
        Font family to use for the plot. If None, uses default font.
        Default is None.
    return_axis : bool, optional
        Whether to return the matplotlib axis object along with the figure.
        Default is False.

    Returns
    -------
    matplotlib.pyplot.Figure or tuple
        If return_axis=False, returns only the figure.
        If return_axis=True, returns (figure, axis) tuple.

    Raises
    ------
    TypeError
        If the estimator is not an NGBoost or LightGBM model.

    Examples
    --------
    >>> from yikit.visualize import get_learning_curve_gb
    >>> from ngboost import NGBRegressor
    >>> import numpy as np
    >>> X = np.random.randn(100, 10)
    >>> y = np.random.randn(100)
    >>> model = NGBRegressor()
    >>> model.fit(X, y)
    >>> fig = get_learning_curve_gb(model)
    """
    # check_estimator
    if isinstance(estimator, (NGBClassifier, NGBRegressor)):
        evals_result = estimator.evals_result
    elif isinstance(estimator, (LGBMClassifier, LGBMRegressor)):
        evals_result = estimator.evals_result_
    else:
        raise TypeError(estimator.__class__.__name__)

    # font
    set_font(fontfamily=fontfamily)

    # generate figure
    fig, ax = plt.subplots(1, 1, facecolor="white", dpi=72)

    # plot
    for data_name, result in evals_result.items():
        if data_name not in COLORS:
            raise ValueError("{} is not in `COLOR`.".format(data_name))
        score_name = list(result.keys())[0]
        score = result[score_name]
        ax.plot(range(len(score)), score, label=data_name, c=COLORS[data_name])

    # label
    ax.set_xlabel("n_iter")
    ax.set_ylabel(score_name)

    # legend
    ax.legend(facecolor="#f0f0f0", edgecolor="None", fontsize=10)
    # tight_layout
    fig.tight_layout()
    if return_axis:
        return fig, ax
    else:
        return fig
