# -*- coding: utf-8 -*-
"""
padding_patch.py — Controla EXPLICITAMENTE a cor de preenchimento (padding) do
Ultralytics, para a ablation black-edge vs white-edge.

Por que e necessario: o Ultralytics 8.4.x codifica 114 (cinza) em TRES lugares
independentes, nenhum exposto no cfg de treino:
  1. Mosaic._mosaic4/_mosaic9 -> np.full(..., 114)          [treino]
  2. RandomPerspective        -> borderValue=(114,114,114)  [treino]
  3. LetterBox                -> padding_value=114          [val/inferencia]

Patchear so o LetterBox contaminaria a ablation (o treino seguiria cinza).
Este modulo intercepta os tres de forma auditavel e verificavel.

Uso:
    import padding_patch; padding_patch.apply("black")   # 0
    import padding_patch; padding_patch.apply("white")   # 255
    import padding_patch; padding_patch.apply("gray")    # 114 (default upstream)
"""
import os

import numpy as np

VALUES = {"black": 0, "white": 255, "gray": 114}
ENV = "WHST_PADDING_MODE"
_APPLIED = None


def worker_init(worker_id: int) -> None:
    """worker_init_fn que REAPLICA o patch dentro de cada worker process.

    ISTO E O CORACAO DA CORRECAO — sem isto a ablation nao acontece.

    No Windows o DataLoader usa SPAWN: cada worker e um processo novo, que nao
    herda monkeypatch nenhum do pai. Diagnosticado em 2026-07-27: os workers
    reportavam `np.full(114) -> 114` enquanto o pai reportava 255, e as 100
    epocas de treino de black e white sairam com loss BIT-IDENTICA.

    O LetterBox escapava do problema por acidente — o padding_value vira atributo
    de instancia no __init__ (executado no pai) e viaja no pickle do dataset. Por
    isso a VALIDACAO diferia e o TREINO nao: exatamente o padrao observado.

    Esta funcao e passada como `worker_init_fn`. Ela e picklada POR REFERENCIA
    (`padding_patch.worker_init`), o que obriga o worker a importar este modulo —
    e entao o patch e aplicado la dentro, antes do primeiro batch. O modo vem da
    variavel de ambiente, que processos filhos herdam.
    """
    modo = os.environ.get(ENV)
    if modo in VALUES:
        apply(modo, _em_worker=True)
    # preserva o comportamento original do Ultralytics
    try:
        import random

        import torch
        s = torch.initial_seed() % 2 ** 32
        np.random.seed(s)
        random.seed(s)
    except Exception:
        pass


def apply(mode: str, _em_worker: bool = False):
    """Aplica o valor de padding globalmente. Retorna o valor numerico usado."""
    global _APPLIED
    if mode not in VALUES:
        raise ValueError(f"modo invalido: {mode} (use {list(VALUES)})")
    fill = VALUES[mode]
    os.environ[ENV] = mode          # herdado pelos processos filhos (spawn)

    import ultralytics.data.augment as A

    # ---- 1) LetterBox: forca o default de padding_value ----
    if not hasattr(A.LetterBox, "_orig_init"):
        A.LetterBox._orig_init = A.LetterBox.__init__

    def _lb_init(self, *args, **kw):
        kw["padding_value"] = fill
        return A.LetterBox._orig_init(self, *args, **kw)

    A.LetterBox.__init__ = _lb_init

    # ---- 2) RandomPerspective: borderValue do warp ----
    cv2 = A.cv2
    if not hasattr(cv2, "_orig_warpAffine"):
        cv2._orig_warpAffine = cv2.warpAffine
        cv2._orig_warpPerspective = cv2.warpPerspective

    # guard POR VALOR (==114), simetrico ao _full: so intercepta o preenchimento
    # cinza da imagem em RandomPerspective. NAO toca o borderValue=255 usado por
    # apply_semantic para a mascara semantica (latente neste pipeline, mas blinda).
    def _wa(src, M, dsize, *a, **kw):
        if kw.get("borderValue") == (114, 114, 114):
            kw["borderValue"] = (fill,) * 3
        return cv2._orig_warpAffine(src, M, dsize, *a, **kw)

    def _wp(src, M, dsize, *a, **kw):
        if kw.get("borderValue") == (114, 114, 114):
            kw["borderValue"] = (fill,) * 3
        return cv2._orig_warpPerspective(src, M, dsize, *a, **kw)

    cv2.warpAffine = _wa
    cv2.warpPerspective = _wp

    # ---- 3) Mosaic: canvas np.full(..., 114) ----
    if not hasattr(A.np, "_orig_full"):
        A.np._orig_full = np.full

    def _full(shape, fill_value, *a, **kw):
        # so intercepta o canvas cinza do mosaico; nao mexe em outros np.full
        if isinstance(fill_value, (int, np.integer)) and int(fill_value) == 114:
            fill_value = fill
        return A.np._orig_full(shape, fill_value, *a, **kw)

    A.np.full = _full

    # ---- 4) worker_init_fn: leva o patch para dentro dos workers ----
    # Sem isto, os itens 2 e 3 nao existem no processo que de fato monta o batch.
    if not _em_worker:
        import ultralytics.data.build as B
        if getattr(B, "_padding_worker_init", None) is not worker_init:
            B.seed_worker = worker_init
            B._padding_worker_init = worker_init

    _APPLIED = mode
    return fill


def applied():
    return _APPLIED