export_static.py

# -*- coding: utf-8 -*-
"""
Builds the zero-install GitHub Pages demo.

The Streamlit app needs Python + scikit-learn + a 23 MB forest. Nobody who just
wants to look at the tool should have to install that, so this script bakes the
frozen models into two artefacts the browser can evaluate on its own:

  1. Etch Rate  - Polynomial2 + StandardScaler + Ridge is collapsed into plain
                  coefficients. JS reproduces predict() EXACTLY (not approximated).
  2. Profile CD - ExtraTrees cannot be expressed in closed form, so it is sampled
                  on a 13^4 grid and shipped as Float32. JS does quadrilinear
                  interpolation. Measured interpolation error is ~0.33 nm, i.e.
                  11% of the model's own 2.90 nm blind-test error - negligible.

Run:  python export_static.py
"""

from __future__ import annotations

import json
import struct
import warnings
from pathlib import Path

import numpy as np

warnings.filterwarnings("ignore")

import etch_core as ec  # noqa: E402

BASE_DIR = Path(__file__).resolve().parent
SITE_DIR = BASE_DIR / "docs"
API_DIR = SITE_DIR / "api"

REPO_URL = "https://github.com/MinJaeJun/dry-etch-ai"

# 13 points per axis -> 28,561 grid nodes. See module docstring for the
# accuracy justification.
GRID_N = 13
VALIDATION_SAMPLES = 3000


def human(n: int) -> str:
    return f"{n / 1024 / 1024:.1f} MB" if n > 1024 * 1024 else f"{n / 1024:.0f} KB"


# --------------------------------------------------------------------------- #
# 1. Etch rate: exact closed form
# --------------------------------------------------------------------------- #
def extract_rate_polynomial(rate_model):
    poly = rate_model.named_steps["poly"]
    scaler = rate_model.named_steps["scale"]
    ridge = rate_model.named_steps["ridge"]

    payload = {
        "powers": poly.powers_.astype(int).tolist(),
        "mean": scaler.mean_.astype(float).tolist(),
        "scale": scaler.scale_.astype(float).tolist(),
        "coef": np.asarray(ridge.coef_, dtype=float).ravel().tolist(),
        "intercept": float(np.asarray(ridge.intercept_).ravel()[0]),
    }

    # Prove the closed form is bit-for-bit equivalent to the sklearn pipeline.
    rng = np.random.default_rng(1)
    probe = np.column_stack([
        rng.uniform(ec.H2_MIN, ec.H2_MAX, 2000),
        rng.uniform(ec.PRESSURE_MIN, ec.PRESSURE_MAX, 2000),
        rng.uniform(ec.BIAS_MIN, ec.BIAS_MAX, 2000),
    ])
    import pandas as pd

    reference = rate_model.predict(pd.DataFrame(probe, columns=ec.RATE_FEATURES))

    powers = np.asarray(payload["powers"], dtype=float)
    features = np.prod(np.power(probe[:, None, :], powers[None, :, :]), axis=2)
    scaled = (features - payload["mean"]) / payload["scale"]
    manual = scaled @ np.asarray(payload["coef"]) + payload["intercept"]

    max_dev = float(np.max(np.abs(manual - reference)))
    if max_dev > 1e-8:
        raise SystemExit(f"rate polynomial mismatch: {max_dev:.3e}")
    print(f"  rate polynomial verified against sklearn (max dev {max_dev:.2e} nm/min)")
    return payload


# --------------------------------------------------------------------------- #
# 2. Profile: 4-D grid -> Float32
# --------------------------------------------------------------------------- #
def build_axes():
    return [
        np.linspace(ec.H2_MIN, ec.H2_MAX, GRID_N),
        np.linspace(ec.PRESSURE_MIN, ec.PRESSURE_MAX, GRID_N),
        np.linspace(ec.BIAS_MIN, ec.BIAS_MAX, GRID_N),
        np.linspace(ec.DEPTH_MIN, ec.DEPTH_MAX, GRID_N),
    ]


def build_profile_grid(rate_model, profile_model, axes):
    mesh = np.meshgrid(*axes, indexing="ij")
    flat = [m.ravel() for m in mesh]

    frame = ec.predict_batch(
        rate_model, profile_model, flat[0], flat[1], flat[2], flat[3],
        with_uncertainty=True,
    )

    # Layout: [node][target] so JS reads 6 contiguous floats per node.
    cd = np.column_stack([frame[name].to_numpy() for name in ec.PROFILE_TARGETS])
    std = np.column_stack([frame[name + "_std"].to_numpy() for name in ec.PROFILE_TARGETS])
    return cd.astype(np.float32), std.astype(np.float32)


