import cv2
import numpy as np
import argparse
import os
# step3ortho_from_perspective.py
# ------------------------------------------------------------
# Configuration (must match forward script exactly)
# ------------------------------------------------------------

ORTHO_W = 2000
ORTHO_H = 2000

PLANE_SIZE_MM = 2000.0

CAM_POS = np.array([2050.0, 2000.0, 4000.0], dtype=float)
TARGET = np.array([PLANE_SIZE_MM / 2.0, PLANE_SIZE_MM / 2.0, 0.0], dtype=float)

FOCAL_LENGTH_MM = 4.25
SENSOR_WIDTH_MM = 7.6
SENSOR_HEIGHT_MM = 5.7

SENSOR_PIXELS_WIDTH = 2000.0
SENSOR_PIXELS_HEIGHT = SENSOR_PIXELS_WIDTH * (SENSOR_HEIGHT_MM / SENSOR_WIDTH_MM)


# ------------------------------------------------------------
# Helper functions
# ------------------------------------------------------------

def build_intrinsics():
    scale_x = SENSOR_PIXELS_WIDTH / SENSOR_WIDTH_MM
    scale_y = SENSOR_PIXELS_HEIGHT / SENSOR_HEIGHT_MM

    fx = FOCAL_LENGTH_MM * scale_x
    fy = FOCAL_LENGTH_MM * scale_y

    K = np.array([
        [fx, 0.0, 0.0],
        [0.0, fy, 0.0],
        [0.0, 0.0, 1.0]
    ], dtype=float)
    return K


def look_at(camera_pos, target_pos, world_up=np.array([0.0, 0.0, 1.0], dtype=float)):
    forward = target_pos - camera_pos
    forward = forward / np.linalg.norm(forward)

    if abs(np.dot(forward, world_up)) > 0.999:
        world_up = np.array([0.0, 1.0, 0.0], dtype=float)

    right = np.cross(world_up, forward)
    right = right / np.linalg.norm(right)

    up = np.cross(forward, right)

    R = np.vstack([right, up, forward])
    return R


# ------------------------------------------------------------
# Main inverse pipeline
# ------------------------------------------------------------

def main():
    # ----------------------------------------
    # Command-line argument for input PNG
    # ----------------------------------------
    parser = argparse.ArgumentParser(description="Recover orthophoto from perspective image.")
    parser.add_argument("input_png", help="Path to the perspective PNG")
    args = parser.parse_args()

    input_path = args.input_png

    # ----------------------------------------
    # Build output filename
    # ----------------------------------------
    base = os.path.splitext(os.path.basename(input_path))[0]

    # Take everything before the FIRST underscore
    if "_" in base:
        prefix = base.split("_")[0]
    else:
        prefix = base  # fallback

    output_path = f"{prefix}_ortho.png"

    # ----------------------------------------
    # Load final perspective image (2000×2000)
    # ----------------------------------------
    img = cv2.imread(input_path, cv2.IMREAD_COLOR)
    if img is None:
        raise FileNotFoundError(f"Could not load: {input_path}")

    # 1. Undo vertical flip
    unflipped = cv2.flip(img, 0)

    # --------------------------------------------------------
    # Recompute the same homography H used in the forward pass
    # --------------------------------------------------------
    K = build_intrinsics()
    R = look_at(CAM_POS, TARGET)
    t = -R @ CAM_POS
    P = K @ np.hstack([R, t.reshape(3, 1)])

    # 3D plane corners
    plane_3d = np.float32([
        [0.0,              0.0,              0.0],
        [PLANE_SIZE_MM,    0.0,              0.0],
        [PLANE_SIZE_MM,    PLANE_SIZE_MM,    0.0],
        [0.0,              PLANE_SIZE_MM,    0.0]
    ])

    # Project corners
    projected = []
    for X in plane_3d:
        X_h = np.append(X, 1.0)
        x = P @ X_h
        x /= x[2]
        projected.append(x[:2])
    projected = np.float32(projected)

    # Compute bounding box
    xs = projected[:, 0]
    ys = projected[:, 1]
    min_x, max_x = np.min(xs), np.max(xs)
    min_y, max_y = np.min(ys), np.max(ys)

    out_w = int(np.ceil(max_x - min_x))
    out_h = int(np.ceil(max_y - min_y))

    shift = np.array([min_x, min_y], dtype=float)
    projected_shifted = projected - shift

    # Source = orthophoto corners
    src_pts = np.float32([
        [0.0,          0.0],
        [ORTHO_W - 1,  0.0],
        [ORTHO_W - 1,  ORTHO_H - 1],
        [0.0,          ORTHO_H - 1]
    ])

    # Destination = projected plane corners
    dst_pts = projected_shifted

    # Forward homography
    H, _ = cv2.findHomography(src_pts, dst_pts)

    # Inverse homography
    H_inv = np.linalg.inv(H)

    # 2. Undo resize
    unscaled = cv2.resize(
        unflipped,
        (out_w, out_h),
        interpolation=cv2.INTER_LINEAR
    )

    # 3. Warp back to orthophoto space
    ortho = cv2.warpPerspective(unscaled, H_inv,     (ORTHO_W, ORTHO_H),
        flags=cv2.INTER_NEAREST)

    # Save
    cv2.imwrite(output_path, ortho)
    print(f"Saved orthophoto to: {output_path}")
    print("Recovered ortho size:", ortho.shape[1], "x", ortho.shape[0])

    # Suggest next step
    print(f"next step: python step4analyzeOrtho.py {base}.png")

    # ------------------------------------------------------------
    # Show images (scaled to 33% for display only)
    # ------------------------------------------------------------
    scale = 0.33

    display_unscaled = cv2.resize(
        unscaled,
        (int(unscaled.shape[1] * scale), int(unscaled.shape[0] * scale)),
        interpolation=cv2.INTER_AREA
    )

    display_ortho = cv2.resize(
        ortho,
        (int(ortho.shape[1] * scale), int(ortho.shape[0] * scale)),
        interpolation=cv2.INTER_AREA
    )

    cv2.imshow("Unscaled synthetic (scaled)", display_unscaled)
    cv2.imshow("Recovered orthophoto (scaled)", display_ortho)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()