// Local folder : ~/RUST/stereo_calib

// make sure create a subfolder ~/RUST/stereo_calib/data, and copy jpg files created by rover_vo to here.

// CRITICAL FIX: The auto-generated type aliases (VectorOfVectorOfPoint3f, etc.) are causing E0432 errors.
// We must fall back to the generic 'Vector' type with the specific Point types.
use opencv::{
    calib3d::{
        self,
        calibrate_camera, find_chessboard_corners, stereo_calibrate,
        CALIB_FIX_K1, CALIB_FIX_K2, CALIB_FIX_K3, CALIB_FIX_K4, CALIB_FIX_K5, CALIB_FIX_K6,
        CALIB_ZERO_TANGENT_DIST,
    },
    core::{
        self,
        // Core items (structs, constants, enums)
        Mat,
        Point2f,
        Point3f,
        Size2i,
        Vector,
        CV_64F,
        TermCriteria_Type,
        TermCriteria,
        Code,
    },
    imgcodecs,
    imgproc::{self, COLOR_BGR2GRAY, corner_sub_pix},
    prelude::*,
    Result,
    // --- NEW IMPORT FOR SAVING DATA ---
    core::{FileStorage, FileStorage_Mode::WRITE}, // Removed FileNode as it was unused
};


// --- Custom Function to Convert Mat (CV_64F) to String ---
// This bypasses unreliable Mat::dump() or Mat::to_string() by directly
// accessing the underlying Mat data pointer. This assumes the Mat is of type CV_64F.
fn mat_to_string(mat: &Mat) -> Result<String, opencv::Error> {
    
    let mat_type = mat.typ(); 
    
    if mat_type != CV_64F {
         return Err(opencv::Error::new(core::StsBadArg, "Mat type is not CV_64F (f64). Cannot print data."));
    }

    let is_continuous = mat.is_continuous(); 
    if !is_continuous {
        return Err(opencv::Error::new(core::StsBadArg, "Mat is not continuous"));
    }
    
    let total_elements = mat.total() as usize; 
    let cols = mat.cols();
    let rows = mat.rows(); 

    // Use data_bytes() which is the correct method for raw data pointer access
    let data_slice: &[f64] = unsafe {
        // data_bytes() returns Result<*const u8>, so we use '?' here to unwrap the Result
        // We then cast the *const u8 pointer to *const f64
        std::slice::from_raw_parts(
            mat.data_bytes()?.as_ptr() as *const f64, 
            total_elements,
        )
    };

    // Format the data into a string
    let mut result = String::new();
    result.push_str("[\n");
    for i in 0..rows {
        result.push_str("  [");
        for j in 0..cols {
            let index = i as usize * cols as usize + j as usize;
            result.push_str(&format!("{:.6}", data_slice[index])); // Use a fixed precision for clean output
            if j < cols - 1 {
                result.push_str(", ");
            }
        }
        result.push_str("]");
        if i < rows - 1 {
            result.push_str(",");
        }
        result.push_str("\n");
    }
    result.push_str("]");

    Ok(result)
}


// FIX E0432: Manually define the aliases that the compiler is failing to resolve from 'opencv::'
type VectorOfVectorOfPoint3f = Vector<Vector<Point3f>>;
type VectorOfVectorOfPoint2f = Vector<Vector<Point2f>>;
type VectorOfPoint2f = Vector<Point2f>;

// --- NEW FUNCTION TO SAVE MATRICES TO YAML ---
fn save_calibration_data(
    filename: &str, 
    rms_error: f64,
    size: Size2i,
    k1: &Mat, d1: &Mat, 
    k2: &Mat, d2: &Mat, 
    r: &Mat, t: &Mat
) -> Result<()> {
    
    // FIX E0308: Convert FileStorage_Mode::WRITE to i32 using .into()
    let mut fs = FileStorage::new(filename, WRITE.into(), "utf-8")?;

    // Save configuration and results
    fs.write_f64("rms_error", rms_error)?;
    // FIX E0599: Use write_i32 instead of write_int
    fs.write_i32("image_width", size.width)?;
    fs.write_i32("image_height", size.height)?;
    
    // Write matrices (FIX E0277: Passed Mat references, not references to Mat references)
    fs.write_mat("K1", k1)?;
    fs.write_mat("D1", d1)?;
    fs.write_mat("K2", k2)?;
    fs.write_mat("D2", d2)?;
    fs.write_mat("R", r)?;
    fs.write_mat("T", t)?;

    println!("\nCalibration data successfully saved to '{}'.", filename);
    Ok(())
}