def write_binary(path: Path, array: np.ndarray) -> int:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = array.astype("<f4").tobytes()
    path.write_bytes(payload)
    return len(payload)


# --------------------------------------------------------------------------- #
# 3. Interpolation accuracy report (shipped to the UI, not just printed)
# --------------------------------------------------------------------------- #
def validate_interpolation(rate_model, profile_model, axes, cd_grid):
    from scipy.interpolate import RegularGridInterpolator

    rng = np.random.default_rng(0)
    q = np.column_stack([
        rng.uniform(ec.H2_MIN, ec.H2_MAX, VALIDATION_SAMPLES),
        rng.uniform(ec.PRESSURE_MIN, ec.PRESSURE_MAX, VALIDATION_SAMPLES),
        rng.uniform(ec.BIAS_MIN, ec.BIAS_MAX, VALIDATION_SAMPLES),
        rng.uniform(ec.DEPTH_MIN, ec.DEPTH_MAX, VALIDATION_SAMPLES),
    ])
    truth = ec.predict_batch(rate_model, profile_model, q[:, 0], q[:, 1], q[:, 2], q[:, 3])

    shape = [GRID_N] * 4
    def _stable(value: float) -> float:
        """
        These are diagnostics measured at runtime, so their last bits differ
        between BLAS builds. Rounding keeps docs/ byte-reproducible in CI while
        preserving far more precision than the UI ever shows.
        """
        return round(float(value), 4)

    report = {}
    for i, name in enumerate(ec.PROFILE_TARGETS):
        volume = cd_grid[:, i].reshape(shape).astype(float)
        interp = RegularGridInterpolator(axes, volume)
        report[name] = _stable(np.mean(np.abs(interp(q) - truth[name].to_numpy())))

    undercut_truth = truth["effective_undercut_nm"].to_numpy()
    cd10 = RegularGridInterpolator(axes, cd_grid[:, 0].reshape(shape).astype(float))(q)
    report["effective_undercut_nm"] = _stable(
        np.mean(np.abs((cd10 - ec.MASK_OPENING_NM) / 2.0 - undercut_truth))
    )
    report["_mean_cd"] = _stable(np.mean([report[n] for n in ec.PROFILE_TARGETS]))
    return report


