import os
import json
import time
import argparse
from datetime import datetime
from pathlib import Path

import cv2
import numpy as np

from core.oak_fcc3_client import OakFcc3Client as MultiSpectralClient


# ============================================================
# Helpers gerais
# ============================================================

def now_str() -> str:
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def ensure_dir(path: str):
    os.makedirs(path, exist_ok=True)


def safe_float(v, default=None):
    try:
        if v is None:
            return default
        return float(v)
    except Exception:
        return default


def safe_int(v, default=None):
    try:
        if v is None:
            return default
        return int(v)
    except Exception:
        return default


def overlay_hud(
    img_bgr,
    lines,
    x=12,
    y=24,
    font_scale=0.58,
    line_step=22,
    color=(255, 255, 255),
    shadow=(0, 0, 0),
):
    yy = int(y)
    for s in lines:
        cv2.putText(img_bgr, str(s), (int(x), yy), cv2.FONT_HERSHEY_SIMPLEX,
                    font_scale, shadow, 3, cv2.LINE_AA)
        cv2.putText(img_bgr, str(s), (int(x), yy), cv2.FONT_HERSHEY_SIMPLEX,
                    font_scale, color, 1, cv2.LINE_AA)
        yy += int(line_step)


def to_bgr_u8_from_rgb01(rgb01: np.ndarray) -> np.ndarray:
    rgb_u8 = np.clip(rgb01 * 255.0, 0, 255).astype(np.uint8)
    return cv2.cvtColor(rgb_u8, cv2.COLOR_RGB2BGR)


def gray_to_bgr_u8(gray01: np.ndarray) -> np.ndarray:
    g = np.clip(gray01 * 255.0, 0, 255).astype(np.uint8)
    return cv2.cvtColor(g, cv2.COLOR_GRAY2BGR)


def resize_if_needed(img: np.ndarray, target_hw: tuple[int, int]) -> np.ndarray:
    if img is None:
        return None
    th, tw = target_hw
    if img.shape[:2] == (th, tw):
        return img
    return cv2.resize(img, (tw, th), interpolation=cv2.INTER_LINEAR)


def get_decoded_by_role(decoded: dict, role: str):
    role = str(role).lower()
    for cam_id, item in (decoded or {}).items():
        if str(item.get("role", "")).lower() == role:
            return cam_id, item
    return None, None


def get_image_by_role(decoded: dict, role: str):
    cam_id, item = get_decoded_by_role(decoded, role)
    if item is None:
        return cam_id, None
    return cam_id, item.get("image")


def validate_module_ready(status: dict, raw_policy: str):
    if not status.get("ok", True):
        raise RuntimeError(f"Status invalido retornado pelo modulo: {status}")

    active_roles = status.get("active_roles", {}) or {}
    active_count = int(status.get("camera_count_active", 0))

    if raw_policy == "require_triple":
        missing = [role for role in ("rgb", "nir", "re") if role not in active_roles]
        if missing:
            raise RuntimeError(
                "RAW_BRUTO com require_triple exige rgb/nir/re ativas. "
                f"Faltando: {missing}. Ativas: {active_roles}"
            )
    elif active_count < 1:
        raise RuntimeError("RAW_BRUTO requer ao menos uma câmera ativa.")


def compute_image_stats(img01: np.ndarray) -> dict:
    if img01 is None:
        return {
            "valid": False,
            "mean": 0.0,
            "std": 0.0,
            "p01": 0.0,
            "p05": 0.0,
            "p50": 0.0,
            "p95": 0.0,
            "p99": 0.0,
            "sat_pct": 0.0,
            "dark_pct": 0.0,
        }

    arr = img01.astype(np.float32).reshape(-1)
    return {
        "valid": True,
        "mean": float(arr.mean()),
        "std": float(arr.std()),
        "p01": float(np.percentile(arr, 1)),
        "p05": float(np.percentile(arr, 5)),
        "p50": float(np.percentile(arr, 50)),
        "p95": float(np.percentile(arr, 95)),
        "p99": float(np.percentile(arr, 99)),
        "sat_pct": float((arr >= 0.985).mean() * 100.0),
        "dark_pct": float((arr <= 0.015).mean() * 100.0),
    }


