from __future__ import annotations

import copy
import math
import random
from dataclasses import asdict, dataclass
from typing import Callable

import numpy as np
import pandas as pd
import torch
from sklearn.metrics import (
    accuracy_score,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
)
from sklearn.model_selection import train_test_split
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from .data import PreparedWorkspace
from .model import BioLMNet


ProgressCallback = Callable[[float, str], None]

# 75% training / 25% validation. The validation partition drives early
# stopping and best-checkpoint selection, so it is not an untouched test set.
DEFAULT_VALIDATION_FRACTION = 0.25

# Activations BioLMNet accepts.
ACTIVATIONS = ("relu", "sigmoid", "tanh")


# Defaults follow the Bozdag Lab grid (J. Rifat, 2026-10-17): the configuration that stayed
# stable across 8 seeds on BRCA (macro F1 0.749 +/- 0.018). The published activations
# (relu/sigmoid/tanh) collapse at this learning rate and L2; see
# diagnostics/2026-10-30-lab-hyperparameter-grid.
@dataclass(frozen=True)
class Hyperparameters:
    epochs: int = 200
    batch_size: int = 32
    learning_rate: float = 0.01
    weight_decay: float = 0.01
    dropout: float = 0.3
    projection_dim: int = 64
    fusion_dim: int = 12
    validation_fraction: float = DEFAULT_VALIDATION_FRACTION
    optimizer: str = "Adam"
    class_weighting: bool = True
    early_stopping_patience: int = 12
    seed: int = 42
    biological_activation: str = "tanh"
    projection_activation: str = "relu"
    fusion_activation: str = "tanh"
    sgd_momentum: float = 0.03

    def validate(self) -> None:
        if not 1 <= self.epochs <= 1000:
            raise ValueError("Epochs must be between 1 and 1,000.")
        if not 2 <= self.batch_size <= 1024:
            raise ValueError("Batch size must be between 2 and 1,024.")
        if not 0 < self.learning_rate <= 1:
            raise ValueError("Learning rate must be in (0, 1].")
        if not 0 <= self.dropout < 1:
            raise ValueError("Dropout must be in [0, 1).")
        if not 0.05 <= self.validation_fraction <= 0.5:
            raise ValueError("Validation fraction must be between 0.05 and 0.5.")
        if self.optimizer.lower() not in {"adam", "sgd"}:
            raise ValueError("Optimizer must be Adam or SGD.")
        for name in ("biological_activation", "projection_activation", "fusion_activation"):
            if getattr(self, name) not in ACTIVATIONS:
                raise ValueError(f"{name.replace('_', ' ').capitalize()} must be one of {', '.join(ACTIVATIONS)}.")
        if not 0 <= self.sgd_momentum < 1:
            raise ValueError("SGD momentum must be in [0, 1).")


@dataclass
class ModelBundle:
    model: BioLMNet
    gene_features: list[str]
    dna_features: list[str]
    label_names: list[str]
    gene_mean: np.ndarray
    gene_scale: np.ndarray
    dna_mean: np.ndarray
    dna_scale: np.ndarray
    config: dict
    metrics: dict
    history: list[dict[str, float]]


@dataclass
class TrainingResult:
    bundle: ModelBundle
    validation_predictions: pd.DataFrame
    confusion: np.ndarray


def set_reproducible_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def split_train_validation(
    labels: np.ndarray, validation_fraction: float, seed: int
) -> tuple[np.ndarray, np.ndarray]:
    """Stratified sample-index split shared by both omics and the labels."""
    classes, class_counts = np.unique(labels, return_counts=True)
    if class_counts.min() < 2:
        raise ValueError(
            "Each class needs at least two samples for a stratified train/validation split."
        )
    sample_count = len(labels)
    validation_count = math.ceil(validation_fraction * sample_count)
    train_count = sample_count - validation_count
    if min(validation_count, train_count) < len(classes):
        raise ValueError(
            f"A validation fraction of {validation_fraction:.2f} on {sample_count} samples "
            f"gives {train_count} training and {validation_count} validation samples; "
            f"each partition needs at least one sample per class ({len(classes)} classes) "
            "for a stratified split. Adjust the validation fraction or add samples."
        )

    indices = np.arange(sample_count)
    train_index, validation_index = train_test_split(
        indices,
        test_size=validation_fraction,
        random_state=seed,
        stratify=labels,
    )
    return train_index, validation_index


