from __future__ import annotations

import gc
from dataclasses import dataclass
from typing import Any

import numpy as np
import pandas as pd
import tensorflow as tf
from sklearn.model_selection import KFold
from tensorflow.keras import regularizers
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.models import Sequential

np.random.seed(42)
tf.random.set_seed(42)

try:
    gpus = tf.config.experimental.list_physical_devices("GPU")
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)
except Exception:
    pass


def hard_free() -> None:
    try:
        tf.keras.backend.clear_session()
    except Exception:
        pass
    try:
        gc.collect()
        gc.collect()
    except Exception:
        pass


@dataclass
class TrainingConfig:
    epochs: int = 50
    batch_size: int = 32
    k_splits: int = 5
    n_repeats: int = 1
    seed_bases: tuple[int, int] = (111, 777)
    l1_strength: float = 0.01
    hidden_1: int = 128
    hidden_2: int = 32


def _coerce_label(v: Any) -> Any:
    if pd.isna(v):
        return v
    if isinstance(v, str):
        s = v.strip()
        try:
            f = float(s)
            if f.is_integer():
                return int(f)
            return f
        except Exception:
            return s
    return v


def _normalize_label_list(values: list[Any] | None) -> list[Any] | None:
    if values is None:
        return None
    return [_coerce_label(v) for v in values]


