hcie-rust-v3.05/hcie-io/examples/tune_mae_bevel_emboss_grid.rs

#![allow(dead_code, unused_imports, unused_variables, unused_macros)]
use rayon::prelude::*;
/// Grid-search Bevel & Emboss MAE tuning against Photoshop ground-truth PNGs.
///
/// Expects `_tmp/psd_tunes/bevel_emboss_png_set/` with transparent PNGs exported
/// from `_tmp/psd_tunes/bevel_emboss_psd_set/`.
///
/// Uses the full `apply_layer_effects` pipeline (including inside/behind split
/// for Emboss/OuterBevel) so the tuning matches real rendering behavior.
///
/// Usage:
///   cargo run --example tune_mae_bevel_emboss_grid -p hcie-io --release
use std::fs;
use std::io::Write;

const GT_DIR: &str = "_tmp/psd_tunes/bevel_emboss_png_set";
const RESULTS_FILE: &str = "bevel_emboss_grid_results_v4.txt";

struct Sample {
    name: String,
    style: String,
    depth: f32,
    size: f32,
    soften: f32,
    direction: String,
    altitude: f32,
    angle: f32,
    ref_pixels: Vec<u8>,
    content_mask: Vec<bool>,
    layer_pixels: Vec<u8>,
    fx_effects: Vec<hcie_fx::types::LayerEffect>,
    fill_opacity: f32,
    w: u32,
    h: u32,
}

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("be_") && 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, style, depth, size, soften, direction, altitude, _technique) = 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 w = img.width();
        let h = img.height();
        let npix = (w * h) as usize;
        let mut content_mask = vec![false; npix];
        for i in 0..npix {
            content_mask[i] = raw[i * 4 + 3] > 0;
        }

        let psd_path_str = format!("_tmp/psd_tunes/bevel_emboss_psd_set/{}.psd", stem);
        let (layer_pixels, fx_effects, fill_opacity, angle) = if let Ok(layers) =
            hcie_psd::import_psd(std::path::Path::new(&psd_path_str))
        {
            let mut found = None;
            for l in layers {
                if !l.effects.is_empty() {
                    let fx_effects: Vec<hcie_fx::types::LayerEffect> = l
                        .effects
                        .iter()
                        .map(|e| hcie_fx::types::protocol_to_hcie_fx_effect(e))
                        .collect();
                    // Extract angle from BevelEmboss effect
                    let mut angle = 120.0f32;
                    for e in &l.effects {
                        if let hcie_protocol::effects::LayerEffect::BevelEmboss {
                            angle: a, ..
                        } = e
                        {
                            angle = *a;
                        }
                    }
                    found = Some((l.pixels, fx_effects, l.fill_opacity, angle));
                    break;
                }
            }
            match found {
                Some(v) => v,
                None => {
                    eprintln!("  SKIP no effects layer: {}", stem);
                    continue;
                }
            }
        } else {
            eprintln!("  SKIP PSD missing for {}: {}", stem, psd_path_str);
            continue;
        };

        samples.push(Sample {
            name: stem,
            style,
            depth,
            size,
            soften,
            direction,
            altitude,
            angle,
            ref_pixels: raw,
            content_mask,
            layer_pixels,
            fx_effects,
            fill_opacity,
            w,
            h,
        });
    }

    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: {}", content_count);

    // Use all samples (inner bevel, outer bevel, emboss) for combined tuning

    // Grid search over lighting_blur, hl_scale, sh_scale, tilt_scale
    // Wide ranges to cover high-depth PSDs (e.g. Sultan depth=350) and
    // various bevel sizes (10, 25+). Each axis spans low→high so the
    // search can find optima for both simple and complex shapes.
    let lighting_blurs: &[i32] = &[0, 2, 5];
    let hl_scales: &[f32] = &[0.05, 0.1, 0.2, 0.5];
    let sh_scales: &[f32] = &[2.0, 8.0, 16.0];
    let tilt_scales: &[f32] = &[0.05, 0.1, 0.2, 0.3, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0, 5.0, 8.0];

    let total_combos = lighting_blurs.len() * hl_scales.len() * sh_scales.len() * tilt_scales.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 = (0i32, 0.0f32, 0.0f32, 0.0f32);
    let mut eval_count = 0usize;

    for &lighting_blur in lighting_blurs {
        for &hl_scale in hl_scales {
            for &sh_scale in sh_scales {
                for &tilt_scale in tilt_scales {
                    let total_mae: f64 = samples
                        .par_iter()
                        .map(|s| {
                            // Use generate_bevel_tuned to get highlight/shadow buffers
                            let style_enum = match s.style.as_str() {
                                "emboss" => hcie_fx::types::BevelStyle::Emboss,
                                "outer" => hcie_fx::types::BevelStyle::OuterBevel,
                                _ => hcie_fx::types::BevelStyle::InnerBevel,
                            };
                            let dir_enum = match s.direction.as_str() {
                                "down" => hcie_fx::types::Direction::Down,
                                _ => hcie_fx::types::Direction::Up,
                            };

                            let (highlight, shadow) = hcie_fx::tuned::generate_bevel_tuned(
                                &s.layer_pixels
                                    .iter()
                                    .skip(3)
                                    .step_by(4)
                                    .copied()
                                    .collect::<Vec<u8>>(),
                                s.w,
                                s.h,
                                s.depth,
                                s.size,
                                s.soften,
                                s.angle,
                                s.altitude,
                                &dir_enum,
                                &hcie_fx::types::Technique::Smooth,
                                &style_enum,
                                lighting_blur,
                                hl_scale,
                                sh_scale,
                                1.0,
                                tilt_scale,
                            );

                            let px = (s.w * s.h) as usize;
                            let mut out = vec![0u8; px * 4];

                            // Base shape: fill with layer pixels (black fill)
                            for i in 0..px {
                                if s.layer_pixels[i * 4 + 3] > 0 {
                                    out[i * 4] = s.layer_pixels[i * 4];
                                    out[i * 4 + 1] = s.layer_pixels[i * 4 + 1];
                                    out[i * 4 + 2] = s.layer_pixels[i * 4 + 2];
                                    out[i * 4 + 3] = s.layer_pixels[i * 4 + 3];
                                }
                            }

                            // For emboss/outerbevel: outside highlight/shadow go to empty canvas (behind)
                            let has_outer = matches!(
                                style_enum,
                                hcie_fx::types::BevelStyle::Emboss
                                    | hcie_fx::types::BevelStyle::OuterBevel
                            );

                            // Get opacity and blend mode from fx_effects
                            let (hl_op, sh_op, hl_bm, sh_bm) = {
                                let mut hl_op = 0.75f32;
                                let mut sh_op = 0.75f32;
                                let mut hl_bm = hcie_blend::BlendMode::Screen;
                                let mut sh_bm = hcie_blend::BlendMode::Multiply;
                                for e in &s.fx_effects {
                                    if let hcie_fx::types::LayerEffect::BevelEmboss {
                                        highlight_opacity,
                                        shadow_opacity,
                                        highlight_blend,
                                        shadow_blend,
                                        ..
                                    } = e
                                    {
                                        hl_op = *highlight_opacity / 100.0;
                                        sh_op = *shadow_opacity / 100.0;
                                        hl_bm = *highlight_blend;
                                        sh_bm = *shadow_blend;
                                    }
                                }
                                (hl_op, sh_op, hl_bm, sh_bm)
                            };

                            // Composite highlight — inside on fill, outside on empty canvas
                            for i in 0..px {
                                let sa = highlight[i * 4 + 3];
                                if sa > 0 {
                                    let is_inner = s.layer_pixels[i * 4 + 3] > 0;
                                    if !has_outer || is_inner {
                                        let dst = [
                                            out[i * 4],
                                            out[i * 4 + 1],
                                            out[i * 4 + 2],
                                            out[i * 4 + 3],
                                        ];
                                        let src = [
                                            highlight[i * 4],
                                            highlight[i * 4 + 1],
                                            highlight[i * 4 + 2],
                                            sa,
                                        ];
                                        let blended =
                                            hcie_blend::blend_pixels(dst, src, hl_bm, hl_op);
                                        out[i * 4..i * 4 + 4].copy_from_slice(&blended);
                                    } else {
                                        // Outside: composite highlight on empty canvas
                                        let dst = [0u8, 0, 0, 0];
                                        let src = [
                                            highlight[i * 4],
                                            highlight[i * 4 + 1],
                                            highlight[i * 4 + 2],
                                            sa,
                                        ];
                                        let blended =
                                            hcie_blend::blend_pixels(dst, src, hl_bm, hl_op);
                                        out[i * 4..i * 4 + 4].copy_from_slice(&blended);
                                    }
                                }
                            }
                            // Composite shadow
                            for i in 0..px {
                                let sa = shadow[i * 4 + 3];
                                if sa > 0 {
                                    let is_inner = s.layer_pixels[i * 4 + 3] > 0;
                                    if !has_outer || is_inner {
                                        let dst = [
                                            out[i * 4],
                                            out[i * 4 + 1],
                                            out[i * 4 + 2],
                                            out[i * 4 + 3],
                                        ];
                                        let src = [
                                            shadow[i * 4],
                                            shadow[i * 4 + 1],
                                            shadow[i * 4 + 2],
                                            sa,
                                        ];
                                        let blended =
                                            hcie_blend::blend_pixels(dst, src, sh_bm, sh_op);
                                        out[i * 4..i * 4 + 4].copy_from_slice(&blended);
                                    } else {
                                        let dst = [
                                            out[i * 4],
                                            out[i * 4 + 1],
                                            out[i * 4 + 2],
                                            out[i * 4 + 3],
                                        ];
                                        let src = [
                                            shadow[i * 4],
                                            shadow[i * 4 + 1],
                                            shadow[i * 4 + 2],
                                            sa,
                                        ];
                                        let blended =
                                            hcie_blend::blend_pixels(dst, src, sh_bm, sh_op);
                                        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 {
                                diff_sum / content_pixels as f64 / 4.0
                            } else {
                                0.0
                            }
                        })
                        .sum();

                    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 = (lighting_blur, hl_scale, sh_scale, tilt_scale);
                        println!("*** [{}/{}] NEW BEST avg MAE = {:.4} (lb={}, hl={:.3}, sh={:.2}, tilt={:.2})",
                        eval_count, total_combos, best_avg_mae, lighting_blur, hl_scale, sh_scale, tilt_scale);
                        std::io::stdout().flush().ok();
                    } else if eval_count % 100 == 0 {
                        println!(
                            "  [{}/{}] avg={:.4}  best={:.4}",
                            eval_count, total_combos, avg_mae, best_avg_mae
                        );
                        std::io::stdout().flush().ok();
                    }
                }
            }
        }
    }

    let (lb, hl, sh, ts) = best_params;
    println!("\n=== RESULT ===");
    println!("Best avg MAE = {:.4}", best_avg_mae);
    println!("  lighting_blur = {}", lb);
    println!("  hl_scale      = {:.3}", hl);
    println!("  sh_scale      = {:.2}", sh);
    println!("  tilt_scale    = {:.2}", ts);

    if let Ok(mut file) = std::fs::OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(RESULTS_FILE)
    {
        let _ = writeln!(
            file,
            "lighting_blur={}  hl_scale={:.3}  sh_scale={:.2}  tilt_scale={:.2}\n\
             avg_mae={:.4}  samples={}  combos={}",
            lb,
            hl,
            sh,
            ts,
            best_avg_mae,
            samples.len(),
            total_combos
        );
    }
}

fn parse_filename(name: &str) -> Option<(String, String, f32, f32, f32, String, f32, String)> {
    // be_{shape}_{style}_d{depth}_sz{size}_sf{soften}_{dir}_alt{altitude}_{technique}
    let parts: Vec<&str> = name.split('_').collect();
    if parts.len() < 9 || parts[0] != "be" {
        return None;
    }
    let shape = parts[1].to_string();
    let style = parts[2].to_string();
    let depth = parts[3].strip_prefix('d')?.parse::<f32>().ok()?;
    let size = parts[4].strip_prefix("sz")?.parse::<f32>().ok()?;
    let soften = parts[5].strip_prefix("sf")?.parse::<f32>().ok()?;
    let direction = parts[6].to_string();
    let altitude = parts[7].strip_prefix("alt")?.parse::<f32>().ok()?;
    let technique = parts[8].to_string();
    Some((
        shape, style, depth, size, soften, direction, altitude, technique,
    ))
}