#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Tests for hit area detection across all plot types.

Ensures that hitmap generation and bbox extraction work correctly
for every supported plot type.
"""

import matplotlib.pyplot as plt
import numpy as np
import pytest

from figrecipe._editor._bbox import extract_bboxes
from figrecipe._editor._hitmap import generate_hitmap


# Simple plot creators that use plain matplotlib (no figrecipe id kwarg)
def create_plot():
    """Create a simple line plot."""
    fig, ax = plt.subplots()
    x = np.linspace(0, 2 * np.pi, 100)
    ax.plot(x, np.sin(x))
    ax.plot(x, np.cos(x))
    return fig, ax


def create_scatter():
    """Create a simple scatter plot."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    x, y = rng.uniform(0, 10, 50), rng.uniform(0, 10, 50)
    ax.scatter(x, y)
    return fig, ax


def create_bar():
    """Create a simple bar chart."""
    fig, ax = plt.subplots()
    categories = ["A", "B", "C", "D"]
    values = [3, 7, 2, 5]
    ax.bar(categories, values)
    return fig, ax


def create_barh():
    """Create a simple horizontal bar chart."""
    fig, ax = plt.subplots()
    categories = ["A", "B", "C", "D"]
    values = [3, 7, 2, 5]
    ax.barh(categories, values)
    return fig, ax


def create_hist():
    """Create a simple histogram."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    data = rng.normal(0, 1, 1000)
    ax.hist(data, bins=30)
    return fig, ax


def create_imshow():
    """Create a simple image plot."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    data = rng.uniform(0, 1, (10, 10))
    ax.imshow(data)
    return fig, ax


def create_contour():
    """Create a simple contour plot."""
    fig, ax = plt.subplots()
    x = np.linspace(-3, 3, 50)
    y = np.linspace(-3, 3, 50)
    X, Y = np.meshgrid(x, y)
    Z = np.sin(X) * np.cos(Y)
    ax.contour(X, Y, Z)
    return fig, ax


def create_contourf():
    """Create a simple filled contour plot."""
    fig, ax = plt.subplots()
    x = np.linspace(-3, 3, 50)
    y = np.linspace(-3, 3, 50)
    X, Y = np.meshgrid(x, y)
    Z = np.sin(X) * np.cos(Y)
    ax.contourf(X, Y, Z)
    return fig, ax


def create_quiver():
    """Create a simple quiver (vector) plot."""
    fig, ax = plt.subplots()
    x = np.linspace(0, 2, 5)
    y = np.linspace(0, 2, 5)
    X, Y = np.meshgrid(x, y)
    U = np.cos(X)
    V = np.sin(Y)
    ax.quiver(X, Y, U, V)
    return fig, ax


def create_pie():
    """Create a simple pie chart."""
    fig, ax = plt.subplots()
    sizes = [15, 30, 45, 10]
    ax.pie(sizes)
    return fig, ax


def create_fill():
    """Create a simple fill plot."""
    fig, ax = plt.subplots()
    x = np.linspace(0, 2 * np.pi, 100)
    y = np.sin(x)
    ax.fill(x, y)
    return fig, ax


def create_fill_between():
    """Create a fill_between plot."""
    fig, ax = plt.subplots()
    x = np.linspace(0, 2 * np.pi, 100)
    y1 = np.sin(x)
    y2 = np.cos(x)
    ax.fill_between(x, y1, y2)
    return fig, ax


def create_boxplot():
    """Create a simple box plot."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    data = [rng.normal(0, 1, 100) for _ in range(3)]
    ax.boxplot(data)
    return fig, ax


def create_violinplot():
    """Create a simple violin plot."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    data = [rng.normal(0, 1, 100) for _ in range(3)]
    ax.violinplot(data)
    return fig, ax


def create_errorbar():
    """Create an error bar plot."""
    fig, ax = plt.subplots()
    x = np.arange(5)
    y = [2, 4, 3, 5, 4]
    yerr = [0.5, 0.3, 0.4, 0.6, 0.2]
    ax.errorbar(x, y, yerr=yerr)
    return fig, ax


def create_step():
    """Create a step plot."""
    fig, ax = plt.subplots()
    x = np.arange(10)
    y = np.random.default_rng(42).integers(0, 10, 10)
    ax.step(x, y)
    return fig, ax


def create_stem():
    """Create a stem plot."""
    fig, ax = plt.subplots()
    x = np.linspace(0, 2 * np.pi, 20)
    y = np.sin(x)
    ax.stem(x, y)
    return fig, ax


