"""
Overlap detection utilities.

Detects identifier overlap between train and test sets to identify
potential group-based data leakage.
"""

from dataclasses import dataclass
from typing import List, Sequence

import pandas as pd

from ..utils.validation import (
    validate_dataframe,
    validate_non_empty_sequence,
    validate_columns_exist,
)

# train/test 데이터에 같은 ID가 얼마나 겹치는지?

@dataclass(frozen=True)
class OverlapResult:
    """
    Result of overlap analysis for a single identifier column.

    Attributes
    ----------
    column : str
        Name of the identifier column.
    unique_train : int
        Number of unique IDs in training set.
    unique_test : int
        Number of unique IDs in test set.
    overlap_count : int
        Number of IDs appearing in both sets.
    overlap_percentage_test : float
        Percentage of test IDs that also appear in training (0-100).
    """

    column: str
    unique_train: int
    unique_test: int
    overlap_count: int
    overlap_percentage_test: float


def compute_id_overlap(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
    id_columns: Sequence[str],
    dropna: bool = True,
) -> pd.DataFrame:
    """
    Compute overlap statistics of identifier columns between train and test splits.

    Parameters
    ----------
    train_df : pd.DataFrame
        Training split.
    test_df : pd.DataFrame
        Test split.
    id_columns : Sequence[str]
        Identifier columns to analyze.
    dropna : bool, default=True
        If True, ignore NaN IDs.

    Returns
    -------
    pd.DataFrame
        Columns:
        - column
        - unique_train
        - unique_test
        - overlap_count
        - overlap_percentage_test
        Sorted by overlap_percentage_test (desc).
    """
    validate_dataframe(train_df, name="train_df")
    validate_dataframe(test_df, name="test_df")
    validate_non_empty_sequence(id_columns, name="id_columns")
    validate_columns_exist(train_df, id_columns, name="train_df")
    validate_columns_exist(test_df, id_columns, name="test_df")

    results: List[OverlapResult] = []

    for col in id_columns:
        train_ids = train_df[col]
        test_ids = test_df[col]

        # Remove NaN
        if dropna:
            train_ids = train_ids.dropna()
            test_ids = test_ids.dropna()

        # Compute unique sets and overlap
        train_set = set(train_ids.unique().tolist())
        test_set = set(test_ids.unique().tolist())
        overlap_set = train_set & test_set

        unique_train = len(train_set)
        unique_test = len(test_set)
        overlap_count = len(overlap_set)

        # Calculate percentage (avoid division by zero)
        if unique_test == 0:
            overlap_pct = 0.0
        else:
            overlap_pct = (overlap_count / unique_test) * 100.0

        results.append(
            OverlapResult(
                column=col,
                unique_train=unique_train,
                unique_test=unique_test,
                overlap_count=overlap_count,
                overlap_percentage_test=overlap_pct,
            )
        )
    df_results = pd.DataFrame([r.__dict__ for r in results])
    return df_results.sort_values(
        "overlap_percentage_test",
        ascending=False
    ).reset_index(drop=True)
