cumulus3d-dense 0.3.1

//! 왜곡 보정: PINHOLE 카메라 계산, 카메라별 재표본 맵 캐시, 영상 재표본, 희소 모델 변환.

use crate::image::{GrayImage, ImageBuffer, Integral};
use rayon::prelude::*;
use cumulus3d_core::interop;
use cumulus3d_core::{Camera, CameraId, CameraModelKind, Error, Image, ImageId, Point3D, Reconstruction, Result, Vec2};
use std::collections::{BTreeMap, HashMap};
use std::hash::{Hash, Hasher};
use std::path::Path;
use std::sync::{Arc, Mutex};

/// 왜곡 보정 옵션(파이프라인은 `max_image_size = 960`).
#[derive(Clone, Debug, PartialEq)]
pub struct UndistortOptions {
    /// 가장자리 빈 픽셀 허용 비율(0 = 빈 픽셀 없음, 1 = 입력 전체 포함).
    pub blank_pixels: f64,
    /// 출력 크기 배율 하한.
    pub min_scale: f64,
    /// 출력 크기 배율 상한.
    pub max_scale: f64,
    /// ≤ 0 이면 제한 없음(관례상 −1).
    pub max_image_size: i64,
    /// 관심 영역 왼쪽(영상 너비 비율).
    pub roi_min_x: f64,
    /// 관심 영역 위쪽(영상 높이 비율).
    pub roi_min_y: f64,
    /// 관심 영역 오른쪽(영상 너비 비율).
    pub roi_max_x: f64,
    /// 관심 영역 아래쪽(영상 높이 비율).
    pub roi_max_y: f64,
    /// patch-match.cfg 에 쓰는 원천(이웃) 영상 수.
    pub num_patch_match_src_images: usize,
    /// JPEG 품질(−1 = 100).
    pub jpeg_quality: i32,
}

impl Default for UndistortOptions {
    fn default() -> Self {
        Self {
            blank_pixels: 0.0,
            min_scale: 0.2,
            max_scale: 2.0,
            max_image_size: -1,
            roi_min_x: 0.0,
            roi_min_y: 0.0,
            roi_max_x: 1.0,
            roi_max_y: 1.0,
            num_patch_match_src_images: 20,
            jpeg_quality: -1,
        }
    }
}

impl UndistortOptions {
    /// 사용자 파이프라인 값(max_image_size 960).
    pub fn pipeline() -> Self {
        Self { max_image_size: 960, ..Self::default() }
    }
    fn roi_enabled(&self) -> bool {
        self.roi_min_x != 0.0 || self.roi_min_y != 0.0 || self.roi_max_x != 1.0 || self.roi_max_y != 1.0
    }
    fn hash_into<H: Hasher>(&self, h: &mut H) {
        for v in [self.blank_pixels, self.min_scale, self.max_scale, self.roi_min_x, self.roi_min_y, self.roi_max_x, self.roi_max_y] {
            v.to_bits().hash(h);
        }
        self.max_image_size.hash(h);
    }
}

fn is_perspective(m: CameraModelKind) -> bool {
    // 이 크레이트가 아는 모델(0..4)은 모두 원근 모델이다.
    matches!(
        m,
        CameraModelKind::SingleFocalPinhole | CameraModelKind::Pinhole | CameraModelKind::SingleFocalRadial | CameraModelKind::Radial | CameraModelKind::OpenCv
    )
}

fn has_no_distortion(cam: &Camera) -> bool {
    cam.model.extra_param_slots().iter().all(|&i| cam.params[i].abs() <= 1e-8)
}

