/*
 *  optical_flow_video.rs
 *  purecv
 *
 *  This file is part of purecv - WebARKit.
 *
 *  purecv is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  purecv is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with purecv.  If not, see <http://www.gnu.org/licenses/>.
 *
 *  As a special exception, the copyright holders of this library give you
 *  permission to link this library with independent modules to produce an
 *  executable, regardless of the license terms of these independent modules, and to
 *  copy and distribute the resulting executable under terms of your choice,
 *  provided that you also meet, for each linked independent module, the terms and
 *  conditions of the license of that module. An independent module is a module
 *  which is neither derived from nor based on this library. If you modify this
 *  library, you may extend this exception to your version of the library, but you
 *  are not obligated to do so. If you do not wish to do so, delete this exception
 *  statement from your version.
 *
 *  Copyright 2026 WebARKit.
 *
 *  Author(s): Walter Perdan @kalwalt https://github.com/kalwalt
 *
 */

//! Video-style optical flow example — multi-frame tracking with pyramidal LK.
//!
//! This example demonstrates a **continuous tracking loop** over a video
//! sequence, the same pattern you would use when processing a real webcam or
//! video file frame by frame.
//!
//! ## Input modes
//!
//! 1. **GIF file** (recommended) — pass a `.gif` file path.  Each animation
//!    frame is decoded and used as a real video frame.
//! 2. **Static image** — pass a `.jpg` / `.png` file.  A synthetic 20-frame
//!    "video" is generated by translating the image by a constant offset per
//!    frame, simulating a smooth camera pan.
//! 3. **Auto-generate** — when run with no arguments and the default test GIF
//!    does not exist, the example creates one from the bundled butterfly image.
//!
//! Because purecv is a zero-FFI library (no C++ backend, no system video
//! codecs), GIF is used as a lightweight "video" format that the `image` crate
//! can decode natively — no extra dependencies needed.
//!
//! ## Pipeline (per-frame pair)
//!
//! ```text
//!   frame[t-1]  ──────►  prev_gray
//!   frame[t]    ──────►  next_gray
//!
//!   prev_gray ──► goodFeaturesToTrack ──► corners  (first frame or re-detect)
//!
//!   (prev_gray, next_gray, corners) ──► calcOpticalFlowPyrLK
//!                                           │
//!                          ┌────────────────┼────────────────┐
//!                       next_pts         status             err
//!
//!   Prune lost points, update tracks, swap frames, repeat.
//! ```
//!
//! ## Usage
//!
//! ```bash
//! # Auto-generate a test GIF and run on it:
//! cargo run --example optical_flow_video
//!
//! # Run on your own GIF:
//! cargo run --example optical_flow_video -- path/to/animation.gif
//!
//! # Run on a static image (synthetic frames):
//! cargo run --example optical_flow_video -- examples/data/butterfly.jpg
//! ```
//!
//! ## Output
//!
//! Written to `examples/data/out/`:
//! * `optical_flow_video_frame_NN.png` — annotated frame with coloured flow trails
//! * `optical_flow_video_summary.csv`  — per-frame tracking statistics

use image::codecs::gif::{GifDecoder, GifEncoder, Repeat};
use image::{AnimationDecoder, DynamicImage, Frame, GenericImageView, ImageBuffer, Rgb, RgbaImage};
use purecv::core::types::{Point2f, Size2i, TermCriteria, TermType};
use purecv::core::Matrix;
use purecv::imgproc::{cvt_color_rgb_to_gray, good_features_to_track};
use purecv::version;
use purecv::video::{calc_optical_flow_pyramid_lk, OPTFLOW_LK_GET_MIN_EIGENVALS};
use std::fs::File;
use std::io::BufReader;

use std::io::Write;
use std::path::Path;

// ---------------------------------------------------------------------------
// Simulation parameters (used when synthesising from a static image)
// ---------------------------------------------------------------------------

/// Number of synthetic video frames to generate.
const SYNTH_NUM_FRAMES: usize = 20;

/// Per-frame translation in pixels for synthetic mode.
const SYNTH_DX: f32 = 3.0;
const SYNTH_DY: f32 = 2.0;

// ---------------------------------------------------------------------------
// Tracking parameters
// ---------------------------------------------------------------------------

