from __future__ import annotations

from typing import Iterable
import numpy as np
import pandas as pd


def train_val_test_split(
    df: pd.DataFrame,
    target: str,
    train_size: float = 0.7,
    val_size: float = 0.15,
    random_state: int | None = None,
):
    """Split dataframe into (X_train,y_train), (X_val,y_val), (X_test,y_test)."""
    if target not in df.columns:
        raise ValueError("target column not found")
    if train_size <= 0 or val_size <= 0 or train_size + val_size >= 1:
        raise ValueError("invalid split sizes")

    rng = np.random.default_rng(random_state)
    idx = np.arange(len(df))
    rng.shuffle(idx)

    n = len(df)
    n_train = int(n * train_size)
    n_val = int(n * val_size)

    i_train = idx[:n_train]
    i_val = idx[n_train : n_train + n_val]
    i_test = idx[n_train + n_val :]

    def _xy(i: np.ndarray):
        part = df.iloc[i]
        return (
            part.drop(columns=[target]).reset_index(drop=True),
            part[target].reset_index(drop=True),
        )

    return _xy(i_train), _xy(i_val), _xy(i_test)


def stratified_split(
    df: pd.DataFrame,
    target: str,
    train_size: float = 0.7,
    val_size: float = 0.15,
    random_state: int | None = None,
):
    """Class-stratified train/val/test split without sklearn dependency."""
    if target not in df.columns:
        raise ValueError("target column not found")
    if train_size <= 0 or val_size <= 0 or train_size + val_size >= 1:
        raise ValueError("invalid split sizes")

    rng = np.random.default_rng(random_state)
    train_idx: list[int] = []
    val_idx: list[int] = []
    test_idx: list[int] = []

    for cls, group in df.groupby(target):
        cls_idx = group.index.to_numpy().copy()
        rng.shuffle(cls_idx)
        n = len(cls_idx)
        n_train = int(n * train_size)
        n_val = int(n * val_size)
        train_idx.extend(cls_idx[:n_train].tolist())
        val_idx.extend(cls_idx[n_train : n_train + n_val].tolist())
        test_idx.extend(cls_idx[n_train + n_val :].tolist())

    rng.shuffle(train_idx)
    rng.shuffle(val_idx)
    rng.shuffle(test_idx)

    def _xy(i: list[int]):
        part = df.loc[i]
        return (
            part.drop(columns=[target]).reset_index(drop=True),
            part[target].reset_index(drop=True),
        )

    return _xy(train_idx), _xy(val_idx), _xy(test_idx)


def check_nulls(df: pd.DataFrame) -> pd.Series:
    """Return null count per column."""
    return df.isna().sum()


def validate_schema(df: pd.DataFrame, schema: dict[str, str]) -> tuple[bool, list[str]]:
    """Validate required columns and dtypes."""
    errors: list[str] = []
    for col, dtype in schema.items():
        if col not in df.columns:
            errors.append(f"missing:{col}")
            continue
        if str(df[col].dtype) != dtype:
            errors.append(f"dtype:{col}:{df[col].dtype}!={dtype}")
    return (len(errors) == 0, errors)


def distribution_summary(df: pd.DataFrame, cols: Iterable[str]) -> pd.DataFrame:
    """Return describe() summary for selected columns."""
    return df[list(cols)].describe(include="all").T


def imbalance_report(y: pd.Series) -> pd.DataFrame:
    """Return class count and ratio table."""
    counts = y.value_counts(dropna=False)
    total = counts.sum()
    out = pd.DataFrame({"count": counts, "ratio": counts / total})
    out.index.name = "class"
    return out.reset_index()
