from __future__ import annotations

import random
from collections.abc import Iterable

import numpy as np
import torch
from torch.utils.data import DataLoader

from .config import ModelConfig
from .dataset import RouterCollator, RouterDataset
from .model import RouterModel
from .schemas import RewardRecord


def seed_everything(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 resolve_device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def make_model(
    model_config: ModelConfig,
    worker_ids: list[str],
    checkpoint: str | None = None,
) -> RouterModel:
    device = resolve_device()
    if checkpoint:
        model = RouterModel.from_checkpoint(checkpoint, device=device)
        if model.worker_ids != worker_ids:
            raise ValueError(
                f"Checkpoint workers {model.worker_ids} do not match dataset workers {worker_ids}"
            )
        return model
    return RouterModel(model_config, worker_ids, device)


def make_loader(
    records: list[RewardRecord],
    model: RouterModel,
    split: str,
    batch_size: int,
    shuffle: bool,
) -> DataLoader:
    dataset = RouterDataset(records, split=split)
    if not dataset:
        raise ValueError(f"No {split} records found")
    generator = torch.Generator().manual_seed(0)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        generator=generator,
        collate_fn=RouterCollator(model.tokenizer, model.router_config.max_length),
        num_workers=0,
        pin_memory=torch.cuda.is_available(),
    )


def move_batch(batch: dict, device: torch.device) -> tuple[dict, torch.Tensor]:
    inputs = {
        "input_ids": batch["input_ids"].to(device, non_blocking=True),
        "attention_mask": batch["attention_mask"].to(device, non_blocking=True),
    }
    return inputs, batch["rewards"].to(device, non_blocking=True)


def optimizer_parameters(model: RouterModel) -> Iterable[torch.nn.Parameter]:
    return (parameter for parameter in model.parameters() if parameter.requires_grad)


@torch.no_grad()
def evaluate_loader(model: RouterModel, loader: DataLoader) -> dict:
    model.eval()
    chosen_rewards = []
    oracle_rewards = []
    oracle_hits = []
    all_rewards = []
    route_counts = torch.zeros(len(model.worker_ids), dtype=torch.long)
    for batch in loader:
        inputs, rewards = move_batch(batch, model.device_ref)
        logits = model(**inputs)
        actions = logits.argmax(dim=-1)
        chosen = rewards.gather(1, actions.unsqueeze(1)).squeeze(1)
        oracle, oracle_actions = rewards.max(dim=1)
        chosen_rewards.append(chosen.cpu())
        oracle_rewards.append(oracle.cpu())
        oracle_hits.append(actions.eq(oracle_actions).float().cpu())
        all_rewards.append(rewards.cpu())
        route_counts += torch.bincount(actions.cpu(), minlength=len(model.worker_ids))
    chosen = torch.cat(chosen_rewards)
    oracle = torch.cat(oracle_rewards)
    rewards = torch.cat(all_rewards)
    fixed_means = rewards.mean(dim=0)
    report = {
        "examples": int(chosen.numel()),
        "router_utility": float(chosen.mean()),
        "oracle_utility": float(oracle.mean()),
        "regret": float((oracle - chosen).mean()),
        "oracle_route_accuracy": float(torch.cat(oracle_hits).mean()),
        "best_fixed_worker": model.worker_ids[int(fixed_means.argmax())],
        "best_fixed_utility": float(fixed_means.max()),
        "random_utility": float(rewards.mean()),
        "route_counts": {
            worker_id: int(count) for worker_id, count in zip(model.worker_ids, route_counts)
        },
    }
    return report