import argparse
import time
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort
from PIL import Image


IMG_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}


def get_providers(force_cpu: bool = False):
    available = ort.get_available_providers()

    if force_cpu:
        return ["CPUExecutionProvider"]

    providers = []

    if "CUDAExecutionProvider" in available:
        providers.append("CUDAExecutionProvider")

    providers.append("CPUExecutionProvider")
    return providers


def print_model_io(session: ort.InferenceSession):
    print("\n[MODEL INPUTS]")
    for i, inp in enumerate(session.get_inputs()):
        print(f"  {i}: name={inp.name} shape={inp.shape} type={inp.type}")

    print("\n[MODEL OUTPUTS]")
    for i, out in enumerate(session.get_outputs()):
        print(f"  {i}: name={out.name} shape={out.shape} type={out.type}")
    print("")


def resolve_hw_from_input_shape(shape, fallback_h: int, fallback_w: int):
    """
    Tenta descobrir H/W do input ONNX.
    Esperado normalmente: [1, 3, H, W] ou [1, H, W, 3].
    Se for dinâmico, usa fallback.
    """
    if shape is None:
        return fallback_h, fallback_w

    dims = list(shape)

    def is_int(x):
        return isinstance(x, int) and x > 0

    # NCHW
    if len(dims) == 4 and dims[1] == 3:
        h = dims[2] if is_int(dims[2]) else fallback_h
        w = dims[3] if is_int(dims[3]) else fallback_w
        return int(h), int(w)

    # NHWC
    if len(dims) == 4 and dims[-1] == 3:
        h = dims[1] if is_int(dims[1]) else fallback_h
        w = dims[2] if is_int(dims[2]) else fallback_w
        return int(h), int(w)

    return fallback_h, fallback_w


def is_nchw_input(shape):
    if shape is None:
        return True
    dims = list(shape)
    return len(dims) == 4 and dims[1] == 3


def preprocess_image(image_path: Path, input_shape, input_h: int, input_w: int, mean, std):
    img_pil = Image.open(image_path).convert("RGB")
    rgb_orig = np.array(img_pil)
    bgr_orig = cv2.cvtColor(rgb_orig, cv2.COLOR_RGB2BGR)

    rgb_resized = cv2.resize(rgb_orig, (input_w, input_h), interpolation=cv2.INTER_AREA)
    x = rgb_resized.astype(np.float32) / 255.0

    mean_arr = np.array(mean, dtype=np.float32).reshape(1, 1, 3)
    std_arr = np.array(std, dtype=np.float32).reshape(1, 1, 3)
    x = (x - mean_arr) / std_arr

    if is_nchw_input(input_shape):
        x = np.transpose(x, (2, 0, 1))  # CHW

    x = np.expand_dims(x, axis=0).astype(np.float32)

    return bgr_orig, x


def robust_normalize_for_view(depth_m: np.ndarray, invert: bool = False):
    """
    Só para visualização colorida.
    O modo bandas usa o depth_m direto.
    """
    d = depth_m.astype(np.float32)
    valid = np.isfinite(d) & (d > 0)

    if np.count_nonzero(valid) < 20:
        return np.zeros_like(d, dtype=np.float32)

    vals = d[valid]
    p2 = np.percentile(vals, 2)
    p98 = np.percentile(vals, 98)

    dn = (d - p2) / (p98 - p2 + 1e-6)
    dn = np.clip(dn, 0.0, 1.0)

    if invert:
        dn = 1.0 - dn

    dn[~valid] = 0.0
    return dn.astype(np.float32)


def depth_to_colormap(depth_norm: np.ndarray):
    u8 = (depth_norm * 255).astype(np.uint8)
    return cv2.applyColorMap(u8, cv2.COLORMAP_TURBO)


def make_metric_bands(depth_m: np.ndarray, near_m: float, mid_m: float, far_m: float):
    """
    Bandas baseadas no valor métrico estimado.

    0: < near_m
    1: near_m até mid_m
    2: mid_m até far_m
    3: >= far_m
    """
    d = depth_m.astype(np.float32)
    valid = np.isfinite(d) & (d > 0)

    bands = np.zeros_like(d, dtype=np.uint8)

    bands[(d >= near_m) & (d < mid_m)] = 1
    bands[(d >= mid_m) & (d < far_m)] = 2
    bands[d >= far_m] = 3
    bands[~valid] = 255

    out = np.zeros((bands.shape[0], bands.shape[1], 3), dtype=np.uint8)

    # BGR
    out[bands == 0] = (40, 40, 255)      # muito perto
    out[bands == 1] = (40, 180, 255)     # perto/médio
    out[bands == 2] = (40, 255, 120)     # médio/longe
    out[bands == 3] = (255, 180, 40)     # longe
    out[bands == 255] = (0, 0, 0)        # inválido

    return out


