#!/usr/bin/env python3
import cv2
import numpy as np
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from surround_view import FisheyeCameraModel
import surround_view.param_settings as settings

def main():
    names = settings.camera_names
    images = [os.path.join(os.getcwd(), "images", name + ".png") for name in names]
    yamls = [os.path.join(os.getcwd(), "yaml", name + ".yaml") for name in names]
    
    camera_models = [FisheyeCameraModel(camera_file, camera_name) for camera_file, camera_name in zip(yamls, names)]
    
    projected = []
    for image_file, camera in zip(images, camera_models):
        img = cv2.imread(image_file)
        img = camera.undistort(img)
        img = camera.project(img)
        img = camera.flip(img)
        projected.append(img)
    
    front, back, left, right = projected
    
    # Generate weights for each corner (matching birdview.py logic)
    from surround_view.utils import get_weight_mask_matrix
    
    G0, M0 = get_weight_mask_matrix(front[:, :settings.xl], left[:settings.yt, :])
    G1, M1 = get_weight_mask_matrix(front[:, settings.xr:], right[:settings.yt, :])
    G2, M2 = get_weight_mask_matrix(back[:, :settings.xl], left[settings.yb:, :])
    G3, M3 = get_weight_mask_matrix(back[:, settings.xr:], right[settings.yb:, :])
    
    # Save as separate PNG files (each is a single-channel float matrix)
    # Convert to 16-bit PNG for precision
    G0_16 = (G0 * 65535).astype(np.uint16)
    G1_16 = (G1 * 65535).astype(np.uint16)
    G2_16 = (G2 * 65535).astype(np.uint16)
    G3_16 = (G3 * 65535).astype(np.uint16)
    
    cv2.imwrite("weights_corner_0_left_top.png", G0_16)
    cv2.imwrite("weights_corner_1_right_top.png", G1_16)
    cv2.imwrite("weights_corner_2_left_bottom.png", G2_16)
    cv2.imwrite("weights_corner_3_right_bottom.png", G3_16)
    
    print(f"Generated weights for 4 corners:")
    print(f"  - weights_corner_0_left_top.png: {G0.shape} (front+left)")
    print(f"  - weights_corner_1_right_top.png: {G1.shape} (front+right)")
    print(f"  - weights_corner_2_left_bottom.png: {G2.shape} (back+left)")
    print(f"  - weights_corner_3_right_bottom.png: {G3.shape} (back+right)")

if __name__ == "__main__":
    main()