def create_stackplot():
    """Create a stack plot."""
    fig, ax = plt.subplots()
    x = np.arange(5)
    y1 = [1, 2, 3, 4, 5]
    y2 = [2, 3, 2, 3, 2]
    y3 = [1, 1, 1, 1, 1]
    ax.stackplot(x, y1, y2, y3)
    return fig, ax


def create_pcolormesh():
    """Create a pcolormesh plot."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    data = rng.uniform(0, 1, (10, 10))
    ax.pcolormesh(data)
    return fig, ax


def create_hexbin():
    """Create a hexbin plot."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    x = rng.normal(0, 1, 1000)
    y = rng.normal(0, 1, 1000)
    ax.hexbin(x, y, gridsize=20)
    return fig, ax


def create_specgram():
    """Create a spectrogram."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    Fs = 1000
    t = np.arange(0, 1, 1 / Fs)
    signal = np.sin(2 * np.pi * 50 * t) + rng.normal(0, 0.1, len(t))
    ax.specgram(signal, Fs=Fs)
    return fig, ax


def create_streamplot():
    """Create a streamplot."""
    fig, ax = plt.subplots()
    x = np.linspace(-3, 3, 20)
    y = np.linspace(-3, 3, 20)
    X, Y = np.meshgrid(x, y)
    U = -Y
    V = X
    ax.streamplot(X, Y, U, V)
    return fig, ax


def create_eventplot():
    """Create an event plot."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    data = [rng.uniform(0, 10, 20) for _ in range(3)]
    ax.eventplot(data)
    return fig, ax


def create_hist2d():
    """Create a 2D histogram."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    x = rng.normal(0, 1, 1000)
    y = rng.normal(0, 1, 1000)
    ax.hist2d(x, y, bins=20)
    return fig, ax


def create_spy():
    """Create a spy plot (sparse matrix visualization)."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    data = rng.choice([0, 1], size=(10, 10), p=[0.7, 0.3])
    ax.spy(data)
    return fig, ax


def create_matshow():
    """Create a matshow plot."""
    fig, ax = plt.subplots()
    rng = np.random.default_rng(42)
    data = rng.uniform(0, 1, (10, 10))
    ax.matshow(data)
    return fig, ax


# Registry of all plot creators
PLOT_CREATORS = {
    "plot": create_plot,
    "scatter": create_scatter,
    "bar": create_bar,
    "barh": create_barh,
    "hist": create_hist,
    "imshow": create_imshow,
    "contour": create_contour,
    "contourf": create_contourf,
    "quiver": create_quiver,
    "pie": create_pie,
    "fill": create_fill,
    "fill_between": create_fill_between,
    "boxplot": create_boxplot,
    "violinplot": create_violinplot,
    "errorbar": create_errorbar,
    "step": create_step,
    "stem": create_stem,
    "stackplot": create_stackplot,
    "pcolormesh": create_pcolormesh,
    "hexbin": create_hexbin,
    "specgram": create_specgram,
    "streamplot": create_streamplot,
    "eventplot": create_eventplot,
    "hist2d": create_hist2d,
    "spy": create_spy,
    "matshow": create_matshow,
}

ALL_PLOT_TYPES = list(PLOT_CREATORS.keys())


class TestHitmapGeneration:
    """Test hitmap generation for all plot types."""

    @pytest.mark.parametrize("plot_type", ALL_PLOT_TYPES)
    def test_hitmap_generates_without_error(self, plot_type):
        """Each plot type should generate a hitmap without errors."""
        creator = PLOT_CREATORS[plot_type]
        fig, ax = creator()

        try:
            hitmap, color_map = generate_hitmap(fig)

            # Basic assertions
            assert hitmap is not None, f"{plot_type}: hitmap should not be None"
            assert isinstance(color_map, dict), f"{plot_type}: color_map should be dict"
            assert len(color_map) > 0, f"{plot_type}: color_map should not be empty"

        finally:
            plt.close(fig)

    @pytest.mark.parametrize("plot_type", ALL_PLOT_TYPES)
    def test_hitmap_has_plot_elements(self, plot_type):
        """Each plot type should have at least one plot element in color_map."""
        creator = PLOT_CREATORS[plot_type]
        fig, ax = creator()

        try:
            hitmap, color_map = generate_hitmap(fig)

            # Filter out text-only elements (title, xlabel, ylabel)
            text_types = {
                "title",
                "xlabel",
                "ylabel",
                "suptitle",
                "supxlabel",
                "supylabel",
            }
            plot_elements = [
                k
                for k, v in color_map.items()
                if isinstance(v, dict) and v.get("type") not in text_types
            ]

            assert len(plot_elements) > 0, (
                f"{plot_type}: should have at least one plot element. "
                f"Found types: {[color_map[k].get('type') for k in color_map if isinstance(color_map[k], dict)]}"
            )

        finally:
            plt.close(fig)


