mod common;

use gif::Frame;
use image::ImageReader;
use nalgebra::SMatrix;
use plotters::prelude::*;

use common::fcns::{BoxConstraints, MultiModalFunction};
use common::img::{
    create_contour_data, find_closest_color, get_color_palette, setup_chart, setup_gif, ChartParams,
};

use non_convex_opt::utils::config::Config;
use non_convex_opt::NonConvexOpt;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let config_json = r#"
    {
        "opt_conf": {
            "max_iter": 10,
            "rtol": "1e-6",
            "atol": "1e-6",
            "rtol_max_iter_fraction": 1.0
        },
        "alg_conf": {
            "LBFGS": {
                "common": {
                    "memory_size": 10
                },
                "line_search": {
                    "Backtracking": {
                        "c1": 0.0001,
                        "rho": 0.5
                    }
                },
                "advanced": {
                    "adaptive_parameters": false,
                    "adaptation_rate": 0.1,
                    "restart_strategy": "None",
                    "stagnation_detection": {
                        "stagnation_window": 50,
                        "improvement_threshold": 1e-6,
                        "gradient_threshold": 1e-6
                    },
                    "memory_adaptation": {
                        "adaptive_memory": false,
                        "min_memory_size": 5,
                        "max_memory_size": 20,
                        "memory_adaptation_rate": 0.1
                    },
                    "numerical_safeguards": {
                        "conditioning_threshold": 1e-12,
                        "curvature_threshold": 1e-8,
                        "use_scaling": false,
                        "scaling_factor": 1.0
                    },
                    "success_history_size": 20,
                    "improvement_history_size": 20
                }
            }
        }
    }"#;

    let config: Config = serde_json::from_str(config_json)?;

    let obj_f = MultiModalFunction;
    let constraints = BoxConstraints;

    let mut opt = NonConvexOpt::new(
        config,
        SMatrix::<f64, 1, 2>::from_vec(vec![4.0, 9.0]),
        obj_f.clone(),
        Some(constraints.clone()),
        42,
    );

    let resolution = 100;
    let (z_values, min_val, max_val) = create_contour_data(&obj_f, resolution);
    let color_palette = get_color_palette();
    let mut encoder = setup_gif("examples/gifs/lbfgs_kbf.gif")?;

    for frame in 0..10 {
        let mut chart = setup_chart(ChartParams {
            frame,
            algorithm_name: "L-BFGS",
            resolution,
            z_values: &z_values,
            min_val,
            max_val,
            constraints: &constraints,
            frame_path: "examples/lbfgs_frame.png",
        })?;

        // Draw best individual in yellow
        let best_x = opt.get_best_individual();
        chart.draw_series(std::iter::once(Circle::new(
            (best_x[0], best_x[1]),
            6,
            RGBColor(255, 0, 0).filled(),
        )))?;

        // Save frame and convert to GIF
        chart.plotting_area().present()?;

        // Convert PNG to GIF frame
        let img = ImageReader::open("examples/lbfgs_frame.png")?
            .decode()?
            .into_rgba8();

        let mut indexed_pixels = Vec::with_capacity((img.width() * img.height()) as usize);
        for pixel in img.pixels() {
            let idx = find_closest_color(pixel[0], pixel[1], pixel[2], &color_palette);
            indexed_pixels.push(idx as u8);
        }

        let frame = Frame::<'_> {
            width: 800,
            height: 800,
            delay: 40,
            buffer: std::borrow::Cow::from(indexed_pixels),
            ..Default::default()
        };
        encoder.write_frame(&frame)?;

        opt.step();
    }

    std::fs::remove_file("examples/lbfgs_frame.png")?;

    Ok(())
}
