use opencv::core::{
    cart_to_polar, merge, no_array, normalize, split, Mat as Matrix, MatTraitConst,
    MatTraitConstManual, Ptr, Size, ToOutputArray, Vector, CV_32F, CV_8U, NORM_MINMAX,
};
use opencv::imgcodecs::imwrite;
use opencv::imgproc::{cvt_color, COLOR_HSV2BGR};
use opencv::optflow::{calc_optical_flow_dense_rlof, InterpolationType, RLOFOpticalFlowParameter};
use opencv::prelude::MatExprTraitConst;
use opencv::types::VectorOfMat;
use opencv::videoio::{
    VideoCapture, VideoCaptureTrait, VideoCaptureTraitConst, VideoWriter, VideoWriterTrait,
    CAP_ANY, CAP_PROP_FOURCC, CAP_PROP_FPS, CAP_PROP_FRAME_HEIGHT, CAP_PROP_FRAME_WIDTH,
};

fn read_frame(
    video: &mut VideoCapture,
    frame: &mut impl ToOutputArray,
) -> Result<bool, anyhow::Error> {
    Ok(video.read(frame)?)
}

fn main() -> Result<(), anyhow::Error> {
    let mut input_video = VideoCapture::from_file("../clips/edit.webm", CAP_ANY)?;
    let mut output_video = VideoWriter::new(
        "output.webm",
        input_video.get(CAP_PROP_FOURCC)? as i32,
        input_video.get(CAP_PROP_FPS)?,
        Size::new(
            input_video.get(CAP_PROP_FRAME_WIDTH)? as i32,
            input_video.get(CAP_PROP_FRAME_HEIGHT)? as i32,
        ),
        true,
    )?;

    let mut frame_count = 0;

    let mut last_frame = Matrix::default();
    input_video
        .read(&mut last_frame)
        .expect("Video cannot be empty");
    let mut current_frame = Matrix::default();

    while read_frame(&mut input_video, &mut current_frame)? {
        println!("Frame {frame_count}");
        let mut computed_flow = unsafe { Matrix::new_size(current_frame.size()?, CV_32F)? };
        calc_optical_flow_dense_rlof(
            &last_frame,
            &current_frame,
            &mut computed_flow,
            Ptr::new(RLOFOpticalFlowParameter::default()?),
            1 as f32,
            Size::new(6, 6),
            InterpolationType::INTERP_EPIC,
            128,
            0.005,
            999 as f32,
            15,
            100,
            true,
            500 as f32,
            1.5,
            false,
        )?;

        let mut flow_parts: VectorOfMat = Vec::with_capacity(2).into();
        split(&computed_flow, &mut flow_parts)?;
        assert_eq!(flow_parts.len(), 2);

        let mut magnitude = Matrix::default();
        let mut normalized_magnitude = Matrix::default();
        let mut angle = Matrix::default();
        cart_to_polar(
            &flow_parts.get(0)?,
            &flow_parts.get(1)?,
            &mut magnitude,
            &mut angle,
            true,
        )?;
        normalize(
            &magnitude,
            &mut normalized_magnitude,
            0 as f64,
            1 as f64,
            NORM_MINMAX,
            -1,
            &no_array(),
        )?;
        angle = (angle * (1_f64 / 360_f64) * (180_f64 / 255_f64))
            .into_result()?
            .to_mat()?;

        let mut hsv_vec = VectorOfMat::with_capacity(3);
        hsv_vec.push(angle.clone());
        hsv_vec.push(Matrix::ones_size(angle.size()?, CV_32F)?.to_mat()?);
        hsv_vec.push(normalized_magnitude);
        assert_eq!(hsv_vec.len(), 3);

        let mut hsv = Matrix::default();
        let mut hsv_8 = Matrix::default();
        let mut bgr = Matrix::default();

        merge(&hsv_vec, &mut hsv)?;
        hsv.convert_to(&mut hsv_8, CV_8U, 255 as f64, 0 as f64)?;
        cvt_color(&hsv_8, &mut bgr, COLOR_HSV2BGR, 0)?;
        // output_video.write(&bgr)?;
        imwrite(
            &format!("frames/frame{frame_count}.jpg"),
            &bgr,
            &Vector::new(),
        )?;

        assert_eq!(current_frame.size()?, last_frame.size()?);
        last_frame = current_frame.clone();

        // if frame_count > 5 {
        //     break;
        // } else {
        // }
        frame_count += 1;
    }

    Ok(())
}