"""Une las fuentes descargadas en un solo conjunto de datos.

    uv run python -m entrenamiento.preparar

1. Lee cada fuente descargada y traduce sus etiquetas a material y defectos.
2. Quita imágenes repetidas (mismo contenido aunque vengan de fuentes distintas).
3. Limita cuántas imágenes aporta cada clase de cada fuente y cada material, para que
   ninguna domine el entrenamiento.
4. Reparte en entrenamiento / validación / prueba (80/10/10) por "grupo": las variantes
   aumentadas de una misma foto (Roboflow las nombra foto_jpg.rf.<hash>.jpg) quedan siempre
   en el mismo lado, para que la evaluación no sea tramposa.
5. Guarda una copia reducida (512 px) de cada imagen en datos/procesados/imagenes/ y el
   listado en datos/procesados/manifiesto.jsonl.
"""

import hashlib
import json
import random
import re
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor
from dataclasses import asdict, dataclass, field
from pathlib import Path

from PIL import Image, ImageOps

from app.ml.estado.defectos import ORDEN_DEFECTOS
from app.ml.materiales import ORDEN_MATERIALES
from entrenamiento.descargar import MARCA
from entrenamiento.fuentes import FUENTES, Fuente
from entrenamiento.mapeo import LIMPIO, defectos_de_clase
from entrenamiento.rutas import (
    CRUDOS,
    IMAGENES,
    LADO_GUARDADO,
    MANIFIESTO,
    asegurar_carpetas,
)

EXTENSIONES = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
MAX_POR_MATERIAL = 4000
SEMILLA = 2026
Image.MAX_IMAGE_PIXELS = 60_000_000


@dataclass
class Registro:
    fuente: str
    ruta: str  # archivo original
    material: str | None
    # Solo los defectos conocidos: 1 = presente, 0 = ausente. Los que faltan, desconocidos.
    defectos: dict[str, int] = field(default_factory=dict)
    grupo: str = ""
    id: str = ""
    split: str = ""
    # Recuadro (x1, y1, x2, y2) si el registro es un objeto recortado de una foto mayor.
    recorte: list[int] | None = None

    @property
    def clave_cupo(self) -> str:
        positivos = sorted(d for d, v in self.defectos.items() if v)
        return f"{self.material}|{'+'.join(positivos) or ('limpio' if self.defectos else '-')}"


def _norm(texto: str) -> str:
    return re.sub(r"[\s_\-]+", " ", texto.lower()).strip()


def _grupo(fuente: str, ruta: Path) -> str:
    # Roboflow: "IMG_123_jpg.rf.abc123.jpg" → todas las variantes comparten "IMG_123_jpg".
    return f"{fuente}/{ruta.name.split('.rf.')[0]}"


def _imagenes(carpeta: Path):
    for ruta in carpeta.rglob("*"):
        if (
            ruta.suffix.lower() in EXTENSIONES
            and "__MACOSX" not in ruta.parts
            and not ruta.name.startswith("._")
        ):
            yield ruta


# ------------------------------------------------------------------------------ lectores


def _leer_carpetas_material(fuente: Fuente, carpeta: Path) -> list[Registro]:
    clases = {_norm(k): v for k, v in fuente.clases.items()}
    ausentes = dict.fromkeys(fuente.defectos_ausentes, 0)
    registros = []
    for ruta in _imagenes(carpeta):
        clase = next((_norm(p) for p in reversed(ruta.parts[:-1]) if _norm(p) in clases), None)
        if clase is None or clases[clase] is None:
            continue
        registros.append(
            Registro(
                fuente.nombre, str(ruta), clases[clase], dict(ausentes), _grupo(fuente.nombre, ruta)
            )
        )
    return registros


def _defectos_conocidos(nombres_clase: list[str]) -> set[str]:
    conocidos: set[str] = set()
    for nombre in nombres_clase:
        conocidos |= defectos_de_clase(nombre) - {LIMPIO}
    return conocidos


def _etiquetar(clases_imagen: set[str], conocidos: set[str]) -> dict[str, int]:
    presentes: set[str] = set()
    for nombre in clases_imagen:
        presentes |= defectos_de_clase(nombre) - {LIMPIO}
    return {d: int(d in presentes) for d in conocidos}