def robust_center(x: np.ndarray, low_pct=10.0, high_pct=90.0) -> float:
    arr = x[np.isfinite(x)].astype(np.float32).reshape(-1)
    if arr.size == 0:
        return 1.0

    lo = np.percentile(arr, low_pct)
    hi = np.percentile(arr, high_pct)
    core = arr[(arr >= lo) & (arr <= hi)]
    if core.size == 0:
        core = arr

    v = float(np.median(core))
    if not np.isfinite(v) or v <= 1e-8:
        return 1.0
    return v


def normalize_for_display(img: np.ndarray) -> np.ndarray:
    if img is None:
        return None

    arr = img.astype(np.float32)
    finite = arr[np.isfinite(arr)]
    if finite.size == 0:
        return np.zeros_like(arr, dtype=np.float32)

    lo = np.percentile(finite, 1)
    hi = np.percentile(finite, 99)
    if hi <= lo:
        hi = lo + 1e-6

    out = (arr - lo) / (hi - lo)
    return np.clip(out, 0.0, 1.0)


def smooth_map_gain(gain: np.ndarray, ksize: int) -> np.ndarray:
    if ksize is None or ksize <= 1:
        return gain.astype(np.float32)

    if ksize % 2 == 0:
        ksize += 1

    return cv2.GaussianBlur(
        gain.astype(np.float32),
        (ksize, ksize),
        sigmaX=0,
        sigmaY=0,
        borderType=cv2.BORDER_REFLECT,
    )


def clip_gain_map(gain: np.ndarray, min_gain: float, max_gain: float) -> np.ndarray:
    return np.clip(gain.astype(np.float32), float(min_gain), float(max_gain)).astype(np.float32)


# ============================================================
# Controle de câmera e extração de canais
# ============================================================

def get_controls_for_role(cam: MultiSpectralClient, role: str) -> dict:
    try:
        ctrl = cam.svc.get_camera_controls(role=role) or {}
        return {
            "ok": True,
            "role": role,
            "ae_enable": bool(ctrl.get("ae_enable", False)),
            "awb_enable": bool(ctrl.get("awb_enable", False)),
            "exposure_time_us": safe_int(ctrl.get("exposure_time_us"), None),
            "analogue_gain": safe_float(ctrl.get("analogue_gain"), None),
            "colour_gains": ctrl.get("colour_gains", None),
            "raw": ctrl,
        }
    except Exception as e:
        return {
            "ok": False,
            "role": role,
            "error": str(e),
            "ae_enable": None,
            "awb_enable": None,
            "exposure_time_us": None,
            "analogue_gain": None,
            "colour_gains": None,
        }


def get_all_controls(cam: MultiSpectralClient) -> dict:
    return {role: get_controls_for_role(cam, role) for role in ("rgb", "re", "nir")}


def exposure_gain_factor(ctrl: dict, fallback_exp=1.0, fallback_gain=1.0) -> float:
    exp = safe_float(ctrl.get("exposure_time_us"), fallback_exp)
    gain = safe_float(ctrl.get("analogue_gain"), fallback_gain)

    if exp is None or exp <= 0:
        exp = fallback_exp
    if gain is None or gain <= 0:
        gain = fallback_gain

    return float(exp * gain)


