"""Symbioza example job: train.

An unattended training run that hands back everything you need afterwards: a small ResNet trained
on Fashion-MNIST (MIT license; downloaded from a pinned commit, 3 attempts, sha256-verified), a
checkpoint after every epoch, metrics.csv, the final model, a confusion matrix (CSV and PNG) and
results.json. Seeded, with PyTorch's deterministic mode requested.

The schedule is sized once, after the first epoch: EPOCHS, or fewer when the first epoch's time
says EPOCHS would not fit in TRAIN_SECONDS, so a slow GPU runs a complete, shorter schedule instead
of being cut off partway through one. Epoch 1 warms the learning rate up; the rest decay it by cosine.
On any failure it writes error.json; SHA256SUMS is written on every exit path.
Local testing only (the published spec sets none of these): CPU_OK=1 runs without CUDA,
LIMIT=n trains and tests on the first n images, WIDTH=n shrinks the network.
"""
import gzip
import hashlib
import json
import math
import os
import re
import shutil
import signal
import struct
import sys
import time
import traceback
import urllib.request
import warnings
import zlib

os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")  # needed for deterministic cuBLAS
warnings.filterwarnings("ignore", message=".*cuda capability")  # torch names the card in this warning
import torch  # noqa: E402
import torch.nn as nn  # noqa: E402
import torch.nn.functional as F  # noqa: E402

OUT = os.environ.get("ARTIFACTS_DIR", "/workspace/artifacts")
WORK = os.environ.get("WORK_DIR", "/tmp/fashion-mnist")
EPOCHS = int(os.environ.get("EPOCHS", "15"))
TRAIN_SECONDS = float(os.environ.get("TRAIN_SECONDS", "540"))
WIDTH = int(os.environ.get("WIDTH", "64"))
LIMIT = int(os.environ.get("LIMIT", "0"))
BATCH = 256
MAX_LR = 0.2
SEED = 1234
MIRROR = ("https://raw.githubusercontent.com/zalandoresearch/fashion-mnist/"
          "b2617bb6d3ffa2e429640350f613e3291e10b141/data/fashion/")
FILES = {
    "train-images-idx3-ubyte.gz": "3aede38d61863908ad78613f6a32ed271626dd12800ba2636569512369268a84",
    "train-labels-idx1-ubyte.gz": "a04f17134ac03560a47e3764e11b92fc97de4d1bfaf8ba1a3aa29af54cc90845",
    "t10k-images-idx3-ubyte.gz": "346e55b948d973a97e58d2351dde16a484bd415d4595297633bb08f03db6a073",
    "t10k-labels-idx1-ubyte.gz": "67da17c76eaffca5446c3361aaab5c3cd6d1c2608764d35dfb1850b086bf8dd5",
}
CLASSES = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
           "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]
MEAN, STD = 0.2860, 0.3530
STAGE = "start"


class Fail(Exception):
    """A failure the job names itself: an exit code and the reason."""

    def __init__(self, code, msg):
        super().__init__(msg)
        self.code = code


def stage(name):
    global STAGE
    STAGE = name
    print(f"== {name}", flush=True)


def sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for block in iter(lambda: fh.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


def write_sha256sums():
    with open(os.path.join(OUT, "SHA256SUMS"), "w") as fh:
        for name in sorted(os.listdir(OUT)):
            path = os.path.join(OUT, name)
            if name != "SHA256SUMS" and os.path.isfile(path):
                fh.write(f"{sha256(path)}  {name}\n")


def redact(text):
    return re.sub(r"\d+(\.\d+)? ?[KMGT]i?B\b", "<size>", text)  # an out-of-memory message states the card's memory


def run(main, job):
    """Run main(). On ANY failure leave error.json; write SHA256SUMS on every exit path."""
    signal.signal(signal.SIGTERM, lambda *_: sys.exit(143))
    code, err = 0, None
    try:
        main()
    except Fail as e:
        code, err = e.code, str(e)
    except SystemExit as e:
        code = e.code if isinstance(e.code, int) else 1
        err = "stopped by SIGTERM" if code == 143 else f"exit {e.code}"
    except BaseException as e:
        code, err = 1, f"{type(e).__name__}: {e}"
        print(redact(traceback.format_exc()), flush=True)
    os.makedirs(OUT, exist_ok=True)
    for name in os.listdir(OUT):
        if name.endswith(".part"):  # a save the failure interrupted
            os.remove(os.path.join(OUT, name))
    if code:
        err = redact(err)[:2000]
        print(f"{job} FAILED during '{STAGE}' (exit {code}): {err}", flush=True)
        with open(os.path.join(OUT, "error.json"), "w") as fh:
            json.dump({"ok": False, "job": job, "stage": STAGE, "error": err, "exit_code": code}, fh, indent=2)
    write_sha256sums()
    sys.exit(code)


def fetch(name, url, digest):
    """Download into WORK (3 attempts, 60 s timeout each); refuse bytes whose sha256 is not the pinned one."""
    path, err = os.path.join(WORK, name), None
    for attempt in range(3):
        time.sleep(10 * attempt)
        try:
            with urllib.request.urlopen(url, timeout=60) as r, open(path, "wb") as fh:
                shutil.copyfileobj(r, fh)
        except Exception as e:
            err = f"{type(e).__name__}: {e}"
            continue
        got = sha256(path)
        if got == digest:
            return path
        err = f"sha256 {got} is not the pinned {digest}"
    raise Fail(2, f"{name} from {url}: {err} (3 attempts)")


def read_idx(path):
    blob = gzip.open(path).read()
    magic, n = struct.unpack(">II", blob[:8])
    if magic == 2051:
        rows, cols = struct.unpack(">II", blob[8:16])
        return torch.frombuffer(bytearray(blob[16:]), dtype=torch.uint8).reshape(n, 1, rows, cols)
    if magic == 2049:
        return torch.frombuffer(bytearray(blob[8:]), dtype=torch.uint8).long()
    raise Fail(2, f"{path}: not an idx file (magic {magic})")


class Block(nn.Module):
    def __init__(self, cin, cout, stride):
        super().__init__()
        self.c1, self.b1 = nn.Conv2d(cin, cout, 3, stride, 1, bias=False), nn.BatchNorm2d(cout)
        self.c2, self.b2 = nn.Conv2d(cout, cout, 3, 1, 1, bias=False), nn.BatchNorm2d(cout)
        self.sc = nn.Identity() if stride == 1 and cin == cout else nn.Sequential(
            nn.Conv2d(cin, cout, 1, stride, bias=False), nn.BatchNorm2d(cout))

    def forward(self, x):
        return F.relu(self.b2(self.c2(F.relu(self.b1(self.c1(x))))) + self.sc(x))


class Net(nn.Module):
    def __init__(self, w):
        super().__init__()
        self.stem = nn.Sequential(nn.Conv2d(1, w, 3, 1, 1, bias=False), nn.BatchNorm2d(w), nn.ReLU())
        self.body = nn.Sequential(Block(w, w, 1), Block(w, w, 1), Block(w, 2 * w, 2),
                                  Block(2 * w, 2 * w, 1), Block(2 * w, 4 * w, 2), Block(4 * w, 4 * w, 1))
        self.head = nn.Linear(4 * w, 10)

    def forward(self, x):
        return self.head(self.body(self.stem(x)).mean((2, 3)))  # mean: deterministic backward


def lr_at(step, steps, planned):
    """Epoch 1 warms up linearly; the remaining planned epochs decay by cosine to 0."""
    if step < steps:
        return MAX_LR * (step + 1) / steps
    t = (step - steps) / max(1, (planned - 1) * steps)
    return MAX_LR * 0.5 * (1 + math.cos(math.pi * min(1.0, t)))


def augment(x, g):
    """Random horizontal flip per image, random shift of up to 2 px per batch."""
    flip = (torch.rand(x.shape[0], generator=g) < 0.5).to(x.device)[:, None, None, None]
    x = torch.where(flip, x.flip(3), x)
    dy, dx = torch.randint(0, 5, (2,), generator=g).tolist()
    return F.pad(x, (2, 2, 2, 2), value=-MEAN / STD)[:, :, dy:dy + 28, dx:dx + 28]


def train_epoch(model, x, y, opt, g, step, lr_of):
    model.train()
    order = torch.randperm(len(x), generator=g).to(x.device)
    loss_sum = torch.zeros((), device=x.device)
    correct = torch.zeros((), device=x.device, dtype=torch.long)
    for i in range(0, len(x), BATCH):
        for group in opt.param_groups:
            group["lr"] = lr_of(step)
        idx = order[i:i + BATCH]
        out = model(augment(x[idx], g))
        loss = F.cross_entropy(out, y[idx])
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()
        step += 1
        loss_sum += loss.detach() * len(idx)
        correct += (out.argmax(1) == y[idx]).sum()
    return loss_sum.item() / len(x), correct.item() / len(x), step


@torch.no_grad()
def evaluate(model, x, y):
    model.eval()
    loss, preds = torch.zeros((), device=x.device), []
    for i in range(0, len(x), 1000):
        out = model(x[i:i + 1000])
        loss += F.cross_entropy(out, y[i:i + 1000], reduction="sum")
        preds.append(out.argmax(1))
    p = torch.cat(preds)
    return loss.item() / len(x), (p == y).float().mean().item(), p


FONT = ["111101101101111", "010110010010111", "111001111100111", "111001111001111", "101101111001001",
        "111100111001111", "111100111101111", "111001001001001", "111101111101111", "111101111001111"]


def write_png(path, w, h, px):
    raw = b"".join(b"\x00" + bytes(px[y * w * 3:(y + 1) * w * 3]) for y in range(h))

    def chunk(tag, data):
        return struct.pack(">I", len(data)) + tag + data + struct.pack(">I", zlib.crc32(tag + data))
    with open(path, "wb") as fh:
        fh.write(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
                 + chunk(b"IDAT", zlib.compress(raw, 9)) + chunk(b"IEND", b""))


def confusion_png(cm, path, cell=40, m=24, s=3):
    """Rows = true class, columns = predicted, 0-9 in CLASSES order; shade = share of the row."""
    size = m + 10 * cell
    px = bytearray(b"\xff" * size * size * 3)

    def fill(x0, y0, x1, y1, rgb):
        for yy in range(y0, y1):
            px[(yy * size + x0) * 3:(yy * size + x1) * 3] = bytes(rgb) * (x1 - x0)
    for i in range(10):
        total = max(1, sum(cm[i]))
        for j in range(10):
            f = cm[i][j] / total
            rgb = [round(255 + (c - 255) * f) for c in (8, 48, 107)]
            fill(m + j * cell + 1, m + i * cell + 1, m + (j + 1) * cell, m + (i + 1) * cell, rgb)
        for ox, oy in ((5, m + i * cell + 12), (m + i * cell + 15, 5)):  # row label, column label
            for k, bit in enumerate(FONT[i]):
                if bit == "1":
                    x, y = ox + (k % 3) * s, oy + (k // 3) * s
                    fill(x, y, x + s, y + s, (0, 0, 0))
    write_png(path, size, size, px)


def main():
    t_start = time.time()
    stage("check CUDA")
    cuda = torch.cuda.is_available()
    if not cuda and os.environ.get("CPU_OK") != "1":
        raise Fail(3, "torch.cuda.is_available() is False; this job requires a GPU")
    dev = torch.device("cuda" if cuda else "cpu")
    torch.manual_seed(SEED)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False
    torch.use_deterministic_algorithms(True, warn_only=True)
    os.makedirs(OUT, exist_ok=True)
    os.makedirs(WORK, exist_ok=True)

    stage("download and verify Fashion-MNIST")
    xtr, ytr, xte, yte = (read_idx(fetch(n, MIRROR + n, d)) for n, d in FILES.items())
    if (len(xtr), len(xte)) != (60000, 10000):
        raise Fail(2, f"unexpected dataset size: {len(xtr)} train / {len(xte)} test")
    if LIMIT:
        xtr, ytr, xte, yte = xtr[:LIMIT], ytr[:LIMIT], xte[:LIMIT], yte[:LIMIT]
    xtr, xte = (((x.float() / 255 - MEAN) / STD).to(dev) for x in (xtr, xte))
    ytr, yte = ytr.to(dev), yte.to(dev)
    t_data = time.time()

    model = Net(WIDTH).to(dev)
    steps = (len(xtr) + BATCH - 1) // BATCH
    opt = torch.optim.SGD(model.parameters(), lr=MAX_LR, momentum=0.9, nesterov=True, weight_decay=5e-4)
    g = torch.Generator().manual_seed(SEED)
    config = {"epochs_max": EPOCHS, "train_seconds": TRAIN_SECONDS, "width": WIDTH, "batch": BATCH,
              "max_lr": MAX_LR, "seed": SEED, "limit": LIMIT}
    ckpt = os.path.join(OUT, "checkpoint.pt")
    planned, epoch, step, history = EPOCHS, 0, 0, []
    while epoch < planned:
        epoch += 1
        stage(f"epoch {epoch}/{planned}")
        t_epoch = time.time()
        lr = lr_at(step, steps, planned)
        train_loss, train_acc, step = train_epoch(model, xtr, ytr, opt, g, step,
                                                  lambda s: lr_at(s, steps, planned))
        test_loss, test_acc, _ = evaluate(model, xte, yte)
        seconds = time.time() - t_epoch
        if epoch == 1:  # size the schedule once, from the first epoch's time
            planned = min(EPOCHS, max(2, int(TRAIN_SECONDS // seconds)))
        row = {"epoch": epoch, "lr": round(lr, 6), "train_loss": round(train_loss, 5),
               "train_acc": round(train_acc, 5), "test_loss": round(test_loss, 5),
               "test_acc": round(test_acc, 5), "epoch_seconds": round(seconds, 2)}
        history.append(row)
        print(json.dumps(row), flush=True)
        torch.save({"epoch": epoch, "config": config, "model": model.state_dict(), "opt": opt.state_dict(),
                    "history": history}, ckpt + ".part")
        os.replace(ckpt + ".part", ckpt)
        with open(os.path.join(OUT, "metrics.csv"), "w") as fh:
            fh.write(",".join(row) + "\n")
            fh.writelines(",".join(str(r[k]) for k in row) + "\n" for r in history)

    stage("final evaluation and outputs")
    _, final_acc, pred = evaluate(model, xte, yte)
    cm = torch.bincount((yte * 10 + pred).cpu(), minlength=100).reshape(10, 10).tolist()
    torch.save(model.state_dict(), os.path.join(OUT, "model.pt"))
    with open(os.path.join(OUT, "confusion_matrix.csv"), "w") as fh:
        fh.write("true\\predicted," + ",".join(CLASSES) + "\n")
        fh.writelines(CLASSES[i] + "," + ",".join(map(str, cm[i])) + "\n" for i in range(10))
    confusion_png(cm, os.path.join(OUT, "confusion_matrix.png"))
    best = max(history, key=lambda r: r["test_acc"])
    results = {
        "dataset": {"name": "Fashion-MNIST", "license": "MIT", "source": MIRROR, "sha256": FILES,
                    "train_images": len(xtr), "test_images": len(xte)},
        "model": {"arch": f"ResNet, 6 basic blocks, width {WIDTH}",
                  "parameters": sum(p.numel() for p in model.parameters()), "file": "model.pt (state_dict)"},
        "config": config,
        "schedule": "SGD, Nesterov momentum 0.9, weight decay 5e-4; epoch 1 warms the learning rate up "
                    "linearly to max_lr, the remaining planned epochs decay it by cosine to 0",
        "epochs_planned": planned,
        "epochs_planned_basis": "epochs_max, or fewer when the first epoch's time says epochs_max would not "
                                "fit in train_seconds",
        "epochs_completed": epoch,
        "final_test_accuracy": round(final_acc, 5), "best_test_accuracy": best["test_acc"],
        "best_epoch": best["epoch"],
        "per_class_test_accuracy": {CLASSES[i]: round(cm[i][i] / max(1, sum(cm[i])), 4) for i in range(10)},
        "seconds": {"download_decode": round(t_data - t_start, 2), "train": round(time.time() - t_data, 2)},
        "device": dev.type, "torch": torch.__version__, "deterministic_mode_requested": True,
    }
    if LIMIT:
        results["LOCAL_TEST_SUBSET"] = LIMIT
    with open(os.path.join(OUT, "results.json"), "w") as fh:
        json.dump(results, fh, indent=2)
    print(json.dumps(results, indent=2), flush=True)


if __name__ == "__main__":
    run(main, "train")