#![allow(dead_code, unused_imports, unused_variables, unused_macros)]
use std::io::Write;

const STROKE_RS: &str = "hcie-fx/src/stroke.rs";
const RESULTS_FILE: &str = "/mnt/extra/00_PROJECTS/hcie-rust-v4/stroke_results.txt";

struct TestCase {
    layer_name: String,
    width: usize,
    height: usize,
    size: f32,
    position: hcie_fx::types::StrokePosition,
    color: [u8; 4],
    blend_mode_str: String,
    opacity: f32,          // stroke effect opacity
    layer_opacity: f32,    // overall layer opacity
    fill_opacity: f32,     // fill opacity (affects only content)
    alpha: Vec<u8>,
    layer_pixels: Vec<u8>,
    ref_pixels: Vec<u8>,
}

struct EdtPrecomputed {
    dist_in_sq: Vec<f32>,
    dist_out_sq: Vec<f32>,
    nearest: Vec<usize>,
}

fn main() {
    let original_code = std::fs::read_to_string(STROKE_RS).expect("read stroke.rs");
    let start_tag = "    // BEGIN_TUNING_ZONE";
    let end_tag = "    // END_TUNING_ZONE";

    let zone_start = original_code.find(start_tag).unwrap();
    let zone_end = original_code.find(end_tag).unwrap();

    let inside_dir = "/mnt/extra/00_PROJECTS/hcie-rust-v4/psd-test-images/stroke/stroke-inside";
    let outside_dir = "/mnt/extra/00_PROJECTS/hcie-rust-v4/psd-test-images/stroke/stroke-outside";

    let individual_psds: &[(&str, &str)] = &[
        (inside_dir, "Shape 1"),
        (inside_dir, "Shape 2"),
        (inside_dir, "Shape 3"),
        (inside_dir, "Shape 4"),
        (outside_dir, "Shape 1"),
        (outside_dir, "Shape 2"),
        (outside_dir, "Shape 3"),
        (outside_dir, "Layer 2"),
    ];

    let composite_psds: &[&str] = &[
        "psd-test-images/stroke/stroke-inside/stroke-inside.psd",
        "psd-test-images/stroke/stroke-outside/stroke-outside.psd",
    ];

    let mut test_cases = Vec::new();
    let mut seen_layers: std::collections::HashSet<String> = std::collections::HashSet::new();

    for &(dir, stem) in individual_psds {
        let psd_path = format!("{}/{}.psd", dir, stem);
        let layers = hcie_psd::import_psd(std::path::Path::new(&psd_path)).expect("import psd");

        for layer in layers {
            let alpha: Vec<u8> = layer.pixels.iter().skip(3).step_by(4).copied().collect();
            if alpha.iter().all(|&a| a == 0) {
                continue;
            }

            let (size, position, color, blend_mode_str, opacity) = extract_stroke_effect(&layer);
            println!("Layer '{}' — size={}, position={:?}, stroke_opacity={:.2}, layer_opacity={:.2}, fill_opacity={:.2}",
                     layer.name, size, position, opacity, layer.opacity, layer.fill_opacity);

            let png_path = format!("{}/{}.png", dir, layer.name);
            let ref_img = image::open(&png_path)
                .unwrap_or_else(|_| panic!("Failed to open reference png: {}", png_path))
                .to_rgba8();

            seen_layers.insert(layer.name.clone());

            test_cases.push(TestCase {
                layer_name: layer.name.clone(),
                width: layer.width as usize,
                height: layer.height as usize,
                size,
                position,
                color,
                blend_mode_str,
                opacity,
                layer_opacity: layer.opacity,
                fill_opacity: layer.fill_opacity,
                alpha,
                layer_pixels: layer.pixels.clone(),
                ref_pixels: ref_img.into_raw(),
            });
            break;
        }
    }

    for psd_path in composite_psds {
        let psd_abs = std::path::Path::new("/mnt/extra/00_PROJECTS/hcie-rust-v4").join(psd_path);
        let dir = psd_abs.parent().unwrap().to_str().unwrap();
        let layers = hcie_psd::import_psd(&psd_abs).expect("import psd");

        for layer in layers {
            if seen_layers.contains(&layer.name) {
                continue;
            }

            let alpha: Vec<u8> = layer.pixels.iter().skip(3).step_by(4).copied().collect();
            if alpha.iter().all(|&a| a == 0) {
                continue;
            }

            let (size, position, color, blend_mode_str, opacity) = extract_stroke_effect(&layer);
            println!("Layer '{}' (composite) — size={}, position={:?}, stroke_opacity={:.2}, layer_opacity={:.2}, fill_opacity={:.2}",
                     layer.name, size, position, opacity, layer.opacity, layer.fill_opacity);

            let png_path = format!("{}/{}.png", dir, layer.name);
            let ref_img = image::open(&png_path)
                .unwrap_or_else(|_| panic!("Failed to open reference png: {}", png_path))
                .to_rgba8();

            test_cases.push(TestCase {
                layer_name: layer.name.clone(),
                width: layer.width as usize,
                height: layer.height as usize,
                size,
                position,
                color,
                blend_mode_str,
                opacity,
                layer_opacity: layer.opacity,
                fill_opacity: layer.fill_opacity,
                alpha,
                layer_pixels: layer.pixels.clone(),
                ref_pixels: ref_img.into_raw(),
            });
        }
    }

    println!("Loaded {} stroke test cases.", test_cases.len());

    // 1. Calculate Baselines using the original O(R^2) algorithm
    let mut baseline_maes = Vec::new();
    for tc in &test_cases {
        let stroke = generate_stroke_original(&tc.alpha, tc.width as u32, tc.height as u32, tc.size, tc.position.clone(), tc.color, tc.opacity);
        let mae = evaluate_composition(tc, &stroke);
        baseline_maes.push(mae);
        println!("  Layer '{}': Baseline MAE = {:.4}", tc.layer_name, mae);
    }
    let avg_baseline_mae: f64 = baseline_maes.iter().sum::<f64>() / baseline_maes.len() as f64;
    println!("Average Baseline MAE = {:.4}", avg_baseline_mae);
    std::io::stdout().flush().ok();

    // Precompute EDT matrices for grid search
    let mut precomputed = Vec::new();
    for tc in &test_cases {
        let dist_in_sq = hcie_fx::tuned::edt_inside(&tc.alpha, tc.width, tc.height);
        let dist_out_sq = hcie_fx::tuned::edt_outside(&tc.alpha, tc.width, tc.height);
        let nearest = hcie_fx::tuned::edt_outside_src(&tc.alpha, tc.width, tc.height);
        precomputed.push(EdtPrecomputed {
            dist_in_sq,
            dist_out_sq,
            nearest,
        });
    }

    // Grid Search
    let aa_list: &[f32] = &[0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0, 4.0];
    let feather_list: &[f32] = &[0.0, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0];
    let normal_step_list: &[f32] = &[0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 4.0, 5.0];

    let mut best_avg_mae = f64::MAX;
    let mut best_aa = 1.0f32;
    let mut best_feather = 0.0f32;
    let mut best_normal_step = 1.5f32;

    let total_combos = aa_list.len() * feather_list.len() * normal_step_list.len();
    println!("Grid search starting with {} combinations...", total_combos);
    std::io::stdout().flush().ok();

    let mut count = 0;
    for &aa in aa_list {
        for &feather in feather_list {
            for &normal_step in normal_step_list {
                count += 1;
                if count % 100 == 0 {
                    println!("  [{}/{}] Current Best Avg MAE = {:.4}", count, total_combos, best_avg_mae);
                    std::io::stdout().flush().ok();
                }

                let mut maes = Vec::new();
                for (idx, tc) in test_cases.iter().enumerate() {
                    let prep = &precomputed[idx];
                    let stroke = render_stroke_tuned(tc, prep, aa, feather, normal_step);
                    let mae = evaluate_composition(tc, &stroke);
                    maes.push(mae);
                }

                let avg_mae: f64 = maes.iter().sum::<f64>() / maes.len() as f64;
                if avg_mae < best_avg_mae {
                    best_avg_mae = avg_mae;
                    best_aa = aa;
                    best_feather = feather;
                    best_normal_step = normal_step;
                    println!("*** NEW BEST: Avg MAE = {:.4} (aa={:.2}, feather={:.2}, normal_step={:.2})", 
                             avg_mae, aa, feather, normal_step);
                    std::io::stdout().flush().ok();
                }
            }
        }
    }

    println!("\nBest parameters found:");
    println!("  aa = {:.2}", best_aa);
    println!("  feather = {:.2}", best_feather);
    println!("  normal_step = {:.2}", best_normal_step);
    println!("  Average MAE = {:.4} (Baseline = {:.4})", best_avg_mae, avg_baseline_mae);

    // Save best parameters to stroke.rs
    let zone_code = format!(
        "    // BEGIN_TUNING_ZONE
    let aa = {:.2}f32;
    let feather = {:.2}f32;
    let normal_step = {:.2}f32;
    // END_TUNING_ZONE",
        best_aa, best_feather, best_normal_step
    );

    let mut patched = String::with_capacity(original_code.len());
    patched.push_str(&original_code[..zone_start]);
    patched.push_str(&zone_code);
    patched.push_str(&original_code[zone_end + end_tag.len()..]);

    std::fs::write(STROKE_RS, &patched).expect("write patched stroke.rs");
    println!("Patched hcie-fx/src/stroke.rs with optimal parameters.");

    // Write final summary to stroke_results.txt
    if let Ok(mut file) = std::fs::OpenOptions::new().create(true).write(true).truncate(true).open(RESULTS_FILE) {
        let _ = writeln!(file, "=== Stroke Tuned ===\n  Optimal Params: aa={:.2}, feather={:.2}, normal_step={:.2}\n  Average Tuned MAE: {:.4}\n  Average Baseline MAE: {:.4}",
            best_aa, best_feather, best_normal_step, best_avg_mae, avg_baseline_mae);
    }
}