def extract_depth_from_outputs(outputs, output_names):
    """
    Tenta achar depth em diferentes formatos:
    - output chamado depth
    - output chamado points/xyz, usando canal Z
    - saída única
    """
    name_to_out = {
        name.lower(): out
        for name, out in zip(output_names, outputs)
    }

    # 1) Procura saída com nome depth
    for name, out in name_to_out.items():
        if "depth" in name:
            return squeeze_depth(out)

    # 2) Procura points/xyz e usa Z
    for name, out in name_to_out.items():
        if "point" in name or "xyz" in name:
            arr = np.array(out)
            return extract_z_from_points(arr)

    # 3) Se tiver só uma saída, usa ela
    if len(outputs) == 1:
        arr = np.array(outputs[0])

        # Se parece points, extrai Z
        if arr.ndim == 4 and (arr.shape[1] == 3 or arr.shape[-1] == 3):
            return extract_z_from_points(arr)

        return squeeze_depth(arr)

    # 4) Fallback: pega a primeira saída 2D/3D/4D plausível
    for out in outputs:
        arr = np.array(out)
        try:
            d = squeeze_depth(arr)
            if d.ndim == 2:
                return d
        except Exception:
            pass

    raise RuntimeError("Não consegui identificar o mapa de depth nas saídas ONNX.")


def squeeze_depth(arr):
    arr = np.array(arr)

    # Remove batch/canal unitário
    arr = np.squeeze(arr)

    if arr.ndim == 2:
        return arr.astype(np.float32)

    if arr.ndim == 3:
        # CHW com 1 canal
        if arr.shape[0] == 1:
            return arr[0].astype(np.float32)

        # HWC com 1 canal
        if arr.shape[-1] == 1:
            return arr[..., 0].astype(np.float32)

        # Se for 3 canais, talvez seja XYZ
        if arr.shape[0] == 3:
            return arr[2].astype(np.float32)

        if arr.shape[-1] == 3:
            return arr[..., 2].astype(np.float32)

    raise RuntimeError(f"Formato de depth não suportado: shape={arr.shape}")


def extract_z_from_points(arr):
    arr = np.array(arr)

    # NCHW: [1, 3, H, W]
    if arr.ndim == 4 and arr.shape[1] == 3:
        return arr[0, 2].astype(np.float32)

    # NHWC: [1, H, W, 3]
    if arr.ndim == 4 and arr.shape[-1] == 3:
        return arr[0, ..., 2].astype(np.float32)

    # CHW: [3, H, W]
    if arr.ndim == 3 and arr.shape[0] == 3:
        return arr[2].astype(np.float32)

    # HWC: [H, W, 3]
    if arr.ndim == 3 and arr.shape[-1] == 3:
        return arr[..., 2].astype(np.float32)

    raise RuntimeError(f"Formato de points/xyz não suportado: shape={arr.shape}")


def maybe_invert_depth_if_needed(depth: np.ndarray, auto_invert: bool):
    """
    Alguns modelos podem devolver inverso/disparity.
    Aqui deixei opcional. Por padrão não mexe.
    """
    if not auto_invert:
        return depth

    d = depth.astype(np.float32)
    valid = np.isfinite(d) & (d > 1e-6)
    out = np.zeros_like(d, dtype=np.float32)
    out[valid] = 1.0 / d[valid]
    return out