/// 왜곡 없는 PINHOLE 카메라 계산. 카메라 id 는 유지.
pub fn undistorted_camera(cam: &Camera, opts: &UndistortOptions) -> Result<Camera> {
    if !is_perspective(cam.model) {
        return Err(Error::Unsupported(format!("왜곡 보정: 모델 {}", cam.model.name())));
    }
    let (w, h) = (cam.width as f64, cam.height as f64);
    let (fx, fy) = (cam.focal_length_x(), cam.focal_length_y());
    let (mut cx, mut cy) = (cam.principal_point_x(), cam.principal_point_y());
    let mut out_w = cam.width;
    let mut out_h = cam.height;

    // ROI
    let (mut x_min, mut y_min, mut x_max, mut y_max) = (0i64, 0i64, cam.width as i64, cam.height as i64);
    let roi = opts.roi_enabled();
    if roi {
        x_min = (opts.roi_min_x * w).round() as i64;
        y_min = (opts.roi_min_y * h).round() as i64;
        x_max = (opts.roi_max_x * w).round() as i64;
        y_max = (opts.roi_max_y * h).round() as i64;
        x_min = x_min.min(cam.width as i64 - 1);
        y_min = y_min.min(cam.height as i64 - 1);
        x_max = x_max.max(x_min + 1);
        y_max = y_max.max(y_min + 1);
        out_w = (x_max - x_min) as u64;
        out_h = (y_max - y_min) as u64;
        cx -= x_min as f64;
        cy -= y_min as f64;
    }

    let pinhole_proj = |uv: &Vec2, cx: f64, cy: f64| Vec2::new(fx * uv.x + cx, fy * uv.y + cy);

    if roi || !matches!(cam.model, CameraModelKind::SingleFocalPinhole | CameraModelKind::Pinhole) {
        let (mut left_min, mut left_max) = (f64::MAX, f64::MIN);
        let (mut right_min, mut right_max) = (f64::MAX, f64::MIN);
        for y in y_min..y_max {
            let yc = y as f64 + 0.5;
            if let Some(uv) = cam.img_to_normalized(&Vec2::new(0.5, yc)) {
                let p = pinhole_proj(&uv, cx, cy);
                left_min = left_min.min(p.x);
                left_max = left_max.max(p.x);
            }
            if let Some(uv) = cam.img_to_normalized(&Vec2::new(w - 0.5, yc)) {
                let p = pinhole_proj(&uv, cx, cy);
                right_min = right_min.min(p.x);
                right_max = right_max.max(p.x);
            }
        }
        let (mut top_min, mut top_max) = (f64::MAX, f64::MIN);
        let (mut bottom_min, mut bottom_max) = (f64::MAX, f64::MIN);
        for x in x_min..x_max {
            let xc = x as f64 + 0.5;
            if let Some(uv) = cam.img_to_normalized(&Vec2::new(xc, 0.5)) {
                let p = pinhole_proj(&uv, cx, cy);
                top_min = top_min.min(p.y);
                top_max = top_max.max(p.y);
            }
            if let Some(uv) = cam.img_to_normalized(&Vec2::new(xc, h - 0.5)) {
                let p = pinhole_proj(&uv, cx, cy);
                bottom_min = bottom_min.min(p.y);
                bottom_max = bottom_max.max(p.y);
            }
        }
        let (w0, h0) = (out_w as f64, out_h as f64);
        let scale_x_min = (cx / (cx - left_min)).min((w0 - 0.5 - cx) / (right_max - cx));
        let scale_y_min = (cy / (cy - top_min)).min((h0 - 0.5 - cy) / (bottom_max - cy));
        let scale_x_max = (cx / (cx - left_max)).max((w0 - 0.5 - cx) / (right_min - cx));
        let scale_y_max = (cy / (cy - top_max)).max((h0 - 0.5 - cy) / (bottom_min - cy));
        let b = opts.blank_pixels;
        let scale_x = (1.0 / (scale_x_min * b + scale_x_max * (1.0 - b))).clamp(opts.min_scale, opts.max_scale);
        let scale_y = (1.0 / (scale_y_min * b + scale_y_max * (1.0 - b))).clamp(opts.min_scale, opts.max_scale);
        let nw = (scale_x * w0).max(1.0).floor() as u64;
        let nh = (scale_y * h0).max(1.0).floor() as u64;
        cx *= nw as f64 / w0;
        cy *= nh as f64 / h0;
        out_w = nw;
        out_h = nh;
    }

    let mut out = Camera::new(cam.camera_id, CameraModelKind::Pinhole, out_w, out_h, vec![fx, fy, cx, cy])?;
    apply_max_image_size(&mut out, opts.max_image_size);
    Ok(out)
}