fn render_stroke_tuned(tc: &TestCase, prep: &EdtPrecomputed, aa: f32, feather: f32, normal_step: f32) -> Vec<u8> {
    let px = tc.width * tc.height;
    let iw = tc.width;
    let ih = tc.height;
    let mut stroke_alpha = vec![0.0f32; px];
    let radius = tc.size.max(1.0);
    let mut buf = vec![0u8; px * 4];

    match tc.position {
        hcie_fx::types::StrokePosition::Inside => {
            for i in 0..px {
                if tc.alpha[i] == 0 { continue; }
                let d = prep.dist_in_sq[i].sqrt();
                if d >= radius { continue; }
                let outer_dist = radius - d;
                let coverage = (outer_dist / aa).clamp(0.0, 1.0);
                stroke_alpha[i] = coverage * tc.alpha[i] as f32;
            }
        }
        hcie_fx::types::StrokePosition::Outside => {
            for i in 0..px {
                let d = prep.dist_out_sq[i].sqrt();
                if d == 0.0 || d >= radius { continue; }
                let x_i = (i % iw) as f32;
                let y_i = (i / iw) as f32;
                let x_n = (prep.nearest[i] % iw) as f32;
                let y_n = (prep.nearest[i] / iw) as f32;
                let dx = x_i - x_n;
                let dy = y_i - y_n;
                let src_a = if normal_step > 0.0 {
                    let sx = (x_n - dx / d * normal_step).round() as i32;
                    let sy = (y_n - dy / d * normal_step).round() as i32;
                    if sx >= 0 && sx < iw as i32 && sy >= 0 && sy < ih as i32 {
                        tc.alpha[sy as usize * iw + sx as usize] as f32
                    } else {
                        tc.alpha[prep.nearest[i]] as f32
                    }
                } else {
                    tc.alpha[prep.nearest[i]] as f32
                };
                let outer_dist = radius - d;
                let coverage = (outer_dist / aa).clamp(0.0, 1.0);
                stroke_alpha[i] = coverage * src_a;
            }
        }
        hcie_fx::types::StrokePosition::Center => {
            let half_r = radius / 2.0;
            for i in 0..px {
                let (d, is_inside) = if tc.alpha[i] > 0 {
                    (prep.dist_in_sq[i].sqrt(), true)
                } else {
                    (prep.dist_out_sq[i].sqrt(), false)
                };
                if d == 0.0 || d >= half_r { continue; }

                let src_a = if is_inside {
                    tc.alpha[i] as f32
                } else {
                    let x_i = (i % iw) as f32;
                    let y_i = (i / iw) as f32;
                    let x_n = (prep.nearest[i] % iw) as f32;
                    let y_n = (prep.nearest[i] / iw) as f32;
                    let dx = x_i - x_n;
                    let dy = y_i - y_n;
                    if normal_step > 0.0 {
                        let sx = (x_n - dx / d * normal_step).round() as i32;
                        let sy = (y_n - dy / d * normal_step).round() as i32;
                        if sx >= 0 && sx < iw as i32 && sy >= 0 && sy < ih as i32 {
                            tc.alpha[sy as usize * iw + sx as usize] as f32
                        } else {
                            tc.alpha[prep.nearest[i]] as f32
                        }
                    } else {
                        tc.alpha[prep.nearest[i]] as f32
                    }
                };

                let outer_dist = half_r - d;
                let coverage = (outer_dist / aa).clamp(0.0, 1.0);
                stroke_alpha[i] = coverage * src_a;
            }
        }
    }

    if feather > 0.0 {
        let feather_radius = (feather.round() as i32).max(1);
        stroke_alpha = hcie_fx::box_blur_f32(&stroke_alpha, tc.width as u32, tc.height as u32, feather_radius);
    }

    for i in 0..px {
        let a = (stroke_alpha[i] * tc.opacity).round() as u8;
        if a > 0 {
            buf[i * 4] = tc.color[0];
            buf[i * 4 + 1] = tc.color[1];
            buf[i * 4 + 2] = tc.color[2];
            buf[i * 4 + 3] = a;
        }
    }

    buf
}

