"""
Tests for plotting module.
"""

import pytest
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend for testing
import matplotlib.pyplot as plt
from medtda.plotting import (
    plot_persistence_diagram,
    plot_barcode,
    plot_betti_curve,
    plot_landscape,
    plot_entropy_summary,
    plot_lifespan,
    plot_silhouette,
    plot_tropical_coordinates,
)


@pytest.fixture
def sample_barcode():
    """Create a simple barcode for testing."""
    return np.array([
        [0.0, 0.5],
        [0.1, 0.8],
        [0.2, 1.2],
        [0.3, 0.7],
    ])


@pytest.fixture
def sample_barcodes_dict():
    """Create sample barcodes dictionary with multiple dimensions."""
    return {
        'H0': np.array([
            [0.0, 0.5],
            [0.1, 0.8],
            [0.2, 1.5],
        ]),
        'H1': np.array([
            [0.3, 0.9],
            [0.4, 1.0],
            [0.5, 1.2],
        ]),
        'H2': np.array([
            [0.6, 1.1],
        ])
    }


@pytest.fixture
def empty_barcode():
    """Create an empty barcode."""
    return np.array([]).reshape(0, 2)


class TestPlotPersistenceDiagram:
    """Tests for plot_persistence_diagram function."""
    
    def test_plot_single_barcode(self, sample_barcode):
        """Test plotting a single barcode array."""
        ax = plot_persistence_diagram(sample_barcode)
        
        assert ax is not None
        assert isinstance(ax, plt.Axes)
        plt.close()
    
    def test_plot_dict_barcodes(self, sample_barcodes_dict):
        """Test plotting barcodes from dictionary."""
        ax = plot_persistence_diagram(sample_barcodes_dict)
        
        assert ax is not None
        assert isinstance(ax, plt.Axes)
        plt.close()
    
    def test_plot_specific_dimensions(self, sample_barcodes_dict):
        """Test plotting specific dimensions only."""
        ax = plot_persistence_diagram(sample_barcodes_dict, dimensions=[0, 1])
        
        assert ax is not None
        plt.close()
    
    def test_custom_color_palette(self, sample_barcodes_dict):
        """Test custom color palette."""
        color_palette = ['red', 'blue', 'green']
        
        ax = plot_persistence_diagram(
            sample_barcodes_dict,
            color_palette=color_palette
        )
        
        assert ax is not None
        plt.close()
    
    def test_no_diagonal(self, sample_barcode):
        """Test plot without diagonal line."""
        ax = plot_persistence_diagram(sample_barcode, show_diagonal=False)
        
        assert ax is not None
        plt.close()
    
    def test_no_legend(self, sample_barcodes_dict):
        """Test plot without legend."""
        ax = plot_persistence_diagram(sample_barcodes_dict, legend=False)
        
        assert ax is not None
        plt.close()
    
    def test_custom_title(self, sample_barcode):
        """Test plot with custom title."""
        ax = plot_persistence_diagram(sample_barcode, title='Custom Title')
        
        assert ax.get_title() == 'Custom Title'
        plt.close()
    
    def test_custom_axes(self, sample_barcode):
        """Test plotting on custom axes."""
        fig, ax = plt.subplots()
        result_ax = plot_persistence_diagram(sample_barcode, ax=ax)
        
        assert result_ax is ax
        plt.close()
    
    def test_empty_barcode(self, empty_barcode):
        """Test plotting empty barcode doesn't crash."""
        ax = plot_persistence_diagram(empty_barcode)
        
        assert ax is not None
        plt.close()
    
    def test_barcode_with_infinities(self):
        """Test handling of infinite death times."""
        barcode = np.array([
            [0.0, 0.5],
            [0.1, np.inf],  # Infinite death
            [0.2, 1.2],
        ])
        
        ax = plot_persistence_diagram(barcode)
        
        assert ax is not None
        plt.close()