def _leer_coco(fuente: Fuente, carpeta: Path) -> list[Registro]:
    anotaciones = list(carpeta.rglob("_annotations.coco.json"))
    categorias_todas: list[str] = []
    datos = []
    for archivo in anotaciones:
        coco = json.loads(archivo.read_text(encoding="utf-8"))
        categorias = {c["id"]: c["name"] for c in coco.get("categories", [])}
        categorias_todas += list(categorias.values())
        datos.append((archivo.parent, coco, categorias))
    conocidos = _defectos_conocidos(categorias_todas)

    registros = []
    material = fuente.material if fuente.usar_para_material else None
    for carpeta_split, coco, categorias in datos:
        por_imagen: dict[int, set[str]] = defaultdict(set)
        for anotacion in coco.get("annotations", []):
            por_imagen[anotacion["image_id"]].add(categorias.get(anotacion["category_id"], ""))
        for imagen in coco.get("images", []):
            ruta = carpeta_split / imagen["file_name"]
            if ruta.is_file():
                registros.append(
                    Registro(
                        fuente.nombre,
                        str(ruta),
                        material,
                        _etiquetar(por_imagen.get(imagen["id"], set()), conocidos),
                        _grupo(fuente.nombre, ruta),
                    )
                )
    return registros


MAX_RECORTES_POR_IMAGEN = 2
LADO_MINIMO_RECORTE = 48
MARGEN_RECORTE = 0.1


def _leer_recortes(fuente: Fuente, carpeta: Path) -> list[Registro]:
    """Un registro por objeto etiquetado (los más grandes de cada foto), con su material."""
    patrones = [(re.compile(p, re.I), m) for p, m in fuente.recortes.items()]
    ausentes = dict.fromkeys(fuente.defectos_ausentes, 0)
    registros = []
    for archivo in carpeta.rglob("_annotations.coco.json"):
        coco = json.loads(archivo.read_text(encoding="utf-8"))
        material_de = {}
        for c in coco.get("categories", []):
            material_de[c["id"]] = next((m for p, m in patrones if p.search(c["name"])), None)
        imagenes = {i["id"]: i for i in coco.get("images", [])}
        por_imagen: dict[int, list[tuple[float, list[int], str]]] = defaultdict(list)
        for a in coco.get("annotations", []):
            material = material_de.get(a["category_id"])
            imagen = imagenes.get(a["image_id"])
            if material is None or imagen is None:
                continue
            x, y, w, h = a["bbox"]
            if min(w, h) < LADO_MINIMO_RECORTE:
                continue
            mx, my = w * MARGEN_RECORTE, h * MARGEN_RECORTE
            caja = [
                max(0, int(x - mx)),
                max(0, int(y - my)),
                min(imagen["width"], int(x + w + mx)),
                min(imagen["height"], int(y + h + my)),
            ]
            por_imagen[a["image_id"]].append((w * h, caja, material))
        for id_imagen, objetos in por_imagen.items():
            ruta = archivo.parent / imagenes[id_imagen]["file_name"]
            if not ruta.is_file():
                continue
            for _, caja, material in sorted(objetos, reverse=True)[:MAX_RECORTES_POR_IMAGEN]:
                registros.append(
                    Registro(
                        fuente.nombre,
                        str(ruta),
                        material,
                        dict(ausentes),
                        _grupo(fuente.nombre, ruta),
                        recorte=caja,
                    )
                )
    return registros


def _leer_clasificacion_rf(fuente: Fuente, carpeta: Path) -> list[Registro]:
    rutas = [r for r in _imagenes(carpeta)]
    conocidos = _defectos_conocidos(sorted({r.parent.name for r in rutas}))
    material = fuente.material if fuente.usar_para_material else None
    registros = []
    for ruta in rutas:
        clase = ruta.parent.name
        if not defectos_de_clase(clase) and conocidos:
            continue  # clase que no dice nada del estado
        registros.append(
            Registro(
                fuente.nombre,
                str(ruta),
                material,
                _etiquetar({clase}, conocidos),
                _grupo(fuente.nombre, ruta),
            )
        )
    return registros


def leer_fuente(fuente: Fuente) -> list[Registro]:
    carpeta = CRUDOS / fuente.nombre
    if not fuente.entrenar or not (carpeta / MARCA).is_file():
        return []
    if fuente.recortes:
        return _leer_recortes(fuente, carpeta)
    if fuente.clases:
        return _leer_carpetas_material(fuente, carpeta)
    if any(carpeta.rglob("_annotations.coco.json")):
        return _leer_coco(fuente, carpeta)
    return _leer_clasificacion_rf(fuente, carpeta)


# ---------------------------------------------------------------------- limpieza y reparto