def prepare_bin_data(
    df: pd.DataFrame,
    bin_value: int | float,
    label_column: str = "target",
    bin_column: str = "bin",
    control_labels: list[Any] | None = None,
    experiment_labels: list[Any] | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    fdf = df[df[bin_column] == bin_value].copy()
    if fdf.empty:
        raise ValueError(f"No rows found for bin={bin_value}")

    fdf[label_column] = fdf[label_column].map(_coerce_label)
    control_labels = _normalize_label_list(control_labels)
    experiment_labels = _normalize_label_list(experiment_labels)

    if control_labels is not None or experiment_labels is not None:
        if not control_labels or not experiment_labels:
            raise ValueError(
                "Both control_labels and experiment_labels must be provided "
                "for binary comparison."
            )

        allowed = set(control_labels) | set(experiment_labels)
        fdf = fdf[fdf[label_column].isin(allowed)].copy()
        if fdf.empty:
            raise ValueError(
                "No rows left for "
                f"bin={bin_value} after filtering control/experiment labels."
            )

        overlap = set(control_labels) & set(experiment_labels)
        if overlap:
            raise ValueError(
                "Overlapping control and experiment labels are not allowed: "
                f"{sorted(overlap)}"
            )

        label_map = {lab: 0 for lab in control_labels}
        label_map.update({lab: 1 for lab in experiment_labels})
        fdf[label_column] = fdf[label_column].map(label_map)

    feature_cols = [c for c in fdf.columns if c not in [bin_column, label_column]]
    if not feature_cols:
        raise ValueError("No feature columns found.")

    for col in feature_cols:
        max_val = float(fdf[col].max()) if len(fdf[col]) else 0.0
        fdf[col] = fdf[col] / (max_val + 1.0)

    y = fdf[label_column].astype(int).to_numpy()
    x = np.nan_to_num(fdf[feature_cols].to_numpy(), copy=False).astype(np.float32)

    if x.shape[0] < 2 or x.shape[1] < 1:
        raise ValueError(
            f"Insufficient data for bin={bin_value}. "
            f"samples={x.shape[0]}, dim={x.shape[1]}"
        )

    unique_labels = np.sort(np.unique(y))
    if unique_labels.size < 2:
        raise ValueError(
            "Need at least two classes after filtering "
            f"for bin={bin_value}. Found: {unique_labels.tolist()}"
        )
    if unique_labels[0] != 0 or not np.array_equal(
        unique_labels, np.arange(unique_labels[-1] + 1)
    ):
        raise ValueError(f"Labels must be contiguous 0..C-1. Found: {unique_labels}")

    return x, y


def build_model(
    input_dim: int,
    num_classes: int,
    cfg: TrainingConfig | None = None,
) -> tf.keras.Model:
    cfg = cfg or TrainingConfig()

    model = Sequential(
        [
            Input(shape=(input_dim,)),
            Dense(
                cfg.hidden_1,
                activation="relu",
                kernel_regularizer=regularizers.l1(cfg.l1_strength),
            ),
            Dense(cfg.hidden_2, activation="relu"),
            Dense(num_classes, activation="softmax"),
        ]
    )
    model.compile(
        optimizer=tf.keras.optimizers.Adam(),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
        metrics=["accuracy"],
    )
    return model


def train_kfold_repeats(
    x: np.ndarray,
    y: np.ndarray,
    seed_base: int,
    cfg: TrainingConfig | None = None,
) -> list[tf.keras.Model]:
    cfg = cfg or TrainingConfig()
    kf = KFold(n_splits=cfg.k_splits, shuffle=True, random_state=42)
    all_models: list[tf.keras.Model] = []
    num_classes = int(np.max(y)) + 1

    for fold, (tr, va) in enumerate(kf.split(x, y), start=1):
        x_tr, y_tr = x[tr], y[tr]
        x_va, y_va = x[va], y[va]

        for repeat_idx in range(cfg.n_repeats):
            seed = seed_base * 1000 + fold * 100 + repeat_idx
            tf.keras.utils.set_random_seed(seed)
            np.random.seed(seed)

            model = build_model(x.shape[1], num_classes, cfg)
            model.fit(
                x_tr,
                y_tr,
                epochs=cfg.epochs,
                batch_size=cfg.batch_size,
                validation_data=(x_va, y_va),
                verbose=0,
            )
            all_models.append(model)

        print(
            f"[Seed base {seed_base}] Fold {fold}/{cfg.k_splits} "
            f"trained {cfg.n_repeats} models (total: {len(all_models)})"
        )

    return all_models


def train_for_bin(
    df: pd.DataFrame,
    bin_value: int | float,
    label_column: str = "target",
    bin_column: str = "bin",
    control_labels: list[Any] | None = None,
    experiment_labels: list[Any] | None = None,
    cfg: TrainingConfig | None = None,
) -> dict:
    cfg = cfg or TrainingConfig()
    x, y = prepare_bin_data(
        df,
        bin_value,
        label_column=label_column,
        bin_column=bin_column,
        control_labels=control_labels,
        experiment_labels=experiment_labels,
    )

    models_a = train_kfold_repeats(x, y, seed_base=cfg.seed_bases[0], cfg=cfg)
    models_b = train_kfold_repeats(x, y, seed_base=cfg.seed_bases[1], cfg=cfg)

    return {
        "bin_value": bin_value,
        "x": x,
        "y": y,
        "models_a": models_a,
        "models_b": models_b,
    }


def run_training(
    csv_path: str,
    label_column: str = "target",
    bin_column: str = "bin",
    bin_value: int | float | None = None,
    control_labels: list[Any] | None = None,
    experiment_labels: list[Any] | None = None,
    cfg: TrainingConfig | None = None,
) -> dict:
    df = pd.read_csv(csv_path)

    if bin_column not in df.columns:
        raise ValueError(f"Missing required column: {bin_column}")
    if label_column not in df.columns:
        raise ValueError(f"Missing required column: {label_column}")

    bins_found = sorted(df[bin_column].dropna().unique().tolist())
    if not bins_found:
        raise ValueError(f"No bin values found in column '{bin_column}'")

    if bin_value is None:
        bin_value = bins_found[0]

    return train_for_bin(
        df=df,
        bin_value=bin_value,
        label_column=label_column,
        bin_column=bin_column,
        control_labels=control_labels,
        experiment_labels=experiment_labels,
        cfg=cfg,
    )