/// Maximum number of features to detect.
const MAX_CORNERS: i32 = 80;

/// Minimum tracked features before re-detecting.
const MIN_CORNERS: usize = 20;

/// Maximum number of trailing positions to draw per track.
const MAX_TRAIL_LEN: usize = 10;

/// Result type for functions that return decoded frames + a base RGB image.
type FrameLoadResult = (Vec<Matrix<u8>>, image::RgbImage);

// ---------------------------------------------------------------------------
// Track — stores the history of a single feature across frames.
// ---------------------------------------------------------------------------

struct Track {
    positions: Vec<Point2f>,
}

impl Track {
    fn new(pt: Point2f) -> Self {
        Self {
            positions: vec![pt],
        }
    }

    fn push(&mut self, pt: Point2f) {
        self.positions.push(pt);
        if self.positions.len() > MAX_TRAIL_LEN {
            self.positions.remove(0);
        }
    }

    fn last(&self) -> Point2f {
        *self.positions.last().unwrap()
    }
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("--- purecv Optical Flow Video Example ---");
    println!("purecv v{}\n", version::get_version());

    let out_dir = "examples/data/out";
    std::fs::create_dir_all(out_dir)?;

    // Determine input path.
    let default_gif = "examples/data/out/butterfly_motion.gif";
    let default_img = "examples/data/butterfly.jpg";
    let input_path = std::env::args().nth(1).unwrap_or_else(|| {
        if Path::new(default_gif).exists() {
            default_gif.to_string()
        } else {
            default_img.to_string()
        }
    });

    if !Path::new(&input_path).exists() {
        eprintln!(
            "Error: '{}' not found. Run from the project root.",
            input_path
        );
        return Ok(());
    }

    // -----------------------------------------------------------------------
    // 1. Load frames — either from a GIF or by synthesising from a static image.
    // -----------------------------------------------------------------------
    let is_gif = input_path.to_lowercase().ends_with(".gif");

    let (gray_frames, rgb_base) = if is_gif {
        load_gif_frames(&input_path)?
    } else {
        // If the default GIF doesn't exist, generate it first so future runs
        // can use the GIF path directly.
        let (gray, rgb) = synthesise_frames(&input_path)?;
        if !Path::new(default_gif).exists() {
            println!("\nGenerating test GIF: {}", default_gif);
            save_test_gif(&input_path, default_gif)?;
            println!(
                "  → saved {} ({} frames). Re-run to use GIF input.\n",
                default_gif, SYNTH_NUM_FRAMES
            );
        }
        (gray, rgb)
    };

    let num_frames = gray_frames.len();
    if num_frames < 2 {
        eprintln!("Error: need at least 2 frames, got {}.", num_frames);
        return Ok(());
    }
    println!("Frames to process: {}\n", num_frames);

    // -----------------------------------------------------------------------
    // 2. Continuous tracking loop.
    // -----------------------------------------------------------------------
    let criteria = TermCriteria::new(TermType::Both, 30, 0.001);
    let win_size = Size2i::new(15, 15);
    let max_level = 3;

    let mut tracks: Vec<Track> = Vec::new();
    let mut csv_rows: Vec<String> = Vec::new();
    csv_rows.push("frame,n_tracked,n_lost,n_total,mean_flow_x,mean_flow_y".to_string());

