"""Preprocessing: transform canonical data into model-ready format."""

from __future__ import annotations

import logging

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


def aggregate_to_weekly(df: pd.DataFrame) -> pd.DataFrame:
    """Aggregate daily data to weekly sums.

    Input: canonical format (date, channel, spend, impressions, clicks, conversions)
    Output: same format, aggregated to ISO weeks
    """
    df = df.copy()
    df["date"] = pd.to_datetime(df["date"])
    # Use Monday as week start
    df["week"] = df["date"].dt.to_period("W-SUN").dt.start_time

    agg = (
        df.groupby(["week", "channel"])
        .agg(
            spend=("spend", "sum"),
            impressions=("impressions", "sum"),
            clicks=("clicks", "sum"),
            conversions=("conversions", "sum"),
        )
        .reset_index()
        .rename(columns={"week": "date"})
    )

    return agg.sort_values(["date", "channel"]).reset_index(drop=True)


def fill_date_gaps(df: pd.DataFrame) -> pd.DataFrame:
    """Fill missing weeks with zero spend for each channel."""
    df = df.copy()
    df["date"] = pd.to_datetime(df["date"])

    all_weeks = sorted(df["date"].unique())
    all_channels = df["channel"].unique()

    full_index = pd.MultiIndex.from_product(
        [all_weeks, all_channels], names=["date", "channel"]
    )

    df = df.set_index(["date", "channel"]).reindex(full_index, fill_value=0).reset_index()
    return df.sort_values(["date", "channel"]).reset_index(drop=True)


def pivot_channels(df: pd.DataFrame, value_col: str = "spend") -> pd.DataFrame:
    """Pivot from long format to wide: one column per channel's spend.

    Returns DataFrame indexed by date with columns like 'google_search_spend', etc.
    """
    pivoted = df.pivot_table(
        index="date",
        columns="channel",
        values=value_col,
        aggfunc="sum",
        fill_value=0,
    )

    # Clean column names: lowercase, underscores, append value_col
    pivoted.columns = [
        f"{c.lower().replace(' ', '_').replace('-', '_')}_{value_col}" for c in pivoted.columns
    ]

    return pivoted.reset_index()


def _build_holiday_column(df: pd.DataFrame) -> pd.Series:
    """Return a float Series (0.0/1.0) indicating whether each row's week contains a US holiday.

    For each date row, the week window is [date, date+6 days].  If any day in
    that window is a recognised US public holiday the value is 1.0.

    Raises ``ImportError`` if the *holidays* package is not installed so the
    caller can degrade gracefully.
    """
    import holidays  # noqa: F811 – intentionally deferred import

    dates = pd.to_datetime(df["date"])
    min_year = int(dates.dt.year.min())
    max_year = int(dates.dt.year.max())
    # Expand by one year to cover week windows that span year boundaries
    us_holidays = holidays.US(years=range(min_year, max_year + 2))

    flags: list[float] = []
    for d in dates:
        hit = False
        for offset in range(7):
            if (d + pd.Timedelta(days=offset)) in us_holidays:
                hit = True
                break
        flags.append(1.0 if hit else 0.0)
    return pd.Series(flags, index=df.index, name="is_holiday_week")


def add_control_columns(df: pd.DataFrame) -> pd.DataFrame:
    """Add trend and any other control columns to the wide-format DataFrame."""
    df = df.copy()
    df["trend"] = np.arange(len(df), dtype=float)

    # Holiday indicator – degrades gracefully if holidays package unavailable
    try:
        df["is_holiday_week"] = _build_holiday_column(df)
    except Exception:  # noqa: BLE001
        logger.warning("Could not build holiday indicator; skipping is_holiday_week column.")
        df["is_holiday_week"] = 0.0

    # Month-of-year cyclical encoding (smooth seasonal patterns without 12 dummy columns)
    month = pd.to_datetime(df["date"]).dt.month
    df["month_sin"] = np.sin(2 * np.pi * month / 12)
    df["month_cos"] = np.cos(2 * np.pi * month / 12)

    # Q4 indicator: sharp DTC retail spike that Fourier terms alone can't capture
    df["is_q4"] = (month >= 10).astype(float)

    return df


def prepare_model_data(
    canonical_df: pd.DataFrame,
    target_series: pd.Series | None = None,
) -> tuple[pd.DataFrame, pd.Series, list[str]]:
    """Full preprocessing pipeline: canonical → model-ready X and y.

    Args:
        canonical_df: Canonical format DataFrame (date, channel, spend, ...)
        target_series: Optional revenue/target series aligned by date.
                      If None, uses conversions summed across channels.

    Returns:
        X: DataFrame with spend columns + controls, indexed by date
        y: Target series
        channel_columns: List of channel spend column names
    """
    # Aggregate to weekly
    weekly = aggregate_to_weekly(canonical_df)

    # Fill gaps
    weekly = fill_date_gaps(weekly)

    # Pivot spend columns
    X = pivot_channels(weekly, value_col="spend")

    # Add controls
    X = add_control_columns(X)

    # Identify channel columns
    channel_columns = [c for c in X.columns if c.endswith("_spend")]

    # Build target
    if target_series is not None:
        target_series = target_series.copy()
        target_series.index = pd.to_datetime(target_series.index)
        y = target_series.reindex(X["date"]).fillna(0)
    else:
        # Sum conversions per week as target
        conv_pivot = pivot_channels(weekly, value_col="conversions")
        conv_cols = [c for c in conv_pivot.columns if c.endswith("_conversions")]
        y = conv_pivot[conv_cols].sum(axis=1)

    # Ensure alignment
    assert len(X) == len(y), f"X has {len(X)} rows but y has {len(y)}"

    return X, y, channel_columns
