import cv2
import numpy as np
import argparse
import os
# step2perspective_from_ortho.py
# ------------------------------------------------------------
# Configuration (constants only)
# ------------------------------------------------------------

PLANE_SIZE_MM = 2000.0

# Camera position in mm (x horizontal, y vertical, z height)
CAM_POS = np.array([2050.0, 2000.0, 4000.0], dtype=float)

# Camera looks at the center of the square
TARGET = np.array([PLANE_SIZE_MM / 2.0, PLANE_SIZE_MM / 2.0, 0.0], dtype=float)

# iPhone-like physical camera parameters
FOCAL_LENGTH_MM = 4.25
SENSOR_WIDTH_MM = 7.6
SENSOR_HEIGHT_MM = 5.7

# Synthetic camera pixel density
SENSOR_PIXELS_WIDTH = 2000.0
SENSOR_PIXELS_HEIGHT = SENSOR_PIXELS_WIDTH * (SENSOR_HEIGHT_MM / SENSOR_WIDTH_MM)

FINAL_SIZE = 2000   # final output resolution


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

def build_intrinsics():
    """Build a pinhole camera intrinsic matrix using physical focal length."""
    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)):
    """Build rotation matrix for a camera looking at a target."""
    forward = target_pos - camera_pos
    forward = forward / np.linalg.norm(forward)

    if np.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 pipeline
# ------------------------------------------------------------

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

    input_path = args.input_png

    # Build output filename automatically
    base = os.path.splitext(os.path.basename(input_path))[0]
    output_path = f"{base}_perspective.png"

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

    h, w = img.shape[:2]

    # Build camera intrinsics
    K = build_intrinsics()

    # Extrinsics
    R = look_at(CAM_POS, TARGET)
    t = -R @ CAM_POS

    # Projection matrix
    P = K @ np.hstack([R, t.reshape(3, 1)])

    # ----------------------------------------
    # Define the 3D plane
    # ----------------------------------------
    plane_3d = np.float32([
        [0.0,              0.0,              0.0],   # top-left
        [PLANE_SIZE_MM,    0.0,              0.0],   # top-right
        [PLANE_SIZE_MM,    PLANE_SIZE_MM,    0.0],   # bottom-right
        [0.0,              PLANE_SIZE_MM,    0.0]    # bottom-left
    ])

    # Project plane 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 so bounding box starts at (0,0)
    shift = np.array([min_x, min_y], dtype=float)
    projected_shifted = projected - shift

    # Source image corners
    src_pts = np.float32([
        [0.0,      0.0],
        [w - 1.0,  0.0],
        [w - 1.0,  h - 1.0],
        [0.0,      h - 1.0]
    ])

    # Destination points
    dst_pts = projected_shifted

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

    # Warp
    synthetic = cv2.warpPerspective(img, H, (out_w, out_h))

    # ----------------------------------------
    # Resize to 2000x2000
    # ----------------------------------------
    synthetic_resized = cv2.resize(
        synthetic,
        (FINAL_SIZE, FINAL_SIZE),
        interpolation=cv2.INTER_LINEAR
    )

    # ----------------------------------------
    # Flip vertically
    # ----------------------------------------
    synthetic_final = cv2.flip(synthetic_resized, 0)

    # Save
    cv2.imwrite(output_path, synthetic_final)
    print(f"Saved perspective view to: {output_path}")
    print(f"Final size: {synthetic_final.shape[1]} x {synthetic_final.shape[0]}")

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

    # Show images
    cv2.imshow("Original bocce field", img)
    cv2.imshow("Synthetic perspective view", synthetic_final)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()