class TestPlotBarcode:
    """Tests for plot_barcode function."""
    
    def test_plot_single_barcode(self, sample_barcode):
        """Test barcode plot with single array."""
        ax = plot_barcode(sample_barcode)
        
        assert ax is not None
        assert isinstance(ax, plt.Axes)
        plt.close()
    
    def test_plot_dict_barcodes(self, sample_barcodes_dict):
        """Test barcode plot with dictionary."""
        ax = plot_barcode(sample_barcodes_dict)
        
        assert ax is not None
        plt.close()
    
    def test_max_bars_limit(self, sample_barcodes_dict):
        """Test limiting number of bars."""
        ax = plot_barcode(sample_barcodes_dict, max_bars=2)
        
        assert ax is not None
        plt.close()
    
    def test_with_linewidth(self, sample_barcode):
        """Test custom linewidth."""
        ax = plot_barcode(sample_barcode, linewidth=2.0)
        
        assert ax is not None
        plt.close()
    
    def test_with_title(self, sample_barcode):
        """Test custom title."""
        ax = plot_barcode(sample_barcode, title='My Barcode')
        
        assert ax is not None
        plt.close()
    
    def test_with_dimensions(self, sample_barcodes_dict):
        """Test plotting specific dimensions."""
        ax = plot_barcode(sample_barcodes_dict, dimensions=[0, 1])
        
        assert ax is not None
        plt.close()
    
    def test_custom_color_palette(self, sample_barcodes_dict):
        """Test custom color palette."""
        color_palette = ['red', 'blue', 'green']
        ax = plot_barcode(sample_barcodes_dict, color_palette=color_palette)
        
        assert ax is not None
        plt.close()
    
    def test_no_legend(self, sample_barcodes_dict):
        """Test plot without legend."""
        ax = plot_barcode(sample_barcodes_dict, legend=False)
        
        assert ax is not None
        plt.close()
    
    def test_custom_axes(self, sample_barcode):
        """Test plotting on custom axes."""
        fig, ax = plt.subplots()
        result_ax = plot_barcode(sample_barcode, ax=ax)
        
        assert result_ax is ax
        plt.close()
    
    def test_empty_barcode(self, empty_barcode):
        """Test plotting empty barcode."""
        ax = plot_barcode(empty_barcode)
        
        assert ax is not None
        plt.close()


class TestPlotBettiCurve:
    """Tests for plot_betti_curve function."""
    
    def test_plot_single_barcode(self, sample_barcode):
        """Test Betti curve with single barcode."""
        ax = plot_betti_curve(sample_barcode)
        
        assert ax is not None
        assert isinstance(ax, plt.Axes)
        plt.close()
    
    def test_plot_dict_barcodes(self, sample_barcodes_dict):
        """Test Betti curve with dictionary."""
        ax = plot_betti_curve(sample_barcodes_dict)
        
        assert ax is not None
        plt.close()
    
    def test_specific_dimensions(self, sample_barcodes_dict):
        """Test plotting specific dimensions."""
        ax = plot_betti_curve(sample_barcodes_dict, dimensions=[0, 1])
        
        assert ax is not None
        plt.close()
    
    def test_custom_resolution(self, sample_barcode):
        """Test custom resolution."""
        ax = plot_betti_curve(sample_barcode, resolution=50)
        
        assert ax is not None
        plt.close()
    
    def test_no_fill(self, sample_barcode):
        """Test without area fill."""
        ax = plot_betti_curve(sample_barcode, fill=False)
        
        assert ax is not None
        plt.close()
    
    def test_custom_color_palette(self, sample_barcodes_dict):
        """Test custom color palette."""
        color_palette = ['red', 'blue', 'green']
        ax = plot_betti_curve(sample_barcodes_dict, color_palette=color_palette)
        
        assert ax is not None
        plt.close()
    
    def test_no_legend(self, sample_barcodes_dict):
        """Test plot without legend."""
        ax = plot_betti_curve(sample_barcodes_dict, legend=False)
        
        assert ax is not None
        plt.close()
    
    def test_custom_axes(self, sample_barcode):
        """Test plotting on custom axes."""
        fig, ax = plt.subplots()
        result_ax = plot_betti_curve(sample_barcode, ax=ax)
        
        assert result_ax is ax
        plt.close()
    
    def test_empty_barcode(self, empty_barcode):
        """Test plotting empty barcode."""
        ax = plot_betti_curve(empty_barcode)
        
        assert ax is not None
        plt.close()