/// max_image_size 규칙: 줄이기만, 반올림은 0 에서 먼 쪽.
fn apply_max_image_size(cam: &mut Camera, max_size: i64) {
    if max_size <= 0 {
        return;
    }
    let m = max_size as f64;
    let s = (m / cam.width as f64).min(m / cam.height as f64);
    if s < 1.0 {
        let nw = (s * cam.width as f64).round() as u64;
        let nh = (s * cam.height as f64).round() as u64;
        cam.rescale_to(nw.max(1), nh.max(1));
    }
}

/// 카메라 하나의 보정 결과: PINHOLE 카메라 + 입력 해상도 재표본 맵.
#[derive(Debug)]
pub struct CameraUndistortion {
    /// 입력(왜곡 있는) 카메라.
    pub source: Camera,
    /// 보정된 PINHOLE 카메라.
    pub pinhole: Camera,
    /// 보조 카메라(입력 크기) 픽셀마다 입력 영상 표본 위치 (sx−0.5, sy−0.5). 실패는 NaN.
    lut: Vec<[f32; 2]>,
}

impl CameraUndistortion {
    /// 카메라 계산 + LUT 구성.
    pub fn new(source: &Camera, opts: &UndistortOptions) -> Result<Self> {
        let pinhole = undistorted_camera(source, opts)?;
        let (sw, sh) = (source.width as usize, source.height as usize);
        let mut aux = pinhole.clone();
        aux.rescale_to(source.width, source.height);
        let (fx, fy, cx, cy) = (aux.focal_length_x(), aux.focal_length_y(), aux.principal_point_x(), aux.principal_point_y());
        let mut lut = vec![[f32::NAN; 2]; sw * sh];
        lut.par_chunks_mut(sw).enumerate().for_each(|(y, row)| {
            for (x, v) in row.iter_mut().enumerate() {
                let uv = Vec2::new((x as f64 + 0.5 - cx) / fx, (y as f64 + 0.5 - cy) / fy);
                let p = source.normalized_to_img(&uv);
                if p.x.is_finite() && p.y.is_finite() {
                    *v = [(p.x - 0.5) as f32, (p.y - 0.5) as f32];
                }
            }
        });
        Ok(Self { source: source.clone(), pinhole, lut })
    }

    /// 영상 재표본: 입력 해상도에서 쌍선형 변형(범위 밖 검정) → 면적 평균 축소.
    pub fn undistort_image(&self, img: &ImageBuffer) -> Result<ImageBuffer> {
        let (sw, sh) = (self.source.width as usize, self.source.height as usize);
        if img.width != sw || img.height != sh {
            return Err(Error::InvalidArgument(format!(
                "왜곡 보정: 영상 크기 {}x{} != 카메라 {}x{}",
                img.width, img.height, sw, sh
            )));
        }
        let ch = img.channels;
        let mut warped = ImageBuffer::new(sw, sh, ch);
        warped.data.par_chunks_mut(sw * ch).enumerate().for_each(|(y, row)| {
            for x in 0..sw {
                let [fx, fy] = self.lut[y * sw + x];
                if !fx.is_finite() {
                    continue;
                }
                let x0f = fx.floor();
                let y0f = fy.floor();
                let (x0, y0) = (x0f as i64, y0f as i64);
                if x0 < 0 || y0 < 0 || x0 + 1 >= sw as i64 || y0 + 1 >= sh as i64 {
                    continue;
                }
                let (ax, ay) = (fx - x0f, fy - y0f);
                let (x0, y0) = (x0 as usize, y0 as usize);
                for c in 0..ch {
                    let p = |xx: usize, yy: usize| img.data[(yy * sw + xx) * ch + c] as f32;
                    let top = p(x0, y0) * (1.0 - ax) + p(x0 + 1, y0) * ax;
                    let bot = p(x0, y0 + 1) * (1.0 - ax) + p(x0 + 1, y0 + 1) * ax;
                    let v = top * (1.0 - ay) + bot * ay;
                    row[x * ch + c] = v.round().clamp(0.0, 255.0) as u8;
                }
            }
        });
        let (ow, oh) = (self.pinhole.width as usize, self.pinhole.height as usize);
        if ow == sw && oh == sh {
            return Ok(warped);
        }
        // 설계 결정: 축소는 정확한 면적 평균으로 한다.
        let mut out = ImageBuffer::new(ow, oh, ch);
        let planes: Vec<Vec<f32>> = (0..ch)
            .into_par_iter()
            .map(|c| {
                let g = GrayImage { width: sw, height: sh, data: (0..sw * sh).map(|i| warped.data[i * ch + c] as f32).collect() };
                Integral::new(&g).resample(ow, oh).data
            })
            .collect();
        for (c, plane) in planes.iter().enumerate() {
            for (i, v) in plane.iter().enumerate() {
                out.data[i * ch + c] = v.round().clamp(0.0, 255.0) as u8;
            }
        }
        Ok(out)
    }
}

