"""Validation utilities for temporal integrity and panel data checks.

Provides guards against look-ahead bias and panel dimension diagnostics.
"""

from __future__ import annotations

import pandas as pd


class TemporalLeakageError(Exception):
    """Raised when training data temporally overlaps with test data."""


def check_no_lookahead(
    train_data: pd.DataFrame,
    test_data: pd.DataFrame,
    date_col: str = "date",
) -> None:
    """Raise TemporalLeakageError if max(train dates) >= min(test dates).

    This is the primary API-boundary guard against look-ahead bias.
    Call this before any fit/predict operation that uses separate
    train and test splits.

    Parameters
    ----------
    train_data : pd.DataFrame
        Training set containing a date column.
    test_data : pd.DataFrame
        Test set containing a date column.
    date_col : str
        Name of the date column in both DataFrames.

    Raises
    ------
    TemporalLeakageError
        If the latest training date is on or after the earliest test date.
    """
    # Extract the boundary dates from each split
    max_train_date = train_data[date_col].max()
    min_test_date = test_data[date_col].min()

    # Strict temporal separation: train must end before test begins
    if max_train_date >= min_test_date:
        raise TemporalLeakageError(
            f"Temporal leakage detected: max train date ({max_train_date}) "
            f">= min test date ({min_test_date}). "
            f"Training data must strictly precede test data."
        )


def check_panel_dims(
    data: pd.DataFrame,
    entity_col: str = "permno",
    date_col: str = "date",
    required_cols: list[str] | None = None,
) -> dict[str, int | float]:
    """Return summary statistics about panel dimensions and completeness.

    Parameters
    ----------
    data : pd.DataFrame
        Long-format panel DataFrame.
    entity_col : str
        Column identifying cross-sectional entities.
    date_col : str
        Column identifying time periods.
    required_cols : list[str] | None
        If provided, raises ValueError when any column is missing.

    Returns
    -------
    dict[str, int | float]
        Keys: n_entities, n_dates, n_obs, pct_missing, unbalanced.
        ``pct_missing`` is the fraction of cells in the full grid that are absent.
        ``unbalanced`` is 1 if the panel is unbalanced, 0 otherwise.

    Raises
    ------
    ValueError
        If required_cols are not all present in the DataFrame.
    """
    # Validate required columns exist
    if required_cols is not None:
        missing = [c for c in required_cols if c not in data.columns]
        if missing:
            raise ValueError(f"Missing required columns: {missing}")

    # Count unique entities and dates
    n_entities = data[entity_col].nunique()
    n_dates = data[date_col].nunique()
    n_obs = len(data)

    # A balanced panel would have exactly n_entities * n_dates observations
    full_grid_size = n_entities * n_dates
    pct_missing = 1.0 - (n_obs / full_grid_size) if full_grid_size > 0 else 0.0

    # Check whether the panel is balanced by comparing per-entity counts
    obs_per_entity = data.groupby(entity_col).size()
    unbalanced = int(obs_per_entity.nunique() > 1)

    return {
        "n_entities": n_entities,
        "n_dates": n_dates,
        "n_obs": n_obs,
        "pct_missing": round(pct_missing, 6),
        "unbalanced": unbalanced,
    }