class TestPlotLandscape:
    """Tests for plot_landscape function."""
    
    def test_plot_from_array(self, sample_barcode):
        """Test landscape plot from array."""
        ax = plot_landscape(sample_barcode, dimension=0)
        
        assert ax is not None
        assert isinstance(ax, plt.Axes)
        plt.close()
    
    def test_plot_from_dict(self, sample_barcodes_dict):
        """Test landscape plot from dictionary."""
        ax = plot_landscape(sample_barcodes_dict, dimension=1)
        
        assert ax is not None
        plt.close()
    
    def test_custom_num_landscapes(self, sample_barcode):
        """Test custom number of landscapes."""
        ax = plot_landscape(sample_barcode, dimension=0, num_landscapes=3)
        
        assert ax is not None
        plt.close()
    
    def test_custom_resolution(self, sample_barcode):
        """Test custom resolution."""
        ax = plot_landscape(sample_barcode, dimension=0, resolution=50)
        
        assert ax is not None
        plt.close()
    
    def test_custom_color_palette(self, sample_barcode):
        """Test custom color palette."""
        from matplotlib.colors import to_rgba
        color_palette = [to_rgba('blue'), to_rgba('orange'), to_rgba('green'),
                         to_rgba('red'), to_rgba('purple')]
        ax = plot_landscape(sample_barcode, dimension=0, color_palette=color_palette)
        
        assert ax is not None
        plt.close()
    
    def test_no_legend(self, sample_barcode):
        """Test plot without legend."""
        ax = plot_landscape(sample_barcode, dimension=0, legend=False)
        
        assert ax is not None
        plt.close()
    
    def test_custom_axes(self, sample_barcode):
        """Test plotting on custom axes."""
        fig, ax = plt.subplots()
        result_ax = plot_landscape(sample_barcode, dimension=0, ax=ax)
        
        assert result_ax is ax
        plt.close()
    
    def test_missing_dimension(self, sample_barcodes_dict):
        """Test error when dimension not found."""
        with pytest.raises(ValueError, match="Dimension .* not found"):
            plot_landscape(sample_barcodes_dict, dimension=5)
    
    def test_empty_barcode(self):
        """Test plotting empty barcode."""
        empty = {'H0': np.array([]).reshape(0, 2)}
        ax = plot_landscape(empty, dimension=0)
        
        assert ax is not None
        plt.close()
    
    def test_barcode_with_infinities(self):
        """Test handling of infinite death times."""
        barcode = {'H1': np.array([
            [0.0, 0.5],
            [0.1, np.inf],  # Infinite death
            [0.2, 1.2],
        ])}
        
        ax = plot_landscape(barcode, dimension=1)
        
        assert ax is not None
        plt.close()


class TestPlotEntropySummary:
    """Tests for plot_entropy_summary function."""
    
    def test_basic_plot(self, sample_barcode):
        """Test basic entropy summary plot."""
        ax = plot_entropy_summary(sample_barcode)
        
        assert ax is not None
        plt.close()
    
    def test_with_dict(self, sample_barcodes_dict):
        """Test with barcodes dictionary."""
        ax = plot_entropy_summary(sample_barcodes_dict)
        
        assert ax is not None
        plt.close()
    
    def test_custom_resolution(self, sample_barcode):
        """Test custom resolution."""
        ax = plot_entropy_summary(sample_barcode, resolution=50)
        
        assert ax is not None
        plt.close()
    
    def test_no_fill(self, sample_barcode):
        """Test without area fill."""
        ax = plot_entropy_summary(sample_barcode, fill=False)
        
        assert ax is not None
        plt.close()