def _limitar(
    registros: list[Registro], fuentes: dict[str, Fuente], rng: random.Random
) -> list[Registro]:
    # 1) cupo por (fuente, material, tipo de defecto)
    por_clave: dict[tuple[str, str], list[Registro]] = defaultdict(list)
    for r in registros:
        por_clave[(r.fuente, r.clave_cupo)].append(r)
    elegidos: list[Registro] = []
    for (fuente, _), lista in por_clave.items():
        rng.shuffle(lista)
        elegidos += lista[: fuentes[fuente].max_por_clase]

    # 2) cupo por material: se priorizan las imágenes que también enseñan defectos
    por_material: dict[str | None, list[Registro]] = defaultdict(list)
    for r in elegidos:
        por_material[r.material].append(r)
    final: list[Registro] = []
    for material, lista in por_material.items():
        if material is None:
            final += lista
            continue
        con_defectos = [r for r in lista if any(r.defectos.values()) or "limpio" in r.clave_cupo]
        resto = [r for r in lista if r not in con_defectos]
        rng.shuffle(con_defectos)
        rng.shuffle(resto)
        final += (con_defectos + resto)[:MAX_POR_MATERIAL]
    return final


def _split(grupo: str) -> str:
    valor = int(hashlib.md5(grupo.encode()).hexdigest(), 16) % 100
    return "entrenamiento" if valor < 80 else "validacion" if valor < 90 else "prueba"


def _guardar_copia(registro: Registro) -> tuple[Registro, str | None]:
    """Guarda la copia reducida. Devuelve (registro, huella del contenido) o None si falla."""
    try:
        with Image.open(registro.ruta) as original:
            if registro.recorte:
                # Las cajas de COCO se refieren a los píxeles tal como están guardados.
                imagen = original.convert("RGB").crop(tuple(registro.recorte))
            else:
                imagen = ImageOps.exif_transpose(original).convert("RGB")
        if min(imagen.size) < 32:
            return registro, None
        imagen.thumbnail((LADO_GUARDADO, LADO_GUARDADO), Image.Resampling.LANCZOS)
        huella = hashlib.md5(imagen.resize((16, 16)).convert("L").tobytes()).hexdigest()
        imagen.save(IMAGENES / f"{registro.id}.jpg", quality=92)
        return registro, huella
    except Exception:
        return registro, None


def main() -> None:
    asegurar_carpetas()
    rng = random.Random(SEMILLA)
    fuentes = {f.nombre: f for f in FUENTES}

    registros: list[Registro] = []
    for fuente in FUENTES:
        leidos = leer_fuente(fuente)
        if leidos:
            print(f"- {fuente.nombre}: {len(leidos)} imágenes")
        registros += leidos
    if not registros:
        raise SystemExit(
            "No hay fuentes descargadas. Ejecuta antes: python -m entrenamiento.descargar"
        )

    registros = _limitar(registros, fuentes, rng)
    for r in registros:
        r.id = hashlib.md5(f"{r.ruta}|{r.recorte}".encode()).hexdigest()[:16]
        r.split = _split(r.grupo)

    print(f"Guardando {len(registros)} copias reducidas en {IMAGENES} ...")
    vistos: set[str] = set()
    validos: list[Registro] = []
    with ThreadPoolExecutor(max_workers=6) as pool:
        for registro, huella in pool.map(_guardar_copia, registros):
            if huella is None:
                continue
            if huella in vistos:  # misma imagen en dos fuentes o repetida
                (IMAGENES / f"{registro.id}.jpg").unlink(missing_ok=True)
                continue
            vistos.add(huella)
            validos.append(registro)

    with MANIFIESTO.open("w", encoding="utf-8") as salida:
        for r in sorted(validos, key=lambda r: r.id):
            salida.write(json.dumps(asdict(r), ensure_ascii=False) + "\n")

    print(f"\nManifiesto: {MANIFIESTO} ({len(validos)} imágenes)")
    print("\nMaterial por reparto:")
    tabla = Counter((r.material, r.split) for r in validos)
    for material in (*ORDEN_MATERIALES, None):
        print(
            f"  {str(material):10} "
            + "  ".join(
                f"{s[:5]}={tabla[(material, s)]:5}"
                for s in ("entrenamiento", "validacion", "prueba")
            )
        )
    print("\nDefectos (positivos / conocidos):")
    for d in ORDEN_DEFECTOS:
        conocidos = sum(1 for r in validos if d in r.defectos)
        positivos = sum(r.defectos.get(d, 0) for r in validos)
        print(f"  {d:12} {positivos:5} / {conocidos:5}")


if __name__ == "__main__":
    main()