from __future__ import annotations

import copy
import json
import random
from pathlib import Path

import numpy as np
import torch
import trackio
from model import FederatedMLP, parameter_count
from safetensors.torch import save_file
from sklearn.datasets import load_digits
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from torch.nn import functional as F
from torch.utils.data import DataLoader, TensorDataset

PROJECT_DIR = Path(__file__).resolve().parent
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "fedavg-global"
CLIENTS = 5


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


def dirichlet_partition(
    labels: np.ndarray,
    clients: int = CLIENTS,
    alpha: float = 0.25,
    seed: int = 2026,
) -> list[np.ndarray]:
    rng = np.random.default_rng(seed)
    assignments = [[] for _ in range(clients)]
    for label in sorted(np.unique(labels)):
        indices = np.flatnonzero(labels == label)
        rng.shuffle(indices)
        proportions = rng.dirichlet(np.full(clients, alpha))
        boundaries = (np.cumsum(proportions)[:-1] * len(indices)).astype(int)
        for client, shard in enumerate(np.split(indices, boundaries)):
            assignments[client].extend(shard.tolist())
    return [np.asarray(sorted(indices), dtype=np.int64) for indices in assignments]


def train_local(
    global_model: FederatedMLP,
    dataset: TensorDataset,
    *,
    local_epochs: int,
    learning_rate: float,
) -> tuple[dict[str, torch.Tensor], float]:
    model = copy.deepcopy(global_model)
    loader = DataLoader(dataset, batch_size=48, shuffle=True)
    optimizer = torch.optim.SGD(
        model.parameters(),
        lr=learning_rate,
        momentum=0.85,
    )
    losses = []
    for _ in range(local_epochs):
        model.train()
        for pixels, labels in loader:
            logits = model(pixels)
            loss = F.cross_entropy(logits, labels)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            losses.append(loss.item())
    return (
        {key: value.detach().clone() for key, value in model.state_dict().items()},
        float(np.mean(losses)),
    )


def fedavg(
    states: list[dict[str, torch.Tensor]],
    weights: list[int],
) -> dict[str, torch.Tensor]:
    total = sum(weights)
    averaged = {}
    for key in states[0]:
        averaged[key] = sum(
            state[key] * (weight / total)
            for state, weight in zip(states, weights, strict=True)
        )
    return averaged


def update_disagreement(
    global_state: dict[str, torch.Tensor],
    client_states: list[dict[str, torch.Tensor]],
) -> float:
    distances = []
    for state in client_states:
        squared = sum(
            torch.sum((state[key] - global_state[key]) ** 2).item() for key in global_state
        )
        distances.append(squared**0.5)
    return float(np.mean(distances))


@torch.inference_mode()
def evaluate(model: FederatedMLP, dataset: TensorDataset) -> float:
    model.eval()
    loader = DataLoader(dataset, batch_size=256, shuffle=False)
    labels, predictions = [], []
    for pixels, targets in loader:
        predictions.extend(model(pixels).argmax(dim=1).tolist())
        labels.extend(targets.tolist())
    return float(accuracy_score(labels, predictions))


def centralized_control(
    dataset: TensorDataset,
    test: TensorDataset,
    epochs: int,
) -> tuple[FederatedMLP, float]:
    seed_everything(2026)
    model = FederatedMLP()
    loader = DataLoader(
        dataset,
        batch_size=48,
        shuffle=True,
        generator=torch.Generator().manual_seed(2026),
    )
    optimizer = torch.optim.SGD(model.parameters(), lr=0.06, momentum=0.85)
    for _ in range(epochs):
        model.train()
        for pixels, labels in loader:
            loss = F.cross_entropy(model(pixels), labels)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
    return model, evaluate(model, test)


def main() -> None:
    seed_everything(2026)
    digits = load_digits()
    pixels = digits.data.astype(np.float32) / 16.0
    labels = digits.target.astype(np.int64)
    train_indices, test_indices = train_test_split(
        np.arange(len(labels)),
        test_size=0.25,
        stratify=labels,
        random_state=2026,
    )
    train_pixels = torch.from_numpy(pixels[train_indices])
    train_labels = torch.from_numpy(labels[train_indices])
    test_dataset = TensorDataset(
        torch.from_numpy(pixels[test_indices]),
        torch.from_numpy(labels[test_indices]),
    )
    client_indices = dirichlet_partition(labels[train_indices])
    client_datasets = [
        TensorDataset(train_pixels[indices], train_labels[indices])
        for indices in client_indices
    ]
    distributions = [
        np.bincount(train_labels[indices].numpy(), minlength=10).tolist()
        for indices in client_indices
    ]

    rounds = 30
    global_model = FederatedMLP()
    history = []
    trackio.init(
        project="federated-digits",
        name="fedavg-dirichlet-alpha-0.25",
        config={
            "parameters": parameter_count(global_model),
            "clients": CLIENTS,
            "rounds": rounds,
            "local_epochs": 2,
            "dirichlet_alpha": 0.25,
            "client_sizes": [len(dataset) for dataset in client_datasets],
        },
    )
    for round_index in range(1, rounds + 1):
        global_state = {
            key: value.detach().clone() for key, value in global_model.state_dict().items()
        }
        client_states = []
        client_losses = []
        for dataset in client_datasets:
            state, loss = train_local(
                global_model,
                dataset,
                local_epochs=2,
                learning_rate=0.06,
            )
            client_states.append(state)
            client_losses.append(loss)
        disagreement = update_disagreement(global_state, client_states)
        global_model.load_state_dict(
            fedavg(client_states, [len(dataset) for dataset in client_datasets])
        )
        accuracy = evaluate(global_model, test_dataset)
        record = {
            "round": round_index,
            "global_test_accuracy": accuracy,
            "mean_client_loss": float(np.mean(client_losses)),
            "client_update_disagreement": disagreement,
        }
        history.append(record)
        trackio.log(record)
    trackio.finish()

    full_train = TensorDataset(train_pixels, train_labels)
    centralized_model, centralized_accuracy = centralized_control(
        full_train,
        test_dataset,
        epochs=rounds * 2,
    )
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    save_file(global_model.state_dict(), ARTIFACT_DIR / "model.safetensors")
    results = {
        "model": "Federated Digits FedAvg",
        "parameters": parameter_count(global_model),
        "clients": CLIENTS,
        "dirichlet_alpha": 0.25,
        "client_sizes": [len(dataset) for dataset in client_datasets],
        "client_class_distributions": distributions,
        "rounds": rounds,
        "local_epochs_per_round": 2,
        "initial_test_accuracy": history[0]["global_test_accuracy"],
        "final_test_accuracy": history[-1]["global_test_accuracy"],
        "best_test_accuracy": max(record["global_test_accuracy"] for record in history),
        "centralized_control_accuracy": centralized_accuracy,
        "federated_centralized_gap": (
            history[-1]["global_test_accuracy"] - centralized_accuracy
        ),
        "history": history,
    }
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()