from __future__ import annotations

import numpy as np
import pytest

from hist import Hist, NamedHist, axis

plt = pytest.importorskip("matplotlib.pyplot")
pytest.importorskip("scipy")


def test_general_plot1d():
    """
    Test general plot1d -- whether 1d-Hist can be plotted properly.
    """

    h = Hist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
    ).fill(np.random.normal(size=10))

    assert h.plot1d(color="green", ls="--", lw=3)
    plt.close("all")

    # dimension error
    h = Hist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(np.random.normal(size=10), np.random.normal(size=10))

    with pytest.raises(Exception):
        h.plot1d()

    # wrong kwargs names
    with pytest.raises(Exception):
        h.project("A").plot1d(abc="red")

    # wrong kwargs type
    with pytest.raises(Exception):
        h.project("B").plot1d(ls="red")

    plt.close("all")


def test_general_plot2d():
    """
    Test general plot2d -- whether 2d-Hist can be plotted properly.
    """

    h = Hist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(np.random.normal(size=10), np.random.normal(size=10))

    assert h.plot2d(cmap="cividis")

    # dimension error
    h = Hist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(np.random.normal(size=10), np.random.normal(size=10))

    with pytest.raises(Exception):
        h.project("A").plot2d()

    # wrong kwargs names
    with pytest.raises(Exception):
        h.plot2d(abc="red")

    # wrong kwargs type
    with pytest.raises(Exception):
        h.plot2d(cmap=0.1)

    plt.close("all")


def test_general_plot2d_full():
    """
    Test general plot2d_full -- whether 2d-Hist can be fully plotted properly.
    """

    h = Hist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(np.random.normal(size=10), np.random.normal(size=10))

    assert h.plot2d_full(
        main_cmap="cividis",
        top_ls="--",
        top_color="orange",
        top_lw=2,
        side_ls="-.",
        side_lw=1,
        side_color="steelblue",
    )

    assert h.plot2d_full(
        main_cmap="cividis",
        top_kw={
            "ls": "--",
            "color": "orange",
            "lw": 2,
        },
        side_kw={
            "ls": "-.",
            "lw": 1,
            "color": "steelblue",
        },
    )

    # dimension error
    h = Hist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(np.random.normal(size=10), np.random.normal(size=10))

    with pytest.raises(Exception):
        h.project("A").plot2d_full()

    # wrong kwargs names
    with pytest.raises(Exception):
        h.plot2d_full(abc="red")

    with pytest.raises(Exception):
        h.plot2d_full(color="red")

    # wrong kwargs type
    with pytest.raises(Exception):
        h.plot2d_full(main_cmap=0.1, side_lw="autumn")

    plt.close("all")


def test_general_plot():
    """
    Test general plot -- whether Hist can be plotted properly.
    """

    h = Hist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
    ).fill(np.random.normal(size=10))

    assert h.plot(color="green", ls="--", lw=3)

    h = Hist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(np.random.normal(size=10), np.random.normal(size=10))

    assert h.plot(cmap="cividis")

    # dimension error
    h = Hist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="C", label="c [units]", underflow=False, overflow=False
        ),
    ).fill(
        np.random.normal(size=10), np.random.normal(size=10), np.random.normal(size=10)
    )

    with pytest.raises(Exception):
        h.plot()

    # wrong kwargs names
    with pytest.raises(Exception):
        h.project("A").plot(abc="red")

    with pytest.raises(Exception):
        h.project("A", "C").plot(abc="red")

    # wrong kwargs type
    with pytest.raises(Exception):
        h.project("B").plot(ls="red")

    with pytest.raises(Exception):
        h.project("A", "C").plot(cmap=0.1)

    plt.close("all")