def infer_depth(session, image_path: Path, args):
    input0 = session.get_inputs()[0]
    input_shape = input0.shape
    fallback_h, fallback_w = args.input_h, args.input_w
    input_h, input_w = resolve_hw_from_input_shape(input_shape, fallback_h, fallback_w)

    preview_bgr, x = preprocess_image(
        image_path=image_path,
        input_shape=input_shape,
        input_h=input_h,
        input_w=input_w,
        mean=args.mean,
        std=args.std,
    )

    feed = {}

    # Alimenta input principal
    feed[input0.name] = x

    # Alguns modelos têm input extra de intrinsics/K.
    # Se existir, manda uma matriz aproximada baseada no tamanho de entrada.
    # Para nosso teste visual, isso é melhor do que travar.
    for inp in session.get_inputs()[1:]:
        name = inp.name.lower()
        shape = inp.shape

        fx = args.fx if args.fx > 0 else input_w * 0.9
        fy = args.fy if args.fy > 0 else input_w * 0.9
        cx = args.cx if args.cx >= 0 else input_w / 2.0
        cy = args.cy if args.cy >= 0 else input_h / 2.0

        K = np.array(
            [
                [fx, 0.0, cx],
                [0.0, fy, cy],
                [0.0, 0.0, 1.0],
            ],
            dtype=np.float32,
        )

        if "k" == name or "intr" in name or "camera" in name:
            if len(shape) == 3:
                feed[inp.name] = K[None, ...]
            else:
                feed[inp.name] = K
        else:
            # Fallback para input extra desconhecido
            # Evita crash, mas imprime para a gente ajustar se precisar.
            print(f"[WARN] Input extra desconhecido: {inp.name}, shape={inp.shape}. Enviando zeros.")
            concrete_shape = []
            for dim in shape:
                concrete_shape.append(dim if isinstance(dim, int) and dim > 0 else 1)
            feed[inp.name] = np.zeros(concrete_shape, dtype=np.float32)

    t0 = time.time()
    outputs = session.run(None, feed)
    infer_ms = (time.time() - t0) * 1000.0

    output_names = [o.name for o in session.get_outputs()]
    depth = extract_depth_from_outputs(outputs, output_names)
    depth = maybe_invert_depth_if_needed(depth, auto_invert=args.inv_depth)

    # Redimensiona para o preview original
    if depth.shape[:2] != preview_bgr.shape[:2]:
        depth = cv2.resize(
            depth.astype(np.float32),
            (preview_bgr.shape[1], preview_bgr.shape[0]),
            interpolation=cv2.INTER_CUBIC,
        )

    # Remove valores absurdos só para visual/estatística
    depth = depth.astype(np.float32)
    depth[~np.isfinite(depth)] = 0.0
    depth[depth < 0] = 0.0

    return preview_bgr, depth, infer_ms


def calc_stats(depth_m: np.ndarray, infer_ms: float):
    valid = np.isfinite(depth_m) & (depth_m > 0)

    if np.count_nonzero(valid) < 20:
        return {
            "p01": 0.0,
            "p05": 0.0,
            "p50": 0.0,
            "p95": 0.0,
            "p99": 0.0,
            "mean": 0.0,
            "std": 0.0,
            "infer_ms": infer_ms,
            "valid_pct": 0.0,
        }

    vals = depth_m[valid]
    return {
        "p01": float(np.percentile(vals, 1)),
        "p05": float(np.percentile(vals, 5)),
        "p50": float(np.percentile(vals, 50)),
        "p95": float(np.percentile(vals, 95)),
        "p99": float(np.percentile(vals, 99)),
        "mean": float(np.mean(vals)),
        "std": float(np.std(vals)),
        "infer_ms": float(infer_ms),
        "valid_pct": float(np.mean(valid) * 100.0),
    }


def resize_to_height(img: np.ndarray, target_h: int):
    h, w = img.shape[:2]
    if h == target_h:
        return img

    scale = target_h / h
    new_w = max(1, int(w * scale))
    return cv2.resize(img, (new_w, target_h), interpolation=cv2.INTER_AREA)


def draw_header(canvas, lines):
    header_h = 24 + 24 * len(lines)
    cv2.rectangle(canvas, (0, 0), (canvas.shape[1], header_h), (0, 0, 0), -1)

    y = 24
    for line in lines:
        cv2.putText(
            canvas,
            line,
            (12, y),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.55,
            (255, 255, 255),
            1,
            cv2.LINE_AA,
        )
        y += 24

    return canvas


def compose_view(preview_bgr, depth_m, image_path, idx, total, mode, invert_view, stats, args):
    depth_norm = robust_normalize_for_view(depth_m, invert=invert_view)
    depth_color = depth_to_colormap(depth_norm)
    bands_color = make_metric_bands(depth_m, args.near_m, args.mid_m, args.far_m)

    right_src = depth_color if mode == "depth" else bands_color

    left = resize_to_height(preview_bgr, args.view_h)
    right = resize_to_height(right_src, args.view_h)

    h = min(left.shape[0], right.shape[0])
    left = left[:h]
    right = right[:h]

    canvas = np.hstack([left, right])

    mode_name = "depth colormap" if mode == "depth" else "metric bands"

    lines = [
        f"{idx + 1}/{total} - {image_path.name}",
        f"modo={mode_name} | p05={stats['p05']:.2f}m p50={stats['p50']:.2f}m p95={stats['p95']:.2f}m valid={stats['valid_pct']:.1f}% infer={stats['infer_ms']:.1f}ms",
        f"bands: <{args.near_m:.2f}m | {args.near_m:.2f}-{args.mid_m:.2f}m | {args.mid_m:.2f}-{args.far_m:.2f}m | >{args.far_m:.2f}m",
        "N/SPACE prox | A ant | M modo | I inverte visual | S salvar | Q sair",
    ]

    return draw_header(canvas, lines)


