"""Clean FRED macro data and build a unified event calendar."""

from __future__ import annotations

from datetime import date

import pandas as pd

from src.config import (
    CPI_RELEASE_DATES,
    DATA_PROCESSED,
    DATA_RAW_FRED,
    ELECTION_KEY_DATES,
    FOMC_DATES,
    NFP_RELEASE_DATES,
)
from src.utils import setup_logging

logger = setup_logging("macro_cleaner")


# ---------------------------------------------------------------------------
# Event calendar
# ---------------------------------------------------------------------------

def build_event_calendar() -> pd.DataFrame:
    """Build a unified event calendar with columns: date, event_type, description."""
    rows: list[dict] = []

    for d in FOMC_DATES:
        rows.append({"date": d, "event_type": "FOMC", "description": f"FOMC meeting {d}"})

    for d in CPI_RELEASE_DATES:
        rows.append({"date": d, "event_type": "CPI", "description": f"CPI release {d}"})

    for d in NFP_RELEASE_DATES:
        rows.append({"date": d, "event_type": "NFP", "description": f"Non-farm payrolls {d}"})

    for label, d in ELECTION_KEY_DATES.items():
        rows.append({"date": d, "event_type": "Election", "description": label})

    df = pd.DataFrame(rows)
    df["date"] = pd.to_datetime(df["date"])
    df = df.sort_values("date").reset_index(drop=True)
    logger.info("Built event calendar with %d events", len(df))
    return df


# ---------------------------------------------------------------------------
# FRED data
# ---------------------------------------------------------------------------

def clean_fred_data() -> dict[str, pd.DataFrame]:
    """Load and clean all FRED CSVs.

    Returns a dict mapping series name to a DataFrame with DatetimeIndex.
    """
    result: dict[str, pd.DataFrame] = {}

    for csv_path in sorted(DATA_RAW_FRED.glob("*.csv")):
        df = pd.read_csv(csv_path, parse_dates=["date"], index_col="date")
        # Coerce all value columns to numeric
        for col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")
        df = df.sort_index()
        name = csv_path.stem  # e.g. "fed_funds_rate", "cpi", "unemployment"
        result[name] = df
        logger.info("Loaded FRED series '%s': %d observations, columns=%s",
                     name, len(df), list(df.columns))

    return result


# ---------------------------------------------------------------------------
# Master cleaning
# ---------------------------------------------------------------------------

def clean_all() -> None:
    """Save macro_events.parquet and fred_data.parquet to ``data/processed/``."""
    DATA_PROCESSED.mkdir(parents=True, exist_ok=True)

    # 1. Event calendar
    calendar = build_event_calendar()
    out = DATA_PROCESSED / "macro_events.parquet"
    calendar.to_parquet(out)
    logger.info("Saved %s (%d rows)", out.name, len(calendar))

    # 2. FRED data -- merge all series into a single wide DataFrame
    fred = clean_fred_data()
    if fred:
        # Outer-join all series on date
        combined = pd.DataFrame()
        for name, df in fred.items():
            for col in df.columns:
                col_name = f"{name}_{col}" if col != name else col
                combined[col_name] = df[col]
        # Re-sort by date
        combined = combined.sort_index()
        # Forward-fill monthly series to daily grid
        combined = combined.ffill()
        out = DATA_PROCESSED / "fred_data.parquet"
        combined.to_parquet(out)
        logger.info("Saved %s (%d rows, %d columns)", out.name, len(combined), len(combined.columns))

    logger.info("=== Macro cleaning complete ===")