#![allow(dead_code, unused_imports, unused_variables, unused_macros)]
/// Grid-search stroke MAE tuning v2 — asymmetric shapes.
///
/// Expects `_tmp/stroke_v2_png_set/` with transparent PNGs exported
/// from `_tmp/stroke_v2_psd_set/`.

use std::io::Write;

const GT_DIR: &str = "_tmp/stroke_v2_png_set";
const CANVAS: u32 = 256;
const RESULTS_FILE: &str = "stroke_v2_grid_results.txt";

struct Sample {
    name: String,
    shape: String,
    position: String,
    size: f32,
    opacity: f32,
    color: [u8; 4],
    ref_pixels: Vec<u8>,
    content_mask: Vec<bool>,
}

fn main() {
    let entries: Vec<_> = std::fs::read_dir(GT_DIR)
        .unwrap_or_else(|e| panic!("cannot read {}: {}", GT_DIR, e))
        .filter_map(|e| e.ok())
        .filter(|e| {
            let s = e.file_name();
            let s = s.to_string_lossy();
            s.starts_with("st_") && s.ends_with(".png")
        })
        .collect();

    println!("Found {} PNGs in {}", entries.len(), GT_DIR);

    let mut samples: Vec<Sample> = Vec::new();

    for entry in &entries {
        let path = entry.path();
        let stem = path.file_stem().unwrap().to_string_lossy().to_string();

        let params = match parse_filename(&stem) {
            Some(p) => p,
            None => { eprintln!("  SKIP: {}", stem); continue; }
        };

        let (shape, position, size, opacity, color) = params;

        let img = match image::open(&path) {
            Ok(i) => i,
            Err(e) => { eprintln!("  SKIP {}: {}", stem, e); continue; }
        };

        let rgba = img.to_rgba8();
        let raw = rgba.into_raw();
        let npix = (CANVAS * CANVAS) as usize;
        let mut content_mask = vec![false; npix];
        for i in 0..npix {
            content_mask[i] = raw[i * 4 + 3] > 0;
        }

        samples.push(Sample {
            name: stem,
            shape, position, size, opacity, color,
            ref_pixels: raw,
            content_mask,
        });
    }

    println!("Loaded {} samples", samples.len());
    if samples.is_empty() {
        eprintln!("No valid samples found.");
        std::process::exit(1);
    }

    let content_count = samples[0].content_mask.iter().filter(|&&b| b).count();
    println!("Content pixels per image: {} / {} ({:.1}%)",
        content_count, npix(), content_count as f64 / npix() as f64 * 100.0);

    let anti_alias_widths: &[f32] = &[0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0];
    let feathers: &[f32] = &[0.0, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0];

    let total_combos = anti_alias_widths.len() * feathers.len();
    println!("Grid: {} combos x {} samples = {} evals", total_combos, samples.len(), total_combos * samples.len());
    std::io::stdout().flush().ok();

    let mut best_avg_mae = f64::MAX;
    let mut best_params = (0.0f32, 0.0f32);
    let mut eval_count = 0usize;

    let mut per_shape_best: std::collections::HashMap<String, (f64, f32, f32)> = std::collections::HashMap::new();

    for &aa in anti_alias_widths {
        for &feather in feathers {
            let mut total_mae = 0.0f64;

            for s in &samples {
                let alpha = make_alpha(&s.shape);
                let pos_enum = match s.position.as_str() {
                    "inside" => hcie_fx::types::StrokePosition::Inside,
                    "center" => hcie_fx::types::StrokePosition::Center,
                    "outside" => hcie_fx::types::StrokePosition::Outside,
                    _ => hcie_fx::types::StrokePosition::Outside,
                };

                let stroke = hcie_fx::tuned::generate_stroke_tuned(
                    &alpha, CANVAS, CANVAS,
                    s.size, pos_enum, s.color, s.opacity,
                    aa, feather,
                );

                let px = npix();
                let mut out = vec![0u8; px * 4];
                for i in 0..px {
                    if alpha[i] > 0 {
                        out[i * 4] = 0;
                        out[i * 4 + 1] = 0;
                        out[i * 4 + 2] = 0;
                        out[i * 4 + 3] = 255;
                    }
                }

                for i in 0..px {
                    let sa = stroke[i * 4 + 3];
                    if sa == 0 { continue; }
                    let dst = [out[i * 4], out[i * 4 + 1], out[i * 4 + 2], out[i * 4 + 3]];
                    let src = [stroke[i * 4], stroke[i * 4 + 1], stroke[i * 4 + 2], sa];
                    let blended = hcie_blend::blend_pixels(dst, src, hcie_blend::BlendMode::Normal, 1.0);
                    out[i * 4..i * 4 + 4].copy_from_slice(&blended);
                }

                let mut diff_sum = 0.0f64;
                let mut content_pixels = 0u64;
                for i in 0..px {
                    if !s.content_mask[i] { continue; }
                    content_pixels += 1;
                    for c in 0..4 {
                        diff_sum += (out[i * 4 + c] as i32 - s.ref_pixels[i * 4 + c] as i32).abs() as f64;
                    }
                }
                if content_pixels > 0 {
                    let mae = diff_sum / content_pixels as f64 / 4.0;
                    total_mae += mae;

                    let entry = per_shape_best.entry(s.shape.clone()).or_insert((f64::MAX, 0.0, 0.0));
                    if mae < entry.0 {
                        *entry = (mae, aa, feather);
                    }
                }
            }

            eval_count += 1;
            let avg_mae = total_mae / samples.len() as f64;

            if avg_mae < best_avg_mae {
                best_avg_mae = avg_mae;
                best_params = (aa, feather);
                println!("*** [{}/{}] NEW BEST avg MAE = {:.4} (aa={:.2}, feather={:.2})",
                    eval_count, total_combos, best_avg_mae, aa, feather);
                std::io::stdout().flush().ok();
            } else if eval_count % 10 == 0 {
                println!("  [{}/{}] avg={:.4}  best={:.4}", eval_count, total_combos, avg_mae, best_avg_mae);
                std::io::stdout().flush().ok();
            }
        }
    }

    let (baa, bf) = best_params;
    println!("\n=== RESULT ===");
    println!("Best avg MAE = {:.4}", best_avg_mae);
    println!("  anti_alias_width = {:.2}", baa);
    println!("  feather          = {:.2}", bf);

    println!("\n=== PER-SHAPE BEST ===");
    for (shape, (mae, aa, feather)) in &per_shape_best {
        println!("  {}: MAE={:.4} aa={:.2} feather={:.2}", shape, mae, aa, feather);
    }

    if let Ok(mut file) = std::fs::OpenOptions::new()
        .create(true).write(true).truncate(true).open(RESULTS_FILE)
    {
        let _ = writeln!(file,
            "anti_alias_width={:.2}  feather={:.2}\n\
             avg_mae={:.4}  samples={}  combos={}",
            baa, bf, best_avg_mae, samples.len(), total_combos);

        let _ = writeln!(file, "\n=== PER-SHAPE BEST ===");
        for (shape, (mae, aa, feather)) in &per_shape_best {
            let _ = writeln!(file, "  {}: MAE={:.4} aa={:.2} feather={:.2}", shape, mae, aa, feather);
        }
    }
}