def extract_channels_from_decoded(decoded: dict) -> dict:
    """
    Retorna canais em float32 0..1 no espaço nativo de cada câmera:
      R/G/B vêm do debayer da role rgb.
      RE vem da role re.
      NIR vem da role nir.
    """
    _, rgb01 = get_image_by_role(decoded, "rgb")
    _, re01 = get_image_by_role(decoded, "re")
    _, nir01 = get_image_by_role(decoded, "nir")

    def assert_not_raw10_packed_image(role, img, expected_w=1280):
        if img is None:
            return

        if img.ndim == 2 and img.shape[1] == int(expected_w * 10 / 8):
            raise RuntimeError(
                f"{role.upper()} parece RAW10_PACKED interpretado como imagem: "
                f"shape={img.shape}. Esperado decodificado com largura {expected_w}."
            )
        
    assert_not_raw10_packed_image("re", re01, expected_w=1280)
    assert_not_raw10_packed_image("nir", nir01, expected_w=1280)

    out = {}

    if rgb01 is not None:
        rgb01 = rgb01.astype(np.float32)
        if rgb01.ndim == 3 and rgb01.shape[2] >= 3:
            out["R"] = rgb01[:, :, 0].copy()
            out["G"] = rgb01[:, :, 1].copy()
            out["B"] = rgb01[:, :, 2].copy()

    if re01 is not None:
        out["RE"] = re01.astype(np.float32).copy()

    if nir01 is not None:
        out["NIR"] = nir01.astype(np.float32).copy()

    return out


def channel_to_role(ch: str) -> str:
    ch = ch.upper()
    if ch in ("R", "G", "B"):
        return "rgb"
    if ch == "RE":
        return "re"
    if ch == "NIR":
        return "nir"
    raise ValueError(f"Canal desconhecido: {ch}")


def apply_exp_gain_correction(channels: dict, controls: dict, enabled: bool) -> dict:
    if not enabled:
        return {k: v.astype(np.float32).copy() for k, v in channels.items()}

    corrected = {}
    for ch, img in channels.items():
        role = channel_to_role(ch)
        ctrl = controls.get(role, {}) or {}
        factor = exposure_gain_factor(ctrl, fallback_exp=1.0, fallback_gain=1.0)
        corrected[ch] = (img.astype(np.float32) / max(factor, 1e-6)).astype(np.float32)

    return corrected


# ============================================================
# UI
# ============================================================

def build_board(decoded, controls, state_lines, progress_lines, preview_scale=1.0):
    rgb_id, rgb01 = get_image_by_role(decoded, "rgb")
    re_id, re01 = get_image_by_role(decoded, "re")
    nir_id, nir01 = get_image_by_role(decoded, "nir")

    if rgb01 is not None:
        rgb_panel = to_bgr_u8_from_rgb01(rgb01)
        base_h, base_w = rgb01.shape[:2]
    else:
        base_h, base_w = 800, 1280
        rgb_panel = np.zeros((base_h, base_w, 3), dtype=np.uint8)
        overlay_hud(rgb_panel, ["RGB", "sem frame"])

    re01 = resize_if_needed(re01, (base_h, base_w)) if re01 is not None else None
    nir01 = resize_if_needed(nir01, (base_h, base_w)) if nir01 is not None else None

    re_panel = gray_to_bgr_u8(re01) if re01 is not None else np.zeros_like(rgb_panel)
    nir_panel = gray_to_bgr_u8(nir01) if nir01 is not None else np.zeros_like(rgb_panel)

    rgb_stats = compute_image_stats(rgb01[:, :, 1] if rgb01 is not None and rgb01.ndim == 3 else None)
    re_stats = compute_image_stats(re01)
    nir_stats = compute_image_stats(nir01)

    overlay_hud(rgb_panel, [
        f"RGB ({rgb_id})",
        f"p50={rgb_stats['p50']:.3f} p95={rgb_stats['p95']:.3f} sat={rgb_stats['sat_pct']:.2f}%",
        f"EXP={controls.get('rgb', {}).get('exposure_time_us')} GAIN={controls.get('rgb', {}).get('analogue_gain')}",
    ])

    overlay_hud(re_panel, [
        f"RE ({re_id})",
        f"p50={re_stats['p50']:.3f} p95={re_stats['p95']:.3f} sat={re_stats['sat_pct']:.2f}%",
        f"EXP={controls.get('re', {}).get('exposure_time_us')} GAIN={controls.get('re', {}).get('analogue_gain')}",
    ])

    overlay_hud(nir_panel, [
        f"NIR ({nir_id})",
        f"p50={nir_stats['p50']:.3f} p95={nir_stats['p95']:.3f} sat={nir_stats['sat_pct']:.2f}%",
        f"EXP={controls.get('nir', {}).get('exposure_time_us')} GAIN={controls.get('nir', {}).get('analogue_gain')}",
    ])

    data_panel = np.zeros_like(rgb_panel)

    lines = []
    lines.extend(state_lines)
    lines.append("")
    lines.extend(progress_lines)
    lines.append("")
    lines.extend([
        "ENTER = iniciar etapa atual",
        "S = pular etapa dark/preto",
        "Q / ESC = sair sem salvar",
        "",
        "Dica: branco/preto devem preencher todo o campo de visao.",
        "Para dark-frame perfeito, tampe as lentes em vez de usar fundo preto.",
    ])

    overlay_hud(data_panel, lines, x=18, y=34, font_scale=0.58, line_step=24)

    top = np.hstack([rgb_panel, re_panel])
    bottom = np.hstack([nir_panel, data_panel])
    board = np.vstack([top, bottom])

    if preview_scale != 1.0:
        board = cv2.resize(
            board,
            (int(board.shape[1] * preview_scale), int(board.shape[0] * preview_scale)),
            interpolation=cv2.INTER_NEAREST,
        )

    return board


