#!/usr/bin/env python3
# YOLOX-Nano exp for form-field coarse3 detection (Text / ChoiceButton / Signature).
# Finetune from COCO yolox_nano.pth. Data: COCO_DIR/{train,val}/*.jpg + annotations/instances_{train,val}.json
# Env: COCO_DIR, MAX_EPOCH, NUM_WORKERS, EVAL_INTERVAL.
import os
import torch.nn as nn
from yolox.exp import Exp as MyExp


class Exp(MyExp):
    def __init__(self):
        super().__init__()
        # ---- model: YOLOX-Nano (depthwise) ----
        self.depth = 0.33
        self.width = 0.25
        self.num_classes = 3
        self.act = "silu"

        # ---- input + multiscale (small form widgets: expose up to ~800px) ----
        self.input_size = (640, 640)
        self.test_size = (640, 640)
        self.multiscale_range = int(os.environ.get("MSR", "5"))  # 5->480..800; 2->576..704 (faster, less VRAM)

        # ---- augmentation (document-safe: NO rotation / shear; forms are axis-aligned) ----
        self.mosaic_prob = 1.0
        self.mosaic_scale = (0.5, 1.5)
        self.enable_mixup = False          # nano default
        self.mixup_prob = 0.0
        self.hsv_prob = 1.0
        self.flip_prob = 0.5
        self.degrees = 0.0
        self.shear = 0.0
        self.translate = 0.1

        # ---- optim / schedule (finetune from COCO-pretrained nano) ----
        self.warmup_epochs = 5
        self.max_epoch = int(os.environ.get("MAX_EPOCH", "50"))
        self.no_aug_epochs = 15            # close-mosaic tail for clean box convergence
        self.basic_lr_per_img = float(os.environ.get("BASIC_LR", 0.01 / 64.0))
        self.scheduler = "yoloxwarmcos"
        self.min_lr_ratio = 0.05
        self.weight_decay = 5e-4
        self.momentum = 0.9
        self.ema = True

        # ---- data ----
        self.data_dir = os.environ.get("COCO_DIR", "/workspace/coco_form")
        self.train_ann = "instances_train.json"
        self.val_ann = "instances_val.json"
        self.data_num_workers = int(os.environ.get("NUM_WORKERS", "14"))
        self.eval_interval = int(os.environ.get("EVAL_INTERVAL", "5"))
        self.print_interval = 50

        # ---- eval / postproc (forms are dense) ----
        self.test_conf = 0.01
        self.nmsthre = 0.65

        self.exp_name = "yolox_nano_formfield"

    # nano depthwise model (mirrors exps/default/yolox_nano.py)
    def get_model(self, sublinear=False):
        def init_yolo(M):
            for m in M.modules():
                if isinstance(m, nn.BatchNorm2d):
                    m.eps = 1e-3
                    m.momentum = 0.03
        if "model" not in self.__dict__:
            from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
            in_channels = [256, 512, 1024]
            backbone = YOLOPAFPN(self.depth, self.width, in_channels=in_channels,
                                 act=self.act, depthwise=True)
            head = YOLOXHead(self.num_classes, self.width, in_channels=in_channels,
                             act=self.act, depthwise=True)
            self.model = YOLOX(backbone, head)
        self.model.apply(init_yolo)
        self.model.head.initialize_biases(1e-2)
        return self.model

    # our COCO images live in data_dir/train and data_dir/val (not train2017/val2017); dense forms -> max_labels 200
    def get_dataset(self, cache: bool = False, cache_type: str = "ram"):
        from yolox.data import COCODataset, TrainTransform
        return COCODataset(
            data_dir=self.data_dir,
            json_file=self.train_ann,
            name="train",
            img_size=self.input_size,
            preproc=TrainTransform(max_labels=200, flip_prob=self.flip_prob, hsv_prob=self.hsv_prob),
            cache=cache,
            cache_type=cache_type,
        )

    def get_eval_dataset(self, **kwargs):
        from yolox.data import COCODataset, ValTransform
        legacy = kwargs.get("legacy", False)
        return COCODataset(
            data_dir=self.data_dir,
            json_file=self.val_ann,
            name="val",
            img_size=self.test_size,
            preproc=ValTransform(legacy=legacy),
        )

    # copy of base get_data_loader with max_labels 120 -> 200 (dense forms)
    def get_data_loader(self, batch_size, is_distributed, no_aug=False, cache_img: str = None):
        from yolox.data import (
            TrainTransform, YoloBatchSampler, DataLoader, InfiniteSampler,
            MosaicDetection, worker_init_reset_seed,
        )
        from yolox.utils import wait_for_the_master

        if self.dataset is None:
            with wait_for_the_master():
                assert cache_img is None, \
                    "cache_img must be None if you didn't create self.dataset before launch"
                self.dataset = self.get_dataset(cache=False, cache_type=cache_img)

        self.dataset = MosaicDetection(
            dataset=self.dataset,
            mosaic=not no_aug,
            img_size=self.input_size,
            preproc=TrainTransform(max_labels=200, flip_prob=self.flip_prob, hsv_prob=self.hsv_prob),
            degrees=self.degrees,
            translate=self.translate,
            mosaic_scale=self.mosaic_scale,
            mixup_scale=self.mixup_scale,
            shear=self.shear,
            enable_mixup=self.enable_mixup,
            mosaic_prob=self.mosaic_prob,
            mixup_prob=self.mixup_prob,
        )

        if is_distributed:
            import torch.distributed as dist
            batch_size = batch_size // dist.get_world_size()

        sampler = InfiniteSampler(len(self.dataset), seed=self.seed if self.seed else 0)
        batch_sampler = YoloBatchSampler(
            sampler=sampler, batch_size=batch_size, drop_last=False, mosaic=not no_aug,
        )
        dataloader_kwargs = {"num_workers": self.data_num_workers, "pin_memory": True,
                             "batch_sampler": batch_sampler, "worker_init_fn": worker_init_reset_seed}
        return DataLoader(self.dataset, **dataloader_kwargs)