from __future__ import annotations

import argparse
import csv
import json
import os
from pathlib import Path

os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.makedirs(os.environ["MPLCONFIGDIR"], exist_ok=True)

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from torch.utils.data import DataLoader, Dataset
from torchvision.models import efficientnet_b0, resnet18
from tqdm import tqdm


class SpectrogramDataset(Dataset):
    def __init__(self, processed_dir: Path, split: str, preload: bool = True):
        self.processed_dir = processed_dir
        self.sample_dir = processed_dir / "samples"
        with (processed_dir / "manifest.csv").open() as f:
            rows = list(csv.DictReader(f))
        self.rows = [row for row in rows if row["split"] == split]
        if not self.rows:
            raise ValueError(f"No rows found for split={split}")
        self.x_cache: torch.Tensor | None = None
        self.y_cache: torch.Tensor | None = None
        if preload:
            xs = []
            ys = []
            for row in tqdm(self.rows, desc=f"preload {split}"):
                data = np.load(self.sample_dir / row["path"])
                x = data["x"].astype(np.float32)
                x = (x - x.mean()) / (x.std() + 1e-6)
                xs.append(x)
                ys.append(int(row["label_id"]))
            self.x_cache = torch.from_numpy(np.stack(xs, axis=0)).unsqueeze(1)
            self.y_cache = torch.tensor(ys, dtype=torch.long)

    def __len__(self) -> int:
        return len(self.rows)

    def __getitem__(self, idx: int):
        if self.x_cache is not None and self.y_cache is not None:
            return self.x_cache[idx], self.y_cache[idx]

        row = self.rows[idx]
        data = np.load(self.sample_dir / row["path"])
        x = data["x"].astype(np.float32)
        x = (x - x.mean()) / (x.std() + 1e-6)
        x = torch.from_numpy(x).unsqueeze(0)
        y = torch.tensor(int(row["label_id"]), dtype=torch.long)
        return x, y


class SmallSpectrogramCNN(nn.Module):
    def __init__(self, num_classes: int):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.classifier = nn.Linear(128, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)


def build_model(name: str, num_classes: int) -> nn.Module:
    if name == "small_cnn":
        return SmallSpectrogramCNN(num_classes)

    if name == "resnet18":
        model = resnet18(weights=None)
        model.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        model.fc = nn.Linear(model.fc.in_features, num_classes)
        return model

    if name == "efficientnet_b0":
        model = efficientnet_b0(weights=None)
        model.features[0][0] = nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1, bias=False)
        model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
        return model

    raise ValueError(f"Unknown model: {name}")


def read_labels(processed_dir: Path) -> list[str]:
    labels = []
    with (processed_dir / "labels.txt").open() as f:
        for line in f:
            _, label = line.rstrip("\n").split("\t", 1)
            labels.append(label)
    return labels


def evaluate(model: nn.Module, loader: DataLoader, device: torch.device):
    model.eval()
    y_true = []
    y_pred = []
    loss_sum = 0.0
    criterion = nn.CrossEntropyLoss()
    use_cuda = device.type == "cuda"
    with torch.no_grad():
        for x, y in loader:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
            if use_cuda:
                x = x.contiguous(memory_format=torch.channels_last)
            with torch.amp.autocast("cuda", enabled=use_cuda):
                logits = model(x)
                loss = criterion(logits, y)
            loss_sum += loss.item() * x.size(0)
            y_true.extend(y.cpu().numpy().tolist())
            y_pred.extend(logits.argmax(dim=1).cpu().numpy().tolist())
    return {
        "loss": loss_sum / len(loader.dataset),
        "accuracy": accuracy_score(y_true, y_pred),
        "y_true": y_true,
        "y_pred": y_pred,
    }


