"""Tests for the plots module."""

import holoviews as hv
import numpy as np
import pandas as pd
import pytest

from holoviz_utils.plots import make_boxplot, make_distributions, make_middle


# Load HoloViews extension for testing
@pytest.fixture(scope="module", autouse=True)
def setup_holoviews():
    """Setup HoloViews with bokeh extension for all tests."""
    hv.extension("bokeh")
    yield


class TestMakeMiddle:
    """Tests for the make_middle function."""

    def test_make_middle_basic(self):
        """Test basic middle calculation."""
        df = pd.DataFrame(
            {"left": [0, 10, 20], "right": [10, 20, 30], "density": [0.5, 1.0, 0.5]}
        )

        result = make_middle(df, yoffset=0, close_ends=False)

        # Check that middle column was created
        assert "middle" in result.columns
        # Check middle values
        assert result["middle"].iloc[0] == 5.0
        assert result["middle"].iloc[1] == 15.0
        assert result["middle"].iloc[2] == 25.0

    def test_make_middle_with_left_inf(self):
        """Test handling of -inf on left side."""
        df = pd.DataFrame(
            {
                "left": [-np.inf, 10, 20],
                "right": [10, 20, 30],
                "density": [0.5, 1.0, 0.5],
            }
        )

        result = make_middle(df, yoffset=0, close_ends=False)

        # When left is -inf, middle should equal right
        assert result["middle"].iloc[0] == 10.0
        # Other values should be normal averages
        assert result["middle"].iloc[1] == 15.0

    def test_make_middle_with_right_inf(self):
        """Test handling of inf on right side."""
        df = pd.DataFrame(
            {"left": [0, 10, 20], "right": [10, 20, np.inf], "density": [0.5, 1.0, 0.5]}
        )

        result = make_middle(df, yoffset=0, close_ends=False)

        # When right is inf, middle should equal left
        assert result["middle"].iloc[-1] == 20.0
        # Other values should be normal averages
        assert result["middle"].iloc[0] == 5.0

    def test_make_middle_with_close_ends(self):
        """Test close_ends parameter adds extra rows."""
        df = pd.DataFrame(
            {"left": [0, 10, 20], "right": [10, 20, 30], "density": [0.5, 1.0, 0.5]}
        )

        result = make_middle(df, yoffset=2.5, close_ends=True)

        # Should have extra rows added for closing ends
        assert len(result) > len(df)
        # Check that density values for added rows match yoffset
        # The added rows should have density equal to yoffset (2.5)
        assert result["density"].iloc[-1] == 2.5

    def test_make_middle_preserves_columns(self):
        """Test that original columns are preserved."""
        df = pd.DataFrame(
            {
                "left": [0, 10, 20],
                "right": [10, 20, 30],
                "density": [0.5, 1.0, 0.5],
                "extra_col": ["a", "b", "c"],
            }
        )

        result = make_middle(df, yoffset=0, close_ends=False)

        # Original columns should still exist
        assert "left" in result.columns
        assert "right" in result.columns
        assert "density" in result.columns


class TestMakeDistributions:
    """Tests for the make_distributions function."""

    def test_make_distributions_single_model(self):
        """Test distribution creation with single model."""
        df = pd.DataFrame(
            {
                "category": ["A", "A", "B", "B"],
                "model": ["M1", "M1", "M1", "M1"],
                "left": [0, 10, 0, 10],
                "right": [10, 20, 10, 20],
                "density": [0.5, 1.0, 0.8, 0.6],
            }
        )

        plot, names = make_distributions(df, spacer=1.25, close_ends=True)

        # Check that plot is an Overlay
        assert isinstance(plot, hv.Overlay)
        # Check that names dict was created
        assert isinstance(names, dict)
        # Should have 2 categories
        assert len(names) == 2

    def test_make_distributions_two_models(self):
        """Test distribution creation with two models."""
        df = pd.DataFrame(
            {
                "category": ["A", "A", "A", "A"],
                "model": ["M1", "M1", "M2", "M2"],
                "left": [0, 10, 0, 10],
                "right": [10, 20, 10, 20],
                "density": [0.5, 1.0, 0.8, 0.6],
            }
        )

        plot, names = make_distributions(df, spacer=1.5, close_ends=False)

        # Check that plot is an Overlay
        assert isinstance(plot, hv.Overlay)
        # Should create areas for both models
        assert len(plot) == 2

    def test_make_distributions_with_colors(self):
        """Test distribution creation with custom colors."""
        df = pd.DataFrame(
            {
                "category": ["A", "A"],
                "model": ["M1", "M1"],
                "left": [0, 10],
                "right": [10, 20],
                "density": [0.5, 1.0],
            }
        )

        colors = {"M1": "blue", "M2": "green"}
        plot, names = make_distributions(df, colors=colors)

        # Should not raise an error
        assert isinstance(plot, hv.Overlay)

    def test_make_distributions_with_middle_column(self):
        """Test with DataFrame that already has middle column."""
        df = pd.DataFrame(
            {
                "category": ["A", "A"],
                "model": ["M1", "M1"],
                "left": [0, 10],
                "right": [10, 20],
                "middle": [5, 15],
                "density": [0.5, 1.0],
            }
        )

        plot, names = make_distributions(df)

        # Should use existing middle column
        assert isinstance(plot, hv.Overlay)


