import argparse
import json
import re
from pathlib import Path

import cv2
import numpy as np


# ============================================================
# RAW10
# ============================================================

def unpack_raw10_packed(raw: bytes, width: int, height: int) -> np.ndarray:
    arr = np.frombuffer(raw, dtype=np.uint8)

    pixel_count = width * height
    expected_bytes = (pixel_count // 4) * 5

    if pixel_count % 4 != 0:
        raise RuntimeError(f"width*height precisa ser múltiplo de 4. Recebido: {pixel_count}")

    if arr.size < expected_bytes:
        raise RuntimeError(
            f"RAW10 menor que esperado. bytes={arr.size}, esperado={expected_bytes}, "
            f"width={width}, height={height}"
        )

    arr = arr[:expected_bytes]
    groups = arr.reshape(-1, 5).astype(np.uint16)

    p0 = (groups[:, 0] << 2) | ((groups[:, 4] >> 0) & 0x03)
    p1 = (groups[:, 1] << 2) | ((groups[:, 4] >> 2) & 0x03)
    p2 = (groups[:, 2] << 2) | ((groups[:, 4] >> 4) & 0x03)
    p3 = (groups[:, 3] << 2) | ((groups[:, 4] >> 6) & 0x03)

    out = np.empty(groups.shape[0] * 4, dtype=np.uint16)
    out[0::4] = p0
    out[1::4] = p1
    out[2::4] = p2
    out[3::4] = p3

    return out.reshape(height, width)


def normalize_to_u8(img: np.ndarray, p_low=1.0, p_high=99.0) -> np.ndarray:
    arr = img.astype(np.float32)
    valid = np.isfinite(arr)

    if np.count_nonzero(valid) < 20:
        return np.zeros(arr.shape[:2], dtype=np.uint8)

    vals = arr[valid]
    lo = np.percentile(vals, p_low)
    hi = np.percentile(vals, p_high)

    out = (arr - lo) / (hi - lo + 1e-6)
    out = np.clip(out, 0.0, 1.0)

    return (out * 255).astype(np.uint8)


def raw10_bin_to_gray(path: Path, width: int, height: int, use_clahe=True) -> np.ndarray:
    raw = path.read_bytes()
    mono10 = unpack_raw10_packed(raw, width, height)
    gray = normalize_to_u8(mono10)

    if use_clahe:
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        gray = clahe.apply(gray)

    return gray


# ============================================================
# Pairing CAM bins
# ============================================================

def clean_stem_for_pair(path: Path, cam_key: str):
    s = path.stem

    variants = [
        cam_key,
        cam_key.lower(),
        cam_key.replace("_", ""),
        cam_key.replace("_", "").lower(),
    ]

    for v in variants:
        s = s.replace(v, "")

    s = re.sub(r"[_\-\s]+", "_", s).strip("_").lower()
    return s


def find_cam_bins(root_dir: Path, cam_key: str):
    return sorted([
        p for p in root_dir.rglob("*.bin")
        if cam_key.lower() in p.name.lower()
    ])


def pair_cam_bins(root_dir: Path, left_cam: str, right_cam: str):
    left_bins = find_cam_bins(root_dir, left_cam)
    right_bins = find_cam_bins(root_dir, right_cam)

    right_by_folder_key = {}
    right_by_key = {}

    for rp in right_bins:
        key = clean_stem_for_pair(rp, right_cam)
        right_by_folder_key[(rp.parent, key)] = rp
        right_by_key.setdefault(key, rp)

    pairs = []

    for lp in left_bins:
        key = clean_stem_for_pair(lp, left_cam)

        rp = right_by_folder_key.get((lp.parent, key))
        if rp is None:
            rp = right_by_key.get(key)

        if rp is None:
            same_folder = [r for r in right_bins if r.parent == lp.parent]
            if len(same_folder) == 1:
                rp = same_folder[0]

        if rp is not None:
            pairs.append((lp, rp))

    if pairs:
        return pairs, left_bins, right_bins, "name/key"

    n = min(len(left_bins), len(right_bins))
    return list(zip(left_bins[:n], right_bins[:n])), left_bins, right_bins, "order"


# ============================================================
# Calibration loading
# ============================================================

def load_stereo_calib(calib_path: Path):
    data = np.load(str(calib_path), allow_pickle=True)

    required = ["map1x", "map1y", "map2x", "map2y", "image_size"]
    for k in required:
        if k not in data:
            raise RuntimeError(f"Calibração sem chave obrigatória: {k}")

    calib = {
        "map1x": data["map1x"],
        "map1y": data["map1y"],
        "map2x": data["map2x"],
        "map2y": data["map2y"],
        "image_size": tuple(data["image_size"].astype(int).tolist()),
    }

    for k in ["left_cam", "right_cam", "rms_left", "rms_right", "rms_stereo", "T"]:
        if k in data:
            calib[k] = data[k]

    return calib


def rectify_pair(left_gray, right_gray, calib):
    image_w, image_h = calib["image_size"]

    if left_gray.shape[::-1] != (image_w, image_h):
        left_gray = cv2.resize(left_gray, (image_w, image_h), interpolation=cv2.INTER_AREA)

    if right_gray.shape[::-1] != (image_w, image_h):
        right_gray = cv2.resize(right_gray, (image_w, image_h), interpolation=cv2.INTER_AREA)

    left_rect = cv2.remap(
        left_gray,
        calib["map1x"],
        calib["map1y"],
        interpolation=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=0,
    )

    right_rect = cv2.remap(
        right_gray,
        calib["map2x"],
        calib["map2y"],
        interpolation=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=0,
    )

    return left_rect, right_rect


# ============================================================
# SGBM
# ============================================================

def make_sgbm(args):
    num_disp = max(16, int(round(args.num_disp / 16)) * 16)

    block_size = max(3, int(args.block_size))
    if block_size % 2 == 0:
        block_size += 1

    matcher = cv2.StereoSGBM_create(
        minDisparity=args.min_disp,
        numDisparities=num_disp,
        blockSize=block_size,
        P1=8 * block_size * block_size,
        P2=32 * block_size * block_size,
        disp12MaxDiff=1,
        uniquenessRatio=args.uniqueness,
        speckleWindowSize=args.speckle_window,
        speckleRange=args.speckle_range,
        preFilterCap=63,
        mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY,
    )

    return matcher, num_disp, block_size


def compute_disparity(left_rect, right_rect, args):
    matcher, num_disp, block_size = make_sgbm(args)

    disp_raw = matcher.compute(left_rect, right_rect).astype(np.float32) / 16.0

    valid = disp_raw > args.min_valid_disp

    disp_vis = disp_raw.copy()
    disp_vis[~valid] = 0.0

    if np.count_nonzero(valid) > 20:
        vals = disp_vis[valid]
        p2 = np.percentile(vals, 2)
        p98 = np.percentile(vals, 98)
        disp_norm = (disp_vis - p2) / (p98 - p2 + 1e-6)
        disp_norm = np.clip(disp_norm, 0.0, 1.0)
    else:
        disp_norm = np.zeros_like(disp_vis, dtype=np.float32)

    disp_color = cv2.applyColorMap((disp_norm * 255).astype(np.uint8), cv2.COLORMAP_TURBO)

    valid_mask = np.zeros_like(disp_color)
    valid_mask[valid] = (255, 255, 255)

    stats = {
        "num_disp": num_disp,
        "block_size": block_size,
        "valid_pct": float(np.mean(valid) * 100.0),
        "disp_p05": float(np.percentile(disp_vis[valid], 5)) if np.count_nonzero(valid) > 20 else 0.0,
        "disp_p50": float(np.percentile(disp_vis[valid], 50)) if np.count_nonzero(valid) > 20 else 0.0,
        "disp_p95": float(np.percentile(disp_vis[valid], 95)) if np.count_nonzero(valid) > 20 else 0.0,
    }

    return disp_raw, disp_color, valid_mask, valid, stats


# ============================================================
# Visualization
# ============================================================

def draw_epipolar_lines(img_bgr, step=40):
    out = img_bgr.copy()
    h, w = out.shape[:2]

    for y in range(0, h, step):
        color = (0, 255, 255) if (y // step) % 2 == 0 else (255, 255, 0)
        cv2.line(out, (0, y), (w, y), color, 1, cv2.LINE_AA)

    return out


def resize_to_height(img, target_h):
    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 make_overlay(left_bgr, disp_color, alpha=0.45):
    disp_resized = cv2.resize(
        disp_color,
        (left_bgr.shape[1], left_bgr.shape[0]),
        interpolation=cv2.INTER_AREA,
    )
    return cv2.addWeighted(left_bgr, 1.0 - alpha, disp_resized, alpha, 0)


def compose_view(left_rect, right_rect, disp_color, valid_mask, pair, idx, total, mode, stats, args):
    left_bgr = cv2.cvtColor(left_rect, cv2.COLOR_GRAY2BGR)
    right_bgr = cv2.cvtColor(right_rect, cv2.COLOR_GRAY2BGR)

    if args.lines:
        left_bgr = draw_epipolar_lines(left_bgr, step=args.line_step)
        right_bgr = draw_epipolar_lines(right_bgr, step=args.line_step)

    if mode == "disp":
        third = disp_color
        mode_name = "disparity"
    elif mode == "mask":
        third = valid_mask
        mode_name = "valid mask"
    elif mode == "overlay":
        third = make_overlay(cv2.cvtColor(left_rect, cv2.COLOR_GRAY2BGR), disp_color)
        mode_name = "overlay"
    else:
        diff = cv2.absdiff(left_rect, right_rect)
        third = cv2.cvtColor(diff, cv2.COLOR_GRAY2BGR)
        mode_name = "rect diff"

    left = resize_to_height(left_bgr, args.view_h)
    right = resize_to_height(right_bgr, args.view_h)
    third = resize_to_height(third, args.view_h)

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

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

    lp, rp = pair

    lines = [
        f"{idx + 1}/{total} | L={lp.name} | R={rp.name}",
        f"mode={mode_name} | valid={stats['valid_pct']:.1f}% | disp p05={stats['disp_p05']:.2f} p50={stats['disp_p50']:.2f} p95={stats['disp_p95']:.2f}",
        f"numDisp={stats['num_disp']} | block={stats['block_size']} | uniqueness={args.uniqueness} | lines={args.lines}",
        "N/SPACE prox | A ant | M modo | L linhas | [ ] numDisp | - + block | S salvar | Q sair",
    ]

    return draw_header(canvas, lines)


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

def main():
    parser = argparse.ArgumentParser()

    parser.add_argument("--root_dir", required=True)
    parser.add_argument("--calib_path", required=True)

    parser.add_argument("--left_cam", default="CAM_C")
    parser.add_argument("--right_cam", default="CAM_B")

    parser.add_argument("--width", type=int, default=1280)
    parser.add_argument("--height", type=int, default=800)

    parser.add_argument("--view_h", type=int, default=480)

    parser.add_argument("--num_disp", type=int, default=128)
    parser.add_argument("--block_size", type=int, default=7)
    parser.add_argument("--min_disp", type=int, default=0)
    parser.add_argument("--min_valid_disp", type=float, default=1.0)

    parser.add_argument("--uniqueness", type=int, default=8)
    parser.add_argument("--speckle_window", type=int, default=80)
    parser.add_argument("--speckle_range", type=int, default=2)

    parser.add_argument("--no_clahe", action="store_true")
    parser.add_argument("--lines", action="store_true")
    parser.add_argument("--line_step", type=int, default=40)

    parser.add_argument("--save_dir", default="stereo_rectified_sgbm_saves")

    args = parser.parse_args()

    root_dir = Path(args.root_dir)
    calib_path = Path(args.calib_path)
    save_dir = Path(args.save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)

    calib = load_stereo_calib(calib_path)

    print(f"[INFO] calib_path: {calib_path}")
    print(f"[INFO] calib image_size: {calib['image_size']}")

    if "rms_left" in calib:
        print(f"[INFO] rms_left: {float(calib['rms_left']):.6f}")
    if "rms_right" in calib:
        print(f"[INFO] rms_right: {float(calib['rms_right']):.6f}")
    if "rms_stereo" in calib:
        print(f"[INFO] rms_stereo: {float(calib['rms_stereo']):.6f}")
    if "T" in calib:
        print(f"[INFO] T: {np.array(calib['T']).ravel()}")

    pairs, left_bins, right_bins, pair_mode = pair_cam_bins(root_dir, args.left_cam, args.right_cam)

    print(f"[INFO] root_dir: {root_dir}")
    print(f"[INFO] left_cam: {args.left_cam} | arquivos: {len(left_bins)}")
    print(f"[INFO] right_cam: {args.right_cam} | arquivos: {len(right_bins)}")
    print(f"[INFO] pares: {len(pairs)}")
    print(f"[INFO] pareamento: {pair_mode}")

    if not pairs:
        raise RuntimeError("Nenhum par encontrado.")

    idx = 0
    mode = "disp"

    cached_key = None
    cached_data = None

    cv2.namedWindow("Stereo RAW10 Rectified SGBM Viewer", cv2.WINDOW_NORMAL)

    while True:
        pair = pairs[idx]
        lp, rp = pair

        key_cache = (
            str(lp),
            str(rp),
            args.num_disp,
            args.block_size,
            args.uniqueness,
            args.speckle_window,
            args.speckle_range,
            args.min_disp,
            args.min_valid_disp,
            args.no_clahe,
            args.width,
            args.height,
            str(args.calib_path),
        )

        if key_cache != cached_key:
            print(f"[RUN] {idx + 1}/{len(pairs)} - L={lp.name} | R={rp.name}")

            try:
                left_gray = raw10_bin_to_gray(
                    lp,
                    width=args.width,
                    height=args.height,
                    use_clahe=not args.no_clahe,
                )

                right_gray = raw10_bin_to_gray(
                    rp,
                    width=args.width,
                    height=args.height,
                    use_clahe=not args.no_clahe,
                )

                left_rect, right_rect = rectify_pair(left_gray, right_gray, calib)

                _, disp_color, valid_mask, valid, stats = compute_disparity(left_rect, right_rect, args)

                cached_data = (left_rect, right_rect, disp_color, valid_mask, stats)
                cached_key = key_cache

            except Exception as e:
                print(f"[ERRO] Falha processando par:")
                print(f"       L={lp}")
                print(f"       R={rp}")
                print(f"       erro={e}")
                idx = min(idx + 1, len(pairs) - 1)
                cached_key = None
                cached_data = None
                continue

        else:
            left_rect, right_rect, disp_color, valid_mask, stats = cached_data

        view = compose_view(
            left_rect=left_rect,
            right_rect=right_rect,
            disp_color=disp_color,
            valid_mask=valid_mask,
            pair=pair,
            idx=idx,
            total=len(pairs),
            mode=mode,
            stats=stats,
            args=args,
        )

        cv2.imshow("Stereo RAW10 Rectified SGBM 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(pairs) - 1)
            cached_key = None

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

        elif key in [ord("m"), ord("M")]:
            if mode == "disp":
                mode = "mask"
            elif mode == "mask":
                mode = "overlay"
            elif mode == "overlay":
                mode = "diff"
            else:
                mode = "disp"

        elif key in [ord("l"), ord("L")]:
            args.lines = not args.lines
            print(f"[PARAM] lines={args.lines}")

        elif key == ord("["):
            args.num_disp = max(16, args.num_disp - 16)
            cached_key = None
            print(f"[PARAM] num_disp={args.num_disp}")

        elif key == ord("]"):
            args.num_disp = min(512, args.num_disp + 16)
            cached_key = None
            print(f"[PARAM] num_disp={args.num_disp}")

        elif key in [ord("-"), ord("_")]:
            args.block_size = max(3, args.block_size - 2)
            if args.block_size % 2 == 0:
                args.block_size -= 1
            cached_key = None
            print(f"[PARAM] block_size={args.block_size}")

        elif key in [ord("+"), ord("=")]:
            args.block_size = min(31, args.block_size + 2)
            if args.block_size % 2 == 0:
                args.block_size += 1
            cached_key = None
            print(f"[PARAM] block_size={args.block_size}")

        elif key in [ord("s"), ord("S")]:
            out_path = save_dir / f"stereo_rectified_{idx:04d}_{mode}.png"
            cv2.imwrite(str(out_path), view)
            print(f"[SAVE] {out_path}")

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()