#!/usr/bin/env python3
"""Train the grid-node detector (models/detector.onnx) on v5 canvases.

    DATASET_REPO=<user>/ecg-canvases-v5 MODEL_REPO=<user>/ecg-models RUN_NAME=detector python3 train_detector.py

Head per token: 4 x 4 cells x (heat logit, offset x, offset y) -> a 28 x 28 cell map of the 224 window.
Loss: CenterNet focal loss on the heat map + L1 of the offsets in the cells holding a node, both
normalised by the number of nodes. The checkpoint with the best validation F1 is kept as best.pt.
"""

import math
import os

import numpy as np
import torch
from torch import nn

import common
import crops
import detector_data as dd

DATASET_REPO = os.environ.get("DATASET_REPO", "")
MODEL_REPO = os.environ.get("MODEL_REPO", "")
RUN_NAME = os.environ.get("RUN_NAME", "detector")
EPOCHS = int(os.environ.get("EPOCHS", "156"))
EVAL_EVERY = 2
VAL_SAMPLES = 1000
SUB = crops.PATCH_SIZE // dd.CELL
HEAD_HIDDEN = 384
FOCAL_ALPHA, FOCAL_BETA = 2.0, 4.0
HEAT_PRIOR = 0.05
THRESHOLDS = (0.2, 0.3, 0.4, 0.5)
MATCH_RADIUS_PX = 4.0


class PointViT(common.WindowTransformer):
    def __init__(self):
        super().__init__()
        self.head = nn.Sequential(nn.Linear(common.EMBED_DIM, HEAD_HIDDEN), nn.GELU(), nn.Linear(HEAD_HIDDEN, SUB * SUB * 3))
        self.init_weights(self.head)
        with torch.no_grad():
            self.head[-1].bias.view(SUB * SUB, 3)[:, 0] = -math.log((1 - HEAT_PRIOR) / HEAT_PRIOR)

    def forward(self, images):
        """-> [B, 28, 28, 3] (heat logit, offset logits)."""
        b, g = images.shape[0], common.GRID_SIZE
        out = self.head(self.tokens(images)).view(b, g, g, SUB, SUB, 3)
        return out.permute(0, 1, 3, 2, 4, 5).reshape(b, dd.CELLS, dd.CELLS, 3)


def loss_fn(output, target):
    logits, offsets = output[..., 0].float(), torch.sigmoid(output[..., 1:3].float())
    heat, positive = target[..., dd.CH_HEAT], target[..., dd.CH_POS]
    count = positive.sum().clamp_min(1.0)
    p = torch.sigmoid(logits)
    pos_term = -((1 - p) ** FOCAL_ALPHA) * nn.functional.logsigmoid(logits) * positive
    neg_term = -((1 - heat) ** FOCAL_BETA) * (p ** FOCAL_ALPHA) * nn.functional.logsigmoid(-logits) * (1 - positive)
    focal = (pos_term.sum() + neg_term.sum()) / count
    offset = (torch.abs(offsets - target[..., dd.CH_OFF_X:dd.CH_OFF_Y + 1]).sum(-1) * positive).sum() / count
    return focal + offset


@torch.no_grad()
def validate(model, loader, device):
    model.eval()
    heat, offsets, targets, kinds = [], [], [], []
    for images, target, kind in loader:
        with common.autocast(device):
            output = model(crops.normalize_images(images.to(device, non_blocking=True))).float()
        heat.append(torch.sigmoid(output[..., 0]).cpu().numpy())
        offsets.append(torch.sigmoid(output[..., 1:3]).cpu().numpy())
        targets.append(target.numpy())
        kinds.append(kind.numpy())
    return dd.f1_scores(np.concatenate(heat), np.concatenate(offsets), np.concatenate(targets), np.concatenate(kinds),
                        THRESHOLDS, MATCH_RADIUS_PX)


def main():
    common.set_seed()
    device = common.device_or_fail()
    splits, data = common.prepare_data(DATASET_REPO)
    train_loader = common.make_loader(dd.PointDataset(data, dd.train_selection(data, splits["train"])), True, device)
    val_loader = common.make_loader(dd.PointDataset(data, splits["val"][:VAL_SAMPLES], train=False), False, device)
    model = PointViT().to(device)
    optimizer, scheduler = common.make_optimizer(model, len(train_loader), EPOCHS)
    run = common.Run(RUN_NAME, MODEL_REPO)
    first, best = run.resume(model, optimizer, scheduler, device)
    best = -1.0 if best is None else best
    for epoch in range(first, EPOCHS + 1):
        model.train()
        loss_sum, batches = 0.0, 0
        for images, target, _ in train_loader:
            images = crops.normalize_images(images.to(device, non_blocking=True))
            target = target.to(device, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            with common.autocast(device):
                output = model(images)
            loss = loss_fn(output, target)
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), common.GRADIENT_CLIP_NORM)
            optimizer.step()
            scheduler.step()
            loss_sum, batches = loss_sum + loss.item(), batches + 1
        row = {"epoch": epoch, "train_loss": loss_sum / batches, "val_f1": ""}
        if epoch % EVAL_EVERY == 0 or epoch == EPOCHS:
            row["val_f1"] = f1 = validate(model, val_loader, device)["best_f1"]
            if f1 > best:
                best = f1
                run.save("best.pt", model, optimizer, scheduler, epoch, best)
        row["best_f1"] = best
        run.save("last.pt", model, optimizer, scheduler, epoch, best)
        common.log(run, row)
        if epoch % 10 == 0 or epoch == EPOCHS:
            run.push("epoch {}".format(epoch))


if __name__ == "__main__":
    main()