    for k in 0..num_frames - 1 {
        let prev_gray = &gray_frames[k];
        let next_gray = &gray_frames[k + 1];

        // -----------------------------------------------------------------
        // 2a. Detect new features if we have too few tracks.
        // -----------------------------------------------------------------
        if tracks.len() < MIN_CORNERS {
            let corners =
                good_features_to_track(prev_gray, MAX_CORNERS, 0.01, 10.0, 3, false, 0.04)?;
            let min_dist_sq = 5.0f32 * 5.0;
            for c in &corners {
                let duplicate = tracks.iter().any(|t| {
                    let last = t.last();
                    let dx = last.x - c.x;
                    let dy = last.y - c.y;
                    dx * dx + dy * dy < min_dist_sq
                });
                if !duplicate {
                    tracks.push(Track::new(*c));
                }
            }
            if k == 0 {
                println!("Frame  0: detected {} initial features", tracks.len());
            } else {
                println!(
                    "Frame {:>2}: re-detected features → {} total tracks",
                    k,
                    tracks.len()
                );
            }
        }

        // -----------------------------------------------------------------
        // 2b. Gather current track heads.
        // -----------------------------------------------------------------
        let prev_pts: Vec<Point2f> = tracks.iter().map(|t| t.last()).collect();

        // -----------------------------------------------------------------
        // 2c. Run pyramidal Lucas-Kanade.
        // -----------------------------------------------------------------
        let (next_pts, status, _err) = calc_optical_flow_pyramid_lk(
            prev_gray,
            next_gray,
            &prev_pts,
            None,
            win_size,
            max_level,
            criteria,
            OPTFLOW_LK_GET_MIN_EIGENVALS,
            1e-4,
        )?;

        // -----------------------------------------------------------------
        // 2d. Update tracks — keep successfully tracked points only.
        // -----------------------------------------------------------------
        let mut new_tracks: Vec<Track> = Vec::with_capacity(tracks.len());
        let mut n_tracked = 0usize;
        let mut n_lost = 0usize;
        let mut sum_dx = 0.0f32;
        let mut sum_dy = 0.0f32;

        for (i, mut track) in tracks.into_iter().enumerate() {
            if status[i] == 1 {
                sum_dx += next_pts[i].x - prev_pts[i].x;
                sum_dy += next_pts[i].y - prev_pts[i].y;
                track.push(next_pts[i]);
                new_tracks.push(track);
                n_tracked += 1;
            } else {
                n_lost += 1;
            }
        }
        tracks = new_tracks;

        let mean_dx = if n_tracked > 0 {
            sum_dx / n_tracked as f32
        } else {
            0.0
        };
        let mean_dy = if n_tracked > 0 {
            sum_dy / n_tracked as f32
        } else {
            0.0
        };

        println!(
            "Frame {:>2}→{:>2}: tracked {}/{}, lost {}, mean flow ({:+.2}, {:+.2})",
            k,
            k + 1,
            n_tracked,
            n_tracked + n_lost,
            n_lost,
            mean_dx,
            mean_dy,
        );

        csv_rows.push(format!(
            "{},{},{},{},{:.4},{:.4}",
            k + 1,
            n_tracked,
            n_lost,
            n_tracked + n_lost,
            mean_dx,
            mean_dy,
        ));

        // -----------------------------------------------------------------
        // 2e. Save annotated frame (every 5th + last).
        // -----------------------------------------------------------------
        if k % 5 == 0 || k == num_frames - 2 {
            let frame_path = format!("{}/optical_flow_video_frame_{:02}.png", out_dir, k + 1);
            save_tracked_frame(&rgb_base, &tracks, &frame_path)?;
        }
    }

    // -----------------------------------------------------------------------
    // 3. Write CSV summary.
    // -----------------------------------------------------------------------
    let csv_path = format!("{}/optical_flow_video_summary.csv", out_dir);
    {
        let mut f = File::create(&csv_path)?;
        for row in &csv_rows {
            writeln!(f, "{}", row)?;
        }
    }

    // -----------------------------------------------------------------------
    // 4. Final statistics.
    // -----------------------------------------------------------------------
    println!("\n=== Final Statistics ===");
    println!("  Total frames processed: {}", num_frames - 1);
    println!("  Active tracks remaining: {}", tracks.len());

    if !tracks.is_empty() {
        let (total_dx, total_dy, count) =
            tracks
                .iter()
                .fold((0.0f32, 0.0f32, 0u32), |(sdx, sdy, cnt), track| {
                    if track.positions.len() >= 2 {
                        let first = track.positions[0];
                        let last = *track.positions.last().unwrap();
                        (sdx + (last.x - first.x), sdy + (last.y - first.y), cnt + 1)
                    } else {
                        (sdx, sdy, cnt)
                    }
                });
        if count > 0 {
            println!(
                "  Avg net displacement:   ({:+.2}, {:+.2}) px  ({} tracks)",
                total_dx / count as f32,
                total_dy / count as f32,
                count,
            );
        }
    }

    println!("\nOutput saved to:");
    println!("  {}/optical_flow_video_frame_*.png", out_dir);
    println!("  {}", csv_path);
    println!("\nDone.");
    Ok(())
}

