# train.py
from pathlib import Path
from ultralytics.models.yolo.segment.train import SegmentationTrainer
from custom_channel_dataset import CustomChannelSegDataset
from patch_model_xch import patch_yolov8_first_conv_to_xch
import yaml

with Path("data.yaml").open("r", encoding="utf-8") as f:
    cfg = yaml.safe_load(f) or {}
CHANNELS = int(cfg.get("channels", 4))
SCALE = cfg.get("scale", "m")
IMG_SZ = int(cfg.get("size", 800))

class CustomChannelSegTrainer(SegmentationTrainer):
    def build_dataset(self, img_path, mode="train", batch=None):
        # img_path é train.txt ou val.txt vindo do data.yaml
        return CustomChannelSegDataset(
            img_path=img_path,
            data=self.data,
            rect=False,
            canais=CHANNELS,
            imgsz=self.args.imgsz,
            augment=(mode == "train"),       # 👈 importante
            hyp=self.args,                   # 👈 importante: passa os hypers
        )

    def get_model(self, cfg=None, weights=None, verbose=True):
        model = super().get_model(cfg=cfg, weights=weights, verbose=verbose)

        patch_yolov8_first_conv_to_xch(model, canais=CHANNELS)

        # debug pra garantir
        seq = model.model if hasattr(model, "model") else model
        print("First conv in_channels:", seq[0].conv.in_channels)

        return model


def main():
    args = {
        "project": f"C:/ZendionInc/agrobot_base/Python/yolov8-seg/runs/segment/ch_{CHANNELS}/sc_{SCALE}/sz_{IMG_SZ}",
        "name": "weed_detector_segformer",
        "model": f"backbones/yolov8-seg-{CHANNELS}ch.yaml",
        #"model": f"backbones/yolov8{SCALE}-seg.pt",
        "pretrained": False,
        "data": "data.yaml",
        "epochs": 1000,
        "patience": 200,
        "imgsz": IMG_SZ,
        "device": 0,
        "batch": 4,
        "workers": 4,

        # zerando todas as augmentações “perigosas” pro debug
        "mosaic": 0.0,
        "copy_paste": 0.0,
        "mixup": 0.0,
        "erasing": 0.0,
        "auto_augment": "none",
        "hsv_h": 0.0,
        "hsv_s": 0.0,
        "hsv_v": 0.0,
        # augmentacoes geometricas
        "fliplr": 0.5,
        "flipud": 0.0,
        "translate": 0.05,
        "scale": 0.2,
        "degrees": 10.0,
        "shear": 0.0,
        "perspective": 0.0
    }

    trainer = CustomChannelSegTrainer(overrides=args)
    trainer.train()


if __name__ == "__main__":
    main()