"""
Splitting utilities

Temporal split preserves chronological order and helps avoid leakage
that random splitting often introduces in time-series problems.
"""

from typing import Tuple, Optional
import pandas as pd
import numpy as np

from ..utils.validation import (
    validate_dataframe,
    validate_time_column,
    validate_ratio,
    validate_timedelta_string,
    validate_columns_exist,
)

def temporal_split(
    df: pd.DataFrame,
    time_column: str,
    train_ratio: float = 0.7,
    gap: str = "0D",
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Split a DataFrame into train/test by time order with an optional time-based gap.

    Purpose
    -------
    Prevents temporal leakage by ensuring training data precedes test data.
    The split is based on the actual number of samples, not dates.

    This means train_ratio=0.7 will use 70% of total rows for training,
    regardless of how the data is distributed across dates.

    Workflow
    --------
    1. Sort entire dataset by time_column (ascending)
    2. Take first `train_ratio` portion as training set
    3. Apply optional time gap (discard samples in gap period)
    4. Assign remaining samples to test set

    Parameters
    ----------
    df : pd.DataFrame
        Input dataset.
    time_column : str
        Name of the timestamp column.
    train_ratio : float, default=0.7
        Fraction of samples for training (0 < ratio < 1).
    gap : str, default="0D"
        Time gap between train and test as pandas Timedelta string.
        Examples: "0D"(no gap), "1D"(1 day), "1h"(1 hour), "1min"(1 minute) .

    Returns
    -------
    train_df : pd.DataFrame
        Training set (earliest samples).
    test_df : pd.DataFrame
        Test set (latest samples, after gap if specified).

    Raises
    ------
    TypeError
        If df is not a DataFrame or gap is not a string.
    ValueError
        If time_column doesn't exist, train_ratio is invalid, or gap format is incorrect.

    Examples
    --------
    >>> import pandas as pd
    >>> df = pd.DataFrame({
    ...     'date': pd.date_range('2026-01-01', periods=100),
    ...     'price': range(100)
    ... })
    >>> train, test = temporal_split(df, 'date', train_ratio=0.7, gap='0D')
    >>> len(train), len(test)
    (70, 30)

    >>> # With 7-day gap
    >>> train, test = temporal_split(df, 'date', train_ratio=0.7, gap='7D')
    >>> # Test set starts 7 days after last training sample
    """

    # Validate inputs
    validate_dataframe(df, name="df")
    validate_time_column(df, time_column, name="df")
    validate_ratio(train_ratio, name="train_ratio", inclusive=False)
    gap_td = validate_timedelta_string(gap)

    sorted_df = df.copy()
    sorted_df[time_column] = pd.to_datetime(sorted_df[time_column])
    sorted_df = sorted_df.sort_values(by=time_column).reset_index(drop=True)

    # Calculate split point (row index, not date)
    total_rows = len(sorted_df)
    split_idx = int(total_rows * train_ratio) # row 기준 분할

    # Simple split by index
    train_df = sorted_df.iloc[:split_idx].copy()
    test_candidates = sorted_df.iloc[split_idx:].copy()

    # Apply gap if specified
    if gap_td > pd.Timedelta(0) and not train_df.empty:
        # Find the last timestamp in training set
        train_end_time = train_df[time_column].max()

        # Calculate cutoff: last train time + gap duration
        gap_cutoff = train_end_time + gap_td

        # Keep only test samples after the gap
        test_df = test_candidates[
            test_candidates[time_column] > gap_cutoff
            ].copy()
    else:
        # No gap - use all test candidates
        test_df = test_candidates.copy()

    # Reset indices for clean output
    train_df = train_df.reset_index(drop=True)
    test_df = test_df.reset_index(drop=True)

    return train_df, test_df

def group_split(
    df: pd.DataFrame,
    group_column: str,
    train_ratio: float = 0.7,
    random_state: Optional[int] = None,
    time_column: Optional[str] = None,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Split dataset by group to prevent group-based leakage.

    Purpose
    -------
    Ensures that all samples belonging to the same group
    appear exclusively in either train or test set. This prevents data leakage when groups represent
    entities that should be kept separate

    Workflow
    --------
    1. Extract all unique group identifiers
    2. Shuffle groups randomly (if random_state provided)
    3. Split groups into train/test based on ratio
    4. Assign all rows of each group to their designated set

    Parameters
    ----------
    df : pd.DataFrame
        Input dataset.
    group_column : str
        Column representing group identity (e.g., patient_id).
    train_ratio : float, default=0.7
        Fraction of groups assigned to training set (0 < ratio < 1).
    random_state : Optional[int]
        Seed for reproducibility.
    time_column : Optional[str]
    If provided, ensures monotonic time order within each group.

    Returns
    -------
    train_df : pd.DataFrame
        Training set (all rows from train groups).
    test_df : pd.DataFrame
        Test set (all rows from test groups).

    Raises
    ------
    TypeError
        If df is not a DataFrame.
    ValueError
        If group_column is missing or ratio invalid.

    Examples
    --------
    >>> df = pd.DataFrame({
    ...     'device_id': ['A', 'A', 'B', 'B', 'C', 'C'],
    ...     'value': [1, 2, 3, 4, 5, 6]
    ... })
    >>> train, test = group_split(df, 'device_id', train_ratio=0.67, random_state=42, time_column="date")
    >>> # Devices A and B in train, device C in test (or similar)
    >>> set(train['device_id']).intersection(set(test['device_id']))
    set()  # Empty - no overlap!
    """

    # 1. Validation
    validate_dataframe(df, name="df")
    validate_columns_exist(df, [group_column], name="df")
    validate_ratio(train_ratio, name="train_ratio", inclusive=False)

    # 2. Extract unique groups
    unique_groups = df[group_column].dropna().unique()

    if len(unique_groups) == 0:
        raise ValueError("No valid groups found in group_column.")

    # 3. Shuffle groups
    unique_groups = np.array(unique_groups)
    rng = np.random.default_rng(random_state)
    rng.shuffle(unique_groups)

    # 4. Split groups
    n_train = int(len(unique_groups) * train_ratio)
    train_groups = unique_groups[:n_train]
    test_groups = unique_groups[n_train:]

    # 5. Assign rows
    train_df = df[df[group_column].isin(train_groups)].reset_index(drop=True)
    test_df = df[df[group_column].isin(test_groups)].reset_index(drop=True)

    if time_column is not None:
        train_df[time_column] = pd.to_datetime(train_df[time_column])
        test_df[time_column] = pd.to_datetime(test_df[time_column])

        train_df = train_df.sort_values(
            by=[group_column, time_column]
        )
        test_df = test_df.sort_values(
            by=[group_column, time_column]
        )

    train_df = train_df.reset_index(drop=True)
    test_df = test_df.reset_index(drop=True)
    return train_df, test_df

def temporal_train_val_test_split(
    df: pd.DataFrame,
    time_column: str,
    train_ratio: float = 0.5,
    val_ratio: float = 0.25,
    gap: str = "0D",
) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """
    Perform temporal train/validation/test split.

    Purpose
    -------
    Extension of temporal_split for model development workflows
    that require a validation set for hyperparameter tuning.

    Parameters
    ----------
    df : pd.DataFrame
        Input dataset.
    time_column : str
        Timestamp column.
    train_ratio : float, default=0.5
        Fraction assigned to training set (0 < ratio < 1).
    val_ratio : float, default=0.25
        Fraction assigned to validation set (0 < ratio < 1).
    gap : str, default="0D"
        Optional time gap between train and validation.

    Returns
    -------
    train_df : pd.DataFrame
        Training set (earliest).
    val_df : pd.DataFrame
        Validation set (middle).
    test_df : pd.DataFrame
        Test set (latest).

    Raises
    ------
    ValueError
        If train_ratio + val_ratio >= 1.0.

    Examples
    --------
    >>> df = pd.DataFrame({
    ...     'date': pd.date_range('2024-01-01', periods=100),
    ...     'value': range(100)
    ... })
    >>> train, val, test = temporal_train_val_test_split(
    ...     df, 'date', train_ratio=0.5, val_ratio=0.25, gap='0D'
    ... )
    >>> len(train), len(val), len(test)
    (60, 20, 20)
    """

    # 1. Validation
    validate_dataframe(df, name="df")
    validate_time_column(df, time_column, name="df")
    validate_ratio(train_ratio, name="train_ratio", inclusive=False)
    validate_ratio(val_ratio, name="val_ratio", inclusive=False)

    if train_ratio + val_ratio >= 1.0:
        raise ValueError(
            "train_ratio + val_ratio must be less than 1."
        )

    gap_td = validate_timedelta_string(gap)

    # 2. Sort by time
    sorted_df = df.copy()
    sorted_df[time_column] = pd.to_datetime(sorted_df[time_column])
    sorted_df = sorted_df.sort_values(by=time_column).reset_index(drop=True)

    total_samples = len(sorted_df)

    # Row-based initial train/validation/test split
    train_end = int(total_samples * train_ratio)
    val_end = int(total_samples * (train_ratio + val_ratio))

    train_df = sorted_df.iloc[:train_end].reset_index(drop=True)
    val_candidates = sorted_df.iloc[train_end:val_end].reset_index(drop=True)
    test_df = sorted_df.iloc[val_end:].reset_index(drop=True)

    # 3. Apply gap between train and validation
    if gap_td > pd.Timedelta(0) and not train_df.empty:
        train_end_time = train_df[time_column].max()
        gap_cutoff = train_end_time + gap_td

        val_df = val_candidates[
            val_candidates[time_column] > gap_cutoff
        ].reset_index(drop=True)
    else:
        val_df = val_candidates

    return train_df, val_df, test_df

def group_train_val_test_split(
    df: pd.DataFrame,
    group_column: str,
    train_ratio: float = 0.5,
    val_ratio: float = 0.25,
    random_state: Optional[int] = None,
    time_column: Optional[str] = None,
) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """
    Perform group-based train/validation/test split.

    Purpose
    -------
    Extension of group_split for workflows requiring validation set.
    Ensures no group appears in multiple sets.

    Parameters
    ----------
    df : pd.DataFrame
        Input dataset.
    group_column : str
        Column defining group identity.
    train_ratio : float, default=0.5
        Fraction of groups assigned to training.
    val_ratio : float, default=0.25
        Fraction of groups assigned to validation.
    random_state : Optional[int], default=None
        Random seed for reproducibility.
    time_column : Optional[str]
        If provided, ensures monotonic time order within each group.

    Returns
    -------
    train_df : pd.DataFrame
        Training set.
    val_df : pd.DataFrame
        Validation set.
    test_df : pd.DataFrame
        Test set.

    Raises
    ------
    ValueError
        If train_ratio + val_ratio >= 1.0 or no groups found.

    Examples
    --------
    >>> df = pd.DataFrame({
    ...     'group': ['A']*10 + ['B']*10 + ['C']*10,
    ...     'value': range(30)
    ... })
    >>> train, val, test = group_train_val_test_split(
    ...     df, 'group', train_ratio=0.5, val_ratio=0.25, random_state=42, time_column="date"
    ... )
    >>> # Groups completely separated across sets
    """

    # Validation
    validate_dataframe(df, name="df")
    validate_columns_exist(df, [group_column], name="df")
    validate_ratio(train_ratio, name="train_ratio", inclusive=False)
    validate_ratio(val_ratio, name="val_ratio", inclusive=False)

    if train_ratio + val_ratio >= 1.0:
        raise ValueError(
            "train_ratio + val_ratio must be less than 1."
        )

    # Get unique groups
    unique_groups = df[group_column].dropna().unique()

    if len(unique_groups) == 0:
        raise ValueError("No valid groups found in group_column.")

    # Shuffle groups
    unique_groups = np.array(unique_groups)
    rng = np.random.default_rng(random_state)
    rng.shuffle(unique_groups)

    # Calculate split points for groups
    total_groups = len(unique_groups)
    train_end_idx = int(total_groups * train_ratio)
    val_end_idx = int(total_groups * (train_ratio + val_ratio))

    # Assign groups to each set
    train_groups = unique_groups[:train_end_idx]
    val_groups = unique_groups[train_end_idx:val_end_idx]
    test_groups = unique_groups[val_end_idx:]

    # Filter rows by group membership
    train_df = df[df[group_column].isin(train_groups)].copy()
    val_df = df[df[group_column].isin(val_groups)].copy()
    test_df = df[df[group_column].isin(test_groups)].copy()

    # Optional time sorting within groups
    if time_column is not None:
        train_df[time_column] = pd.to_datetime(train_df[time_column])
        val_df[time_column] = pd.to_datetime(val_df[time_column])
        test_df[time_column] = pd.to_datetime(test_df[time_column])

        train_df = train_df.sort_values(
            by=[group_column, time_column]
        )
        val_df = val_df.sort_values(
            by=[group_column, time_column]
        )
        test_df = test_df.sort_values(
            by=[group_column, time_column]
        )

    # Reset index
    train_df = train_df.reset_index(drop=True)
    val_df = val_df.reset_index(drop=True)
    test_df = test_df.reset_index(drop=True)

    return train_df, val_df, test_df
