"""Shared training code of the window networks (offsets, detector): canvas data from a Hugging Face
dataset repo, the conv + transformer encoder, the LR schedule and resumable checkpoints."""

import json
import math
import os
import random
import tarfile
import time
from pathlib import Path

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

GRID_SIZE = 7
EMBED_DIM = 192
NUM_HEADS = 6
NUM_LAYERS = 6
MLP_RATIO = 4
DROPOUT = 0.10
CONV_STEM_CHANNELS = 24
CONV_STAGE_CHANNELS = (32, 48, 64, 96, 128)
GROUP_NORM_GROUPS = 8
BATCH_SIZE = 128
WARMUP_EPOCHS = 3
LEARNING_RATE = 6.0e-4
WEIGHT_DECAY = 1.0e-4
GRADIENT_CLIP_NORM = 1.0
SEED = int(os.environ.get("SEED", "12345"))
WORK_DIR = Path("work")


def set_seed(seed=SEED):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def prepare_data(dataset_repo):
    """Download and unpack the canvas shards. -> {split: [indices]}, data directory."""
    from huggingface_hub import snapshot_download
    name = dataset_repo.replace("/", "--")
    root = Path(snapshot_download(dataset_repo, repo_type="dataset", allow_patterns=["manifest.json", "shards/*"],
                                  local_dir=str(WORK_DIR / "download" / name)))
    data = WORK_DIR / "data" / name
    data.mkdir(parents=True, exist_ok=True)
    for shard in sorted((root / "shards").glob("*.tar")):
        marker = data / (shard.name + ".done")
        if not marker.is_file():
            with tarfile.open(str(shard)) as archive:
                archive.extractall(str(data), filter="data")
            marker.touch()
    manifest = json.loads((root / "manifest.json").read_text())
    splits = {name: [i for i in range(begin, end) if (data / "{:07d}.png".format(i)).is_file()]
              for name, (begin, end) in manifest["splits"].items()}
    return splits, data


def make_loader(dataset, shuffle, device):
    workers = max(1, len(os.sched_getaffinity(0)) - 1)
    return DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=shuffle, drop_last=shuffle, num_workers=workers,
                      pin_memory=device.type == "cuda", persistent_workers=True,
                      worker_init_fn=lambda _: cv2.setNumThreads(1), prefetch_factor=4)


def autocast(device):
    return torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda")


def conv_block(in_channels, out_channels, stride):
    return nn.Sequential(nn.Conv2d(in_channels, out_channels, 3, stride, 1, bias=False),
                         nn.GroupNorm(GROUP_NORM_GROUPS, out_channels), nn.GELU())


class ConvEncoder(nn.Module):
    """224 x 224 -> EMBED_DIM x 7 x 7: a full-resolution 3x3 stem (grid lines are 1-2 px wide),
    then five stages of a stride-2 and a stride-1 3x3 convolution."""

    def __init__(self):
        super().__init__()
        layers = [conv_block(3, CONV_STEM_CHANNELS, 1)]
        channels = CONV_STEM_CHANNELS
        for stage_channels in CONV_STAGE_CHANNELS:
            layers += [conv_block(channels, stage_channels, 2), conv_block(stage_channels, stage_channels, 1)]
            channels = stage_channels
        layers.append(nn.Conv2d(channels, EMBED_DIM, 1))
        self.layers = nn.Sequential(*layers)

    def forward(self, images):
        return self.layers(images)


class WindowTransformer(nn.Module):
    """Conv encoder -> 7 x 7 tokens with learned row / column positions -> transformer -> [B, 49, EMBED_DIM].
    Subclasses add their heads as attributes, so checkpoint keys stay flat."""

    def __init__(self):
        super().__init__()
        self.patch_embedding = ConvEncoder()
        self.row_position = nn.Parameter(torch.zeros(1, GRID_SIZE, 1, EMBED_DIM))
        self.column_position = nn.Parameter(torch.zeros(1, 1, GRID_SIZE, EMBED_DIM))
        layer = nn.TransformerEncoderLayer(d_model=EMBED_DIM, nhead=NUM_HEADS, dim_feedforward=EMBED_DIM * MLP_RATIO,
                                           dropout=DROPOUT, activation="gelu", batch_first=True, norm_first=True)
        self.encoder = nn.TransformerEncoder(layer, NUM_LAYERS, enable_nested_tensor=False)
        self.output_norm = nn.LayerNorm(EMBED_DIM)

    def init_weights(self, *heads):
        """Called by subclasses after their heads exist (keeps the original initialisation order)."""
        nn.init.trunc_normal_(self.row_position, std=0.02)
        nn.init.trunc_normal_(self.column_position, std=0.02)
        for head in heads:
            for module in head:
                if isinstance(module, nn.Linear):
                    nn.init.trunc_normal_(module.weight, std=0.02)
                    nn.init.zeros_(module.bias)

    def tokens(self, images):
        tokens = self.patch_embedding(images).permute(0, 2, 3, 1) + self.row_position + self.column_position
        return self.output_norm(self.encoder(tokens.reshape(images.shape[0], GRID_SIZE * GRID_SIZE, EMBED_DIM)))