fn camera_key(cam: &Camera, opts: &UndistortOptions) -> u64 {
    let mut h = std::collections::hash_map::DefaultHasher::new();
    cam.model.id().hash(&mut h);
    cam.width.hash(&mut h);
    cam.height.hash(&mut h);
    for p in &cam.params {
        p.to_bits().hash(&mut h);
    }
    opts.hash_into(&mut h);
    h.finish()
}

/// 카메라 매개변수 해시 → 보정 결과(LUT 포함) 캐시. 구역 사이에 공유한다.
#[derive(Default)]
pub struct UndistortCache {
    map: Mutex<HashMap<u64, Arc<CameraUndistortion>>>,
}

impl UndistortCache {
    /// 빈 캐시.
    pub fn new() -> Self {
        Self::default()
    }
    /// 캐시에서 찾거나 새로 계산.
    pub fn get(&self, cam: &Camera, opts: &UndistortOptions) -> Result<Arc<CameraUndistortion>> {
        let key = camera_key(cam, opts);
        if let Some(v) = self.map.lock().unwrap_or_else(|e| e.into_inner()).get(&key) {
            return Ok(v.clone());
        }
        let v = Arc::new(CameraUndistortion::new(cam, opts)?);
        self.map.lock().unwrap_or_else(|e| e.into_inner()).insert(key, v.clone());
        Ok(v)
    }
    /// 캐시된 카메라 수.
    pub fn len(&self) -> usize {
        self.map.lock().unwrap_or_else(|e| e.into_inner()).len()
    }
    /// 비어 있는지.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// 희소 모델 변환: 카메라를 PINHOLE 로, 2D 관측을 새 카메라 좌표로(실패 NaN).
/// 자세·3D 점·트랙·이름·id 는 그대로.
pub fn undistort_reconstruction(rec: &Reconstruction, opts: &UndistortOptions, cache: &UndistortCache) -> Result<Reconstruction> {
    let mut new_cams: BTreeMap<CameraId, Camera> = BTreeMap::new();
    for (id, cam) in rec.cameras() {
        let nc = if has_no_distortion(cam) && opts.max_image_size <= 0 {
            cam.clone()
        } else {
            // 캐시는 매개변수 해시로 공유되므로(같은 값의 다른 카메라) id 는 이 카메라 것으로 되돌린다.
            let mut c = cache.get(cam, opts)?.pinhole.clone();
            c.camera_id = *id;
            c
        };
        new_cams.insert(*id, nc);
    }
    let mut out = Reconstruction::new();
    for c in new_cams.values() {
        out.add_camera(c.clone())?;
    }
    for r in rec.rigs().values() {
        out.add_rig(r.clone())?;
    }
    for f in rec.frames().values() {
        out.add_frame(f.clone())?;
    }
    let images: Vec<&Image> = rec.images().collect();
    let converted: Vec<Image> = images
        .par_iter()
        .map(|im| {
            let src = &rec.cameras()[&im.camera_id];
            let dst = &new_cams[&im.camera_id];
            let same = src == dst;
            let pts = im.points2d().iter().map(|p| {
                if same {
                    return p.xy;
                }
                match src.img_to_normalized(&p.xy) {
                    Some(uv) => dst.normalized_to_img(&uv),
                    None => Vec2::new(f64::NAN, f64::NAN),
                }
            });
            let mut ni = Image::new(im.image_id, im.name.clone(), im.camera_id, pts);
            ni.frame_id = im.frame_id;
            ni
        })
        .collect();
    for im in converted {
        out.add_image(im)?;
    }
    for (id, p) in rec.points3d() {
        let np = Point3D { xyz: p.xyz, color: p.color, error: p.error, track: p.track.clone() };
        out.add_point3d_with_id(id, np)?;
    }
    for fid in rec.registered_frames() {
        out.register_frame(*fid)?;
    }
    Ok(out)
}

/// 메모리 내 왜곡 보정 결과.
pub struct UndistortResult {
    /// 변환된 희소 모델(PINHOLE).
    pub reconstruction: Reconstruction,
    /// 등록 영상별 보정 영상. 읽기 실패한 영상은 없다.
    pub images: BTreeMap<ImageId, Arc<ImageBuffer>>,
    /// 처리 실패(영상 읽기 실패 등)한 영상 id.
    pub failed: Vec<ImageId>,
}

/// 등록 영상 전부를 보정한다. `load` 는 입력 영상을 돌려준다(없으면 실패 처리).
pub fn undistort<F>(rec: &Reconstruction, opts: &UndistortOptions, cache: &UndistortCache, load: F) -> Result<UndistortResult>
where
    F: Fn(&Image) -> Option<Arc<ImageBuffer>> + Sync,
{
    let reconstruction = undistort_reconstruction(rec, opts, cache)?;
    let ids = rec.registered_images();
    // 카메라 LUT 를 먼저 채워 영상 병렬 처리 중 중복 계산을 막는다.
    for im in ids.iter().filter_map(|i| rec.image(*i)) {
        cache.get(&rec.cameras()[&im.camera_id], opts)?;
    }
    let results: Vec<(ImageId, Option<Arc<ImageBuffer>>)> = ids
        .par_iter()
        .map(|&iid| {
            let Some(im) = rec.image(iid) else { return (iid, None) };
            let cam = &rec.cameras()[&im.camera_id];
            let Some(src) = load(im) else { return (iid, None) };
            let r = cache.get(cam, opts).and_then(|u| u.undistort_image(&src)).ok().map(Arc::new);
            (iid, r)
        })
        .collect();
    let mut images = BTreeMap::new();
    let mut failed = Vec::new();
    for (iid, r) in results {
        match r {
            Some(b) => {
                images.insert(iid, b);
            }
            None => failed.push(iid),
        }
    }
    Ok(UndistortResult { reconstruction, images, failed })
}

/// 입력 영상 폴더에서 읽어 보정.
pub fn undistort_from_dir(rec: &Reconstruction, image_dir: impl AsRef<Path>, opts: &UndistortOptions, cache: &UndistortCache) -> Result<UndistortResult> {
    let dir = image_dir.as_ref().to_path_buf();
    undistort(rec, opts, cache, |im| ImageBuffer::load(dir.join(&im.name)).ok().map(Arc::new))
}

/// 보정 결과를 조밀 복원 작업 폴더로 쓴다: images/(하위 폴더 포함), sparse/(이진 모델),
/// stereo/{depth_maps,normal_maps,consistency_graphs}/ 빈 폴더와 설정 파일.
/// 폴더 구성·파일 형식은 [`cumulus3d_core::interop`] 가 정한다.
pub fn write_undistorted_workspace(result: &UndistortResult, out_dir: impl AsRef<Path>, opts: &UndistortOptions) -> Result<()> {
    let out = out_dir.as_ref();
    let rec = &result.reconstruction;
    std::fs::create_dir_all(out.join("images"))?;
    std::fs::create_dir_all(out.join("sparse"))?;
    interop::create_stereo_dirs(out, Path::new(""))?;
    let q = if opts.jpeg_quality < 0 { 100 } else { opts.jpeg_quality.clamp(1, 100) as u8 };
    let names: Vec<(ImageId, String)> = rec.registered_images().into_iter().filter_map(|i| rec.image(i).map(|im| (i, im.name.clone()))).collect();
    names
        .par_iter()
        .filter_map(|(iid, name)| result.images.get(iid).map(|b| (name, b)))
        .try_for_each(|(name, buf)| -> Result<()> {
            let p = out.join("images").join(name);
            if let Some(parent) = p.parent() {
                std::fs::create_dir_all(parent)?;
            }
            if let Some(parent) = Path::new(name).parent() {
                interop::create_stereo_dirs(out, parent)?;
            }
            buf.save(p, q)
        })?;
    interop::write_model_binary(rec, out.join("sparse"), interop::ImageOrder::Registration)?;
    let written = names.iter().filter(|(iid, _)| result.images.contains_key(iid)).map(|(_, n)| n.as_str());
    interop::write_stereo_configs(out, written, opts.num_patch_match_src_images)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn front_cam() -> Camera {
        Camera::new(1, CameraModelKind::OpenCv, 2048, 1152, vec![1609.22, 1608.21, 1024.0, 576.0, 0.0824, -0.0849, 0.0022, 0.0007]).unwrap()
    }

    #[test]
    fn spec_example_camera() {
        let cam = front_cam();
        let mut o = UndistortOptions::pipeline();
        o.max_image_size = -1;
        let mid = undistorted_camera(&cam, &o).unwrap();
        assert_eq!((mid.width, mid.height), (2002, 1123));
        assert!((mid.principal_point_x() - 1001.0).abs() < 1e-9);
        assert!((mid.principal_point_y() - 561.5).abs() < 1e-9);
        let c = undistorted_camera(&cam, &UndistortOptions::pipeline()).unwrap();
        assert_eq!(c.model, CameraModelKind::Pinhole);
        assert_eq!((c.width, c.height), (960, 539));
        let fx = 1609.22 * 960.0 / 2002.0;
        let fy = 1608.21 * 539.0 / 1123.0;
        assert!((c.params[0] - fx).abs() / fx < 1e-12);
        assert!((c.params[1] - fy).abs() / fy < 1e-12);
        assert!((c.params[0] - 771.654).abs() / 771.654 < 1e-6, "fx {}", c.params[0]);
        assert!((c.params[1] - 771.884).abs() / 771.884 < 1e-6, "fy {}", c.params[1]);
        assert!((c.params[2] - 480.0).abs() < 1e-9 && (c.params[3] - 269.5).abs() < 1e-9);
    }

    #[test]
    fn lut_maps_pinhole_to_source() {
        // 작은 카메라에서 보정 영상 위치를 입력 카메라로 다시 투영하면 LUT 와 같아야 한다.
        let cam = Camera::new(1, CameraModelKind::OpenCv, 200, 120, vec![160.0, 160.0, 100.0, 60.0, 0.05, -0.02, 0.001, 0.0005]).unwrap();
        let u = CameraUndistortion::new(&cam, &UndistortOptions::default()).unwrap();
        // 입력 크기의 체커보드 → 보정 영상의 기하 정합 확인: 중심 근처 픽셀값 일치.
        let mut img = ImageBuffer::new(200, 120, 1);
        for y in 0..120 {
            for x in 0..200 {
                img.data[y * 200 + x] = (((x as f64 * 0.3).sin() * 60.0) + 128.0) as u8;
            }
        }
        let out = u.undistort_image(&img).unwrap();
        assert_eq!((out.width as u64, out.height as u64), (u.pinhole.width, u.pinhole.height));
        // 보정 영상의 중심 픽셀 → 입력 위치
        let (x, y) = (u.pinhole.width as f64 / 2.0, u.pinhole.height as f64 / 2.0);
        let uv = u.pinhole.img_to_normalized(&Vec2::new(x, y)).unwrap();
        let p = cam.normalized_to_img(&uv);
        let expect = img.get((p.x - 0.5).round() as usize, (p.y - 0.5).round() as usize, 0) as i32;
        let got = out.get(x as usize, y as usize, 0) as i32;
        assert!((expect - got).abs() <= 25, "{expect} vs {got}");
    }
}