"""Train YOLOv8x on 100% of data (no validation split). For competition only."""
import sys
import torch
from pathlib import Path
from ultralytics import YOLO
import json
import shutil

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

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

print(f"Training: seed={SEED} imgsz={IMGSZ} epochs={EPOCHS} opt={OPTIMIZER} lr={LR}")
print("MODE: 100% training data (NO validation split)")

work_dir = Path(f"train_full_s{SEED}_i{IMGSZ}_e{EPOCHS}")
work_dir.mkdir(exist_ok=True)

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

print(f"Total images: {len(images)} (ALL used for training)")

# Use ALL images for training - copy train as val too (ultralytics requires val)
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_map = {img["id"]: img for img in images}
img_anns = {}
for a in annotations:
    img_anns.setdefault(a["image_id"], []).append(a)

# ALL images go to train AND val (val is just a dummy to satisfy ultralytics)
for img in images:
    src = Path("input/train/images") / img["file_name"]
    if not src.exists():
        continue

    # Symlink to train
    dst_train = train_img_dir / img["file_name"]
    if not dst_train.exists():
        shutil.copy2(src, dst_train)

    # Also copy a small subset to val (just 5 images to satisfy ultralytics)
    # We don't care about val metrics - just need it to not crash

    iw, ih = img["width"], img["height"]
    label_lines = []
    for a in img_anns.get(img["id"], []):
        x, y, w, h = a["bbox"]
        cx = (x + w / 2) / iw
        cy = (y + h / 2) / ih
        nw = w / iw
        nh = h / ih
        cx = max(0, min(1, cx))
        cy = max(0, min(1, cy))
        nw = max(0, min(1, nw))
        nh = max(0, min(1, nh))
        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))

# Copy first 5 images to val (dummy)
val_count = 0
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)
        val_count += 1

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()}\n"
    f"train: images/train\n"
    f"val: images/val\n"
    f"nc: {nc}\n"
    f"names: {names_list}\n"
)

print(f"Train images: {len(images)}, Val images: {val_count} (dummy)")
print(f"Categories: {nc}")

model = YOLO("yolov8x.pt")
model.train(
    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=1.0,
    copy_paste=0.3,
    mixup=0.2,
    degrees=10,
    translate=0.2,
    scale=0.9,
    fliplr=0.0,
    optimizer=OPTIMIZER,
    lr0=LR,
    lrf=0.01,
    warmup_epochs=5,
    cls=4.0,
    label_smoothing=0.1,
    save=True,
    save_period=25,
)

print(f"Done! Best: {work_dir}/run/weights/best.pt")