def test_general_plot_pull():
    """
    Test general plot_pull -- whether 1d-Hist can be plotted pull properly.
    """
    pytest.importorskip("iminuit")

    np.random.seed(42)

    h = Hist(
        axis.Regular(
            50, -4, 4, name="S", label="s [units]", underflow=False, overflow=False
        )
    ).fill(np.random.normal(size=10))

    def pdf(x, a=1 / np.sqrt(2 * np.pi), x0=0, sigma=1, offset=0):
        return a * np.exp(-((x - x0) ** 2) / (2 * sigma**2)) + offset

    assert h.plot_pull(
        pdf,
        eb_ecolor="crimson",
        eb_mfc="crimson",
        eb_mec="crimson",
        eb_fmt="o",
        eb_ms=6,
        eb_capsize=1,
        eb_capthick=2,
        eb_alpha=0.8,
        fp_c="chocolate",
        fp_ls="-",
        fp_lw=3,
        fp_alpha=1.0,
        bar_fc="orange",
        pp_num=6,
        pp_fc="orange",
        pp_alpha=0.618,
        pp_ec=None,
    )

    pdf_str = "a * np.exp(-((x - x0) ** 2) / (2 * sigma ** 2)) + offset"

    assert h.plot_pull(pdf_str)

    assert h.plot_pull("gauss")

    assert h.plot_pull("gauss", likelihood=True)

    # dimension error
    hh = Hist(
        axis.Regular(
            50, -4, 4, name="X", label="s [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="Y", label="s [units]", underflow=False, overflow=False
        ),
    ).fill(np.random.normal(size=10), np.random.normal(size=10))

    with pytest.raises(Exception):
        hh.plot_pull(pdf)

    # no eval-able variable
    with pytest.raises(Exception):
        h.plot_pull("1")

    with pytest.raises(Exception):
        h.plot_pull(1)

    with pytest.raises(Exception):
        h.plot_pull(0.1)

    with pytest.raises(Exception):
        h.plot_pull((1, 2))

    with pytest.raises(Exception):
        h.plot_pull([1, 2])

    with pytest.raises(Exception):
        h.plot_pull({"a": 1})

    # wrong kwargs names
    with pytest.raises(Exception):
        h.plot_pull(pdf, abc="crimson", xyz="crimson")

    with pytest.raises(Exception):
        h.plot_pull(pdf, ecolor="crimson", mfc="crimson")

    # not disabled params
    h.plot_pull(pdf, eb_label="value")

    h.plot_pull(pdf, fp_label="value")

    h.plot_pull(pdf, ub_label="value")

    h.plot_pull(pdf, bar_label="value")

    h.plot_pull(pdf, pp_label="value")

    # disabled params
    with pytest.raises(Exception):
        h.plot_pull(pdf, bar_width="value")

    # wrong kwargs types
    with pytest.raises(Exception):
        h.plot_pull(pdf, eb_ecolor=1.0, eb_mfc=1.0)  # kwargs should be str

    plt.close("all")


def test_named_plot1d():
    """
    Test named plot1d -- whether 1d-NamedHist can be plotted properly.
    """

    h = NamedHist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
    ).fill(A=np.random.normal(size=10))

    assert h.plot1d(color="green", ls="--", lw=3)
    plt.close("all")

    # dimension error
    h = NamedHist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(B=np.random.normal(size=10), A=np.random.normal(size=10))

    with pytest.raises(Exception):
        h.plot1d()

    # wrong kwargs names
    with pytest.raises(Exception):
        h.project("A").plot1d(abc="red")

    # wrong kwargs type
    with pytest.raises(Exception):
        h.project("B").plot1d(ls="red")

    plt.close("all")


def test_named_plot2d():
    """
    Test named plot2d -- whether 2d-NamedHist can be plotted properly.
    """

    h = NamedHist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(B=np.random.normal(size=10), A=np.random.normal(size=10))

    assert h.plot2d(cmap="cividis")
    plt.close("all")

    # dimension error
    h = NamedHist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(B=np.random.normal(size=10), A=np.random.normal(size=10))

    with pytest.raises(Exception):
        h.project("A").plot2d()

    # wrong kwargs names
    with pytest.raises(Exception):
        h.plot2d(abc="red")

    # wrong kwargs type
    with pytest.raises(Exception):
        h.plot2d(cmap=0.1)

    plt.close("all")


def test_named_plot2d_full():
    """
    Test named plot2d_full -- whether 2d-NamedHist can be fully plotted properly.
    """

    h = NamedHist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(B=np.random.normal(size=10), A=np.random.normal(size=10))

    assert h.plot2d_full(
        main_cmap="cividis",
        top_ls="--",
        top_color="orange",
        top_lw=2,
        side_ls="-.",
        side_lw=1,
        side_color="steelblue",
    )
    plt.close("all")

    # dimension error
    h = NamedHist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(B=np.random.normal(size=10), A=np.random.normal(size=10))

    with pytest.raises(Exception):
        h.project("A").plot2d_full()

    # wrong kwargs names
    with pytest.raises(Exception):
        h.plot2d_full(abc="red")

    with pytest.raises(Exception):
        h.plot2d_full(color="red")

    # wrong kwargs type
    with pytest.raises(Exception):
        h.plot2d_full(main_cmap=0.1, side_lw="autumn")

    plt.close("all")