class TestMakeBoxplot:
    """Tests for the make_boxplot function."""

    def test_make_boxplot_basic(self):
        """Test basic boxplot creation."""
        df = pd.DataFrame(
            {
                "category": ["A", "B", "C"],
                "y1": [1.0, 2.0, 1.5],  # min
                "y2": [2.0, 3.0, 2.5],  # Q1
                "y3": [3.0, 4.0, 3.5],  # median
                "y4": [4.0, 5.0, 4.5],  # Q3
                "y5": [5.0, 6.0, 5.5],  # max
            }
        )

        mapping = {
            "x": "category",
            "y1": "y1",
            "y2": "y2",
            "y3": "y3",
            "y4": "y4",
            "y5": "y5",
        }

        plot = make_boxplot(df, mapping)

        # Check that plot is an Overlay
        assert isinstance(plot, hv.Overlay)
        # Should contain 3 elements: Bars, ErrorBars, Points
        assert len(plot) == 3

    def test_make_boxplot_with_colors(self):
        """Test boxplot with colors parameter (when model is x-axis)."""
        # When using colors, model should be the x-axis variable
        df = pd.DataFrame(
            {
                "model": ["M1", "M2"],
                "y1": [1.0, 2.0],
                "y2": [2.0, 3.0],
                "y3": [3.0, 4.0],
                "y4": [4.0, 5.0],
                "y5": [5.0, 6.0],
            }
        )

        mapping = {
            "x": "model",  # model is the x-axis
            "y1": "y1",
            "y2": "y2",
            "y3": "y3",
            "y4": "y4",
            "y5": "y5",
        }

        colors = {"M1": "blue", "M2": "red"}
        plot = make_boxplot(df, mapping, colors=colors)

        # Should not raise an error
        assert isinstance(plot, hv.Overlay)

    def test_make_boxplot_with_nan_values(self):
        """Test boxplot handles NaN values."""
        df = pd.DataFrame(
            {
                "category": ["A", "B"],
                "y1": [1.0, 2.0],
                "y2": [2.0, 3.0],
                "y3": [np.nan, np.nan],  # All NaN in median
                "y4": [4.0, 5.0],
                "y5": [5.0, 6.0],
            }
        )

        mapping = {
            "x": "category",
            "y1": "y1",
            "y2": "y2",
            "y3": "y3",
            "y4": "y4",
            "y5": "y5",
        }

        # Should handle NaN values gracefully (fillna converts NaN to 0.0)
        plot = make_boxplot(df, mapping)
        assert isinstance(plot, hv.Overlay)

    def test_make_boxplot_overlay_elements(self):
        """Test that boxplot contains expected HoloViews elements."""
        df = pd.DataFrame(
            {
                "category": ["A"],
                "y1": [1.0],
                "y2": [2.0],
                "y3": [3.0],
                "y4": [4.0],
                "y5": [5.0],
            }
        )

        mapping = {
            "x": "category",
            "y1": "y1",
            "y2": "y2",
            "y3": "y3",
            "y4": "y4",
            "y5": "y5",
        }

        plot = make_boxplot(df, mapping)

        # Extract element types from overlay
        element_types = [type(elem).__name__ for elem in plot]

        # Should contain Bars, ErrorBars, and Points
        assert "Bars" in element_types
        assert "ErrorBars" in element_types
        assert "Points" in element_types


class TestIntegration:
    """Integration tests combining multiple functions."""

    def test_distribution_to_boxplot_workflow(self):
        """Test creating both distribution and boxplot from same data structure."""
        # Create sample data
        df_dist = pd.DataFrame(
            {
                "category": ["A", "A", "B", "B"],
                "model": ["M1", "M1", "M1", "M1"],
                "left": [0, 10, 0, 10],
                "right": [10, 20, 10, 20],
                "density": [0.5, 1.0, 0.8, 0.6],
            }
        )

        df_box = pd.DataFrame(
            {
                "category": ["A", "B"],
                "y1": [1.0, 2.0],
                "y2": [2.0, 3.0],
                "y3": [3.0, 4.0],
                "y4": [4.0, 5.0],
                "y5": [5.0, 6.0],
            }
        )

        # Create both plots
        dist_plot, names = make_distributions(df_dist)
        box_mapping = {
            "x": "category",
            "y1": "y1",
            "y2": "y2",
            "y3": "y3",
            "y4": "y4",
            "y5": "y5",
        }
        box_plot = make_boxplot(df_box, box_mapping)

        # Both should be valid HoloViews elements
        assert isinstance(dist_plot, hv.Overlay)
        assert isinstance(box_plot, hv.Overlay)