// ===========================================================================
// Frame loading
// ===========================================================================

/// Decode a GIF file into a sequence of grayscale `Matrix<u8>` frames.
///
/// Returns `(gray_frames, rgb_base)` where `rgb_base` is the first frame
/// as an `RgbImage` (used as the background for annotated output).
fn load_gif_frames(path: &str) -> Result<FrameLoadResult, Box<dyn std::error::Error>> {
    println!("Loading GIF: {}", path);
    let file = BufReader::new(File::open(path)?);
    let decoder = GifDecoder::new(file)?;
    let frames: Vec<Frame> = decoder.into_frames().collect::<Result<Vec<_>, _>>()?;

    if frames.is_empty() {
        return Err("GIF contains no frames".into());
    }

    let first_rgba: &RgbaImage = frames[0].buffer();
    let rgb_base = DynamicImage::ImageRgba8(first_rgba.clone()).to_rgb8();
    let (w, h) = rgb_base.dimensions();
    println!("  Frame size: {}×{}, {} frame(s)", w, h, frames.len());

    let mut gray_frames: Vec<Matrix<u8>> = Vec::with_capacity(frames.len());
    for frame in &frames {
        let rgba: &RgbaImage = frame.buffer();
        let rgb = DynamicImage::ImageRgba8(rgba.clone()).to_rgb8();
        let (fw, fh) = rgb.dimensions();
        let mat_rgb = Matrix::from_vec(fh as usize, fw as usize, 3, rgb.into_raw());
        let gray = cvt_color_rgb_to_gray(&mat_rgb)?;
        gray_frames.push(gray);
    }

    Ok((gray_frames, rgb_base))
}

/// Synthesise frames from a static image by translating it.
fn synthesise_frames(img_path: &str) -> Result<FrameLoadResult, Box<dyn std::error::Error>> {
    println!("Loading static image: {}", img_path);
    println!(
        "Synthesising {} frames with per-frame shift ({:+.1}, {:+.1}) px",
        SYNTH_NUM_FRAMES, SYNTH_DX, SYNTH_DY,
    );

    let img = image::open(img_path)?;
    let (width, height) = img.dimensions();
    println!("  Image size: {}×{}", width, height);

    let rgb_base = img.to_rgb8();
    let mat_rgb = Matrix::from_vec(
        height as usize,
        width as usize,
        3,
        rgb_base.clone().into_raw(),
    );
    let gray_base = cvt_color_rgb_to_gray(&mat_rgb)?;
    let rows = gray_base.rows;
    let cols = gray_base.cols;

    let mut frames: Vec<Matrix<u8>> = Vec::with_capacity(SYNTH_NUM_FRAMES);
    for k in 0..SYNTH_NUM_FRAMES {
        let shift_x = (k as f32 * SYNTH_DX).round() as usize;
        let shift_y = (k as f32 * SYNTH_DY).round() as usize;
        let mut data = vec![0u8; rows * cols];
        for r in 0..rows {
            for c in 0..cols {
                let src_r = r.saturating_sub(shift_y);
                let src_c = c.saturating_sub(shift_x);
                data[r * cols + c] = gray_base.data[src_r * cols + src_c];
            }
        }
        frames.push(Matrix::<u8>::from_vec(rows, cols, 1, data));
    }

    Ok((frames, rgb_base))
}

// ===========================================================================
// Test GIF generation — create a small animated GIF from a static image
// ===========================================================================

