import astropy.units as u
import numpy as np
import pytest
from astropy.table import Table
from astropy.time import Time
from qualpipe.core.descriptor import (
    EventRateCounter,
    PercentageDescriptor,
    RateCounter,
    TrivialDescriptor,
)
from traitlets.config import Config


@pytest.fixture()
def sample_table():
    """Create a sample table with time, event_id, and data columns."""
    return Table(
        {
            "time": Time(
                [
                    "2020-01-01T00:00:00",
                    "2020-01-01T00:00:01",
                    "2020-01-01T00:00:02",
                    "2020-01-01T00:00:03",
                    "2020-01-01T00:00:04",
                ]
            ),
            "event_id": np.array([0, 1, 2, 3, 4]),
            "data": np.array([0, 1, 2, 3, 4], dtype=float),
        }
    )


@pytest.mark.verifies_usecase("UC-140-1.2.1")
def test_trivial_descriptor_dl0_monitoring(sample_table):
    """Test TrivialDescriptor with DL0_MONITORING data (raw monitoring with time column)."""
    descriptor = TrivialDescriptor()

    # Test with DL0_MONITORING format (time + data)
    result = descriptor(sample_table, "data")

    # Check output columns
    assert "time_start" in result.colnames
    assert "time_end" in result.colnames
    assert "data" in result.colnames

    # Check that time_start and time_end equal the input time
    assert np.array_equal(result["time_start"], sample_table["time"])
    assert np.array_equal(result["time_end"], sample_table["time"])

    # Check that data is preserved
    assert np.array_equal(result["data"], sample_table["data"])


@pytest.mark.verifies_usecase("UC-140-1.2.1")
def test_trivial_descriptor_dl1_monitoring():
    """Test TrivialDescriptor with DL1_MONITORING data (pre-aggregated with time windows)."""
    from astropy.table import Table
    from astropy.time import Time

    # Create pre-aggregated monitoring data
    pre_aggregated = Table(
        {
            "time_start": Time(
                ["2020-01-01T00:00:00", "2020-01-01T00:01:00"],
                format="isot",
                scale="tai",
            ),
            "time_end": Time(
                ["2020-01-01T00:01:00", "2020-01-01T00:02:00"],
                format="isot",
                scale="tai",
            ),
            "event_id_start": np.array([0, 100]),
            "event_id_end": np.array([99, 199]),
            "temperature": np.array([20.5, 21.0]),
        }
    )

    descriptor = TrivialDescriptor()
    result = descriptor(pre_aggregated, "temperature")

    # Should return the table as-is
    assert result.colnames == pre_aggregated.colnames
    assert np.array_equal(result["temperature"], pre_aggregated["temperature"])
    assert np.array_equal(result["event_id_start"], pre_aggregated["event_id_start"])


@pytest.mark.verifies_usecase("UC-140-1.2.1")
def test_percentage_descriptor():
    """Test that PercentageDescriptor computes percentage of non-zero values."""
    table = Table(
        {
            "time": Time(
                [
                    "2020-01-01T00:00:00",
                    "2020-01-01T00:00:01",
                    "2020-01-01T00:00:02",
                    "2020-01-01T00:00:03",
                    "2020-01-01T00:00:04",
                ]
            ),
            "event_id": np.array([0, 1, 2, 3, 4]),
            "data": np.array([0, 0, 0, 1, 0], dtype=float),  # 20% non-zero
        }
    )

    # Configure for single chunk using SizeChunking with chunk_size=None
    config = Config()
    config.PercentageDescriptor.chunking_type = "SizeChunking"
    config.SizeChunking.chunk_size = None  # Use entire table as one chunk

    descriptor = PercentageDescriptor(config=config)
    result = descriptor(table, col_name="data")

    # Check that percentage column is present
    assert "percentage" in result.colnames

    # Check that percentage is calculated correctly (20%)
    assert result["percentage"][0] == pytest.approx(20.0)


