// Local folder : ~/RUST/stereo_rectify

use opencv::{
    // CRITICAL: Import all traits and methods needed.
    prelude::*,
    // We don't need FileNode or FileNodeTrait for this serde-based approach
    core::{self, Mat, Size2i, Rect}, 
    calib3d,
    imgcodecs,
    imgproc,
    highgui,
    Result,
    Error,
};
use serde::{Deserialize, Serialize};
use std::fs;

// --- Configuration ---
const CALIBRATION_FILE: &str = "../main_data/stereo_params.yml";
const TEST_IMAGE_INDEX: u32 = 0;
const DATA_DIR: &str = "../rover_vo";

// --- Data Structure for OpenCV Mat ---
#[derive(Debug, Deserialize, Serialize)]
struct MatData {
    dt: String,
    rows: i32,
    cols: i32,
    data: Vec<f64>,
}

// --- Data Structure for the ENTIRE YAML file ---
// **FIX IS HERE:** We rename fields to snake_case and use `serde(rename)`
// to map them to the YAML's PascalCase, fixing all 6 warnings.
#[derive(Debug, Deserialize, Serialize)]
struct StereoParams {
    image_width: i32,
    image_height: i32,
    rms_error: f64,
    
    #[serde(rename = "K1")]
    k1: MatData,
    #[serde(rename = "D1")]
    d1: MatData,
    #[serde(rename = "K2")]
    k2: MatData,
    #[serde(rename = "D2")]
    d2: MatData,
    #[serde(rename = "R")]
    r: MatData,
    #[serde(rename = "T")]
    t: MatData,
}

// --- Helper function to convert MatData struct to an OpenCV Mat ---
fn mat_from_data(data: &MatData) -> Result<Mat> {
    if data.dt != "d" { // Check that data type is f64
        return Err(Error::new(core::Code::StsError.into(), "Data type must be f64 ('d')"));
    }

    // Get the number of columns (width of each chunk)
    let cols = data.cols as usize;
    if cols == 0 {
        return Err(Error::new(core::Code::StsError.into(), "Matrix data has 0 columns."));
    }

    // Use .chunks() to create an iterator of row-slices (&[f64])
    // And .collect() them into a Vec<&[f64]>
    let rows: Vec<&[f64]> = data.data.chunks(cols).collect();

    // Pass a slice of our new Vec (&[&[f64]]) to from_slice_2d
    let mat = Mat::from_slice_2d(&rows)?;
    
    Ok(mat)
}