fn evaluate_composition(tc: &TestCase, stroke: &[u8]) -> f64 {
    let px = tc.width * tc.height;
    let bm_parsed = hcie_fx::types::blend_mode_from_string(&tc.blend_mode_str);

    let mut out = vec![0u8; px * 4];
    for i in 0..px {
        let sa = stroke[i * 4 + 3];
        let stroke_px = [stroke[i * 4], stroke[i * 4 + 1], stroke[i * 4 + 2], sa];
        let orig_px = [
            tc.layer_pixels[i * 4],
            tc.layer_pixels[i * 4 + 1],
            tc.layer_pixels[i * 4 + 2],
            tc.layer_pixels[i * 4 + 3],
        ];

        let blended = match tc.position {
            hcie_fx::types::StrokePosition::Outside => {
                let behind = if sa == 0 {
                    [0, 0, 0, 0]
                } else {
                    hcie_blend::blend_pixels([0, 0, 0, 0], stroke_px, bm_parsed, 1.0)
                };
                if orig_px[3] == 0 {
                    behind
                } else if orig_px[3] == 255 {
                    orig_px
                } else {
                    hcie_blend::blend_pixels(behind, orig_px, hcie_blend::BlendMode::Normal, 1.0)
                }
            }
            hcie_fx::types::StrokePosition::Inside | hcie_fx::types::StrokePosition::Center => {
                if sa == 0 {
                    orig_px
                } else {
                    // Opacity already baked into stroke alpha
                    let eff_sa = sa as f32 / 255.0;
                    let keep_alpha = orig_px[3];
                    if bm_parsed == hcie_blend::BlendMode::Normal {
                        if eff_sa >= 1.0 {
                            [stroke_px[0], stroke_px[1], stroke_px[2], keep_alpha]
                        } else {
                            let dst_r = orig_px[0] as f32;
                            let dst_g = orig_px[1] as f32;
                            let dst_b = orig_px[2] as f32;
                            let src_r = stroke_px[0] as f32;
                            let src_g = stroke_px[1] as f32;
                            let src_b = stroke_px[2] as f32;
                            [
                                (dst_r * (1.0 - eff_sa) + src_r * eff_sa).round() as u8,
                                (dst_g * (1.0 - eff_sa) + src_g * eff_sa).round() as u8,
                                (dst_b * (1.0 - eff_sa) + src_b * eff_sa).round() as u8,
                                keep_alpha,
                            ]
                        }
                    } else {
                        let tmp = hcie_blend::blend_pixels(orig_px, stroke_px, bm_parsed, 1.0);
                        [tmp[0], tmp[1], tmp[2], keep_alpha]
                    }
                }
            }
        };
        out[i * 4..i * 4 + 4].copy_from_slice(&blended);
    }

    let diff_sum: f64 = (0..px).map(|i| {
        (0..4).map(|c| (out[i * 4 + c] as i32 - tc.ref_pixels[i * 4 + c] as i32).abs() as f64).sum::<f64>()
    }).sum();

    diff_sum / px as f64 / 4.0
}