# ============================================================
# Captura e processamento
# ============================================================

def capture_stage(
    cam: MultiSpectralClient,
    stage_name: str,
    frames_count: int,
    discard_frames: int,
    exp_gain_correct: bool,
    preview_scale: float,
    window_name: str,
):
    """
    Captura frames decodificados e retorna:
      channel_stack: dict canal -> list[np.ndarray]
      controls_log: lista dos controles reais lidos por frame
      meta_log: lista de metadados básicos por frame
    """
    channel_stack = {}
    controls_log = []
    meta_log = []

    total_target = int(frames_count)
    captured = 0
    seen_frame_ids = set()
    last_decoded = {}
    last_controls = {}

    while captured < total_target:
        frame, meta, decoded = cam.get_next_decoded(timeout=2.0)
        if meta is None or frame is None:
            continue

        frame_id = meta.get("frame_id")
        if frame_id in seen_frame_ids:
            continue
        seen_frame_ids.add(frame_id)

        last_decoded = decoded
        last_controls = get_all_controls(cam)

        if len(seen_frame_ids) <= discard_frames:
            progress_lines = [
                f"Etapa: {stage_name}",
                f"Descartando frames iniciais: {len(seen_frame_ids)}/{discard_frames}",
                "Aguardando estabilizacao de exposicao/stream...",
            ]
            board = build_board(
                decoded=last_decoded,
                controls=last_controls,
                state_lines=[f"CALIBRACAO FLAT-FIELD - {stage_name.upper()}"],
                progress_lines=progress_lines,
                preview_scale=preview_scale,
            )
            cv2.imshow(window_name, board)
            cv2.waitKey(1)
            continue

        channels = extract_channels_from_decoded(decoded)
        channels = apply_exp_gain_correction(channels, last_controls, enabled=exp_gain_correct)

        for ch, img in channels.items():
            channel_stack.setdefault(ch, []).append(img.astype(np.float32).copy())

        meta_log.append({
            "frame_id": frame_id,
            "sync_ok": meta.get("sync_ok"),
            "sync_dt_ms": meta.get("sync_dt_ms"),
            "timestamps": meta.get("timestamps"),
        })
        controls_log.append(last_controls)

        captured += 1

        progress_lines = [
            f"Etapa: {stage_name}",
            f"Capturando: {captured}/{total_target}",
            f"exp_gain_correction={'ON' if exp_gain_correct else 'OFF'}",
        ]

        board = build_board(
            decoded=last_decoded,
            controls=last_controls,
            state_lines=[f"CALIBRACAO FLAT-FIELD - {stage_name.upper()}"],
            progress_lines=progress_lines,
            preview_scale=preview_scale,
        )

        # Barra de progresso simples.
        h, w = board.shape[:2]
        pct = captured / max(total_target, 1)
        cv2.rectangle(board, (30, h - 38), (w - 30, h - 18), (80, 80, 80), -1)
        cv2.rectangle(board, (30, h - 38), (30 + int((w - 60) * pct), h - 18), (0, 220, 0), -1)

        cv2.imshow(window_name, board)

        k = cv2.waitKey(1) & 0xFF
        if k in (ord("q"), ord("Q"), 27):
            raise KeyboardInterrupt("Captura cancelada pelo usuario.")

    return channel_stack, controls_log, meta_log