/// Generate a test animated GIF by translating the source image across
/// multiple frames — this produces a file that can be fed back as GIF input.
fn save_test_gif(img_path: &str, gif_path: &str) -> Result<(), Box<dyn std::error::Error>> {
    let img = image::open(img_path)?;
    let (width, height) = img.dimensions();
    let rgb_base = img.to_rgb8();

    let file = File::create(gif_path)?;
    let mut encoder = GifEncoder::new(file);
    encoder.set_repeat(Repeat::Infinite)?;

    for k in 0..SYNTH_NUM_FRAMES {
        let shift_x = (k as f32 * SYNTH_DX).round() as i32;
        let shift_y = (k as f32 * SYNTH_DY).round() as i32;

        let mut rgba = RgbaImage::new(width, height);
        for y in 0..height {
            for x in 0..width {
                let src_x = (x as i32 - shift_x).clamp(0, width as i32 - 1) as u32;
                let src_y = (y as i32 - shift_y).clamp(0, height as i32 - 1) as u32;
                let px = rgb_base.get_pixel(src_x, src_y);
                rgba.put_pixel(x, y, image::Rgba([px[0], px[1], px[2], 255]));
            }
        }

        let delay = image::Delay::from_numer_denom_ms(100, 1); // 10 fps
        let frame = Frame::from_parts(rgba, 0, 0, delay);
        encoder.encode_frame(frame)?;
    }

    Ok(())
}

// ===========================================================================
// Drawing helpers
// ===========================================================================

/// Save an annotated copy of the base image showing tracked feature trails.
fn save_tracked_frame(
    base: &image::RgbImage,
    tracks: &[Track],
    path: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut out: ImageBuffer<Rgb<u8>, Vec<u8>> = base.clone();

    for (ti, track) in tracks.iter().enumerate() {
        let color = track_color(ti);
        let pts = &track.positions;

        // Draw trail segments.
        for w in pts.windows(2) {
            draw_line(
                &mut out,
                w[0].x as i32,
                w[0].y as i32,
                w[1].x as i32,
                w[1].y as i32,
                color,
            );
        }

        // Draw circle at head position.
        if let Some(head) = pts.last() {
            draw_circle(&mut out, head.x as i32, head.y as i32, 3, color);
        }
    }

    DynamicImage::ImageRgb8(out).save(path)?;
    Ok(())
}

/// Assign a visually distinct colour per track via golden-ratio hue cycling.
fn track_color(index: usize) -> Rgb<u8> {
    let hue = (index as f32 * 137.508) % 360.0;
    hsv_to_rgb(hue, 0.9, 0.95)
}

/// HSV → RGB conversion.
fn hsv_to_rgb(h: f32, s: f32, v: f32) -> Rgb<u8> {
    let c = v * s;
    let x = c * (1.0 - ((h / 60.0) % 2.0 - 1.0).abs());
    let m = v - c;
    let (r1, g1, b1) = match (h / 60.0).floor() as i32 % 6 {
        0 => (c, x, 0.0),
        1 => (x, c, 0.0),
        2 => (0.0, c, x),
        3 => (0.0, x, c),
        4 => (x, 0.0, c),
        _ => (c, 0.0, x),
    };
    Rgb([
        ((r1 + m) * 255.0) as u8,
        ((g1 + m) * 255.0) as u8,
        ((b1 + m) * 255.0) as u8,
    ])
}

/// Bresenham line drawing.
fn draw_line(
    img: &mut image::RgbImage,
    mut x0: i32,
    mut y0: i32,
    x1: i32,
    y1: i32,
    color: Rgb<u8>,
) {
    let (w, h) = (img.width() as i32, img.height() as i32);
    let dx = (x1 - x0).abs();
    let dy = -(y1 - y0).abs();
    let sx = if x0 < x1 { 1 } else { -1 };
    let sy = if y0 < y1 { 1 } else { -1 };
    let mut err = dx + dy;
    loop {
        if x0 >= 0 && x0 < w && y0 >= 0 && y0 < h {
            img.put_pixel(x0 as u32, y0 as u32, color);
        }
        if x0 == x1 && y0 == y1 {
            break;
        }
        let e2 = 2 * err;
        if e2 >= dy {
            err += dy;
            x0 += sx;
        }
        if e2 <= dx {
            err += dx;
            y0 += sy;
        }
    }
}

/// Filled circle drawing.
fn draw_circle(img: &mut image::RgbImage, cx: i32, cy: i32, radius: i32, color: Rgb<u8>) {
    let (w, h) = (img.width() as i32, img.height() as i32);
    for dy in -radius..=radius {
        for dx in -radius..=radius {
            if dx * dx + dy * dy <= radius * radius {
                let px = cx + dx;
                let py = cy + dy;
                if px >= 0 && px < w && py >= 0 && py < h {
                    img.put_pixel(px as u32, py as u32, color);
                }
            }
        }
    }
}
