import sys
import pytest
import numpy as np
from pathlib import Path
from matplotlib.ticker import FuncFormatter
from matplotlib.ticker import PercentFormatter
from matplotlib.ticker import StrMethodFormatter
from unittest.mock import Mock, patch


sys.path.append(Path(__file__).parent.resolve()) # type: ignore

from daspi import AxesFacets
from daspi import flat_unique
from daspi.plotlib.facets import LabelFacets


def test_flat_unique_numpy() -> None:
    nested = np.array([[1, 1], [2, 3]])
    result = flat_unique(nested)
    assert result == [1, 2, 3]

    nested = np.array([[1, 1, 1], [2, 2, 3], [2, 2, 3]])
    result = flat_unique(nested)
    assert result == [1, 2, 3]

    nested = np.array([[2, 2, 3], [2, 2, 3], [1, 1, 1]])
    result = flat_unique(nested)
    assert result == [2, 3, 1]

def test_flat_unique_list() -> None:
    nested = [[1, 1], [2, 3]]
    result = flat_unique(nested)
    assert result == [1, 2, 3]

    nested = [[1, 1, 1], [2, 2, 3], [2, 2, 3]]
    result = flat_unique(nested)
    assert result == [1, 2, 3]

    nested = [[2, 2, 3], [2, 2, 3], [1, 1, 1]]
    result = flat_unique(nested)
    assert result == [2, 3, 1]

def test_flat_unique_empty() -> None:
    result = flat_unique([])
    assert result == []


class TestAxesFacets:

    @pytest.fixture
    def rc(self) -> AxesFacets:
        axes = AxesFacets(
            nrows=2, ncols=2, sharex='col', sharey='row', 
            width_ratios=[4, 1], height_ratios=[1, 4])
        return axes
    
    @pytest.fixture
    def s_mosaic(self) -> AxesFacets:
        mosaic="""
            AAA
            BBC
            BBC"""
        axes = AxesFacets(
            mosaic=mosaic, sharex='col', sharey='row', 
            width_ratios=[2, 2, 1], height_ratios=[1, 2, 2])
        return axes
    
    @pytest.fixture
    def l_mosaic(self) -> AxesFacets:
        mosaic=[
            ['Horizontal', 'Horizontal', '.'],
            ['Vertical', 'Qubic big', 'Qubic big'],
            ['Vertical', 'Qubic big', 'Qubic big']]
        axes = AxesFacets(
            mosaic=mosaic, sharex='col', sharey='row', 
            width_ratios=[2, 2, 1], height_ratios=[1, 2, 2])
        return axes
    
    def test_axes_grid_rc(self,rc: AxesFacets) -> None:
        assert rc.shape == (2, 2)
        assert rc.nrows == 2
        assert rc.ncols == 2
        assert len(rc) == 4

    def test_axes_grid_string(self, s_mosaic: AxesFacets) -> None:
        assert s_mosaic.shape == (3, 3)
        assert s_mosaic.nrows == 3
        assert s_mosaic.ncols == 3
        assert len(s_mosaic) == 3
        assert all(s_mosaic[0] == a for a in s_mosaic[0, :])
        assert not any(s_mosaic[0] == a for a in s_mosaic[1:, :].flat)
        assert all(s_mosaic[1] == a for a in s_mosaic[1:, :2].flat)
        assert not any(s_mosaic[1] == a for a in s_mosaic[1:, 2])
        assert all(s_mosaic[2] == a for a in s_mosaic[1:, 2])

    def test_axes_grid_lists(self, l_mosaic: AxesFacets) -> None:
        assert l_mosaic.shape == (3, 3)
        assert l_mosaic.nrows == 3
        assert l_mosaic.ncols == 3
        assert l_mosaic[0, -1] is None
        assert len(l_mosaic) == 3
        assert all(l_mosaic[0] == a for a in l_mosaic[0, :2])
        assert l_mosaic[0] != l_mosaic[0, -1]
        assert l_mosaic[0] not in l_mosaic[1:, :].flat
        assert all(l_mosaic[1] == a for a in l_mosaic[1:, 0])
        assert l_mosaic[1] not in l_mosaic[1:, 1:].flat
        assert all(l_mosaic[2] == a for a in l_mosaic[1:, 1:].flat)

    def test_iteration(self, l_mosaic: AxesFacets) -> None:
        assert l_mosaic.ax is None
        assert next(iter(l_mosaic)) == l_mosaic[0]
        for i, ax in enumerate(l_mosaic):
            assert l_mosaic.ax == ax
            assert ax == l_mosaic[i]