def test_named_plot():
    """
    Test named plot -- whether NamedHist can be plotted properly.
    """

    h = NamedHist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
    ).fill(A=np.random.normal(size=10))

    assert h.plot(color="green", ls="--", lw=3)

    h = NamedHist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
    ).fill(B=np.random.normal(size=10), A=np.random.normal(size=10))

    assert h.plot(cmap="cividis")
    plt.close("all")

    # dimension error
    h = NamedHist(
        axis.Regular(
            50, -5, 5, name="A", label="a [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="B", label="b [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="C", label="c [units]", underflow=False, overflow=False
        ),
    ).fill(
        A=np.random.normal(size=10),
        B=np.random.normal(size=10),
        C=np.random.normal(size=10),
    )

    with pytest.raises(Exception):
        h.plot()

    # wrong kwargs names
    with pytest.raises(Exception):
        h.project("A").plot(abc="red")

    with pytest.raises(Exception):
        h.project("A", "C").plot(abc="red")

    # wrong kwargs type
    with pytest.raises(Exception):
        h.project("B").plot(ls="red")

    with pytest.raises(Exception):
        h.project("A", "C").plot(cmap=0.1)

    plt.close("all")


def test_named_plot_pull():
    """
    Test named plot_pull -- whether 1d-NamedHist can be plotted pull properly.
    """

    np.random.seed(42)

    h = NamedHist(
        axis.Regular(
            50, -4, 4, name="S", label="s [units]", underflow=False, overflow=False
        )
    ).fill(S=np.random.normal(size=10))

    def pdf(x, a=1 / np.sqrt(2 * np.pi), x0=0, sigma=1, offset=0):
        return a * np.exp(-((x - x0) ** 2) / (2 * sigma**2)) + offset

    assert h.plot_pull(
        pdf,
        eb_ecolor="crimson",
        eb_mfc="crimson",
        eb_mec="crimson",
        eb_fmt="o",
        eb_ms=6,
        eb_capsize=1,
        eb_capthick=2,
        eb_alpha=0.8,
        fp_c="chocolate",
        fp_ls="-",
        fp_lw=3,
        fp_alpha=1.0,
        bar_fc="orange",
        pp_num=6,
        pp_fc="orange",
        pp_alpha=0.618,
        pp_ec=None,
    )

    pdf_str = "a * np.exp(-((x - x0) ** 2) / (2 * sigma ** 2)) + offset"

    assert h.plot_pull(pdf_str)

    # dimension error
    hh = NamedHist(
        axis.Regular(
            50, -4, 4, name="X", label="s [units]", underflow=False, overflow=False
        ),
        axis.Regular(
            50, -4, 4, name="Y", label="s [units]", underflow=False, overflow=False
        ),
    ).fill(X=np.random.normal(size=10), Y=np.random.normal(size=10))

    with pytest.raises(Exception):
        hh.plot_pull(pdf)

    # no eval-able variable
    with pytest.raises(Exception):
        h.plot_pull("1")

    with pytest.raises(Exception):
        h.plot_pull(1)

    with pytest.raises(Exception):
        h.plot_pull(0.1)

    with pytest.raises(Exception):
        h.plot_pull((1, 2))

    with pytest.raises(Exception):
        h.plot_pull([1, 2])

    with pytest.raises(Exception):
        h.plot_pull({"a": 1})
    plt.close("all")

    # wrong kwargs names
    with pytest.raises(Exception):
        h.plot_pull(pdf, abc="crimson", xyz="crimson")

    with pytest.raises(Exception):
        h.plot_pull(pdf, ecolor="crimson", mfc="crimson")

    # not disabled params
    h.plot_pull(pdf, eb_label="value")

    h.plot_pull(pdf, fp_label="value")

    h.plot_pull(pdf, ub_label="value")

    h.plot_pull(pdf, bar_label="value")

    h.plot_pull(pdf, pp_label="value")

    # disabled params
    with pytest.raises(Exception):
        h.plot_pull(pdf, bar_width="value")

    # wrong kwargs types
    with pytest.raises(Exception):
        h.plot_pull(pdf, eb_ecolor=1.0, eb_mfc=1.0)  # kwargs should be str

    plt.close("all")


@pytest.mark.parametrize("str_alias", ["normal", "gauss", "gaus"])
@pytest.mark.parametrize("use_likelihood", [True, False])
def test_ratiolike_str_alias(str_alias, use_likelihood):
    """
    Test str alias for callable in plot_ratio and plot_pull
    """
    pytest.importorskip("iminuit")

    np.random.seed(42)

    h = NamedHist(
        axis.Regular(
            50, -4, 4, name="S", label="s [units]", underflow=False, overflow=False
        )
    ).fill(S=np.random.normal(size=10))

    assert h.plot_ratio(str_alias, likelihood=use_likelihood)
    assert h.plot_pull(str_alias, likelihood=use_likelihood)


@pytest.mark.mpl_image_compare(baseline_dir="baseline", savefig_kwargs={"dpi": 50})
def test_image_plot_pull():
    """
    Test plot_pull by comparing against a reference image generated via
    `pytest --mpl-generate-path=tests/baseline`
    """

    np.random.seed(42)

    h = Hist(
        axis.Regular(
            50, -4, 4, name="S", label="s [units]", underflow=False, overflow=False
        )
    ).fill(np.random.normal(size=100))

    def pdf(x, a=1 / np.sqrt(2 * np.pi), x0=0, sigma=1, offset=0):
        return a * np.exp(-((x - x0) ** 2) / (2 * sigma**2)) + offset

    fig = plt.figure()

    assert h.plot_pull(
        pdf,
        eb_color="black",
        fp_color="blue",
        ub_color="lightblue",
        fit_fmt=r"{name} = {value:.3g} $\pm$ {error:.3g}",
    )

    return fig


@pytest.mark.mpl_image_compare(baseline_dir="baseline", savefig_kwargs={"dpi": 50})
def test_image_plot_ratio_hist():
    """
    Test plot_pull by comparing against a reference image generated via
    `pytest --mpl-generate-path=tests/baseline`
    """

    np.random.seed(42)

    hist_1 = Hist(
        axis.Regular(
            50, -5, 5, name="X", label="x [units]", underflow=False, overflow=False
        )
    ).fill(np.random.normal(size=1000))
    hist_2 = Hist(
        axis.Regular(
            50, -5, 5, name="X", label="x [units]", underflow=False, overflow=False
        )
    ).fill(np.random.normal(size=1700))

    fig = plt.figure()

    assert hist_1.plot_ratio(
        hist_2, rp_num_label="numerator", rp_denom_label="denominator"
    )

    return fig


@pytest.mark.mpl_image_compare(baseline_dir="baseline", savefig_kwargs={"dpi": 50})
def test_image_plot_ratio_callable():
    """
    Test plot_pull by comparing against a reference image generated via
    `pytest --mpl-generate-path=tests/baseline`
    """

    np.random.seed(42)

    hist_1 = Hist(
        axis.Regular(
            50, -5, 5, name="X", label="x [units]", underflow=False, overflow=False
        )
    ).fill(np.random.normal(size=1000))

    def model(x, a=1 / np.sqrt(2 * np.pi), x0=0, sigma=1, offset=0):
        return a * np.exp(-((x - x0) ** 2) / (2 * sigma**2)) + offset

    fig = plt.figure()

    assert hist_1.plot_ratio(
        model, eb_color="black", fp_color="blue", ub_color="lightblue"
    )

    return fig


@pytest.mark.mpl_image_compare(baseline_dir="baseline", savefig_kwargs={"dpi": 50})
def test_plot1d_auto_handling():
    """
    Test plot() by comparing against a reference image generated via
    `pytest --mpl-generate-path=tests/baseline`
    """

    np.random.seed(42)

    h = Hist(
        axis.Regular(10, 0, 10, name="variable", label="variable"),
        axis.StrCategory("", name="dataset", growth=True),
    )

    h_nameless = Hist(
        axis.Regular(10, 0, 10),
        axis.StrCategory("", growth=True),
    )

    h.fill(dataset="A", variable=np.random.normal(3, 2, 100))
    h.fill(dataset="B", variable=np.random.normal(5, 2, 100))
    h.fill(dataset="C", variable=np.random.normal(7, 2, 100))

    h_nameless.fill(np.random.normal(3, 2, 1000), "A")
    h_nameless.fill(np.random.normal(5, 2, 1000), "B")
    h_nameless.fill(np.random.normal(7, 2, 1000), "C")

    fig, (ax1, ax2) = plt.subplots(2, 2, figsize=(14, 10))

    assert h.plot(ax=ax1[0])
    assert h_nameless.plot(ax=ax2[0])

    # Discrete axis plotting not yet implemented
    # assert h.plot(ax=ax1[1], overlay='variable')
    # assert h.plot(ax=ax2[1], overlay=1)

    return fig


def test_plot1d_legend_functionality():
    """
    Test plot1d legend functionality including:
    - Default legend behavior
    - Legend with axis label as title
    - Legend opt-out
    """
    np.random.seed(42)

    # Create histogram with labeled axis for stacked plots
    h = Hist(
        axis.Regular(10, 0, 10, name="variable", label="Variable [units]"),
        axis.StrCategory("", name="dataset", label="Dataset Type", growth=True),
    )

    h.fill(dataset="Signal", variable=np.random.normal(5, 1, 100))
    h.fill(dataset="Background", variable=np.random.normal(3, 2, 100))

    # Test with default legend (should be True)
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

    # Test with legend=True (default)
    artists1 = h.plot1d(ax=ax1)
    legend1 = ax1.get_legend()
    assert legend1 is not None, "Legend should be created by default"
    assert legend1.get_title().get_text() == "Dataset Type", (
        "Legend title should be axis label"
    )

    # Test with legend=False
    artists2 = h.plot1d(ax=ax2, legend=False)
    legend2 = ax2.get_legend()
    assert legend2 is None, "Legend should not be created when legend=False"

    # Test return type is consistent
    assert type(artists1) is type(artists2), "Return type should be consistent"

    plt.close(fig)


def test_plot1d_legend_without_axis_label():
    """
    Test plot1d legend functionality when axis has no explicit label.
    """
    np.random.seed(42)

    # Create histogram without explicit axis label (will default to name)
    h = Hist(
        axis.Regular(10, 0, 10, name="variable"),
        axis.StrCategory("", name="dataset", growth=True),
    )

    h.fill(dataset="A", variable=np.random.normal(5, 1, 100))
    h.fill(dataset="B", variable=np.random.normal(3, 2, 100))

    fig, ax = plt.subplots()

    # Test with legend=True but no explicit axis label
    h.plot1d(ax=ax)
    legend = ax.get_legend()
    assert legend is not None, "Legend should still be created"
    # Title should be the axis name when no explicit label is provided
    title = legend.get_title().get_text()
    assert title == "dataset", "Legend title should be axis name when no explicit label"

    plt.close(fig)


def test_plot1d_legend_1d_histogram():
    """
    Test that 1D histograms don't get legends (since they don't have categories to legend).
    """
    np.random.seed(42)

    # Create simple 1D histogram
    h = Hist(axis.Regular(10, 0, 10, name="variable", label="Variable [units]"))
    h.fill(np.random.normal(5, 1, 100))

    fig, ax = plt.subplots()

    # Test 1D histogram (should not create legend since no categories)
    artists = h.plot1d(ax=ax)
    # For 1D histograms, the legend parameter doesn't apply since there are no categories
    # The function should still work and return artists
    assert artists is not None

    plt.close(fig)


def test_plot_ratio_misalignment():
    import numpy as np

    import hist
    from hist.plot import plot_ratio_array

    h = hist.new.Int(0, 3, name="x").Double()
    h.fill([0, 1, 2])

    ratio = np.ones(3)
    ratio_uncert = np.zeros((2, 3))

    captured = {}

    def record_x(_self, x, *a, **k):
        captured["x"] = np.asarray(x)

    class Ax:
        errorbar = record_x
        bar = record_x

        def axhline(self, *a, **k):
            return None

        def set_xlim(self, *a, **k):
            return None

        def set_ylim(self, *a, **K):
            return None

        def set_xlabel(self, *a, **k):
            return None

        def set_ylabel(self, *a, **K):
            return None

    plot_ratio_array(h, ratio, ratio_uncert, ax=Ax(), uncert_draw_type="line")

    assert np.allclose(captured["x"], h.axes[0].edges[:-1])