fn run_rectification() -> Result<()> {
    // --- 1. Load Calibration Data ---
    println!("--- 1. Loading Calibration Data from {} ---", CALIBRATION_FILE);

    // Read the entire YAML file into a string
    let yaml_content = fs::read_to_string(CALIBRATION_FILE)
        .map_err(|e| Error::new(core::Code::StsError.into(), &format!("Failed to read calibration file: {}", e)))?;
        
    // Deserialize the *entire* string into our master StereoParams struct
    let params: StereoParams = serde_yaml::from_str(&yaml_content)
        .map_err(|e| Error::new(core::Code::StsError.into(), &format!("YAML Deserialize Error: {}. Check stereo_params.yml format.", e)))?;
    
    let image_size = Size2i::new(params.image_width, params.image_height);

    // We now convert our deserialized structs into OpenCV Mats
    println!("Converting loaded data into OpenCV Mats...");
    let k1 = mat_from_data(&params.k1)?;
    let d1 = mat_from_data(&params.d1)?;
    let k2 = mat_from_data(&params.k2)?;
    let d2 = mat_from_data(&params.d2)?;
    let r_stereo = mat_from_data(&params.r)?;
    let t_stereo = mat_from_data(&params.t)?;
    
    println!("Calibration Data Loaded. RMS Error: {:.4}, Image size: {}x{}", params.rms_error, image_size.width, image_size.height);


    // --- 2. Stereo Rectification ---
    let mut r1 = Mat::default();
    let mut r2 = Mat::default();
    let mut p1 = Mat::default();
    let mut p2 = Mat::default();
    let mut q = Mat::default();
    let mut roi1 = Rect::default(); 
    let mut roi2 = Rect::default(); 

    println!("--- 2. Calculating Rectification Maps ---");

    calib3d::stereo_rectify(
        &k1, &d1, &k2, &d2, 
        image_size, 
        &r_stereo, &t_stereo, 
        &mut r1, &mut r2, &mut p1, &mut p2, &mut q, 
        calib3d::CALIB_ZERO_DISPARITY,
        -1.0, // Alpha
        image_size, 
        &mut roi1,
        &mut roi2,
    )?;

    // --- 3. Compute Undistortion and Remapping Maps ---
    let (mut map1x, mut map1y) = (Mat::default(), Mat::default());
    let (mut map2x, mut map2y) = (Mat::default(), Mat::default());
    
    calib3d::init_undistort_rectify_map(
        &k1, &d1, &r1, &p1, image_size, 
        core::CV_32FC1, // Output map type
        &mut map1x, &mut map1y,
    )?;

    calib3d::init_undistort_rectify_map(
        &k2, &d2, &r2, &p2, image_size, 
        core::CV_32FC1, 
        &mut map2x, &mut map2y,
    )?;

    println!("Rectification maps generated successfully.");
    
    // --- 4. Load Test Image Pair ---
    let left_path = format!("{}/left_frame_{:05}.jpg", DATA_DIR, TEST_IMAGE_INDEX);
    let right_path = format!("{}/right_frame_{:05}.jpg", DATA_DIR, TEST_IMAGE_INDEX);

    println!("--- 4. Loading Test Images: {} and {} ---", left_path, right_path);
    
    let src_l = imgcodecs::imread(&left_path, imgcodecs::IMREAD_COLOR)?;
    let src_r = imgcodecs::imread(&right_path, imgcodecs::IMREAD_COLOR)?;
    
    if src_l.empty() || src_r.empty() {
        return Err(Error::new(core::Code::StsError.into(), &format!("Failed to load test image pair index {}. Ensure images exist at the path: {}", TEST_IMAGE_INDEX, DATA_DIR)));
    }

    // --- 5. Apply Remapping to Rectify Images ---
    let (mut rectified_l, mut rectified_r) = (Mat::default(), Mat::default());

    println!("--- 5. Applying Undistortion and Rectification via Remap ---");

    imgproc::remap(
        &src_l, &mut rectified_l, 
        &map1x, &map1y, 
        imgproc::INTER_LINEAR,
        core::BORDER_CONSTANT,
        core::Scalar::default(),
    )?;

    imgproc::remap(
        &src_r, &mut rectified_r, 
        &map2x, &map2y, 
        imgproc::INTER_LINEAR, 
        core::BORDER_CONSTANT, 
        core::Scalar::default(), 
    )?;

    // --- 6. Display Results ---
    highgui::named_window("01 Original Left Image", highgui::WINDOW_AUTOSIZE)?;
    highgui::named_window("02 Rectified Left Image", highgui::WINDOW_AUTOSIZE)?;
    highgui::named_window("03 Original Right Image", highgui::WINDOW_AUTOSIZE)?;
    highgui::named_window("04 Rectified Right Image", highgui::WINDOW_AUTOSIZE)?;

    highgui::imshow("01 Original Left Image", &src_l)?;
    highgui::imshow("03 Original Right Image", &src_r)?;
    highgui::imshow("02 Rectified Left Image", &rectified_l)?;
    highgui::imshow("04 Rectified Right Image", &rectified_r)?;

    println!("\nRectification complete. Check the displayed windows!");
    println!("Press any key to exit...");

    highgui::wait_key(0)?; // Wait indefinitely for a key press
    Ok(())
}

fn main() {
    if let Err(e) = run_rectification() {
        eprintln!("Application Error: {:?}", e);
    }
}