def parse_mean_std(text):
    vals = [float(x.strip()) for x in text.split(",")]
    if len(vals) != 3:
        raise argparse.ArgumentTypeError("Use 3 valores separados por vírgula. Ex: 0.485,0.456,0.406")
    return vals


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_dir", required=True, help="Pasta com previews PNG/JPG")
    parser.add_argument("--model_path", required=True, help="Caminho do arquivo .onnx")
    parser.add_argument("--cpu", action="store_true", help="Força CPUExecutionProvider")

    # Usado se o modelo tiver tamanho dinâmico
    parser.add_argument("--input_h", type=int, default=384)
    parser.add_argument("--input_w", type=int, default=512)

    # Normalização padrão ImageNet. Se o repo do ONNX pedir outra, ajustamos.
    parser.add_argument("--mean", type=parse_mean_std, default=[0.485, 0.456, 0.406])
    parser.add_argument("--std", type=parse_mean_std, default=[0.229, 0.224, 0.225])

    # Intrínsecos aproximados se o ONNX pedir K/intrinsics
    parser.add_argument("--fx", type=float, default=-1.0)
    parser.add_argument("--fy", type=float, default=-1.0)
    parser.add_argument("--cx", type=float, default=-1.0)
    parser.add_argument("--cy", type=float, default=-1.0)

    # Se o output for inverso/disparity, ativa isto
    parser.add_argument("--inv_depth", action="store_true")

    # Bandas métricas para visual
    parser.add_argument("--near_m", type=float, default=0.6)
    parser.add_argument("--mid_m", type=float, default=1.1)
    parser.add_argument("--far_m", type=float, default=1.8)

    parser.add_argument("--view_h", type=int, default=560)
    parser.add_argument("--save_dir", default="unidepth_onnx_viewer_saves")

    args = parser.parse_args()

    input_dir = Path(args.input_dir)
    model_path = Path(args.model_path)
    save_dir = Path(args.save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)

    if not model_path.exists():
        raise FileNotFoundError(f"Modelo ONNX não encontrado: {model_path}")

    image_paths = sorted([
        p for p in input_dir.rglob("*")
        if p.suffix.lower() in IMG_EXTS
    ])

    if not image_paths:
        raise RuntimeError(f"Nenhuma imagem encontrada em: {input_dir}")

    providers = get_providers(force_cpu=args.cpu)

    print(f"[INFO] imagens: {len(image_paths)}")
    print(f"[INFO] model_path: {model_path}")
    print(f"[INFO] providers: {providers}")
    print(f"[INFO] ONNX Runtime providers disponiveis: {ort.get_available_providers()}")

    session = ort.InferenceSession(str(model_path), providers=providers)
    print_model_io(session)

    idx = 0
    mode = "depth"
    invert_view = False

    cached_path = None
    cached_data = None

    cv2.namedWindow("UniDepth ONNX Metric Viewer", cv2.WINDOW_NORMAL)

    while True:
        image_path = image_paths[idx]

        if cached_path != image_path or cached_data is None:
            print(f"[RUN] {idx + 1}/{len(image_paths)} - {image_path.name}")

            try:
                preview_bgr, depth_m, infer_ms = infer_depth(session, image_path, args)
                stats = calc_stats(depth_m, infer_ms)
                cached_data = (preview_bgr, depth_m, stats)
                cached_path = image_path
            except Exception as e:
                print(f"[ERRO] Falha em {image_path}: {e}")
                idx = min(idx + 1, len(image_paths) - 1)
                cached_path = None
                cached_data = None
                continue
        else:
            preview_bgr, depth_m, stats = cached_data

        view = compose_view(
            preview_bgr=preview_bgr,
            depth_m=depth_m,
            image_path=image_path,
            idx=idx,
            total=len(image_paths),
            mode=mode,
            invert_view=invert_view,
            stats=stats,
            args=args,
        )

        cv2.imshow("UniDepth ONNX Metric Viewer", view)
        key = cv2.waitKey(0) & 0xFF

        if key in [27, ord("q"), ord("Q")]:
            break

        elif key in [ord("n"), ord("N"), 32]:
            idx = min(idx + 1, len(image_paths) - 1)
            cached_path = None

        elif key in [ord("a"), ord("A")]:
            idx = max(idx - 1, 0)
            cached_path = None

        elif key in [ord("m"), ord("M")]:
            mode = "bands" if mode == "depth" else "depth"

        elif key in [ord("i"), ord("I")]:
            invert_view = not invert_view
            print(f"[INFO] invert_view={invert_view}")

        elif key in [ord("s"), ord("S")]:
            out_path = save_dir / f"{image_path.stem}_unidepth_onnx_{mode}.png"
            cv2.imwrite(str(out_path), view)
            print(f"[SAVE] {out_path}")

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()