fn npix() -> usize { (CANVAS * CANVAS) as usize }

fn make_alpha(shape: &str) -> Vec<u8> {
    let n = npix();
    let mut a = vec![0u8; n];
    let cx = CANVAS / 2;
    let cy = CANVAS / 2;

    for y in 0..CANVAS {
        for x in 0..CANVAS {
            let i = (y * CANVAS + x) as usize;
            let inside = match shape {
                "rect" => {
                    x >= 64 && x < CANVAS - 64 && y >= 64 && y < CANVAS - 64
                }
                "tall" => {
                    x >= 96 && x < CANVAS - 96 && y >= 32 && y < CANVAS - 32
                }
                "wide" => {
                    x >= 32 && x < CANVAS - 32 && y >= 96 && y < CANVAS - 96
                }
                "L" => {
                    (x >= 64 && x < 96 && y >= 64 && y < CANVAS - 64)
                    || (x >= 64 && x < CANVAS - 64 && y >= CANVAS - 96 && y < CANVAS - 64)
                }
                "T" => {
                    (y >= 64 && y < 96 && x >= 64 && x < CANVAS - 64)
                    || (x >= 112 && x < 144 && y >= 64 && y < CANVAS - 64)
                }
                "cross" => {
                    (x >= 112 && x < 144 && y >= 48 && y < CANVAS - 48)
                    || (y >= 112 && y < 144 && x >= 48 && x < CANVAS - 48)
                }
                "triangle" => {
                    let dx = (x as i32 - cx as i32).abs();
                    let top = 64i32;
                    let bottom = (CANVAS - 64) as i32;
                    if (y as i32) < top || (y as i32) > bottom { false } else {
                        let y_span = ((y as i32) - top) as f32 / (bottom - top) as f32;
                        let half_width = (CANVAS - 128) as f32 / 2.0 * (1.0 - y_span);
                        (dx as f32) <= half_width
                    }
                }
                "star" => {
                    let dx = x as f32 - cx as f32;
                    let dy = y as f32 - cy as f32;
                    let r = (dx * dx + dy * dy).sqrt();
                    if r > 96.0 || r < 20.0 { false } else {
                        let angle = dy.atan2(dx);
                        let sector = (angle * 5.0f32).cos();
                        let limit = 64.0 + sector * 32.0;
                        r <= limit
                    }
                }
                _ => false,
            };
            if inside {
                a[i] = 255;
            }
        }
    }
    a
}

fn parse_filename(name: &str) -> Option<(String, String, f32, f32, [u8; 4])> {
    // st_{shape}_{position}_sz{size}_op{opacity}_{color}
    let parts: Vec<&str> = name.split('_').collect();
    if parts.len() < 5 || parts[0] != "st" { return None; }
    let shape = parts[1].to_string();
    let position = parts[2].to_string();
    let size = parts[3].strip_prefix("sz")?.parse::<f32>().ok()?;
    let opacity = parts[4].strip_prefix("op")?.parse::<f32>().ok()? / 100.0;
    let color = color_from_name(parts[5]);
    Some((shape, position, size, opacity, color))
}

fn color_from_name(name: &str) -> [u8; 4] {
    match name {
        "red"     => [220, 30, 30, 255],
        "green"   => [30, 180, 30, 255],
        "blue"    => [30, 60, 220, 255],
        "white"   => [220, 220, 220, 255],
        "yellow"  => [255, 220, 30, 255],
        "cyan"    => [30, 220, 220, 255],
        "magenta" => [255, 30, 180, 255],
        _         => [220, 30, 30, 255],
    }
}