@pytest.mark.verifies_usecase("UC-140-1.2.1")
def test_rate_counter_with_chunking():
    """Test RateCounter with time-based chunking."""
    table = Table(
        {
            "time": Time(
                [
                    "2020-01-01T00:00:00",
                    "2020-01-01T00:00:01",
                    "2020-01-01T00:00:02",
                    "2020-01-01T00:00:03",
                    "2020-01-01T00:00:04",
                ]
            ),
            "event_id": np.array([0, 1, 2, 3, 4]),
            "data": np.array([0, 1, 2, 3, 4], dtype=float),
        }
    )

    # Configure RateCounter with 2-second chunks (TimeChunking is automatic)
    config = Config()
    config.TimeChunking.chunk_duration = 2 * u.s

    descriptor = RateCounter(config=config)
    result = descriptor(table, col_name="data")

    # Should have multiple chunks
    assert len(result) > 1
    # Each chunk should have rate column with measurements/second
    assert "rate" in result.colnames
    # Check that rates are reasonable - with 2s chunks and our test data:
    # Chunk 1: 3 events in 2s = 1.5 Hz
    # Chunk 2: 1 event in 2s = 0.5 Hz
    # Chunk 3: 2 events in 2s = 1.0 Hz
    for row in result:
        assert row["rate"] > 0.25
        # Rate should be measurements per second (count / chunk_duration)
        # With 2s chunks, max rate is 1.5 Hz for this test data
        assert row["rate"] <= 1.5  # Reasonable upper bound


@pytest.mark.verifies_usecase("UC-140-1.2.1")
def test_event_rate_counter_with_filtering(caplog):
    """Test EventRateCounter with event type filtering."""
    from ctapipe.containers import EventType

    # Create test data with mixed event types
    table = Table(
        {
            "time": Time(
                [
                    "2020-01-01T00:00:00",
                    "2020-01-01T00:00:01",
                    "2020-01-01T00:00:02",
                    "2020-01-01T00:00:03",
                    "2020-01-01T00:00:04",
                    "2020-01-01T00:00:05",
                ]
            ),
            "event_id": np.array([0, 1, 2, 3, 4, 5]),
            "event_type": np.array(
                [
                    EventType.SUBARRAY.value,  # 0
                    EventType.SUBARRAY.value,  # 1
                    EventType.PEDESTAL.value,  # 2 - calibration event
                    EventType.SUBARRAY.value,  # 3
                    EventType.MUON.value,  # 4 - muon event
                    EventType.SUBARRAY.value,  # 5
                ],
                dtype=np.int32,
            ),
        }
    )

    # Configure EventRateCounter for SUBARRAY events with 2-second chunks
    config = Config()
    config.TimeChunking.chunk_duration = 2 * u.s

    descriptor = EventRateCounter(config=config)
    result = descriptor(table, col_name="event_type")

    # Should have multiple chunks
    assert len(result) > 1
    # Each chunk should have rate column with SUBARRAY events/second
    assert "rate" in result.colnames

    # Count total SUBARRAY events: 4 out of 6 events
    total_subarray_events = np.sum(table["event_type"] == EventType.SUBARRAY.value)
    assert total_subarray_events == 4

    # Test filtering for different event type
    config2 = Config()
    config2.TimeChunking.chunk_duration = 2 * u.s
    config2.EventRateCounter.event_type_filter = "MUON"

    descriptor2 = EventRateCounter(config=config2)
    result2 = descriptor2(table, col_name="event_type")

    # Count total MUON events: 1 out of 6 events
    total_muon_rate = sum(row["rate"] for row in result2) * 2  # rate * 2s = count
    assert total_muon_rate == pytest.approx(1.0)  # Should find exactly 1 muon event

    # Test automatic correction with warning
    import logging

    with caplog.at_level(logging.WARNING):
        result_corrected = descriptor(table, col_name="data")
        assert "EventRateCounter requires col_name='event_type'" in caplog.text
        assert "Automatically using 'event_type'" in caplog.text
        # Should still work correctly after auto-correction
        assert len(result_corrected) > 0

    # Test missing event_type column
    bad_table = Table({"time": table["time"], "data": [1] * len(table)})
    with pytest.raises(ValueError, match="must have 'event_type' column"):
        descriptor(bad_table, col_name="event_type")