def median_stack(channel_stack: dict) -> dict:
    med = {}
    for ch, frames in channel_stack.items():
        if not frames:
            continue
        arr = np.stack(frames, axis=0).astype(np.float32)
        med[ch] = np.median(arr, axis=0).astype(np.float32)
    return med


def build_gain_maps(
    white_med: dict,
    dark_med: dict | None,
    epsilon: float,
    smooth_ksize: int,
    min_gain: float,
    max_gain: float,
):
    gain_maps = {}
    flat_norm_maps = {}
    corrected_white = {}
    dark_used = {}

    for ch, white in white_med.items():
        white = white.astype(np.float32)

        if dark_med is not None and ch in dark_med:
            dark = resize_if_needed(dark_med[ch].astype(np.float32), white.shape[:2])
        else:
            dark = np.zeros_like(white, dtype=np.float32)

        signal = white - dark
        signal = np.maximum(signal, float(epsilon)).astype(np.float32)

        center = robust_center(signal, low_pct=10.0, high_pct=90.0)
        flat_norm = signal / max(center, epsilon)

        gain = center / np.maximum(signal, epsilon)
        gain = smooth_map_gain(gain, smooth_ksize)
        gain = clip_gain_map(gain, min_gain=min_gain, max_gain=max_gain)

        gain_maps[ch] = gain.astype(np.float32)
        flat_norm_maps[ch] = flat_norm.astype(np.float32)
        corrected_white[ch] = signal.astype(np.float32)
        dark_used[ch] = dark.astype(np.float32)

    return gain_maps, flat_norm_maps, corrected_white, dark_used


def save_preview_maps(out_dir: str, gain_maps: dict, corrected_white: dict):
    ensure_dir(out_dir)

    for ch, gain in gain_maps.items():
        gain_vis = normalize_for_display(gain)
        cv2.imwrite(os.path.join(out_dir, f"gain_{ch}.png"), gray_to_bgr_u8(gain_vis))

    for ch, white in corrected_white.items():
        white_vis = normalize_for_display(white)
        cv2.imwrite(os.path.join(out_dir, f"white_signal_{ch}.png"), gray_to_bgr_u8(white_vis))