def make_optimizer(model, steps_per_epoch, epochs):
    """AdamW, linear warm-up for WARMUP_EPOCHS, then a per-step cosine."""
    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
    warmup, total = WARMUP_EPOCHS * steps_per_epoch, epochs * steps_per_epoch

    def factor(step):
        if step < warmup:
            return (step + 1) / warmup
        return 0.5 * (1.0 + math.cos(math.pi * min(1.0, (step - warmup) / max(1, total - warmup))))
    return optimizer, torch.optim.lr_scheduler.LambdaLR(optimizer, factor)


class Run:
    """Checkpoints in work/runs/<name>; with MODEL_REPO they are uploaded to runs/<name> and a
    stopped job resumes from the uploaded last.pt."""

    def __init__(self, name, model_repo):
        self.name, self.repo = name, model_repo
        self.dir = WORK_DIR / "runs" / name
        self.dir.mkdir(parents=True, exist_ok=True)
        self.api = None
        if model_repo:
            from huggingface_hub import HfApi, hf_hub_download
            self.api = HfApi()
            self.api.create_repo(model_repo, repo_type="model", private=True, exist_ok=True)
            remote = set(self.api.list_repo_files(model_repo))
            for file in ("last.pt", "best.pt"):
                path = "runs/{}/{}".format(name, file)
                if path in remote and not (self.dir / file).is_file():
                    local = hf_hub_download(model_repo, path, local_dir=str(WORK_DIR / "resume"))
                    (self.dir / file).write_bytes(Path(local).read_bytes())

    def resume(self, model, optimizer, scheduler, device):
        """-> (first epoch, best score) from last.pt, or (1, None)."""
        path = self.dir / "last.pt"
        if not path.is_file():
            return 1, None
        checkpoint = torch.load(str(path), map_location=device, weights_only=False)
        model.load_state_dict(checkpoint["model"])
        optimizer.load_state_dict(checkpoint["optimizer"])
        scheduler.load_state_dict(checkpoint["scheduler"])
        return int(checkpoint["epoch"]) + 1, checkpoint["best"]

    def save(self, file, model, optimizer, scheduler, epoch, best):
        payload = {"epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict(),
                   "scheduler": scheduler.state_dict(), "best": best}
        torch.save(payload, str(self.dir / (file + ".tmp")))
        os.replace(str(self.dir / (file + ".tmp")), str(self.dir / file))

    def push(self, message):
        if self.api is None:
            return
        for attempt in range(5):
            try:
                self.api.upload_folder(folder_path=str(self.dir), path_in_repo="runs/" + self.name, repo_id=self.repo,
                                       repo_type="model", allow_patterns=["*.pt", "*.csv"], commit_message=message)
                return
            except Exception as error:
                print("upload failed ({}), retry {}".format(error, attempt + 1), flush=True)
                time.sleep(20 * (attempt + 1))


def log(run, row):
    """Print a row and append it to history.csv."""
    path = run.dir / "history.csv"
    new = not path.is_file()
    with path.open("a", encoding="utf-8") as output:
        if new:
            output.write(",".join(row) + "\n")
        output.write(",".join("{:.6g}".format(v) if isinstance(v, float) else str(v) for v in row.values()) + "\n")
    print(" ".join("{}={}".format(k, "{:.4g}".format(v) if isinstance(v, float) else v) for k, v in row.items()), flush=True)


def device_or_fail():
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is not available")
    return torch.device("cuda")