import copy
import csv
import os
import warnings
from argparse import ArgumentParser
from typing import cast

import torch
import tqdm
import yaml
from torch.utils import data
from torch.amp.autocast_mode import autocast

from nets import nn
from utils import util
from utils.dataset import Dataset

warnings.filterwarnings("ignore")

data_dir = "/home/image1325/ssd1/dataset/coco"


def train(args, params):
    # Model
    model = nn.yolo_v11_n(len(params["names"]))
    model.cuda()

    # Optimizer
    accumulate = max(round(64 / (args.batch_size * args.world_size)), 1)
    params["weight_decay"] *= args.batch_size * args.world_size * accumulate / 64

    optimizer = torch.optim.SGD(
        util.set_params(model, params["weight_decay"]), params["min_lr"], params["momentum"], nesterov=True
    )

    # EMA
    ema = util.EMA(model) if args.local_rank == 0 else None

    filenames = []
    with open(f"{data_dir}/train2017.txt") as f:
        for filename in f.readlines():
            filename = os.path.basename(filename.rstrip())
            filenames.append(f"{data_dir}/images/train2017/" + filename)

    sampler = None
    dataset = Dataset(filenames, args.input_size, params, augment=True)

    if args.distributed:
        sampler = data.DistributedSampler(dataset)

    loader = data.DataLoader(
        dataset,
        args.batch_size,
        sampler is None,
        sampler,
        num_workers=8,
        pin_memory=True,
        collate_fn=Dataset.collate_fn,
    )

    # Scheduler
    num_steps = len(loader)
    scheduler = util.LinearLR(args, params, num_steps)

    if args.distributed:
        # DDP mode
        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
        model = torch.nn.parallel.DistributedDataParallel(
            module=model, device_ids=[args.local_rank], output_device=args.local_rank
        )

    best = 0
    amp_scale = torch.amp.grad_scaler.GradScaler()
    criterion = util.ComputeLoss(model, params)

    with open("weights/step.csv", "w") as log:
        if args.local_rank == 0:
            logger = csv.DictWriter(
                log, fieldnames=["epoch", "box", "cls", "dfl", "Recall", "Precision", "mAP@50", "mAP"]
            )
            logger.writeheader()

        for epoch in range(args.epochs):
            model.train()
            if args.distributed and sampler:
                sampler.set_epoch(epoch)
            if args.epochs - epoch == 10:
                ds = cast(Dataset, loader.dataset)
                ds.mosaic = False

            p_bar = enumerate(loader)

            if args.local_rank == 0:
                print(("\n" + "%10s" * 5) % ("epoch", "memory", "box", "cls", "dfl"))
                p_bar = tqdm.tqdm(p_bar, total=num_steps, ascii=" >-")

            optimizer.zero_grad()
            avg_box_loss = util.AverageMeter()
            avg_cls_loss = util.AverageMeter()
            avg_dfl_loss = util.AverageMeter()
            for i, (samples, targets) in p_bar:
                step = i + num_steps * epoch
                scheduler.step(step, optimizer)

                samples = samples.cuda().float() / 255

                # Forward
                with autocast("cuda"):
                    outputs = model(samples)  # forward
                    loss_box, loss_cls, loss_dfl = criterion(outputs, targets)

                avg_box_loss.update(loss_box.item(), samples.size(0))
                avg_cls_loss.update(loss_cls.item(), samples.size(0))
                avg_dfl_loss.update(loss_dfl.item(), samples.size(0))

                loss_box *= args.batch_size  # loss scaled by batch_size
                loss_cls *= args.batch_size  # loss scaled by batch_size
                loss_dfl *= args.batch_size  # loss scaled by batch_size
                loss_box *= args.world_size  # gradient averaged between devices in DDP mode
                loss_cls *= args.world_size  # gradient averaged between devices in DDP mode
                loss_dfl *= args.world_size  # gradient averaged between devices in DDP mode

                # Backward
                amp_scale.scale(loss_box + loss_cls + loss_dfl).backward()

                # Optimize
                if step % accumulate == 0:
                    # amp_scale.unscale_(optimizer)  # unscale gradients
                    # util.clip_gradients(model)  # clip gradients
                    amp_scale.step(optimizer)  # optimizer.step
                    amp_scale.update()
                    optimizer.zero_grad()
                    if ema:
                        ema.update(model)

                torch.cuda.synchronize()

                # Log
                if args.local_rank == 0:
                    memory = f"{torch.cuda.memory_reserved() / 1e9:.4g}G"  # (GB)
                    s = ("%10s" * 2 + "%10.3g" * 3) % (
                        f"{epoch + 1}/{args.epochs}",
                        memory,
                        avg_box_loss.avg,
                        avg_cls_loss.avg,
                        avg_dfl_loss.avg,
                    )
                    p_bar = cast(tqdm.tqdm, p_bar)
                    p_bar.set_description(s)

            if args.local_rank == 0:
                # mAP
                last = test(args, params, ema.ema if ema else None)

                logger.writerow(
                    {
                        "epoch": str(epoch + 1).zfill(3),
                        "box": str(f"{avg_box_loss.avg:.3f}"),
                        "cls": str(f"{avg_cls_loss.avg:.3f}"),
                        "dfl": str(f"{avg_dfl_loss.avg:.3f}"),
                        "mAP": str(f"{last[0]:.3f}"),
                        "mAP@50": str(f"{last[1]:.3f}"),
                        "Recall": str(f"{last[2]:.3f}"),
                        "Precision": str(f"{last[3]:.3f}"),
                    }
                )
                log.flush()

                # Update best mAP
                if last[0] > best:
                    best = last[0]

                # Save model
                save = {"epoch": epoch + 1, "model": copy.deepcopy(ema.ema if ema else None)}

                # Save last, best and delete
                torch.save(save, f="./weights/last.pt")
                if best == last[0]:
                    torch.save(save, f="./weights/best.pt")
                del save

    if args.local_rank == 0:
        util.strip_optimizer("./weights/best.pt")  # strip optimizers
        util.strip_optimizer("./weights/last.pt")  # strip optimizers