def plot_confusion(cm: np.ndarray, labels: list[str], out_path: Path) -> None:
    fig, ax = plt.subplots(figsize=(8, 7))
    im = ax.imshow(cm, interpolation="nearest", cmap="Blues")
    fig.colorbar(im, ax=ax)
    ax.set_xticks(np.arange(len(labels)), labels=labels, rotation=45, ha="right")
    ax.set_yticks(np.arange(len(labels)), labels=labels)
    ax.set_ylabel("True label")
    ax.set_xlabel("Predicted label")
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            ax.text(j, i, str(cm[i, j]), ha="center", va="center", color="black")
    fig.tight_layout()
    fig.savefig(out_path, dpi=160)
    plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser(description="Train a spectrogram CNN for RFUAV classification.")
    parser.add_argument("--processed-dir", default="/data/RFUAV_processed")
    parser.add_argument("--out-dir", default="/data/checkpoints")
    parser.add_argument("--results-dir", default="/data/results")
    parser.add_argument("--model", choices=["small_cnn", "resnet18", "efficientnet_b0"], default="small_cnn")
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--num-workers", type=int, default=0)
    parser.add_argument("--no-amp", action="store_true", help="Disable CUDA automatic mixed precision.")
    parser.add_argument("--no-preload", action="store_true", help="Read .npz files lazily instead of preloading tensors into RAM.")
    args = parser.parse_args()

    processed_dir = Path(args.processed_dir)
    out_dir = Path(args.out_dir)
    results_dir = Path(args.results_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    results_dir.mkdir(parents=True, exist_ok=True)

    labels = read_labels(processed_dir)
    preload = not args.no_preload
    train_ds = SpectrogramDataset(processed_dir, "train", preload=preload)
    val_ds = SpectrogramDataset(processed_dir, "val", preload=preload)
    test_ds = SpectrogramDataset(processed_dir, "test", preload=preload)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    use_cuda = device.type == "cuda"
    if use_cuda:
        torch.backends.cudnn.benchmark = True
    loader_kwargs = {
        "batch_size": args.batch_size,
        "num_workers": args.num_workers,
        "pin_memory": use_cuda,
    }
    if args.num_workers > 0:
        loader_kwargs["persistent_workers"] = True
        loader_kwargs["prefetch_factor"] = 4

    train_loader = DataLoader(train_ds, shuffle=True, **loader_kwargs)
    val_loader = DataLoader(val_ds, shuffle=False, **loader_kwargs)
    test_loader = DataLoader(test_ds, shuffle=False, **loader_kwargs)

    model = build_model(args.model, len(labels)).to(device)
    if use_cuda:
        model = model.to(memory_format=torch.channels_last)
    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
    criterion = nn.CrossEntropyLoss()
    use_amp = use_cuda and not args.no_amp
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)

    best_val_acc = -1.0
    history = []
    best_path = out_dir / "best_model.pt"

    for epoch in range(1, args.epochs + 1):
        model.train()
        train_loss = 0.0
        for x, y in tqdm(train_loader, desc=f"epoch {epoch}/{args.epochs}"):
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
            if use_cuda:
                x = x.contiguous(memory_format=torch.channels_last)
            optimizer.zero_grad(set_to_none=True)
            with torch.amp.autocast("cuda", enabled=use_amp):
                logits = model(x)
                loss = criterion(logits, y)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            train_loss += loss.item() * x.size(0)

        train_loss /= len(train_loader.dataset)
        val = evaluate(model, val_loader, device)
        row = {"epoch": epoch, "train_loss": train_loss, "val_loss": val["loss"], "val_accuracy": val["accuracy"]}
        history.append(row)
        print(row)

        if val["accuracy"] > best_val_acc:
            best_val_acc = val["accuracy"]
            torch.save({"model": model.state_dict(), "labels": labels, "model_name": args.model}, best_path)

    checkpoint = torch.load(best_path, map_location=device)
    model.load_state_dict(checkpoint["model"])
    test = evaluate(model, test_loader, device)
    cm = confusion_matrix(test["y_true"], test["y_pred"], labels=list(range(len(labels))))
    report = classification_report(test["y_true"], test["y_pred"], target_names=labels, output_dict=True, zero_division=0)

    metrics = {
        "model": args.model,
        "best_val_accuracy": best_val_acc,
        "test_accuracy": test["accuracy"],
        "history": history,
        "classification_report": report,
    }
    (results_dir / "metrics.json").write_text(json.dumps(metrics, indent=2))
    plot_confusion(cm, labels, results_dir / "confusion_matrix.png")

    print(json.dumps({"best_val_accuracy": best_val_acc, "test_accuracy": test["accuracy"]}, indent=2))
    print(f"Saved checkpoint: {best_path}")
    print(f"Saved metrics: {results_dir / 'metrics.json'}")


if __name__ == "__main__":
    main()