class TestLabelFacets:
    
    @pytest.fixture
    def simple_axes(self) -> AxesFacets:
        return AxesFacets(nrows=1, ncols=2)
    
    def percentage_formatter(self, x, pos):
        """Format tick values as percentages."""
        return f'{x:.1f}%'
    
    def test_label_facets_initialization(self, simple_axes: AxesFacets) -> None:
        """Test that LabelFacets initializes correctly with new parameters."""
        label_facets = LabelFacets(
            simple_axes,
            xlabel='Test X Label',
            ylabel='Test Y Label',
            xlabel_formatter=FuncFormatter(self.percentage_formatter),
            ylabel_formatter=FuncFormatter(self.percentage_formatter),
            xlabel_angle=45,
            ylabel_angle=30
        )
        
        assert label_facets.xlabel == 'Test X Label'
        assert label_facets.ylabel == 'Test Y Label'
        assert label_facets.xlabel_formatter is not None
        assert label_facets.ylabel_formatter is not None
        assert label_facets.xlabel_angle == 45
        assert label_facets.ylabel_angle == 30
    
    def test_label_facets_default_values(self, simple_axes: AxesFacets) -> None:
        """Test that LabelFacets has correct default values for new parameters."""
        label_facets = LabelFacets(simple_axes)
        
        assert label_facets.xlabel_formatter is None
        assert label_facets.ylabel_formatter is None
        assert label_facets.xlabel_angle == 0
        assert label_facets.ylabel_angle == 0
    
    def test_label_facets_draw_with_formatters_and_angles(self, simple_axes: AxesFacets) -> None:
        """Test that formatters and angles are applied correctly when drawing."""
        # Add some data to the axes
        for ax in simple_axes:
            ax.plot([1, 2, 3], [10, 20, 30])
        
        label_facets = LabelFacets(
            simple_axes,
            xlabel_formatter=FuncFormatter(self.percentage_formatter),
            ylabel_formatter=FuncFormatter(self.percentage_formatter),
            xlabel_angle=45,
            ylabel_angle=15
        )
        
        # Draw the labels
        label_facets.draw()
        
        # Verify that angles were applied
        for ax in simple_axes:
            x_tick_labels = ax.get_xticklabels()
            y_tick_labels = ax.get_yticklabels()
            
            if x_tick_labels:
                assert x_tick_labels[0].get_rotation() == 45
            if y_tick_labels:
                assert y_tick_labels[0].get_rotation() == 15
    
    def test_label_facets_alignment_features(self, simple_axes: AxesFacets) -> None:
        """Test that alignment features work correctly."""
        # Add some data to the axes
        for ax in simple_axes:
            ax.plot([1, 2, 3], [10, 20, 30])
        
        label_facets = LabelFacets(
            simple_axes,
            xlabel_align='left',
            ylabel_align='top',
            xlabel_angle=30,
            ylabel_angle=45
        )
        
        # Draw the labels
        label_facets.draw()
        
        # Verify alignment was applied
        for ax in simple_axes:
            x_tick_labels = ax.get_xticklabels()
            y_tick_labels = ax.get_yticklabels()
            
            if x_tick_labels:
                assert x_tick_labels[0].get_horizontalalignment() == 'left'
                assert x_tick_labels[0].get_rotation() == 30
            if y_tick_labels:
                assert y_tick_labels[0].get_verticalalignment() == 'top'
                assert y_tick_labels[0].get_rotation() == 45
    
    def test_label_facets_alignment_with_default_values(self, simple_axes: AxesFacets) -> None:
        """Test alignment with default center values."""
        # Add some data to the axes
        for ax in simple_axes:
            ax.plot([1, 2, 3], [10, 20, 30])
        
        label_facets = LabelFacets(simple_axes)
        
        # Verify default alignment values
        assert label_facets.xlabel_align == 'center'
        assert label_facets.ylabel_align == 'center'
        
        # Draw and verify no changes to alignment when center (default matplotlib behavior)
        label_facets.draw()