// --- Calibration Configuration (ADJUST THESE) ---
// Number of inner corners (width, height of the checkerboard pattern)
const CHESSBOARD_SIZE: Size2i = Size2i::new(9, 6); 
// Size of a single square in real-world units (e.g., millimeters)
const SQUARE_SIZE_MM: f32 = 25.0;
// Total number of frames/images to use for calibration
const MAX_FRAMES: usize = 30;
// Only process every Nth image index (e.g., 5 means process 0, 5, 10, ...)
const FRAME_SKIP: i32 = 5;
// Max index of the image files 
const MAX_IMAGE_INDEX: i32 = 49; 
// Image size (adjust this to your actual image resolution)
const IMAGE_SIZE: Size2i = Size2i::new(640, 480);
// Prefix for image file paths. 
const DATA_PATH: &str = "data/";


// Helper function to generate 3D world coordinates for the checkerboard
fn generate_object_points() -> Result<Vector<Point3f>> {
    let mut object_points = Vector::new();
    for i in 0..CHESSBOARD_SIZE.height {
        for j in 0..CHESSBOARD_SIZE.width {
            let x = j as f32 * SQUARE_SIZE_MM;
            let y = i as f32 * SQUARE_SIZE_MM;
            object_points.push(Point3f::new(x, y, 0.0));
        }
    }
    Ok(object_points)
}