class TestPlotLifespan:
    """Tests for plot_lifespan function."""
    
    def test_basic_plot(self, sample_barcode):
        """Test basic lifespan plot."""
        ax = plot_lifespan(sample_barcode)
        
        assert ax is not None
        plt.close()
    
    def test_with_dict(self, sample_barcodes_dict):
        """Test with barcodes dictionary."""
        ax = plot_lifespan(sample_barcodes_dict)
        
        assert ax is not None
        plt.close()
    
    def test_custom_resolution(self, sample_barcode):
        """Test custom resolution."""
        ax = plot_lifespan(sample_barcode, resolution=50)
        
        assert ax is not None
        plt.close()


class TestPlotSilhouette:
    """Tests for plot_silhouette function."""
    
    def test_basic_plot(self, sample_barcode):
        """Test basic silhouette plot."""
        ax = plot_silhouette(sample_barcode)
        
        assert ax is not None
        plt.close()
    
    def test_with_dict(self, sample_barcodes_dict):
        """Test with barcodes dictionary."""
        ax = plot_silhouette(sample_barcodes_dict)
        
        assert ax is not None
        plt.close()
    
    def test_custom_weight(self, sample_barcode):
        """Test custom weight parameter."""
        ax = plot_silhouette(sample_barcode, weight=2.0)
        
        assert ax is not None
        plt.close()


class TestPlotTropicalCoordinates:
    """Tests for plot_tropical_coordinates function."""
    
    def test_basic_plot(self, sample_barcode):
        """Test basic tropical coordinates plot."""
        ax = plot_tropical_coordinates(sample_barcode)
        
        assert ax is not None
        plt.close()
    
    def test_with_dict(self, sample_barcodes_dict):
        """Test with barcodes dictionary."""
        ax = plot_tropical_coordinates(sample_barcodes_dict)
        
        assert ax is not None
        plt.close()
    
    def test_no_log_scale(self, sample_barcode):
        """Test without logarithmic scale."""
        ax = plot_tropical_coordinates(sample_barcode, log_scale=False)
        
        assert ax is not None
        plt.close()
    
    def test_custom_r_parameter(self, sample_barcode):
        """Test custom r parameter."""
        ax = plot_tropical_coordinates(sample_barcode, r=50)
        
        assert ax is not None
        plt.close()


class TestPlottingEdgeCases:
    """Test edge cases and error handling."""
    
    def test_all_infinite_deaths(self):
        """Test barcode with all infinite death times."""
        barcode = np.array([
            [0.0, np.inf],
            [0.1, np.inf],
        ])
        
        # Should not crash
        ax = plot_persistence_diagram(barcode)
        assert ax is not None
        plt.close()
        
        ax = plot_barcode(barcode)
        assert ax is not None
        plt.close()
    
    def test_single_point(self):
        """Test barcode with single point."""
        barcode = np.array([[0.1, 0.5]])
        
        ax = plot_persistence_diagram(barcode)
        assert ax is not None
        plt.close()
        
        ax = plot_barcode(barcode)
        assert ax is not None
        plt.close()
    
    def test_very_large_barcode(self):
        """Test with large number of features."""
        large_barcode = np.random.rand(1000, 2)
        large_barcode[:, 1] += large_barcode[:, 0]  # Ensure death > birth
        
        ax = plot_barcode({'H0': large_barcode}, max_bars=100)
        assert ax is not None
        plt.close()
    
    def test_negative_values(self):
        """Test barcode with negative values."""
        barcode = np.array([
            [-1.0, -0.5],
            [-0.5, 0.0],
            [0.0, 0.5],
        ])
        
        ax = plot_persistence_diagram(barcode)
        assert ax is not None
        plt.close()
