"""Tests for preprocessing module."""

import numpy as np
import pandas as pd
import pytest

from mixlift.modeling.preprocessing import (
    add_control_columns,
    aggregate_to_weekly,
    fill_date_gaps,
    pivot_channels,
    prepare_model_data,
)


class TestAggregateToWeekly:
    def test_aggregates_daily_to_weekly(self, sample_canonical_df):
        result = aggregate_to_weekly(sample_canonical_df)
        # Should have fewer rows than daily
        assert len(result) < len(sample_canonical_df)
        # Each week should have entries for each channel
        channels_per_week = result.groupby("date")["channel"].nunique()
        assert (channels_per_week == 3).all()

    def test_preserves_totals(self, sample_canonical_df):
        result = aggregate_to_weekly(sample_canonical_df)
        # Total spend should match
        assert abs(result["spend"].sum() - sample_canonical_df["spend"].sum()) < 0.01


class TestFillDateGaps:
    def test_fills_missing_weeks(self):
        df = pd.DataFrame({
            "date": pd.to_datetime(["2024-01-01", "2024-01-01", "2024-01-15", "2024-01-15"]),
            "channel": ["A", "B", "A", "B"],
            "spend": [100, 200, 150, 250],
            "impressions": [1000, 2000, 1500, 2500],
            "clicks": [50, 100, 75, 125],
            "conversions": [5, 10, 8, 13],
        })
        result = fill_date_gaps(df)
        # Should have entries for intermediate weeks
        assert len(result) >= len(df)


class TestPivotChannels:
    def test_pivot_spend(self, sample_canonical_df):
        weekly = aggregate_to_weekly(sample_canonical_df)
        filled = fill_date_gaps(weekly)
        pivoted = pivot_channels(filled)

        assert "date" in pivoted.columns
        spend_cols = [c for c in pivoted.columns if c.endswith("_spend")]
        assert len(spend_cols) == 3  # 3 channels


class TestAddControlColumns:
    def test_adds_trend(self):
        df = pd.DataFrame({"date": pd.date_range("2024-01-01", periods=10)})
        result = add_control_columns(df)
        assert "trend" in result.columns
        assert list(result["trend"]) == list(range(10))


class TestPrepareModelData:
    def test_full_pipeline(self, sample_canonical_df):
        X, y, channel_columns = prepare_model_data(sample_canonical_df)

        assert len(X) == len(y)
        assert len(channel_columns) == 3
        assert all(c.endswith("_spend") for c in channel_columns)
        assert "trend" in X.columns
        assert "date" in X.columns

    def test_with_target_series(self, sample_canonical_df):
        # Create a fake target
        weekly = aggregate_to_weekly(sample_canonical_df)
        dates = weekly["date"].unique()
        target = pd.Series(
            np.random.default_rng(42).uniform(10000, 50000, len(dates)),
            index=pd.to_datetime(dates),
        )

        X, y, channel_columns = prepare_model_data(sample_canonical_df, target)
        assert len(X) == len(y)
