"""
Data leakage analysis orchestrator.

This module provides a high-level interface for analyzing datasets
for potential data leakage using the ICPE measure and supporting tools.
"""

from dataclasses import dataclass
from typing import Optional, Dict, Any, List
import pandas as pd

from ..utils.validation import (
    validate_dataframe,
    validate_time_column,
    validate_time_no_nat,
)

from .duplicates import (
    find_exact_duplicates,
    find_near_duplicates,
    ExactDuplicateSummary,
    NearDuplicateSummary,
)

from .split import (
    temporal_split,
    temporal_train_val_test_split,
    group_split,
    group_train_val_test_split,
)
from .overlap import compute_id_overlap
from .measure import compute_data_leakage


@dataclass
class LeakageAnalysisResult:
    """
    Container for all leakage analysis results.

    Attributes
    ----------
    leakage_score : float
        Normalized leakage score (L̃ ∈ [0,1])
    exact_duplicate_summary : ExactDuplicateSummary
        Summary statistics for exact duplicate detection,
        including count and percentage of exact duplicate rows.
    near_duplicate_summary : NearDuplicateSummary
        Summary statistics for near-duplicate detection,
        including number of duplicate pairs and affected rows.
    overlap_info : Optional[pd.DataFrame]
        DataFrame containing identifier overlap statistics between
        training and test sets. None if overlap analysis was not performed.
    split_info : Dict[str, Any]
        Information about the dataset split, including:
        - split mode (temporal or group)
        - split type (train/test or train/val/test)
        - sizes of each subset
    metadata : Dict[str, Any]
        Additional contextual metadata such as:
        - total number of rows
        - time column name
        - alpha parameter
        - dataset date range
    """
    leakage_score: float
    exact_duplicate_summary: ExactDuplicateSummary
    near_duplicate_summary: NearDuplicateSummary
    overlap_info: pd.DataFrame
    split_info: Dict[str, Any]
    metadata: Dict[str, Any]