@torch.no_grad()
def test(args, params, model=None):
    filenames = []
    with open(f"{data_dir}/val2017.txt") as f:
        for filename in f.readlines():
            filename = os.path.basename(filename.rstrip())
            filenames.append(f"{data_dir}/images/val2017/" + filename)

    dataset = Dataset(filenames, args.input_size, params, augment=False)
    loader = data.DataLoader(
        dataset, batch_size=4, shuffle=False, num_workers=4, pin_memory=True, collate_fn=Dataset.collate_fn
    )

    plot = False
    if not model:
        plot = True
        model = torch.load(f="./weights/best.pt", map_location="cuda", weights_only=False)
        model = model["model"].float().fuse()

    model.half()
    model.eval()

    # Configure
    iou_v = torch.linspace(start=0.5, end=0.95, steps=10).cuda()  # iou vector for mAP@0.5:0.95
    n_iou = iou_v.numel()

    m_pre = 0
    m_rec = 0
    map50 = 0
    mean_ap = 0
    metrics = []
    p_bar = tqdm.tqdm(loader, desc=("%10s" * 5) % ("", "precision", "recall", "mAP50", "mAP"), ascii=" >-")
    for samples, targets in p_bar:
        samples = samples.cuda()
        samples = samples.half()  # uint8 to fp16/32
        samples = samples / 255.0  # 0 - 255 to 0.0 - 1.0
        _, _, h, w = samples.shape  # batch-size, channels, height, width
        scale = torch.tensor((w, h, w, h)).cuda()
        # Inference
        outputs = model(samples)
        # NMS
        outputs = util.non_max_suppression(outputs)
        # Metrics
        for i, output in enumerate(outputs):
            # Ensure idx is a 1D boolean mask (squeeze any trailing dimension) to match cls/box shapes
            idx = targets["idx"]
            if idx.dim() > 1:
                idx = idx.squeeze(-1)
            idx = idx == i

            # XXX: initially, the code was like below, which caused shape mismatch when idx has extra dimension
            # idx = targets["idx"] == i
            cls = targets["cls"][idx]
            box = targets["box"][idx]

            cls = cls.cuda()
            box = box.cuda()

            metric = torch.zeros(output.shape[0], n_iou, dtype=torch.bool).cuda()

            if output.shape[0] == 0:
                if cls.shape[0]:
                    metrics.append((metric, *torch.zeros((2, 0)).cuda(), cls.squeeze(-1)))
                continue
            # Evaluate
            if cls.shape[0]:
                target = torch.cat(tensors=(cls, util.wh2xy(box) * scale), dim=1)
                metric = util.compute_metric(output[:, :6], target, iou_v)
            # Append
            metrics.append((metric, output[:, 4], output[:, 5], cls.squeeze(-1)))

    # Compute metrics
    metrics = [torch.cat(x, dim=0).cpu().numpy() for x in zip(*metrics)]  # to numpy
    if len(metrics) and metrics[0].any():
        tp, fp, m_pre, m_rec, map50, mean_ap = util.compute_ap(*metrics, plot=plot, names=params["names"])
    # Print results
    print(("%10s" + "%10.3g" * 4) % ("", m_pre, m_rec, map50, mean_ap))
    # Return results
    model.float()  # for training
    return mean_ap, map50, m_rec, m_pre


def profile(args, params):
    import thop

    shape = (1, 3, args.input_size, args.input_size)
    model = nn.yolo_v11_n(len(params["names"])).fuse()

    model.eval()
    model(torch.zeros(shape))

    x = torch.empty(shape)
    flops, num_params = thop.profile(model, inputs=[x], verbose=False)
    flops, num_params = thop.clever_format(nums=[2 * flops, num_params], format="%.3f")

    if args.local_rank == 0:
        print(f"Number of parameters: {num_params}")
        print(f"Number of FLOPs: {flops}")


def main():
    parser = ArgumentParser()
    parser.add_argument("--input-size", default=640, type=int)
    parser.add_argument("--batch-size", default=32, type=int)
    parser.add_argument("--local-rank", default=0, type=int)
    parser.add_argument("--local_rank", default=0, type=int)
    parser.add_argument("--epochs", default=600, type=int)
    parser.add_argument("--train", action="store_true")
    parser.add_argument("--test", action="store_true")

    args = parser.parse_args()

    args.local_rank = int(os.getenv("LOCAL_RANK", 0))
    args.world_size = int(os.getenv("WORLD_SIZE", 1))
    args.distributed = int(os.getenv("WORLD_SIZE", 1)) > 1

    if args.distributed:
        torch.cuda.set_device(device=args.local_rank)
        torch.distributed.init_process_group(backend="nccl", init_method="env://")

    if args.local_rank == 0:
        if not os.path.exists("weights"):
            os.makedirs("weights")

    with open("utils/args.yaml", errors="ignore") as f:
        params = yaml.safe_load(f)

    util.setup_seed()
    util.setup_multi_processes()

    profile(args, params)

    if args.train:
        train(args, params)
    if args.test:
        test(args, params)

    # Clean
    if args.distributed:
        torch.distributed.destroy_process_group()
    torch.cuda.empty_cache()


if __name__ == "__main__":
    main()