fn main() -> Result<()> {
    // --- 1. Setup ---

    // Initialize containers for calibration data
    let object_points_model = generate_object_points()?;
    
    // The manually defined types are used here
    let mut object_points_all = VectorOfVectorOfPoint3f::new();
    let mut image_points_l = VectorOfVectorOfPoint2f::new();
    let mut image_points_r = VectorOfVectorOfPoint2f::new();

    let mut calibration_frame_count = 0;

    // --- 2. Corner Detection Loop (Iterating over image files) ---
    println!("Starting image processing (will use files up to index {})", MAX_IMAGE_INDEX);

    for i in 0..=MAX_IMAGE_INDEX {
        if calibration_frame_count >= MAX_FRAMES {
            break; // Stop when the desired number of frames is captured
        }

        // Apply frame skip logic
        if i % FRAME_SKIP != 0 {
            continue;
        }

        // FIX: Corrected format! macro to include the placeholder for the index `i`
        let left_path = format!("{}/left_frame_{:05}.jpg", DATA_PATH, i);
        let right_path = format!("{}/right_frame_{:05}.jpg", DATA_PATH, i);

        // Read images from disk
        let frame_l = imgcodecs::imread(&left_path, imgcodecs::IMREAD_COLOR)?;
        let frame_r = imgcodecs::imread(&right_path, imgcodecs::IMREAD_COLOR)?;

        // Check if images were loaded successfully
        if frame_l.size()?.width <= 0 || frame_r.size()?.width <= 0 {
            if i % FRAME_SKIP == 0 {
                println!("Warning: Could not read image pair index {}. Skipping. (Tried paths: {} and {})", i, left_path, right_path);
            }
            continue;
        }

        let mut gray_l = Mat::default();
        let mut gray_r = Mat::default();
        imgproc::cvt_color(&frame_l, &mut gray_l, COLOR_BGR2GRAY, 0)?;
        imgproc::cvt_color(&frame_r, &mut gray_r, COLOR_BGR2GRAY, 0)?;

        let mut corners_l = VectorOfPoint2f::new();
        let mut corners_r = VectorOfPoint2f::new();
        
        let found_l; 
        let found_r;

        // Find chessboard corners
        found_l = find_chessboard_corners(
            &gray_l,
            CHESSBOARD_SIZE,
            &mut corners_l,
            calib3d::CALIB_CB_ADAPTIVE_THRESH + calib3d::CALIB_CB_NORMALIZE_IMAGE,
        )?;
        
        found_r = find_chessboard_corners(
            &gray_r,
            CHESSBOARD_SIZE,
            &mut corners_r,
            calib3d::CALIB_CB_ADAPTIVE_THRESH + calib3d::CALIB_CB_NORMALIZE_IMAGE,
        )?;

        if found_l && found_r {
            // Refine corner positions for better accuracy
            let criteria = TermCriteria::new(
                TermCriteria_Type::EPS as i32 | TermCriteria_Type::COUNT as i32,
                40, // max_iter
                0.001, // epsilon
            )?;

            // Use corner_sub_pix with the corrected criteria
            corner_sub_pix(&gray_l, &mut corners_l, Size2i::new(11, 11), Size2i::new(-1, -1), criteria)?;
            corner_sub_pix(&gray_r, &mut corners_r, Size2i::new(11, 11), Size2i::new(-1, -1), criteria)?;

            // Save the detected points
            object_points_all.push(object_points_model.clone());
            // Clone the vectors before moving them into the main calibration collection
            image_points_l.push(corners_l.clone());
            image_points_r.push(corners_r.clone());

            calibration_frame_count += 1;
            println!("Captured {}/{} stereo pairs (Image index {}).", calibration_frame_count, MAX_FRAMES, i);

            // Optionally draw corners for verification
            calib3d::draw_chessboard_corners(
                &mut frame_l.clone(),
                CHESSBOARD_SIZE,
                &corners_l, 
                found_l,
            )?;
            calib3d::draw_chessboard_corners(
                &mut frame_r.clone(),
                CHESSBOARD_SIZE,
                &corners_r, 
                found_r,
            )?;
        }
    }

    if calibration_frame_count < 10 {
        // FIX E0308: Use .into() to convert Code enum to i32
        return Err(opencv::Error::new(Code::StsError.into(), "Too few successful calibration pairs found. Need at least 10."));
    }

    // --- 3. Camera Calibration (Individual) ---

    // Initialize camera matrices directly with Mat::eye()
    let mut camera_matrix_l = Mat::eye(3, 3, CV_64F)?.to_mat()?;
    let mut dist_coeffs_l = Mat::default();

    let mut camera_matrix_r = Mat::eye(3, 3, CV_64F)?.to_mat()?;
    let mut dist_coeffs_r = Mat::default();
    
    // FIX: Simplified the assignment of rvecs_r, tvecs_r for cleaner syntax.
    let (mut rvecs_l, mut tvecs_l): (Vector<Mat>, Vector<Mat>) = (Vector::new(), Vector::new());
    // FIX E0283: Explicitly define the type as Vector<Mat>
    let (mut rvecs_r, mut tvecs_r): (Vector<Mat>, Vector<Mat>) = (Vector::new(), Vector::new());


    // Left Camera Calibration
    let criteria = TermCriteria::new(
        TermCriteria_Type::EPS as i32 | TermCriteria_Type::COUNT as i32,
        30,
        1e-6,
    )?;
    let _ = calibrate_camera(
        &object_points_all,
        &image_points_l,
        IMAGE_SIZE,
        &mut camera_matrix_l,
        &mut dist_coeffs_l,
        &mut rvecs_l,
        &mut tvecs_l,
        calib3d::CALIB_FIX_PRINCIPAL_POINT | CALIB_FIX_K4 | CALIB_FIX_K5 | CALIB_FIX_K6,
        criteria,
    )?;

    // Right Camera Calibration
    let criteria = TermCriteria::new(
        TermCriteria_Type::EPS as i32 | TermCriteria_Type::COUNT as i32,
        30,
        1e-6,
    )?;
    let _ = calibrate_camera(
        &object_points_all,
        &image_points_r,
        // REMOVED DUPLICATE ARGUMENT: `&image_points_r`
        IMAGE_SIZE,
        &mut camera_matrix_r,
        &mut dist_coeffs_r,
        &mut rvecs_r,
        &mut tvecs_r,
        calib3d::CALIB_FIX_PRINCIPAL_POINT | CALIB_FIX_K4 | CALIB_FIX_K5 | CALIB_FIX_K6,
        criteria,
    )?;

    // --- 4. Stereo Calibration ---
    let (mut r_stereo, mut t_stereo) = (Mat::default(), Mat::default());
    let (mut e_matrix, mut f_matrix) = (Mat::default(), Mat::default());

    let criteria = TermCriteria::new(
        TermCriteria_Type::EPS as i32 | TermCriteria_Type::COUNT as i32,
        100,
        1e-5,
    )?;

    println!("\nPerforming stereo calibration...");

    let rms = stereo_calibrate(
        &object_points_all,
        &image_points_l,
        &image_points_r,
        &mut camera_matrix_l,
        &mut dist_coeffs_l,
        &mut camera_matrix_r,
        &mut dist_coeffs_r,
        IMAGE_SIZE,
        &mut r_stereo,
        &mut t_stereo,
        &mut e_matrix,
        &mut f_matrix,
        CALIB_FIX_K1
            + CALIB_FIX_K2
            + CALIB_FIX_K3
            + CALIB_FIX_K4
            + CALIB_FIX_K5
            + CALIB_FIX_K6
            + CALIB_ZERO_TANGENT_DIST,
        criteria,
    )?;

    println!("Stereo Calibration RMS Error: {:.6}", rms);

    // --- 5. Results (Print & Save) ---
    // Print results to console
    println!("\nLeft Camera Matrix (K1):\n{}", mat_to_string(&camera_matrix_l)?);
    println!("\nLeft Distortion Coefficients (D1):\n{}", mat_to_string(&dist_coeffs_l)?);
    println!("\nRight Camera Matrix (K2):\n{}", mat_to_string(&camera_matrix_r)?);
    println!("\nRight Distortion Coefficients (D2):\n{}", mat_to_string(&dist_coeffs_r)?);

    println!("\nRotation Matrix (R):\n{}", mat_to_string(&r_stereo)?);
    println!("\nTranslation Vector (T):\n{}", mat_to_string(&t_stereo)?);

    // Save results to YAML file
    save_calibration_data(
        "stereo_params.yml", 
        rms, 
        IMAGE_SIZE, 
        &camera_matrix_l, &dist_coeffs_l, 
        &camera_matrix_r, &dist_coeffs_r, 
        &r_stereo, &t_stereo
    )?;


    println!("\nCalibration complete. The parameters are saved and ready for the next step: Rectification!");
    Ok(())
}