def analyze_dataset(
    df: pd.DataFrame,
    split_mode: str = "temporal",
    split_type: str = "train_test",
    time_column: str = "timestamp",
    group_column: Optional[str] = None,
    id_columns: Optional[List[str]] = None,
    train_ratio: float = 0.7,
    val_ratio: float = 0.15,
    time_gap: Optional[str] = None,
    alpha: float = 1.0,
    similarity_threshold: float = 0.95,
    remove_exact_dups: bool = True,
    remove_near_dups: bool = False,
) -> LeakageAnalysisResult:
    """
    Perform full leakage analysis on a dataset.

    Purpose
    -------
    High-level orchestrator that runs the complete leakage detection pipeline:
    input validation → duplicate detection → splitting → overlap analysis → leakage score.

    Workflow
    --------
    1. Validate inputs
    2. Detect duplicates on the FULL original dataset
    3. Optionally remove duplicates from a working copy
    4. Split the working copy into train/test
    5. Analyze identifier overlap between splits
    6. Compute normalized leakage score

    Parameters
    ----------
    df : pd.DataFrame
        Input dataset. Must contain time_column.
    split_mode : str, default="temporal"
        Splitting strategy: "temporal" or "group".
    split_type : str, default="train_test"
        Number of splits: "train_test" or "train_val_test".
    time_column : str, default="timestamp"
        Name of the timestamp column.
    group_column : Optional[str], default=None
        Required when split_mode="group". Column defining group identity.
    id_columns : Optional[List[str]], default=None
        Columns to analyze for identifier overlap.
        If None, overlap analysis is skipped.
    train_ratio : float, default=0.7
        Fraction of data assigned to training set.
    val_ratio : float, default=0.15
        Fraction of data assigned to validation set (only for train_val_test).
    time_gap : Optional[str], default=None
        Optional time gap between train and test (e.g., "7D" for 7 days).
    alpha : float, default=1.0
        Decay parameter for the leakage measure.
        Higher alpha = stronger penalty for time-distant duplicates.
    similarity_threshold : float, default=0.95
        Cosine similarity threshold for near-duplicate detection (0.0–1.0).
    remove_exact_dups : bool, default=True
        If True, remove exact duplicate rows before splitting.
    remove_near_dups : bool, default=False
        If True, remove near-duplicate rows before splitting.
        Use with caution: may remove legitimate similar records.

    Returns
    -------
    LeakageAnalysisResult
        Dataclass containing leakage score, duplicate summaries,
        overlap info, split info, and metadata.

    Raises
    ------
    TypeError
        If df is not a DataFrame.
    ValueError
        If split_mode/split_type is invalid, time_column is missing,
        or group_column is missing when split_mode="group".

    Examples
    --------
    >>> result = analyze_dataset(
    ...     df,
    ...     split_mode="group",
    ...     time_column="date",
    ...     group_column="serial_number",
    ...     id_columns=["serial_number"],
    ...     alpha=1.0,
    ...     similarity_threshold=0.999,
    ... )
    >>> print(result.leakage_score)
    0.000123
    """

    # STEP 1: Validation
    validate_dataframe(df, name="df")
    validate_time_column(df, time_column, name="df")
    validate_time_no_nat(df, time_column, name="df")

    if split_mode not in ["temporal", "group"]:
        raise ValueError("split_mode must be 'temporal' or 'group'")

    if split_type not in ["train_test", "train_val_test"]:
        raise ValueError("split_type must be 'train_test' or 'train_val_test'")

    if split_mode == "group":
        if group_column is None:
            raise ValueError("group_column required for group-based split")
        if group_column not in df.columns:
            raise ValueError(f"group_column '{group_column}' not found in df")

    # STEP 2: Duplicate Detection
    # find_exact_duplicates() runs on the full df (fast, no memory issue)
    _, exact_summary = find_exact_duplicates(
        df=df,
        subset=None,
        keep="first",
    )

    # find_near_duplicates() internally samples large datasets (max_rows=10_000)
    # df_used is the actual dataframe analyzed (sampled or full)
    # We use df_used as df_working to keep things consistent and computationally feasible
    pairs, near_summary, df_used = find_near_duplicates(
        df=df,
        threshold=similarity_threshold,
        subset=None,
    )

    # STEP 3: Build working copy from df_used (sampled), then remove duplicates
    # Note: df_used may be a subset of df if sampling was triggered
    df_working = df_used.copy()

    if remove_exact_dups:
        # Use drop_duplicates() (value-based) instead of drop(index)
        # because exact_df.index refers to original df, not df_working
        df_working = df_working.drop_duplicates(keep="first").reset_index(drop=True)

    if remove_near_dups:
        # Re-run find_near_duplicates on df_working so pair indices
        # match df_working's current index (not the original df's index)
        pairs_working, _, _ = find_near_duplicates(
            df=df_working,
            threshold=similarity_threshold,
            subset=None,
        )
        if pairs_working:
            # keep="first" → remove j (higher index) from each pair
            indices_to_remove = {j for _, j in pairs_working}
            mask = ~df_working.index.isin(indices_to_remove)
            df_working = df_working[mask].reset_index(drop=True)

    # STEP 4: Splitting
    val_df = None

    if split_mode == "temporal":
        if split_type == "train_test":
            train_df, test_df = temporal_split(
                df=df_working,
                time_column=time_column,
                train_ratio=train_ratio,
                gap=time_gap or "0D",
            )
        else:
            train_df, val_df, test_df = temporal_train_val_test_split(
                df=df_working,
                time_column=time_column,
                train_ratio=train_ratio,
                val_ratio=val_ratio,
                gap=time_gap or "0D",
            )

    else:  # group
        if split_type == "train_test":
            train_df, test_df = group_split(
                df=df_working,
                group_column=group_column,
                train_ratio=train_ratio,
                time_column=time_column,  # ensures time order within groups
            )
        else:
            train_df, val_df, test_df = group_train_val_test_split(
                df=df_working,
                group_column=group_column,
                train_ratio=train_ratio,
                val_ratio=val_ratio,
                time_column=time_column,
            )

    split_info = {
        "mode": split_mode,
        "type": split_type,
        "train_size": len(train_df),
        "val_size": len(val_df) if val_df is not None else 0,
        "test_size": len(test_df),
    }

    # STEP 5: Overlap Analysis
    # If no identifier columns provided, create a global dummy ID
    if id_columns is None:
        train_df = train_df.copy()
        test_df = test_df.copy()

        train_df["_global_id"] = 0
        test_df["_global_id"] = 0

        id_columns = ["_global_id"]

    overlap_info = compute_id_overlap(
        train_df=train_df,
        test_df=test_df,
        id_columns=id_columns,
    )

    # STEP 6: Leakage Score
    leakage_score = compute_data_leakage(
        train_df=train_df,
        test_df=test_df,
        time_column=time_column,
        alpha=alpha,
        group_column=group_column if split_mode == "group" else None,
    )

    # STEP 7: Metadata
    metadata = {
        "total_rows": len(df),
        "analyzed_rows": len(df_used),
        "sampling_applied": len(df) > len(df_used),
        "time_column": time_column,
        "alpha": alpha,
        "date_range": f"{df[time_column].min()} to {df[time_column].max()}",
    }

    if split_mode == "group":
        metadata["group_column"] = group_column

    return LeakageAnalysisResult(
        leakage_score=leakage_score,
        exact_duplicate_summary=exact_summary,
        near_duplicate_summary=near_summary,
        overlap_info=overlap_info,
        split_info=split_info,
        metadata=metadata,
    )