def show_final_preview(window_name: str, gain_maps: dict, preview_scale: float):
    order = ["R", "G", "B", "RE", "NIR"]
    panels = []

    # RGB e mono podem sair com tamanhos diferentes dependendo do decode/preview.
    # Para o painel final, padronizamos tudo para o maior H/W encontrado.
    shapes = [gain_maps[ch].shape[:2] for ch in order if ch in gain_maps]
    if not shapes:
        return

    target_h = max(s[0] for s in shapes)
    target_w = max(s[1] for s in shapes)

    def fit_panel(img_bgr: np.ndarray) -> np.ndarray:
        if img_bgr.shape[:2] == (target_h, target_w):
            return img_bgr
        return cv2.resize(img_bgr, (target_w, target_h), interpolation=cv2.INTER_NEAREST)

    for ch in order:
        if ch not in gain_maps:
            panel = np.zeros((target_h, target_w, 3), dtype=np.uint8)
            overlay_hud(panel, [ch, "sem mapa"])
        else:
            vis = normalize_for_display(gain_maps[ch])
            panel = gray_to_bgr_u8(vis)
            panel = fit_panel(panel)
            g = gain_maps[ch]
            overlay_hud(panel, [
                f"GAIN MAP {ch}",
                f"shape={list(g.shape)}",
                f"min={float(np.min(g)):.3f} max={float(np.max(g)):.3f}",
                f"mean={float(np.mean(g)):.3f} std={float(np.std(g)):.3f}",
            ])
        panels.append(panel)

    blank = np.zeros((target_h, target_w, 3), dtype=np.uint8)
    overlay_hud(blank, [
        "Flat-field salvo com sucesso.",
        "ENTER/qualquer tecla = fechar",
        "",
        "Use estes mapas antes da fusao geometrica.",
        "",
        "Obs: RGB e mono podem ter shapes diferentes;",
        "isso e normal se o decode gerar resolucoes distintas.",
    ], x=18, y=36)

    top = np.hstack([panels[0], panels[1], panels[2]])
    bottom = np.hstack([panels[3], panels[4], blank])
    board = np.vstack([top, bottom])

    if preview_scale != 1.0:
        board = cv2.resize(
            board,
            (int(board.shape[1] * preview_scale), int(board.shape[0] * preview_scale)),
            interpolation=cv2.INTER_NEAREST,
        )

    cv2.imshow(window_name, board)
    cv2.waitKey(0)


def wait_for_enter_or_skip(
    cam: MultiSpectralClient,
    window_name: str,
    title: str,
    instruction_lines: list[str],
    preview_scale: float,
    allow_skip=False,
):
    while True:
        frame, meta, decoded = cam.get_next_decoded(timeout=2.0)
        if meta is None or frame is None:
            continue

        controls = get_all_controls(cam)

        progress_lines = list(instruction_lines)
        if allow_skip:
            progress_lines.append("")
            progress_lines.append("S = pular esta etapa")

        board = build_board(
            decoded=decoded,
            controls=controls,
            state_lines=[title],
            progress_lines=progress_lines,
            preview_scale=preview_scale,
        )
        cv2.imshow(window_name, board)

        k = cv2.waitKey(1) & 0xFF
        if k in (13, 10):
            return "start"
        if allow_skip and k in (ord("s"), ord("S")):
            return "skip"
        if k in (ord("q"), ord("Q"), 27):
            raise KeyboardInterrupt("Cancelado pelo usuario.")


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

