"""
Core leakage measurement implementation.

Implements the data leakage quantification formula from:
"Quantifying Data Leakage in Failure Prediction Tasks" (ICPE 2025)
"""

import pandas as pd
import numpy as np

from ..utils.validation import (
    validate_dataframe,
    validate_time_column,
    validate_time_no_nat,
    validate_group_column,
    validate_monotonic_time_per_group,
    validate_alpha,
    validate_non_empty_dataframe,
)

def compute_data_leakage(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
    time_column: str,
    alpha: float = 1.0,
    group_column: str = None,
) -> float:
    """
    Compute normalized data leakage

    Parameters
    ----------
    train_df : pd.DataFrame
        Training dataset with time column.
    test_df : pd.DataFrame
        Test dataset with time column.
    time_column : str
        Name of the column containing timestamps.
        Must exist in both dataframes and be parseable as datetime.
    alpha : float, default=1.0
        Weight for samples where train data precedes test data.

        - alpha=0: Only future leakage counts (temporal-only task)
        - alpha=1: Past and future equally weighted (group-aware task)
        - alpha>1: Emphasize group contamination over temporal leakage
    group_column : str, optional
        Column defining groups (e.g., 'serial_number').
        If provided, only pairs within same group contribute to leakage.
        If None, all data treated as single group.

    Returns
    -------
    float
        Normalized leakage value in range [0, 1].

        - 0.0: No leakage (clean split)
        - 1.0: Maximum leakage (train == test)
        - (0, 1): Partial leakage

    Raises
    ------
    TypeError
        If train_df or test_df are not DataFrames.
    ValueError
        If time_column is missing, unparseable, or alpha is invalid.
    ZeroDivisionError
        If test set is empty (cannot compute normalization).

    Notes
    -----
    >>> # For large datasets, downsample before splitting
    >>> data_sample = data.sample(n=10000, random_state=42)
    >>> train, test = temporal_split(data_sample, ...)
    >>> leakage = compute_data_leakage(train, test, 'date')

    Examples
    --------
    >>> # Small dataset - no sampling
    >>> leakage = compute_data_leakage(train, test, 'date', alpha=1.0)

    >>> # Group-aware measurement
    >>> leakage = compute_data_leakage(
    ...     train, test, 'date',
    ...     alpha=1.0, group_column='serial_number'
    ... )

    """
    # validate inputs
    validate_dataframe(train_df, name="train_df")
    validate_dataframe(test_df, name="test_df")

    validate_time_column(train_df, time_column, name="train_df")
    validate_time_column(test_df, time_column, name="test_df")

    validate_time_no_nat(train_df, time_column, name="train_df")
    validate_time_no_nat(test_df, time_column, name="test_df")

    validate_alpha(alpha)

    validate_non_empty_dataframe(train_df, name="train_df")
    validate_non_empty_dataframe(test_df, name="test_df")

    # processing single-group (내부적으로만 group_id 생성)
    train = train_df.copy()
    test = test_df.copy()

    if group_column is not None:
        validate_group_column(train_df, group_column, name="train_df")
        validate_group_column(test_df, group_column, name="test_df")

        validate_monotonic_time_per_group(train_df, time_column, group_column, name="train_df")
        validate_monotonic_time_per_group(test_df, time_column, group_column, name="test_df")

        # Use actual group IDs
        merge_on = group_column
    else:
        # Single group -> create dummy group ID
        train["_group_id"] = 0
        test["_group_id"] = 0
        merge_on = "_group_id"

    # test × train 모든 pair 생성
    test_train_pairs = test.merge(
        train,
        on=merge_on,
        suffixes=("_test", "_train")
    )

    if len(test_train_pairs) == 0:
        # No same-group pairs (e.g., group-based split)
        return 0.0

    # 시간 차이(day 단위)
    time_diff = (
        test_train_pairs[f"{time_column}_train"]
        - test_train_pairs[f"{time_column}_test"]
    ).dt.days.astype(float)

    # Exponential decay 적용
    leakage_values = np.where(
        time_diff >= 0,
        np.exp(-time_diff),
        alpha * np.exp(time_diff)
    )

    leakage_test_train = leakage_values.sum()

    # Normalization: test–test leakage
    test_test_pairs = test.merge(
        test,
        on=merge_on,
        suffixes=("_t1", "_t2")
    )

    time_diff_tt = (
        test_test_pairs[f"{time_column}_t2"]
        - test_test_pairs[f"{time_column}_t1"]
    ).dt.days.astype(float)

    leakage_tt = np.where(
        time_diff_tt >= 0,
        np.exp(-time_diff_tt),
        alpha * np.exp(time_diff_tt)
    )

    leakage_test_test = leakage_tt.sum()

    # Check for zero normalization
    if leakage_test_test == 0:
        raise ZeroDivisionError(
            "Normalization factor L(test,test) is zero. "
            "This may occur with single-sample test sets."
        )

    # Normalized leakage
    return leakage_test_train / leakage_test_test