def _fit_scaler(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    mean = values.mean(axis=0, dtype=np.float64).astype(np.float32)
    scale = values.std(axis=0, dtype=np.float64).astype(np.float32)
    scale[scale < 1e-8] = 1.0
    return mean, scale


def _scale(values: np.ndarray, mean: np.ndarray, scale: np.ndarray) -> np.ndarray:
    return ((values.astype(np.float32) - mean) / scale).astype(np.float32)


def _make_model(
    workspace: PreparedWorkspace, hyperparameters: Hyperparameters
) -> BioLMNet:
    gene = workspace.gene_branch
    dna = workspace.dna_branch
    if (
        gene.embeddings is None
        or dna.embeddings is None
        or gene.pathway_mask is None
        or dna.pathway_mask is None
    ):
        raise ValueError("Workspace priors are incomplete; run data preparation first.")

    return BioLMNet(
        gene_biological_mask=torch.from_numpy(gene.biological_mask),
        dna_biological_mask=torch.from_numpy(dna.biological_mask),
        gene_embeddings=torch.from_numpy(gene.embeddings),
        dna_embeddings=torch.from_numpy(dna.embeddings),
        gene_pathway_mask=torch.from_numpy(gene.pathway_mask),
        dna_pathway_mask=torch.from_numpy(dna.pathway_mask),
        n_classes=len(workspace.label_names),
        projection_dim=hyperparameters.projection_dim,
        fusion_dim=hyperparameters.fusion_dim,
        dropout=hyperparameters.dropout,
        biological_activation=hyperparameters.biological_activation,
        projection_activation=hyperparameters.projection_activation,
        fusion_activation=hyperparameters.fusion_activation,
    )


def _evaluate(
    model: BioLMNet,
    gene_values: np.ndarray,
    dna_values: np.ndarray,
    labels: np.ndarray,
    device: torch.device,
) -> tuple[float, np.ndarray, np.ndarray]:
    model.eval()
    with torch.no_grad():
        logits = model(
            torch.from_numpy(gene_values).to(device),
            torch.from_numpy(dna_values).to(device),
        )
        loss = nn.functional.cross_entropy(
            logits, torch.from_numpy(labels).to(device)
        ).item()
        probabilities = torch.softmax(logits, dim=1).cpu().numpy()
    predictions = probabilities.argmax(axis=1)
    return float(loss), probabilities, predictions


def train(
    workspace: PreparedWorkspace,
    hyperparameters: Hyperparameters,
    progress: ProgressCallback | None = None,
) -> TrainingResult:
    hyperparameters.validate()
    set_reproducible_seed(hyperparameters.seed)
    labels = workspace.labels
    train_index, validation_index = split_train_validation(
        labels, hyperparameters.validation_fraction, hyperparameters.seed
    )
    gene_mean, gene_scale = _fit_scaler(
        workspace.gene_expression[train_index]
    )
    dna_mean, dna_scale = _fit_scaler(workspace.dna_methylation[train_index])
    gene_train = _scale(
        workspace.gene_expression[train_index], gene_mean, gene_scale
    )
    gene_validation = _scale(
        workspace.gene_expression[validation_index], gene_mean, gene_scale
    )
    dna_train = _scale(
        workspace.dna_methylation[train_index], dna_mean, dna_scale
    )
    dna_validation = _scale(
        workspace.dna_methylation[validation_index], dna_mean, dna_scale
    )
    y_train = labels[train_index]
    y_validation = labels[validation_index]

    dataset = TensorDataset(
        torch.from_numpy(gene_train),
        torch.from_numpy(dna_train),
        torch.from_numpy(y_train),
    )
    generator = torch.Generator().manual_seed(hyperparameters.seed)
    loader = DataLoader(
        dataset,
        batch_size=min(hyperparameters.batch_size, len(dataset)),
        shuffle=True,
        generator=generator,
    )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = _make_model(workspace, hyperparameters).to(device)
    if hyperparameters.optimizer.lower() == "sgd":
        optimizer = torch.optim.SGD(
            model.parameters(),
            lr=hyperparameters.learning_rate,
            momentum=hyperparameters.sgd_momentum,
            weight_decay=hyperparameters.weight_decay,
        )
    else:
        optimizer = torch.optim.Adam(
            model.parameters(),
            lr=hyperparameters.learning_rate,
            weight_decay=hyperparameters.weight_decay,
        )

    class_weights: torch.Tensor | None = None
    if hyperparameters.class_weighting:
        count_by_class = np.bincount(
            y_train, minlength=len(workspace.label_names)
        ).astype(np.float32)
        weights = len(y_train) / (len(count_by_class) * count_by_class)
        class_weights = torch.from_numpy(weights).to(device)
    criterion = nn.CrossEntropyLoss(weight=class_weights)

    history: list[dict[str, float]] = []
    best_state: dict[str, torch.Tensor] | None = None
    best_loss = float("inf")
    patience = 0
    for epoch in range(hyperparameters.epochs):
        model.train()
        running_loss = 0.0
        seen = 0
        for gene_batch, dna_batch, label_batch in loader:
            gene_batch = gene_batch.to(device)
            dna_batch = dna_batch.to(device)
            label_batch = label_batch.to(device)
            optimizer.zero_grad(set_to_none=True)
            logits = model(gene_batch, dna_batch)
            loss = criterion(logits, label_batch)
            loss.backward()
            optimizer.step()
            running_loss += loss.item() * len(label_batch)
            seen += len(label_batch)

        validation_loss, _, validation_predictions = _evaluate(
            model,
            gene_validation,
            dna_validation,
            y_validation,
            device,
        )
        epoch_row = {
            "epoch": float(epoch + 1),
            "training_loss": float(running_loss / max(seen, 1)),
            "validation_loss": validation_loss,
            "validation_accuracy": float(
                accuracy_score(y_validation, validation_predictions)
            ),
            "validation_f1_macro": float(
                f1_score(
                    y_validation,
                    validation_predictions,
                    average="macro",
                    zero_division=0,
                )
            ),
        }
        history.append(epoch_row)
        if progress:
            progress(
                (epoch + 1) / hyperparameters.epochs,
                (
                    f"Epoch {epoch + 1}/{hyperparameters.epochs} · "
                    f"validation F1 {epoch_row['validation_f1_macro']:.3f}"
                ),
            )

        if validation_loss < best_loss - 1e-5:
            best_loss = validation_loss
            best_state = copy.deepcopy(model.state_dict())
            patience = 0
        else:
            patience += 1
            if patience >= hyperparameters.early_stopping_patience:
                break

    if best_state is not None:
        model.load_state_dict(best_state)
    validation_loss, probabilities, predictions = _evaluate(
        model,
        gene_validation,
        dna_validation,
        y_validation,
        device,
    )
    confusion = confusion_matrix(
        y_validation,
        predictions,
        labels=np.arange(len(workspace.label_names)),
    )
    metrics = {
        "validation_loss": validation_loss,
        "accuracy": float(accuracy_score(y_validation, predictions)),
        "f1_macro": float(
            f1_score(
                y_validation, predictions, average="macro", zero_division=0
            )
        ),
        "f1_weighted": float(
            f1_score(
                y_validation, predictions, average="weighted", zero_division=0
            )
        ),
        "precision_macro": float(
            precision_score(
                y_validation, predictions, average="macro", zero_division=0
            )
        ),
        "recall_macro": float(
            recall_score(
                y_validation, predictions, average="macro", zero_division=0
            )
        ),
        "epochs_completed": len(history),
        "training_samples": int(len(train_index)),
        "validation_samples": int(len(validation_index)),
        "device": str(device),
    }
    predictions_frame = pd.DataFrame(
        {
            "sample_row": validation_index,
            "observed": [
                workspace.label_names[value] for value in y_validation
            ],
            "predicted": [
                workspace.label_names[value] for value in predictions
            ],
        }
    )
    for index, label in enumerate(workspace.label_names):
        predictions_frame[f"P({label})"] = probabilities[:, index]

    gene = workspace.gene_branch
    dna = workspace.dna_branch
    config = {
        "format_version": 1,
        "source_name": workspace.source_name,
        "gene_features": gene.input_genes,
        "dna_features": dna.input_genes,
        "label_names": workspace.label_names,
        "gene_hidden_genes": gene.hidden_genes,
        "dna_hidden_genes": dna.hidden_genes,
        "gene_pathways": gene.pathways,
        "dna_pathways": dna.pathways,
        "hyperparameters": asdict(hyperparameters),
        "architecture": {
            "projection_dim": hyperparameters.projection_dim,
            "fusion_dim": hyperparameters.fusion_dim,
            "dropout": hyperparameters.dropout,
            "biological_activation": hyperparameters.biological_activation,
            "projection_activation": hyperparameters.projection_activation,
            "fusion_activation": hyperparameters.fusion_activation,
        },
    }
    bundle = ModelBundle(
        model=model.cpu(),
        gene_features=gene.input_genes,
        dna_features=dna.input_genes,
        label_names=workspace.label_names,
        gene_mean=gene_mean,
        gene_scale=gene_scale,
        dna_mean=dna_mean,
        dna_scale=dna_scale,
        config=config,
        metrics=metrics,
        history=history,
    )
    return TrainingResult(
        bundle=bundle,
        validation_predictions=predictions_frame,
        confusion=confusion,
    )


def validate_prediction_frames(
    gene_frame: pd.DataFrame,
    dna_frame: pd.DataFrame,
    bundle: ModelBundle,
) -> tuple[np.ndarray, np.ndarray]:
    if len(gene_frame) != len(dna_frame):
        raise ValueError(
            "Prediction gene-expression and DNA-methylation files must have "
            "the same number of rows."
        )
    missing_gene = sorted(set(bundle.gene_features) - set(gene_frame.columns))
    missing_dna = sorted(set(bundle.dna_features) - set(dna_frame.columns))
    if missing_gene or missing_dna:
        details = []
        if missing_gene:
            details.append(
                "gene-expression: " + ", ".join(missing_gene[:8])
            )
        if missing_dna:
            details.append("DNA-methylation: " + ", ".join(missing_dna[:8]))
        raise ValueError(
            "Prediction files are missing trained features (" + "; ".join(details) + ")."
        )
    gene_values = gene_frame.loc[:, bundle.gene_features].apply(
        pd.to_numeric, errors="coerce"
    )
    dna_values = dna_frame.loc[:, bundle.dna_features].apply(
        pd.to_numeric, errors="coerce"
    )
    if gene_values.isna().any().any() or dna_values.isna().any().any():
        raise ValueError("Prediction inputs contain missing or non-numeric values.")
    return (
        _scale(gene_values.to_numpy(), bundle.gene_mean, bundle.gene_scale),
        _scale(dna_values.to_numpy(), bundle.dna_mean, bundle.dna_scale),
    )


def predict(
    gene_frame: pd.DataFrame,
    dna_frame: pd.DataFrame,
    bundle: ModelBundle,
) -> pd.DataFrame:
    gene_values, dna_values = validate_prediction_frames(
        gene_frame, dna_frame, bundle
    )
    bundle.model.eval()
    with torch.no_grad():
        logits = bundle.model(
            torch.from_numpy(gene_values), torch.from_numpy(dna_values)
        )
        probabilities = torch.softmax(logits, dim=1).numpy()
    predicted = probabilities.argmax(axis=1)
    output = pd.DataFrame(
        {
            "sample_row": np.arange(len(gene_frame)),
            "predicted_class": [
                bundle.label_names[index] for index in predicted
            ],
            "confidence": probabilities.max(axis=1),
        }
    )
    for index, label in enumerate(bundle.label_names):
        output[f"P({label})"] = probabilities[:, index]
    return output


def pathway_importance(bundle: ModelBundle) -> pd.DataFrame:
    bundle.model.eval()
    with torch.no_grad():
        attention = bundle.model.pathway_attention()
    records: list[dict[str, float | str]] = []
    for branch_name, config_key in (
        ("Gene expression", "gene_pathways"),
        ("DNA methylation", "dna_pathways"),
    ):
        key = "gene_expression" if branch_name == "Gene expression" else "dna_methylation"
        weights = attention[key].cpu().numpy()
        pathways = bundle.config[config_key]
        peak_attention = weights.max(axis=0)
        entropy = -(
            weights * np.log(np.clip(weights, 1e-12, None))
        ).sum(axis=0)
        for pathway, peak, entropy_value in zip(
            pathways, peak_attention, entropy, strict=True
        ):
            records.append(
                {
                    "branch": branch_name,
                    "pathway": pathway,
                    "peak_gene_attention": float(peak),
                    "attention_entropy": float(entropy_value),
                }
            )
    return (
        pd.DataFrame(records)
        .sort_values("peak_gene_attention", ascending=False)
        .reset_index(drop=True)
    )