class TestLabelFacetsFormatting:
    """Test the new formatting features of LabelFacets."""

    @pytest.fixture
    def axes_with_data(self) -> AxesFacets:
        """Create AxesFacets with some sample data."""
        axes = AxesFacets(nrows=1, ncols=1)
        # Add sample data to create ticks
        # Access the actual axes object from the numpy array
        actual_ax = axes.axes.flat[0]  # This gets the actual matplotlib Axes
        actual_ax.plot([1, 2, 3, 4, 5], [10.5, 20.3, 30.7, 40.2, 50.9])
        return axes

    def test_prepare_formatter_with_none(self, axes_with_data: AxesFacets) -> None:
        """Test _prepare_formatter with None input."""
        label_facets = LabelFacets(axes_with_data)
        result = label_facets._prepare_formatter(None)
        assert result is None

    def test_prepare_formatter_with_existing_formatter(self, axes_with_data: AxesFacets) -> None:
        """Test _prepare_formatter with existing Formatter."""
        existing_formatter = PercentFormatter()
        label_facets = LabelFacets(axes_with_data)
        result = label_facets._prepare_formatter(existing_formatter)
        assert result is existing_formatter

    def test_prepare_formatter_with_single_arg_callable(self, axes_with_data: AxesFacets) -> None:
        """Test _prepare_formatter with single-argument callable."""
        def single_arg_formatter(x):
            return f"{x:.1f}°C"
        
        label_facets = LabelFacets(axes_with_data)
        result = label_facets._prepare_formatter(single_arg_formatter)
        
        assert isinstance(result, FuncFormatter)
        # Test that the wrapped formatter works
        formatted_value = result(25.6, 0)
        assert formatted_value == "25.6°C"

    def test_prepare_formatter_with_two_arg_callable(self, axes_with_data: AxesFacets) -> None:
        """Test _prepare_formatter with two-argument callable."""
        def two_arg_formatter(x, pos):
            return f"#{pos}: {x:.2f}"
        
        label_facets = LabelFacets(axes_with_data)
        result = label_facets._prepare_formatter(two_arg_formatter)
        
        assert isinstance(result, FuncFormatter)
        # Test that the formatter works
        formatted_value = result(25.6, 2)
        assert formatted_value == "#2: 25.60"

    def test_prepare_formatter_with_string_positional(self, axes_with_data: AxesFacets) -> None:
        """Test _prepare_formatter with positional string formats."""
        test_cases = [
            ('{:.2e}', 1234.567, '1.23e+03'),
            ('{:.1f}', 1234.567, '1234.6'),
            ('{:.0f}', 1234.567, '1235'),
            ('${:.2f}', 1234.567, '$1234.57'),
            ('{:.2%}', 0.1234, '12.34%'),
            ('${:.0f}k', 1234.567, '$1235k'),
        ]
        
        label_facets = LabelFacets(axes_with_data)
        
        for format_string, test_value, expected in test_cases:
            result = label_facets._prepare_formatter(format_string)
            assert isinstance(result, FuncFormatter)
            formatted_value = result(test_value, 0)
            assert formatted_value == expected, f"Format '{format_string}' failed: got '{formatted_value}', expected '{expected}'"

    def test_prepare_formatter_with_string_keyword(self, axes_with_data: AxesFacets) -> None:
        """Test _prepare_formatter with keyword string formats."""
        label_facets = LabelFacets(axes_with_data)
        
        # Test keyword format (should use StrMethodFormatter)
        result = label_facets._prepare_formatter('{x:.2e}')
        assert isinstance(result, StrMethodFormatter)
        
        # Test that it works
        formatted_value = result(1234.567, 0)
        assert formatted_value == '1.23e+03'

    def test_formatter_integration_with_axes(self, axes_with_data: AxesFacets) -> None:
        """Test that formatters are properly applied to axes."""
        def temp_formatter(x):
            return f"{x:.1f}°C"
        
        def pressure_formatter(x):
            return f"{x:.2f} bar"
        
        label_facets = LabelFacets(
            axes_with_data,
            xlabel_formatter=pressure_formatter,
            ylabel_formatter=temp_formatter
        )
        label_facets.draw()
        
        # Verify formatters are applied
        actual_ax = axes_with_data.axes.flat[0]
        x_formatter = actual_ax.xaxis.get_major_formatter()
        y_formatter = actual_ax.yaxis.get_major_formatter()
        
        assert isinstance(x_formatter, FuncFormatter)
        assert isinstance(y_formatter, FuncFormatter)

    def test_rotation_angles_applied(self, axes_with_data: AxesFacets) -> None:
        """Test that rotation angles are properly applied."""
        label_facets = LabelFacets(
            axes_with_data,
            xlabel_angle=45,
            ylabel_angle=90
        )
        label_facets.draw()
        
        # Check that rotation is applied to tick labels
        actual_ax = axes_with_data.axes.flat[0]
        x_labels = actual_ax.get_xticklabels()
        y_labels = actual_ax.get_yticklabels()
        
        if x_labels:
            assert x_labels[0].get_rotation() == 45
        if y_labels:
            assert y_labels[0].get_rotation() == 90

    def test_alignment_options_applied(self, axes_with_data: AxesFacets) -> None:
        """Test that alignment options are properly applied."""
        label_facets = LabelFacets(
            axes_with_data,
            xlabel_align='right',
            ylabel_align='top'
        )
        label_facets.draw()
        
        # Check that alignment is applied
        actual_ax = axes_with_data.axes.flat[0]
        x_labels = actual_ax.get_xticklabels()
        y_labels = actual_ax.get_yticklabels()
        
        if x_labels:
            assert x_labels[0].get_horizontalalignment() == 'right'
        if y_labels:
            assert y_labels[0].get_verticalalignment() == 'top'

    def test_rotation_margin_estimation(self, axes_with_data: AxesFacets) -> None:
        """Test the rotation margin estimation."""
        label_facets = LabelFacets(axes_with_data)
        
        # Test with no rotation
        margin_0 = label_facets.estimate_rotation_margin(0)
        assert margin_0 == 0
        
        # Test with small rotation (should be 0)
        margin_small = label_facets.estimate_rotation_margin(4)
        assert margin_small == 0
        
        # Test with 45 degree rotation
        margin_45 = label_facets.estimate_rotation_margin(45, base_margin=20)
        expected_45 = int(20 * np.sin(45 * np.pi / 180))
        assert margin_45 == expected_45
        
        # Test with 90 degree rotation
        margin_90 = label_facets.estimate_rotation_margin(90, base_margin=20)
        expected_90 = int(20 * np.sin(90 * np.pi / 180))
        assert margin_90 == expected_90

    def test_combined_formatting_features(self, axes_with_data: AxesFacets) -> None:
        """Test combining all formatting features together."""
        def custom_formatter(x):
            return f"${x:,.0f}"
        
        label_facets = LabelFacets(
            axes_with_data,
            xlabel='Cost',
            ylabel='Revenue',
            xlabel_formatter=custom_formatter,
            ylabel_formatter=lambda x: f"{x:.1f}%",
            xlabel_angle=30,
            ylabel_angle=15,
            xlabel_align='right',
            ylabel_align='bottom'
        )
        label_facets.draw()
        
        # Verify all features are applied
        actual_ax = axes_with_data.axes.flat[0]
        x_formatter = actual_ax.xaxis.get_major_formatter()
        y_formatter = actual_ax.yaxis.get_major_formatter()
        
        assert isinstance(x_formatter, FuncFormatter)
        assert isinstance(y_formatter, FuncFormatter)
        
        x_labels = actual_ax.get_xticklabels()
        y_labels = actual_ax.get_yticklabels()
        
        if x_labels:
            assert x_labels[0].get_rotation() == 30
            assert x_labels[0].get_horizontalalignment() == 'right'
        if y_labels:
            assert y_labels[0].get_rotation() == 15
            assert y_labels[0].get_verticalalignment() == 'bottom'