# --------------------------------------------------------------------------- #
# 4. Manifest
# --------------------------------------------------------------------------- #
def build_manifest(axes, rate_poly, interp_report, metadata, sizes):
    return {
        "generated_from": "frozen surrogate models (etch_rate_model / etch_profile_model)",
        "repo": REPO_URL,
        "grid": {
            "n": GRID_N,
            "nodes": GRID_N ** 4,
            "order": ["h2_fraction_pct", "pressure_mTorr", "bias_W", "etch_depth_nm"],
            "axes": [a.tolist() for a in axes],
            "targets": ec.PROFILE_TARGETS,
            "files": {"cd": "api/profile_cd.f32", "std": "api/profile_std.f32"},
            "bytes": sizes,
        },
        "rate_model": {"name": "Polynomial2 + Ridge", "exact": True, **rate_poly},
        "profile_model": {
            "name": "ExtraTrees (500 trees)",
            "exact": False,
            "interpolation_mae_nm": interp_report,
        },
        "domain": {
            "h2_fraction_pct": [ec.H2_MIN, ec.H2_MAX],
            "pressure_mTorr": [ec.PRESSURE_MIN, ec.PRESSURE_MAX],
            "bias_W": [ec.BIAS_MIN, ec.BIAS_MAX],
            "etch_depth_nm": [ec.DEPTH_MIN, ec.DEPTH_MAX],
            "depth_confident_max_nm": ec.DEPTH_CONFIDENT_MAX,
        },
        "geometry": {
            "mask_opening_nm": ec.MASK_OPENING_NM,
            "initial_oxide_nm": ec.INITIAL_OXIDE_NM,
            "initial_pr_nm": ec.INITIAL_PR_NM,
            "source_power_w": ec.SOURCE_POWER_W,
            "temperature_c": ec.TEMPERATURE_C,
            "depth_fractions": ec.DEPTH_FRACTIONS.tolist(),
        },
        "pr_model": {
            "h2_pct": ec.PR_RATE_H2.tolist(),
            "pr_rate_nm_min": ec.PR_RATE_NM_MIN.tolist(),
            "note": "문헌 기반 현상학적 모델. AI 학습 결과가 아님.",
        },
        "spec": {
            key: {
                "label": rule["label"],
                "unit": rule["unit"],
                "direction": rule["direction"],
                "green": rule["green"],
                "yellow": rule["yellow"],
                "gate": rule["gate"],
            }
            for key, rule in ec.SPEC.items()
        },
        "cd_spec": {
            "target_nm": ec.CD_TARGET_NM,
            "green_tol_nm": ec.CD_GREEN_TOL_NM,
            "yellow_tol_nm": ec.CD_YELLOW_TOL_NM,
            "window_targets": ec.CD_WINDOW_TARGETS,
            "bottom_target": ec.CD_BOTTOM_TARGET,
            "bottom_green_min_nm": ec.CD95_GREEN_MIN_NM,
            "bottom_yellow_min_nm": ec.CD95_YELLOW_MIN_NM,
        },
        "blind_test": {
            "recipes": 30,
            "profile_states": 90,
            "etch_rate": {"mae_nm_min": 1.7394, "rmse": 1.9066, "r2": 0.9965},
            "profile_mean_cd_mae_nm": 2.9008,
            "per_cd": {
                "cd_d10_nm": {"mae": 3.3493, "r2": 0.6200},
                "cd_d25_nm": {"mae": 2.8683, "r2": 0.7090},
                "cd_d50_nm": {"mae": 2.2303, "r2": 0.8140},
                "cd_d75_nm": {"mae": 1.5062, "r2": 0.8508},
                "cd_d90_nm": {"mae": 2.1388, "r2": 0.9480},
                "cd_d95_nm": {"mae": 5.3118, "r2": 0.7667},
            },
            "derived": {
                "effective_undercut_nm": {"mae": 1.6746, "r2": 0.6200},
                "sidewall_angle_deg": {"mae": 0.2935, "r2": 0.9754},
                "aspect_ratio": {"mae": 0.0180, "r2": 0.9944},
            },
            "by_depth": {
                "1.0": {"cd10_mae": 1.015, "undercut_mae": 0.508, "cd10_r2": 0.889},
                "2.0": {"cd10_mae": 3.801, "undercut_mae": 1.900, "cd10_r2": 0.534},
                "2.5": {"cd10_mae": 5.232, "undercut_mae": 2.616, "cd10_r2": 0.351},
            },
            "uncertainty": {
                "corr_abs_error_vs_tree_std": 0.730,
                "coverage_1sigma": 0.844,
                "coverage_2sigma": 0.978,
            },
        },
        "dataset": {
            "recipes": 200,
            "split": {"train": 140, "validation": 30, "test": 30},
            "profile_states": {"train": 220, "validation": 50, "test": 90},
            "depth_levels_min": 1,
            "note": "270 train/val state 중 깊이 3수준 Recipe는 50개뿐 (Train 40 / Val 10).",
        },
        "metadata": {k: (v if not isinstance(v, np.generic) else v.item())
                     for k, v in (metadata or {}).items()},
    }


# --------------------------------------------------------------------------- #
def main() -> int:
    print("=" * 70)
    print("Dry Etch AI - static demo export")
    print("=" * 70)

    rate_model, profile_model, metadata = ec.load_models()
    print("  models loaded")

    rate_poly = extract_rate_polynomial(rate_model)

    axes = build_axes()
    print(f"  sampling profile model on {GRID_N}^4 = {GRID_N ** 4:,} nodes ...")
    cd_grid, std_grid = build_profile_grid(rate_model, profile_model, axes)

    cd_bytes = write_binary(API_DIR / "profile_cd.f32", cd_grid)
    std_bytes = write_binary(API_DIR / "profile_std.f32", std_grid)
    print(f"  profile_cd.f32  {human(cd_bytes)}")
    print(f"  profile_std.f32 {human(std_bytes)}")

    print("  validating quadrilinear interpolation ...")
    report = validate_interpolation(rate_model, profile_model, axes, cd_grid)
    print(f"    mean CD interpolation MAE {report['_mean_cd']:.3f} nm "
          f"(model blind-test MAE 2.901 nm)")
    print(f"    undercut interpolation MAE {report['effective_undercut_nm']:.3f} nm "
          f"(model blind-test MAE 1.675 nm)")
    if report["_mean_cd"] > 0.5:
        raise SystemExit("interpolation error too large - raise GRID_N")

    manifest = build_manifest(axes, rate_poly, report, metadata,
                              {"cd": cd_bytes, "std": std_bytes})
    manifest_path = API_DIR / "manifest.json"
    manifest_path.write_text(
        json.dumps(manifest, ensure_ascii=False, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(f"  manifest.json   {human(manifest_path.stat().st_size)}")

    (SITE_DIR / ".nojekyll").write_text("", encoding="utf-8")

    total = cd_bytes + std_bytes + manifest_path.stat().st_size
    print("-" * 70)
    print(f"  total payload {human(total)}  ->  {SITE_DIR}")
    print("=" * 70)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())