fn extract_stroke_effect(layer: &hcie_protocol::Layer) -> (f32, hcie_fx::types::StrokePosition, [u8; 4], String, f32) {
    let mut size = 0.0f32;
    let mut position = hcie_protocol::effects::StrokePosition::Outside;
    let mut color = [0u8; 4];
    let mut blend_mode_str = String::new();
    let mut opacity = 0.0f32;
    for e in &layer.effects {
        if let hcie_protocol::effects::LayerEffect::Stroke { enabled: true, size: s, position: p, color: c, blend_mode: bm, opacity: o, .. } = e {
            size = *s;
            position = p.clone();
            color = *c;
            blend_mode_str = bm.clone();
            opacity = *o;
        }
    }
    let fx_pos = match position {
        hcie_protocol::effects::StrokePosition::Outside => hcie_fx::types::StrokePosition::Outside,
        hcie_protocol::effects::StrokePosition::Inside => hcie_fx::types::StrokePosition::Inside,
        hcie_protocol::effects::StrokePosition::Center => hcie_fx::types::StrokePosition::Center,
    };
    (size, fx_pos, color, blend_mode_str, opacity)
}

fn generate_stroke_original(
    alpha: &[u8],
    w: u32,
    h: u32,
    size: f32,
    position: hcie_fx::types::StrokePosition,
    color: [u8; 4],
    opacity: f32,
) -> Vec<u8> {
    let px = (w * h) as usize;
    let mut buf = vec![0u8; px * 4];
    let radius = size.max(1.0).round() as i32;

    match position {
        hcie_fx::types::StrokePosition::Inside => {
            for y in 0..h as i32 {
                for x in 0..w as i32 {
                    let i = (y as u32 * w + x as u32) as usize;
                    if alpha[i] == 0 { continue; }

                    let mut min_dist_sq = f32::INFINITY;
                    for dy in -radius..=radius {
                        for dx in -radius..=radius {
                            let dist_sq = (dx * dx + dy * dy) as f32;
                            if dist_sq > radius as f32 * radius as f32 { continue; }
                            let nx = x + dx;
                            let ny = y + dy;
                            if nx < 0 || nx >= w as i32 || ny < 0 || ny >= h as i32 {
                                min_dist_sq = min_dist_sq.min(dist_sq);
                                continue;
                            }
                            let ni = (ny as u32 * w + nx as u32) as usize;
                            if alpha[ni] == 0 {
                                min_dist_sq = min_dist_sq.min(dist_sq);
                            }
                        }
                    }

                    if min_dist_sq < f32::INFINITY {
                        let d = min_dist_sq.sqrt();
                        let a = if d <= radius as f32 - 1.0 {
                            alpha[i] as f32 * opacity
                        } else {
                            let t = (radius as f32 - d).clamp(0.0, 1.0);
                            alpha[i] as f32 * t * opacity
                        };
                        let a = a.round() as u8;
                        if a > 0 {
                            buf[i * 4] = color[0];
                            buf[i * 4 + 1] = color[1];
                            buf[i * 4 + 2] = color[2];
                            buf[i * 4 + 3] = a;
                        }
                    }
                }
            }
        }
        hcie_fx::types::StrokePosition::Outside => {
            for y in 0..h as i32 {
                for x in 0..w as i32 {
                    let i = (y as u32 * w + x as u32) as usize;

                    let mut min_dist_sq = f32::INFINITY;
                    let mut max_alpha = 0u8;
                    for dy in -radius..=radius {
                        for dx in -radius..=radius {
                            let dist_sq = (dx * dx + dy * dy) as f32;
                            if dist_sq > radius as f32 * radius as f32 { continue; }
                            let nx = x + dx;
                            let ny = y + dy;
                            if nx < 0 || nx >= w as i32 || ny < 0 || ny >= h as i32 { continue; }
                            let ni = (ny as u32 * w + nx as u32) as usize;
                            if alpha[ni] > 0 {
                                if dist_sq < min_dist_sq {
                                    min_dist_sq = dist_sq;
                                    max_alpha = max_alpha.max(alpha[ni]);
                                }
                            }
                        }
                    }

                    if min_dist_sq < f32::INFINITY {
                        let d = min_dist_sq.sqrt();
                        let a = if d <= radius as f32 - 1.0 {
                            max_alpha as f32 * opacity
                        } else {
                            let t = (radius as f32 - d).clamp(0.0, 1.0);
                            max_alpha as f32 * t * opacity
                        };

                        let a = a.round() as u8;
                        if a > 0 {
                            buf[i * 4] = color[0];
                            buf[i * 4 + 1] = color[1];
                            buf[i * 4 + 2] = color[2];
                            buf[i * 4 + 3] = a;
                        }
                    }
                }
            }
        }
        hcie_fx::types::StrokePosition::Center => {
            let half_r = radius as f32 / 2.0;
            for y in 0..h as i32 {
                for x in 0..w as i32 {
                    let i = (y as u32 * w + x as u32) as usize;

                    let mut min_dist_sq = f32::INFINITY;
                    let mut max_alpha = 0u8;
                    let is_inside = alpha[i] > 0;

                    for dy in -radius..=radius {
                        for dx in -radius..=radius {
                            let dist_sq = (dx * dx + dy * dy) as f32;
                            if dist_sq > half_r * half_r { continue; }
                            let nx = x + dx;
                            let ny = y + dy;

                            if nx < 0 || nx >= w as i32 || ny < 0 || ny >= h as i32 {
                                if is_inside {
                                    min_dist_sq = min_dist_sq.min(dist_sq);
                                }
                                continue;
                            }

                            let ni = (ny as u32 * w + nx as u32) as usize;
                            if is_inside {
                                if alpha[ni] == 0 {
                                    min_dist_sq = min_dist_sq.min(dist_sq);
                                }
                            } else {
                                if alpha[ni] > 0 {
                                    if dist_sq < min_dist_sq {
                                        min_dist_sq = dist_sq;
                                        max_alpha = max_alpha.max(alpha[ni]);
                                    }
                                }
                            }
                        }
                    }

                    if min_dist_sq < f32::INFINITY {
                        let d = min_dist_sq.sqrt();
                        let base_a = if is_inside { alpha[i] } else { max_alpha };
                        let a = if d <= half_r - 1.0 {
                            base_a as f32 * opacity
                        } else {
                            let t = (half_r - d).clamp(0.0, 1.0);
                            base_a as f32 * t * opacity
                        };

                        let a = a.round() as u8;
                        if a > 0 {
                            buf[i * 4] = color[0];
                            buf[i * 4 + 1] = color[1];
                            buf[i * 4 + 2] = color[2];
                            buf[i * 4 + 3] = a;
                        }
                    }
                }
            }
        }
    }

    buf
}