class TestBboxExtraction:
    """Test bounding box extraction for all plot types."""

    @pytest.mark.parametrize("plot_type", ALL_PLOT_TYPES)
    def test_bbox_extracts_without_error(self, plot_type):
        """Each plot type should have bbox extraction work without errors."""
        creator = PLOT_CREATORS[plot_type]
        fig, ax = creator()

        try:
            bboxes = extract_bboxes(fig, 800, 600)

            # Basic assertions
            assert isinstance(bboxes, dict), f"{plot_type}: bboxes should be dict"
            assert len(bboxes) > 0, f"{plot_type}: bboxes should not be empty"

        finally:
            plt.close(fig)

    @pytest.mark.parametrize("plot_type", ALL_PLOT_TYPES)
    def test_bbox_has_valid_coordinates(self, plot_type):
        """Each bbox should have valid x, y, width, height coordinates."""
        creator = PLOT_CREATORS[plot_type]
        fig, ax = creator()

        try:
            bboxes = extract_bboxes(fig, 800, 600)

            for key, bbox in bboxes.items():
                if key == "_meta":
                    continue  # Skip metadata

                if "x" in bbox:  # Some elements may not have bbox
                    assert bbox["x"] >= 0, f"{plot_type}/{key}: x should be >= 0"
                    assert bbox["y"] >= 0, f"{plot_type}/{key}: y should be >= 0"
                    assert bbox["width"] > 0, f"{plot_type}/{key}: width should be > 0"
                    assert bbox["height"] > 0, (
                        f"{plot_type}/{key}: height should be > 0"
                    )

        finally:
            plt.close(fig)


class TestHitmapBboxConsistency:
    """Test that hitmap and bbox extraction are consistent."""

    @pytest.mark.parametrize("plot_type", ALL_PLOT_TYPES)
    def test_plot_elements_have_bboxes(self, plot_type):
        """Plot elements in colorMap should have corresponding bboxes."""
        creator = PLOT_CREATORS[plot_type]
        fig, ax = creator()

        try:
            hitmap, color_map = generate_hitmap(fig)
            bboxes = extract_bboxes(fig, 800, 600)

            # Get plot element types (excluding text elements)
            text_types = {
                "title",
                "xlabel",
                "ylabel",
                "suptitle",
                "supxlabel",
                "supylabel",
            }
            plot_elements = [
                k
                for k, v in color_map.items()
                if isinstance(v, dict) and v.get("type") not in text_types
            ]

            # Count elements with bboxes
            elements_with_bbox = 0
            elements_without_bbox = []
            for elem_key in plot_elements:
                # Find matching bbox key (may have different naming)
                has_bbox = any(
                    elem_key in bbox_key or bbox_key in elem_key
                    for bbox_key in bboxes.keys()
                    if bbox_key != "_meta"
                )
                if has_bbox:
                    elements_with_bbox += 1
                else:
                    elements_without_bbox.append(elem_key)

            # At least some elements should have bboxes
            # (Some elements like fills may not always have bboxes)
            # Note: boxplot creates many thin lines (whiskers, caps) that have
            # zero-height bboxes, so we use a lower threshold of 30%
            if len(plot_elements) > 0:
                coverage = elements_with_bbox / len(plot_elements) * 100
                assert coverage >= 30, (
                    f"{plot_type}: Only {coverage:.0f}% of plot elements have bboxes. "
                    f"Missing: {elements_without_bbox[:5]}"
                )

        finally:
            plt.close(fig)


