from __future__ import annotations

import json
import random
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import trackio
from model import TeacherCNN, TinyStudentCNN, parameter_count
from safetensors.torch import save_file
from sklearn.metrics import accuracy_score, confusion_matrix, f1_score
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, TensorDataset

PROJECT_DIR = Path(__file__).resolve().parent
DATA_DIR = PROJECT_DIR / "data"
ARTIFACT_DIR = PROJECT_DIR / "artifacts"
DEVICE = torch.device("cpu")


@dataclass
class TrainResult:
    best_epoch: int
    best_validation_accuracy: float
    history: list[dict[str, float]]


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def load_split(name: str, *, shuffle: bool, batch_size: int) -> DataLoader:
    frame = pd.read_parquet(DATA_DIR / f"{name}.parquet")
    pixels = np.stack(frame["image"].to_numpy()).astype(np.float32) / 16.0
    images = torch.from_numpy(pixels.reshape(-1, 1, 8, 8))
    labels = torch.from_numpy(frame["label"].to_numpy(dtype=np.int64, copy=True))
    generator = torch.Generator().manual_seed(2026)
    return DataLoader(
        TensorDataset(images, labels),
        batch_size=batch_size,
        shuffle=shuffle,
        generator=generator,
    )


@torch.inference_mode()
def evaluate(model: nn.Module, loader: DataLoader) -> dict:
    model.eval()
    labels, predictions, losses = [], [], []
    for images, targets in loader:
        logits = model(images.to(DEVICE))
        losses.append(F.cross_entropy(logits, targets.to(DEVICE), reduction="sum").item())
        labels.extend(targets.tolist())
        predictions.extend(logits.argmax(dim=1).cpu().tolist())
    return {
        "loss": float(sum(losses) / len(labels)),
        "accuracy": float(accuracy_score(labels, predictions)),
        "macro_f1": float(f1_score(labels, predictions, average="macro")),
        "confusion_matrix": confusion_matrix(labels, predictions).tolist(),
    }


def train_model(
    model: nn.Module,
    train_loader: DataLoader,
    validation_loader: DataLoader,
    *,
    run_name: str,
    epochs: int,
    learning_rate: float,
    teacher: nn.Module | None = None,
    temperature: float = 3.0,
    soft_weight: float = 0.65,
) -> TrainResult:
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    best_accuracy = -1.0
    best_epoch = 0
    best_state = None
    history = []

    trackio.init(
        project="tiny-vision-foundry",
        name=run_name,
        config={
            "parameters": parameter_count(model),
            "epochs": epochs,
            "learning_rate": learning_rate,
            "distilled": teacher is not None,
            "temperature": temperature if teacher else None,
            "soft_weight": soft_weight if teacher else None,
        },
    )
    if teacher is not None:
        teacher.eval()
        for parameter in teacher.parameters():
            parameter.requires_grad_(False)

    for epoch in range(1, epochs + 1):
        model.train()
        running_loss = 0.0
        examples = 0
        for images, targets in train_loader:
            images = images.to(DEVICE)
            targets = targets.to(DEVICE)
            optimizer.zero_grad(set_to_none=True)
            logits = model(images)
            hard_loss = F.cross_entropy(logits, targets)
            if teacher is None:
                loss = hard_loss
            else:
                with torch.no_grad():
                    teacher_logits = teacher(images)
                soft_loss = (
                    F.kl_div(
                        F.log_softmax(logits / temperature, dim=1),
                        F.softmax(teacher_logits / temperature, dim=1),
                        reduction="batchmean",
                    )
                    * temperature**2
                )
                loss = (1 - soft_weight) * hard_loss + soft_weight * soft_loss
            loss.backward()
            optimizer.step()
            running_loss += loss.item() * len(targets)
            examples += len(targets)
        scheduler.step()
        validation = evaluate(model, validation_loader)
        record = {
            "epoch": float(epoch),
            "train_loss": running_loss / examples,
            "validation_loss": validation["loss"],
            "validation_accuracy": validation["accuracy"],
            "learning_rate": scheduler.get_last_lr()[0],
        }
        history.append(record)
        trackio.log(record)
        if validation["accuracy"] > best_accuracy:
            best_accuracy = validation["accuracy"]
            best_epoch = epoch
            best_state = {
                key: value.detach().cpu().clone()
                for key, value in model.state_dict().items()
            }
    trackio.finish()
    assert best_state is not None
    model.load_state_dict(best_state)
    return TrainResult(best_epoch, best_accuracy, history)


def save_variant(name: str, model: nn.Module, result: TrainResult, metrics: dict) -> None:
    output = ARTIFACT_DIR / name
    output.mkdir(parents=True, exist_ok=True)
    save_file(model.state_dict(), output / "model.safetensors")
    report = {
        "variant": name,
        "parameters": parameter_count(model),
        "best_epoch": result.best_epoch,
        "best_validation_accuracy": result.best_validation_accuracy,
        "test": metrics,
    }
    (output / "evaluation.json").write_text(
        json.dumps(report, indent=2),
        encoding="utf-8",
    )


def main() -> None:
    seed_everything(2026)
    train_loader = load_split("train", shuffle=True, batch_size=64)
    validation_loader = load_split("validation", shuffle=False, batch_size=256)
    test_loader = load_split("test", shuffle=False, batch_size=256)

    teacher = TeacherCNN().to(DEVICE)
    teacher_result = train_model(
        teacher,
        train_loader,
        validation_loader,
        run_name="teacher-cnn",
        epochs=80,
        learning_rate=3e-3,
    )
    teacher_metrics = evaluate(teacher, test_loader)
    save_variant("teacher-cnn", teacher, teacher_result, teacher_metrics)

    seed_everything(2027)
    scratch_student = TinyStudentCNN().to(DEVICE)
    scratch_result = train_model(
        scratch_student,
        train_loader,
        validation_loader,
        run_name="tiny-student-scratch",
        epochs=100,
        learning_rate=4e-3,
    )
    scratch_metrics = evaluate(scratch_student, test_loader)
    save_variant(
        "tiny-student-scratch",
        scratch_student,
        scratch_result,
        scratch_metrics,
    )

    seed_everything(2027)
    distilled_student = TinyStudentCNN().to(DEVICE)
    distilled_result = train_model(
        distilled_student,
        train_loader,
        validation_loader,
        run_name="tiny-student-distilled",
        epochs=100,
        learning_rate=4e-3,
        teacher=teacher,
        temperature=3.0,
        soft_weight=0.65,
    )
    distilled_metrics = evaluate(distilled_student, test_loader)
    save_variant(
        "tiny-student-distilled",
        distilled_student,
        distilled_result,
        distilled_metrics,
    )

    summary = {
        "teacher": {
            "parameters": parameter_count(teacher),
            "test": teacher_metrics,
        },
        "student_scratch": {
            "parameters": parameter_count(scratch_student),
            "test": scratch_metrics,
        },
        "student_distilled": {
            "parameters": parameter_count(distilled_student),
            "test": distilled_metrics,
        },
        "compression_ratio": parameter_count(teacher) / parameter_count(distilled_student),
        "distillation_accuracy_delta": (
            distilled_metrics["accuracy"] - scratch_metrics["accuracy"]
        ),
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    (ARTIFACT_DIR / "comparison.json").write_text(
        json.dumps(summary, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()