def main():
    parser = argparse.ArgumentParser(
        description="Calibrador automatico de flat-field/dark-frame para o módulo RGB/RE/NIR OAK-FCC-3.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument("--fps", type=int, default=20)
    parser.add_argument("--width", type=int, default=1280)
    parser.add_argument("--height", type=int, default=800)
    parser.add_argument("--bayer", default="BGGR", choices=["GBRG", "GRBG", "RGGB", "BGGR"])
    parser.add_argument("--capture_mode", default="AUTO", choices=["AUTO", "SINGLE", "DOUBLE", "TRIPLE"])
    parser.add_argument("--raw_policy", default="require_triple", choices=["allow_single", "require_triple"])
    parser.add_argument("--module_calibration_json", default="calibration/module_params.json")

    parser.add_argument("--frames", type=int, default=60, help="Frames úteis capturados por etapa.")
    parser.add_argument("--discard_frames", type=int, default=15, help="Frames descartados antes de cada etapa.")
    parser.add_argument("--preview_scale", type=float, default=0.65)

    parser.add_argument("--out_npz", default="calibration/flatfield_maps_v1.npz")
    parser.add_argument("--out_json", default="calibration/flatfield_maps_v1.json")
    parser.add_argument("--preview_dir", default="calibration/flatfield_previews")

    parser.add_argument("--smooth_ksize", type=int, default=31, help="Kernel gaussiano para suavizar mapa de ganho. Use 1 para desligar.")
    parser.add_argument("--min_gain", type=float, default=0.25)
    parser.add_argument("--max_gain", type=float, default=4.0)
    parser.add_argument("--epsilon", type=float, default=1e-8)

    parser.add_argument("--exp_gain_correct", action="store_true",
                        help="Divide frames por exposure_time_us*analogue_gain antes de calcular o flat.")
    parser.add_argument("--skip_dark", action="store_true",
                        help="Pula etapa dark/preto. O mapa será calculado sem subtração dark.")
    parser.add_argument("--notes", default="")

    args = parser.parse_args()

    ensure_dir(os.path.dirname(args.out_npz) or ".")
    ensure_dir(os.path.dirname(args.out_json) or ".")
    ensure_dir(args.preview_dir)

    window_name = "Flat Field Calibration Tool"
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)

    session_meta = {
        "schema": "multispec_flatfield_v1",
        "created_at": now_str(),
        "sensor_width": args.width,
        "sensor_height": args.height,
        "bayer_pattern": args.bayer,
        "fps": args.fps,
        "capture_mode_requested": args.capture_mode,
        "raw_policy": args.raw_policy,
        "module_calibration_json": args.module_calibration_json,
        "frames_per_stage": args.frames,
        "discard_frames": args.discard_frames,
        "exp_gain_correct": bool(args.exp_gain_correct),
        "smooth_ksize": args.smooth_ksize,
        "min_gain": args.min_gain,
        "max_gain": args.max_gain,
        "epsilon": args.epsilon,
        "channels": ["R", "G", "B", "RE", "NIR"],
        "notes": args.notes,
        "stages": {},
        "outputs": {
            "npz": args.out_npz,
            "json": args.out_json,
            "preview_dir": args.preview_dir,
        },
    }

    try:
        with MultiSpectralClient(
            width=args.width,
            height=args.height,
            bayer=args.bayer,
            fps=args.fps,
            frame_type="RAW_BRUTO",
            output_dtype="uint8",
            capture_mode=args.capture_mode,
            raw_policy=args.raw_policy,
            module_calibration_json=args.module_calibration_json or None
        ) as cam:

            validate_module_ready(cam.get_status(), args.raw_policy)

            wait_for_enter_or_skip(
                cam=cam,
                window_name=window_name,
                title="ETAPA 1/2 - WHITE / FLAT FIELD",
                instruction_lines=[
                    "Posicione o modulo na altura real de operacao.",
                    "Aponte para uma superficie branca/cinza fosca, uniforme e sem textura.",
                    "Evite reflexos, sombras laterais e saturacao.",
                    "A superficie deve preencher todo o campo de visao.",
                    "Pressione ENTER para comecar a captura WHITE.",
                ],
                preview_scale=args.preview_scale,
                allow_skip=False,
            )

            white_stack, white_controls, white_meta = capture_stage(
                cam=cam,
                stage_name="white",
                frames_count=args.frames,
                discard_frames=args.discard_frames,
                exp_gain_correct=args.exp_gain_correct,
                preview_scale=args.preview_scale,
                window_name=window_name,
            )

            session_meta["stages"]["white"] = {
                "controls_log": white_controls,
                "meta_log": white_meta,
                "captured_channels": {ch: len(v) for ch, v in white_stack.items()},
            }

            dark_stack = None
            dark_controls = []
            dark_meta = []

            if not args.skip_dark:
                action = wait_for_enter_or_skip(
                    cam=cam,
                    window_name=window_name,
                    title="ETAPA 2/2 - DARK / PRETO",
                    instruction_lines=[
                        "Agora faca a captura dark/preto.",
                        "Melhor opcao: tampe as lentes completamente.",
                        "Alternativa: use fundo preto fosco preenchendo todo o frame.",
                        "Mantenha exposicao/ganho iguais aos da etapa anterior, se possivel.",
                        "Pressione ENTER para capturar DARK/PRETO.",
                    ],
                    preview_scale=args.preview_scale,
                    allow_skip=True,
                )

                if action == "start":
                    dark_stack, dark_controls, dark_meta = capture_stage(
                        cam=cam,
                        stage_name="dark",
                        frames_count=args.frames,
                        discard_frames=args.discard_frames,
                        exp_gain_correct=args.exp_gain_correct,
                        preview_scale=args.preview_scale,
                        window_name=window_name,
                    )

                    session_meta["stages"]["dark"] = {
                        "controls_log": dark_controls,
                        "meta_log": dark_meta,
                        "captured_channels": {ch: len(v) for ch, v in dark_stack.items()},
                    }
                else:
                    session_meta["stages"]["dark"] = {"skipped": True}
            else:
                session_meta["stages"]["dark"] = {"skipped": True}

            # Processamento robusto.
            processing_panel = np.zeros((720, 1280, 3), dtype=np.uint8)
            overlay_hud(processing_panel, [
                "Processando calibracao flat-field...",
                "Calculando medianas robustas por canal.",
                "Gerando mapas de ganho e previews.",
            ], x=40, y=80, font_scale=0.8, line_step=34)
            cv2.imshow(window_name, processing_panel)
            cv2.waitKey(1)

            white_med = median_stack(white_stack)
            dark_med = median_stack(dark_stack) if dark_stack is not None else None

            gain_maps, flat_norm_maps, corrected_white, dark_used = build_gain_maps(
                white_med=white_med,
                dark_med=dark_med,
                epsilon=args.epsilon,
                smooth_ksize=args.smooth_ksize,
                min_gain=args.min_gain,
                max_gain=args.max_gain,
            )

            # Salva .npz.
            save_payload = {}
            for ch, arr in gain_maps.items():
                save_payload[f"gain_{ch}"] = arr.astype(np.float32)
            for ch, arr in flat_norm_maps.items():
                save_payload[f"flat_norm_{ch}"] = arr.astype(np.float32)
            for ch, arr in white_med.items():
                save_payload[f"white_median_{ch}"] = arr.astype(np.float32)
            if dark_med is not None:
                for ch, arr in dark_med.items():
                    save_payload[f"dark_median_{ch}"] = arr.astype(np.float32)

            np.savez_compressed(args.out_npz, **save_payload)

            # Estatísticas finais.
            session_meta["maps"] = {}
            for ch, gain in gain_maps.items():
                session_meta["maps"][ch] = {
                    "shape": list(gain.shape),
                    "gain_key": f"gain_{ch}",
                    "flat_norm_key": f"flat_norm_{ch}",
                    "white_median_key": f"white_median_{ch}",
                    "dark_median_key": f"dark_median_{ch}" if dark_med is not None and ch in dark_med else None,
                    "gain_min": float(np.min(gain)),
                    "gain_max": float(np.max(gain)),
                    "gain_mean": float(np.mean(gain)),
                    "gain_std": float(np.std(gain)),
                    "white_signal_stats": compute_image_stats(normalize_for_display(corrected_white[ch])),
                }

            save_preview_maps(args.preview_dir, gain_maps, corrected_white)

            with open(args.out_json, "w", encoding="utf-8") as f:
                json.dump(session_meta, f, ensure_ascii=False, indent=2)

            show_final_preview(window_name, gain_maps, args.preview_scale)

            print("")
            print("[OK] Flat-field gerado com sucesso.")
            print(f"[OK] NPZ : {args.out_npz}")
            print(f"[OK] JSON: {args.out_json}")
            print(f"[OK] PNGs: {args.preview_dir}")

    except KeyboardInterrupt as e:
        print(f"[CANCELADO] {e}")

    finally:
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()