class TestSpecificPlotTypes:
    """Test specific plot types that had known issues."""

    def test_imshow_has_image_element(self):
        """imshow should create an 'image' type element."""
        fig, ax = create_imshow()

        try:
            hitmap, color_map = generate_hitmap(fig)

            image_elements = [
                k
                for k, v in color_map.items()
                if isinstance(v, dict) and v.get("type") == "image"
            ]
            assert len(image_elements) > 0, "imshow should create image elements"

        finally:
            plt.close(fig)

    def test_hist_has_bar_or_hist_element(self):
        """hist should create 'hist' or 'bar' type elements."""
        fig, ax = create_hist()

        try:
            hitmap, color_map = generate_hitmap(fig)

            # Without recording, hist may show as bar
            bar_or_hist = [
                k
                for k, v in color_map.items()
                if isinstance(v, dict) and v.get("type") in ("bar", "hist")
            ]
            assert len(bar_or_hist) > 0, "hist should create bar/hist elements"

        finally:
            plt.close(fig)

    def test_quiver_has_quiver_element(self):
        """quiver should create a 'quiver' type element."""
        fig, ax = create_quiver()

        try:
            hitmap, color_map = generate_hitmap(fig)

            quiver_elements = [
                k
                for k, v in color_map.items()
                if isinstance(v, dict) and v.get("type") == "quiver"
            ]
            assert len(quiver_elements) > 0, "quiver should create quiver elements"

        finally:
            plt.close(fig)

    def test_pie_has_pie_element(self):
        """pie should create 'pie' type elements."""
        fig, ax = create_pie()

        try:
            hitmap, color_map = generate_hitmap(fig)

            pie_elements = [
                k
                for k, v in color_map.items()
                if isinstance(v, dict) and v.get("type") == "pie"
            ]
            assert len(pie_elements) > 0, "pie should create pie elements"

        finally:
            plt.close(fig)

    def test_contourf_has_elements(self):
        """contourf should create fill or image elements."""
        fig, ax = create_contourf()

        try:
            hitmap, color_map = generate_hitmap(fig)

            # contourf may create fill or image elements
            # Note: contourf may not create easily selectable elements
            # Just verify hitmap generation doesn't error
            assert hitmap is not None
            assert color_map is not None

        finally:
            plt.close(fig)

    def test_boxplot_has_elements(self):
        """boxplot should create 'boxplot' or 'line' type elements."""
        fig, ax = create_boxplot()

        try:
            hitmap, color_map = generate_hitmap(fig)

            # With recording, boxplot elements should be identified
            # Without recording, they may appear as lines
            elements = [
                k
                for k, v in color_map.items()
                if isinstance(v, dict) and v.get("type") in ("boxplot", "line")
            ]
            assert len(elements) > 0, "boxplot should create boxplot or line elements"

        finally:
            plt.close(fig)

    def test_violinplot_has_elements(self):
        """violinplot should create 'violin' type elements."""
        fig, ax = create_violinplot()

        try:
            hitmap, color_map = generate_hitmap(fig)

            # With recording, violin elements should be identified
            elements = [
                k
                for k, v in color_map.items()
                if isinstance(v, dict)
                and v.get("type") in ("violin", "fill", "linecollection")
            ]
            assert len(elements) > 0, "violinplot should create violin elements"

        finally:
            plt.close(fig)


class TestBboxTypes:
    """Test that specific element types have valid bboxes."""

    def test_scatter_bbox(self):
        """scatter should have valid bbox."""
        fig, ax = create_scatter()

        try:
            bboxes = extract_bboxes(fig, 800, 600)

            scatter_bboxes = [k for k in bboxes.keys() if "scatter" in k]
            assert len(scatter_bboxes) > 0, "scatter should have bbox"

        finally:
            plt.close(fig)

    def test_bar_bbox(self):
        """bar should have valid bbox."""
        fig, ax = create_bar()

        try:
            bboxes = extract_bboxes(fig, 800, 600)

            bar_bboxes = [k for k in bboxes.keys() if "bar" in k]
            assert len(bar_bboxes) > 0, "bar should have bbox"

        finally:
            plt.close(fig)

    def test_line_bbox(self):
        """plot (line) should have valid bbox."""
        fig, ax = create_plot()

        try:
            bboxes = extract_bboxes(fig, 800, 600)

            line_bboxes = [k for k in bboxes.keys() if "line" in k]
            assert len(line_bboxes) > 0, "plot should have line bbox"

        finally:
            plt.close(fig)

    def test_image_bbox(self):
        """imshow should have valid bbox."""
        fig, ax = create_imshow()

        try:
            bboxes = extract_bboxes(fig, 800, 600)

            image_bboxes = [k for k in bboxes.keys() if "image" in k]
            assert len(image_bboxes) > 0, "imshow should have image bbox"

        finally:
            plt.close(fig)

    def test_quiver_bbox(self):
        """quiver should have valid bbox."""
        fig, ax = create_quiver()

        try:
            bboxes = extract_bboxes(fig, 800, 600)

            quiver_bboxes = [k for k in bboxes.keys() if "quiver" in k]
            assert len(quiver_bboxes) > 0, "quiver should have bbox"

        finally:
            plt.close(fig)

    def test_pie_bbox(self):
        """pie should have valid bbox."""
        fig, ax = create_pie()

        try:
            bboxes = extract_bboxes(fig, 800, 600)

            pie_bboxes = [k for k in bboxes.keys() if "wedge" in k or "pie" in k]
            assert len(pie_bboxes) > 0, "pie should have wedge bbox"

        finally:
            plt.close(fig)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
