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

import matplotlib.pyplot as plt
import numpy as np
from ngboost.distns import Distn
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 is_installed("tqdm"):
    from tqdm.auto import tqdm
else:
    from yikit.helpers import dummy_tqdm as tqdm

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


if is_installed("ngboost"):
    from ngboost import NGBClassifier, NGBRegressor
    from ngboost.distns import Distn
else:
    NGBClassifier = None  # type: ignore[misc,assignment]
    NGBRegressor = None  # type: ignore[misc,assignment]
    Distn = None  # 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
):
    # 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
