#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Augmentação por grupos para o módulo multiespectral (opção 1: bins como canais sincronizados).

Entrada:
  dataset/original/group/<grupo>/previews
  dataset/original/group/<grupo>/metas
  dataset/original/group/<grupo>/bins
  dataset/original/group/<grupo>/masks
  (opcional) dataset/original/group/<grupo>/masks2

Saída:
  dataset/augmented/group/<grupo>/previews
  dataset/augmented/group/<grupo>/metas
  dataset/augmented/group/<grupo>/bins
  dataset/augmented/group/<grupo>/masks
  (opcional) dataset/augmented/group/<grupo>/masks2

Amostra esperada:
  <base>.png               # preview
  <base>.json              # meta
  <base>_cam0.bin          # bin câmera 0
  <base>_cam1.bin          # bin câmera 1
  <base>_cam2.bin          # bin câmera 2 (opcional)
  <base>.png               # mask

Estratégia:
- Geometria sincronizada em preview + masks + todos os bins.
- Blur / ruído / ganho apenas nos bins.
- Preview de saída recebe a mesma geometria; não recebe blur pesado para continuar útil como inspeção visual.
- Meta é copiado e marcado como augmentado.

Importante:
- Este script assume que os .bin são RAW10 packed, um arquivo por câmera.
- A largura/altura do bin é lida do meta.json quando possível; se não existir, cai para config['raw_size'].
"""

import os
import re
import cv2
import json
import math
import shutil
import argparse
import csv
import random
from copy import deepcopy
from typing import Dict, List, Optional, Tuple

import numpy as np
from PIL import Image

from pi.raw_processor_core import RawProcessorCore


with open("config.json", "r", encoding="utf-8") as f:
    config = json.load(f)

MODELO = config.get("camera", ".")
USE_MASKS2 = config.get("dual_head", False)
RAW_SIZE = config.get("raw_size", [1296, 1028])  # [W, H]

DATASET_BASE = os.path.join("dataset")
ORIG_GROUP_ROOT = os.path.join(DATASET_BASE, "original", "group")
AUG_GROUP_ROOT = os.path.join(DATASET_BASE, "augmented", "group")

PREVIEW_EXTS = (".jpg", ".jpeg", ".png")
MASK_EXTS = (".png", ".jpg", ".jpeg")
MASK2_EXTS = (".png", ".jpg", ".jpeg")
META_EXTS = (".json",)
BIN_RE = re.compile(r"^(?P<base>.+)_cam(?P<cam>\d+)\.bin$", re.IGNORECASE)
MANIFESTO_DEFAULT = "manifest_aug.csv"


# ============================================================
# Helpers básicos
# ============================================================

def garantir_dir(p: str):
    os.makedirs(p, exist_ok=True)


def save_rgb(path: str, arr_rgb: np.ndarray):
    Image.fromarray(arr_rgb).save(path)


def load_rgb(path: str) -> np.ndarray:
    im = cv2.imread(path, cv2.IMREAD_COLOR)
    if im is None:
        raise FileNotFoundError(path)
    return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)


def load_mask_any(path: str) -> np.ndarray:
    m = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    if m is None:
        raise FileNotFoundError(path)
    if m.ndim == 2:
        return m
    if m.shape[2] == 1:
        return m[:, :, 0]
    return cv2.cvtColor(m, cv2.COLOR_BGR2RGB)


def save_mask_any(path: str, mask: np.ndarray):
    if mask.ndim == 2:
        cv2.imwrite(path, mask)
    else:
        bgr = cv2.cvtColor(mask, cv2.COLOR_RGB2BGR)
        cv2.imwrite(path, bgr)


def list_groups(root: str) -> List[str]:
    if not os.path.isdir(root):
        return []
    grupos = []
    for name in sorted(os.listdir(root)):
        gdir = os.path.join(root, name)
        if not os.path.isdir(gdir):
            continue
        has_prev = os.path.isdir(os.path.join(gdir, "previews"))
        has_meta = os.path.isdir(os.path.join(gdir, "metas"))
        has_bins = os.path.isdir(os.path.join(gdir, "bins"))
        has_masks = os.path.isdir(os.path.join(gdir, "masks"))
        if has_prev and has_meta and has_bins and has_masks:
            grupos.append(name)
    return grupos


def map_by_base_priorizando_png(folder: str, exts: Tuple[str, ...]) -> Dict[str, str]:
    if not os.path.isdir(folder):
        return {}
    by_base = {}
    for fname in os.listdir(folder):
        if not fname.lower().endswith(exts):
            continue
        base, ext = os.path.splitext(fname)
        full = os.path.join(folder, fname)
        if base not in by_base:
            by_base[base] = full
        else:
            cur_ext = os.path.splitext(by_base[base])[1].lower()
            if cur_ext != ".png" and ext.lower() == ".png":
                by_base[base] = full
    return by_base


def map_bins_by_base(folder: str) -> Dict[str, List[str]]:
    by_base = {}
    if not os.path.isdir(folder):
        return by_base
    for fname in os.listdir(folder):
        m = BIN_RE.match(fname)
        if not m:
            continue
        base = m.group("base")
        cam = int(m.group("cam"))
        by_base.setdefault(base, []).append((cam, os.path.join(folder, fname)))
    for base in list(by_base.keys()):
        by_base[base] = [p for _, p in sorted(by_base[base], key=lambda x: x[0])]
    return by_base


def ensure_aug_dirs(group_name: str, use_masks2: bool):
    base = os.path.join(AUG_GROUP_ROOT, group_name)
    prev_out = os.path.join(base, "previews")
    meta_out = os.path.join(base, "metas")
    bins_out = os.path.join(base, "bins")
    mask_out = os.path.join(base, "masks")
    mask2_out = os.path.join(base, "masks2") if use_masks2 else None

    garantir_dir(prev_out)
    garantir_dir(meta_out)
    garantir_dir(bins_out)
    garantir_dir(mask_out)
    if use_masks2 and mask2_out:
        garantir_dir(mask2_out)

    return prev_out, meta_out, bins_out, mask_out, mask2_out


def make_raw_core() -> RawProcessorCore:
    return RawProcessorCore(sensor_width=RAW_SIZE[0], sensor_height=RAW_SIZE[1])


def load_all_bins(meta_path, bins_paths):
    import json

    with open(meta_path, "r", encoding="utf-8") as f:
        meta = json.load(f)

    core = make_raw_core()

    bins_data = []
    bins_meta = []

    for path in bins_paths:
        filename = os.path.basename(path)
        cam_id = filename.split("_")[-1].replace(".bin", "")

        cam_meta = core.extract_camera_meta(meta, cam_id)
        data = core.load_native_bin(path, cam_meta)

        bins_data.append(data)
        bins_meta.append(cam_meta)

    return bins_data, bins_meta


def save_all_bins(core, bins_data, bins_meta, base_name, out_dir):
    paths = []

    for data, meta in zip(bins_data, bins_meta):
        cam_id = meta["camera_id"]
        out_path = os.path.join(out_dir, f"{base_name}_{cam_id}.bin")

        core.save_native_bin(out_path, data, meta)
        paths.append(out_path)

    return paths


# ============================================================
# Meta / resolução
# ============================================================

def infer_bin_hw_from_meta(meta_path: str) -> Tuple[int, int]:
    with open(meta_path, "r", encoding="utf-8") as f:
        meta = json.load(f)

    # tenta campos mais prováveis
    width = None
    height = None

    for k in ("sensor_width", "width", "raw_width"):
        if k in meta:
            width = int(meta[k])
            break
    for k in ("sensor_height", "height", "raw_height"):
        if k in meta:
            height = int(meta[k])
            break

    if (width is None or height is None) and "raw_size" in meta and isinstance(meta["raw_size"], (list, tuple)) and len(meta["raw_size"]) == 2:
        width, height = int(meta["raw_size"][0]), int(meta["raw_size"][1])

    if width is None or height is None:
        width, height = int(RAW_SIZE[0]), int(RAW_SIZE[1])

    return width, height


def build_augmented_meta(meta_path: str, source_group: str, source_base: str,
                         aug_base: str, aug_index: int, params: Dict) -> Dict:
    with open(meta_path, "r", encoding="utf-8") as f:
        meta = json.load(f)

    meta_aug = deepcopy(meta)
    meta_aug["augmented"] = True
    meta_aug["augmentation"] = {
        "source_group": source_group,
        "source_base": source_base,
        "aug_base": aug_base,
        "aug_index": aug_index,
        "params": params,
    }
    return meta_aug


# ============================================================
# Geometria sincronizada
# ============================================================

def sample_geom_params() -> Dict:
    do_hflip = np.random.rand() < 0.5
    shift_x_frac = float(np.random.uniform(-0.01, 0.01))
    shift_y_frac = float(np.random.uniform(-0.01, 0.01))
    scale = float(np.random.uniform(0.92, 1.08))
    angle = float(np.random.uniform(-5.0, 5.0))
    return {
        "hflip": do_hflip,
        "shift_x_frac": shift_x_frac,
        "shift_y_frac": shift_y_frac,
        "scale": scale,
        "angle": angle,
    }


def build_affine_matrix(width: int, height: int, params: Dict) -> np.ndarray:
    cx = (width - 1) / 2.0
    cy = (height - 1) / 2.0
    M = cv2.getRotationMatrix2D((cx, cy), params["angle"], params["scale"])
    M[0, 2] += params["shift_x_frac"] * width
    M[1, 2] += params["shift_y_frac"] * height
    return M


def apply_geom_to_image(img: np.ndarray, params: Dict, is_mask: bool = False) -> np.ndarray:
    out = img
    if params["hflip"]:
        out = cv2.flip(out, 1)

    h, w = out.shape[:2]
    M = build_affine_matrix(w, h, params)

    interp = cv2.INTER_NEAREST if is_mask else cv2.INTER_LINEAR
    if out.ndim == 2:
        warped = cv2.warpAffine(out, M, (w, h), flags=interp, borderMode=cv2.BORDER_REFLECT_101)
    else:
        warped = cv2.warpAffine(out, M, (w, h), flags=interp, borderMode=cv2.BORDER_REFLECT_101)
    return warped


def apply_effects_rgb(img):
    out = img.astype(np.float32)

    if np.random.rand() < 0.2:
        out = cv2.GaussianBlur(out, (3,3), 0)

    gain = np.random.uniform(0.97, 1.03)
    out *= gain

    noise = np.random.normal(0, 2, out.shape)
    out += noise

    return np.clip(out, 0, 255).astype(np.uint8)


# ============================================================
# Blur / ruído / ganho coerentes nos bins
# ============================================================

def motion_blur_kernel(ksize=5, angle=0.0):
    ksize = int(ksize)
    if ksize < 3:
        ksize = 3
    if ksize % 2 == 0:
        ksize += 1

    kernel = np.zeros((ksize, ksize), dtype=np.float32)
    kernel[ksize // 2, :] = 1.0
    center = (ksize / 2.0 - 0.5, ksize / 2.0 - 0.5)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    kernel = cv2.warpAffine(kernel, M, (ksize, ksize))
    s = kernel.sum()
    if s > 0:
        kernel /= s
    return kernel


def sample_bin_effects() -> Dict:
    effect = {"kind": "none"}
    r = np.random.rand()
    if r < 0.15:
        effect["kind"] = "motion"
        effect["ksize"] = int(np.random.choice([3, 5, 7]))
        effect["angle"] = float(np.random.uniform(-20.0, 20.0))
    elif r < 0.30:
        effect["kind"] = "gaussian"
        effect["ksize"] = int(np.random.choice([3, 5, 7]))

    # leves variações por bin (mantendo coerência física e sem enlouquecer)
    effect["gain_min"] = float(np.random.uniform(0.97, 0.995))
    effect["gain_max"] = float(np.random.uniform(1.005, 1.03))
    effect["noise_sigma"] = float(np.random.uniform(0.0, 2.0))  # escala raw10
    return effect


def apply_effects_to_bin(bin_img: np.ndarray, effect: Dict, gain: float) -> np.ndarray:
    out = bin_img.astype(np.float32)

    if effect["kind"] == "motion":
        kernel = motion_blur_kernel(effect["ksize"], effect["angle"])
        out = cv2.filter2D(out, ddepth=-1, kernel=kernel, borderType=cv2.BORDER_REFLECT_101)
    elif effect["kind"] == "gaussian":
        k = effect["ksize"]
        if k % 2 == 0:
            k += 1
        out = cv2.GaussianBlur(out, (k, k), 0, borderType=cv2.BORDER_REFLECT_101)

    out *= gain

    sigma = effect.get("noise_sigma", 0.0)
    if sigma > 0:
        noise = np.random.normal(0.0, sigma, out.shape).astype(np.float32)
        out += noise

    out = np.clip(out, 0.0, 1023.0)
    return np.round(out).astype(np.uint16)


# ============================================================
# Núcleo da augmentação
# ============================================================

def make_preview_from_augmented_preview(preview_geom: np.ndarray) -> np.ndarray:
    """
    Por enquanto, o preview final é o preview original com a mesma geometria.
    Mantemos isso simples nesta etapa para inspeção humana.
    """
    return preview_geom


def augment_sample(group_name: str, base: str,
                   preview_path: str, meta_path: str, bins_paths: List[str], mask_path: str,
                   preview_out_dir: str, meta_out_dir: str, bins_out_dir: str, mask_out_dir: str,
                   copies: int, mask2_path: Optional[str] = None, mask2_out_dir: Optional[str] = None,
                   aug_suffix: str = "aug") -> int:

    preview_ext = os.path.splitext(preview_path)[1].lower()
    meta_ext = os.path.splitext(meta_path)[1].lower()
    mask_ext = os.path.splitext(mask_path)[1].lower()
    mask2_ext = os.path.splitext(mask2_path)[1].lower() if mask2_path else None

    preview = load_rgb(preview_path)
    mask = load_mask_any(mask_path)
    mask2 = load_mask_any(mask2_path) if mask2_path else None

    raw_w, raw_h = infer_bin_hw_from_meta(meta_path)
    bins_imgs, bins_meta = load_all_bins(meta_path, bins_paths)
    core = make_raw_core()

    generated = 0
    for i in range(copies):
        params = sample_geom_params()
        effects = sample_bin_effects()

        preview_g = apply_geom_to_image(preview, params, is_mask=False)
        mask_g = apply_geom_to_image(mask, params, is_mask=True)
        mask2_g = apply_geom_to_image(mask2, params, is_mask=True) if mask2 is not None else None

        bins_g = [apply_geom_to_image(b, params, is_mask=False) for b in bins_imgs]

        gains = [float(np.random.uniform(effects["gain_min"], effects["gain_max"])) for _ in bins_g]
        bins_aug = []

        for b, meta in zip(bins_g, bins_meta):
            if meta["channels"] == 3:
                # RGB (cam2)
                out = apply_effects_rgb(b)
            else:
                # RAW mono (cam0, cam1)
                gain = float(np.random.uniform(effects["gain_min"], effects["gain_max"]))
                out = apply_effects_to_bin(b, effects, gain)

            bins_aug.append(out)

        preview_aug = make_preview_from_augmented_preview(preview_g)

        aug_base = f"{base}_{aug_suffix}_{i:02d}"
        out_preview = os.path.join(preview_out_dir, aug_base + preview_ext)
        out_meta = os.path.join(meta_out_dir, aug_base + meta_ext)
        out_mask = os.path.join(mask_out_dir, aug_base + mask_ext)

        save_rgb(out_preview, preview_aug)
        save_mask_any(out_mask, mask_g)

        meta_aug = build_augmented_meta(meta_path, group_name, base, aug_base, i, {
            "geometry": params,
            "bin_effects": effects,
            "bin_gains": gains,
        })
        with open(out_meta, "w", encoding="utf-8") as f:
            json.dump(meta_aug, f, ensure_ascii=False, indent=2)

        out_bins = save_all_bins(core, bins_aug, bins_meta, aug_base, bins_out_dir)

        if mask2_g is not None and mask2_out_dir:
            out_mask2 = os.path.join(mask2_out_dir, aug_base + mask2_ext)
            save_mask_any(out_mask2, mask2_g)

        generated += 1

    return generated


# ============================================================
# Processamento por grupo
# ============================================================

def process_group(group_name: str, copies: int, limit: Optional[int] = None,
                  seed: int = 42, aug_suffix: str = "aug") -> Tuple[int, List[List[str]]]:
    gdir = os.path.join(ORIG_GROUP_ROOT, group_name)
    previews_dir = os.path.join(gdir, "previews")
    metas_dir = os.path.join(gdir, "metas")
    bins_dir = os.path.join(gdir, "bins")
    masks_dir = os.path.join(gdir, "masks")
    masks2_dir = os.path.join(gdir, "masks2")

    if not (os.path.isdir(previews_dir) and os.path.isdir(metas_dir) and os.path.isdir(bins_dir) and os.path.isdir(masks_dir)):
        print(f"[WARN] Grupo '{group_name}' inválido. Precisa de previews/metas/bins/masks.")
        return 0, []

    use_masks2 = USE_MASKS2 and os.path.isdir(masks2_dir)

    previews_map = map_by_base_priorizando_png(previews_dir, PREVIEW_EXTS)
    metas_map = map_by_base_priorizando_png(metas_dir, META_EXTS)
    masks_map = map_by_base_priorizando_png(masks_dir, MASK_EXTS)
    bins_map = map_bins_by_base(bins_dir)
    masks2_map = map_by_base_priorizando_png(masks2_dir, MASK2_EXTS) if use_masks2 else {}

    bases = sorted(set(previews_map.keys()) & set(metas_map.keys()) & set(masks_map.keys()) & set(bins_map.keys()))
    if limit is not None and limit > 0 and limit < len(bases):
        rng = np.random.default_rng(seed)
        idx = sorted(rng.choice(len(bases), size=limit, replace=False).tolist())
        bases = [bases[i] for i in idx]

    preview_out, meta_out, bins_out, mask_out, mask2_out = ensure_aug_dirs(group_name, use_masks2)

    registros = []
    count = 0
    for base in bases:
        try:
            gen = augment_sample(
                group_name=group_name,
                base=base,
                preview_path=previews_map[base],
                meta_path=metas_map[base],
                bins_paths=bins_map[base],
                mask_path=masks_map[base],
                preview_out_dir=preview_out,
                meta_out_dir=meta_out,
                bins_out_dir=bins_out,
                mask_out_dir=mask_out,
                copies=copies,
                mask2_path=masks2_map.get(base),
                mask2_out_dir=mask2_out,
                aug_suffix=aug_suffix,
            )
            count += gen
            registros.append([
                group_name,
                base,
                previews_map[base],
                metas_map[base],
                json.dumps(bins_map[base], ensure_ascii=False),
                masks_map[base],
                gen,
            ])
        except Exception as e:
            print(f"[ERRO] [{group_name}] {base}: {e}")

    print(f"[OK] Grupo '{group_name}' -> {count} amostras geradas.")
    return count, registros


# ============================================================
# Main
# ============================================================

def main(copies: int = 5, groups_csv: Optional[str] = None,
         limit: Optional[int] = None, seed: int = 42,
         suffix: str = "aug", manifesto: str = MANIFESTO_DEFAULT):
    total = 0
    all_records = []

    grupos = list_groups(ORIG_GROUP_ROOT)
    if groups_csv:
        want = {g.strip() for g in groups_csv.split(",") if g.strip()}
        grupos = [g for g in grupos if g in want]
        if not grupos:
            print("[WARN] Nenhum grupo válido encontrado após filtro.")
            return

    if not grupos:
        print("[WARN] Nenhum grupo encontrado em dataset/original/group.")
        return

    print(f"Grupos encontrados: {', '.join(grupos)}")
    for g in grupos:
        count, records = process_group(g, copies, limit=limit, seed=seed, aug_suffix=suffix)
        total += count
        all_records.extend(records)

    if manifesto:
        with open(manifesto, "w", newline="", encoding="utf-8") as f:
            w = csv.writer(f)
            w.writerow([
                "grupo",
                "base",
                "src_preview",
                "src_meta",
                "src_bins_json",
                "src_mask",
                "generated_copies",
            ])
            w.writerows(all_records)

    print(f"\nAugmentation completed! Total: {total} amostras geradas.")


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Augmentação por grupos para o módulo multiespectral usando bins sincronizados.")
    ap.add_argument("--copies", type=int, default=5, help="Número de cópias augmentadas por amostra (default=5).")
    ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula.")
    ap.add_argument("--limit", type=int, default=None, help="Quantidade máxima de amostras originais do grupo a augmentar.")
    ap.add_argument("--seed", type=int, default=42, help="Seed para seleção reproduzível quando usar --limit.")
    ap.add_argument("--suffix", type=str, default="aug", help="Sufixo usado no nome dos arquivos gerados.")
    ap.add_argument("--manifest", type=str, default=MANIFESTO_DEFAULT, help="CSV de manifesto.")
    args = ap.parse_args()

    random.seed(args.seed)
    np.random.seed(args.seed)

    main(
        copies=args.copies,
        groups_csv=args.groups,
        limit=args.limit,
        seed=args.seed,
        suffix=args.suffix,
        manifesto=args.manifest,
    )