"""Train YOLOv8x with experimental settings on 100% data."""
import sys
import torch
import json
import shutil
from pathlib import Path
from ultralytics import YOLO

_orig = torch.load
def _safe(*a, **kw): kw["weights_only"] = False; return _orig(*a, **kw)
torch.load = _safe

# Parse args: seed imgsz epochs optimizer lr experiment_name [extra_args...]
SEED = int(sys.argv[1]) if len(sys.argv) > 1 else 42
IMGSZ = int(sys.argv[2]) if len(sys.argv) > 2 else 1280
EPOCHS = int(sys.argv[3]) if len(sys.argv) > 3 else 100
OPTIMIZER = sys.argv[4] if len(sys.argv) > 4 else "AdamW"
LR = float(sys.argv[5]) if len(sys.argv) > 5 else 0.0005
EXP_NAME = sys.argv[6] if len(sys.argv) > 6 else "default"

# Experiment presets
PRESETS = {
    "light_aug": {"mosaic": 0.3, "copy_paste": 0.0, "mixup": 0.0, "degrees": 5, "scale": 0.5, "cls": 4.0},
    "no_aug": {"mosaic": 0.0, "copy_paste": 0.0, "mixup": 0.0, "degrees": 0, "scale": 0.0, "cls": 4.0},
    "high_cls": {"mosaic": 1.0, "copy_paste": 0.3, "mixup": 0.2, "degrees": 10, "scale": 0.9, "cls": 8.0},
    "freeze10": {"mosaic": 1.0, "copy_paste": 0.3, "mixup": 0.2, "degrees": 10, "scale": 0.9, "cls": 4.0, "freeze": 10},
    "cosine_low": {"mosaic": 1.0, "copy_paste": 0.3, "mixup": 0.2, "degrees": 10, "scale": 0.9, "cls": 4.0, "cos_lr": True, "lr0": 0.0001},
    "default": {"mosaic": 1.0, "copy_paste": 0.3, "mixup": 0.2, "degrees": 10, "scale": 0.9, "cls": 4.0},
}

preset = PRESETS.get(EXP_NAME, PRESETS["default"])
print(f"Experiment: {EXP_NAME}")
print(f"Settings: seed={SEED} imgsz={IMGSZ} epochs={EPOCHS} opt={OPTIMIZER} lr={LR}")
print(f"Preset: {preset}")

work_dir = Path(f"train_exp_{EXP_NAME}")
work_dir.mkdir(exist_ok=True)

ann = json.load(open("input/train/annotations.json"))
categories = ann["categories"]
images = ann["images"]
annotations = ann["annotations"]

train_img_dir = work_dir / "images" / "train"
val_img_dir = work_dir / "images" / "val"
train_lbl_dir = work_dir / "labels" / "train"
val_lbl_dir = work_dir / "labels" / "val"
for d in [train_img_dir, val_img_dir, train_lbl_dir, val_lbl_dir]:
    d.mkdir(parents=True, exist_ok=True)

img_anns = {}
for a in annotations:
    img_anns.setdefault(a["image_id"], []).append(a)

for img in images:
    src = Path("input/train/images") / img["file_name"]
    if not src.exists(): continue
    dst = train_img_dir / img["file_name"]
    if not dst.exists(): shutil.copy2(src, dst)
    iw, ih = img["width"], img["height"]
    label_lines = []
    for a in img_anns.get(img["id"], []):
        x, y, w, h = a["bbox"]
        cx = max(0, min(1, (x + w/2) / iw))
        cy = max(0, min(1, (y + h/2) / ih))
        nw = max(0, min(1, w / iw))
        nh = max(0, min(1, h / ih))
        label_lines.append(f"{a['category_id']} {cx} {cy} {nw} {nh}")
    lbl_name = img["file_name"].rsplit(".", 1)[0] + ".txt"
    (train_lbl_dir / lbl_name).write_text("\n".join(label_lines))

for img in images[:5]:
    src = Path("input/train/images") / img["file_name"]
    dst = val_img_dir / img["file_name"]
    if src.exists() and not dst.exists():
        shutil.copy2(src, dst)
        lbl_name = img["file_name"].rsplit(".", 1)[0] + ".txt"
        lbl_src = train_lbl_dir / lbl_name
        if lbl_src.exists(): shutil.copy2(lbl_src, val_lbl_dir / lbl_name)

nc = len(categories)
cat_names = {c["id"]: c["name"] for c in categories}
names_list = [cat_names.get(i, f"class_{i}") for i in range(nc)]
data_yaml = work_dir / "data.yaml"
data_yaml.write_text(f"path: {work_dir.resolve()}\ntrain: images/train\nval: images/val\nnc: {nc}\nnames: {names_list}\n")

train_args = dict(
    data=str(data_yaml), epochs=EPOCHS, imgsz=IMGSZ,
    batch=2 if IMGSZ <= 1280 else 1, workers=0,
    device=0 if torch.cuda.is_available() else "cpu",
    seed=SEED, close_mosaic=10,
    mosaic=preset.get("mosaic", 1.0),
    copy_paste=preset.get("copy_paste", 0.3),
    mixup=preset.get("mixup", 0.2),
    degrees=preset.get("degrees", 10),
    translate=0.2, scale=preset.get("scale", 0.9),
    fliplr=0.0, optimizer=OPTIMIZER,
    lr0=preset.get("lr0", LR), lrf=0.01,
    warmup_epochs=5, cls=preset.get("cls", 4.0),
    label_smoothing=0.1, save=True, save_period=25,
    cos_lr=preset.get("cos_lr", False),
)

if "freeze" in preset:
    train_args["freeze"] = preset["freeze"]

model = YOLO("yolov8x.pt")
model.train(**train_args)
print(f"Done! Experiment: {EXP_NAME}")