"""Task 01 - Laboratory configuration.

Writes ./out/config/lab_config.yaml, the authoritative machine-readable
laboratory model (volumes, deposition matrix, physical-model coefficients,
random seed) that all downstream tasks read, plus the ./out/.done/01.ok
completion marker. Spec: docs/specs/01_lab_configuration.md.
"""

from __future__ import annotations

import logging
import sys
from pathlib import Path

import yaml

from common.pipeline import parse_task_args, remove_stale_marker, write_marker

logger = logging.getLogger(__name__)

TASK_ID = "01"

# Transcription of docs/specs/01_lab_configuration.md. This literal is the
# ONLY sanctioned home of these values in code; every other task reads the
# YAML this script writes (docs/rules/code-config-yaml.md).
LAB_CONFIG: dict = {
	"project": "LabDataStorageEvaluation",
	"description": "Simulated tribology laboratory (Sandia Pt-Au LDRD context)",
	"spec": "docs/specs/01_lab_configuration.md",
	"seed": 20260711,
	"material_system": {
		"substrate": {"material": "Ti-6Al-4V", "dimensions_mm": [10, 10, 3]},
		"adhesion_layer": {"material": "Cr", "process": "sputtered"},
		"coating": {
			"material": "Pt-Au",
			"deposition": "composition gradient across each wafer",
			"thickness_um_range": [0.3, 1.1],
		},
	},
	"instruments": [
		{
			"instrument_id": "rapid",
			"role": "friction",
			"name": "RAPID custom high-throughput parallelized 6-probe tribometer",
			"data_produced": "COF vs cycle per track; test conditions",
		},
		{
			"instrument_id": "ti980",
			"role": "nanoindentation",
			"name": "Bruker TI980 TriboIndenter",
			"data_produced": "hardness, reduced modulus (25 indents per coupon)",
		},
		{
			"instrument_id": "m4_tornado",
			"role": "composition",
			"name": "Bruker M4 Tornado micro-XRF",
			"data_produced": "Pt/Au wt% map per coupon (20x20 grid)",
		},
		{
			"instrument_id": "pvd200",
			"role": "deposition",
			"name": "Kurt J. Lesker PVD 200 sputter-down",
			"data_produced": "batch deposition parameters",
		},
		{
			"instrument_id": "afm",
			"role": "surface_roughness",
			"name": "AFM",
			"data_produced": "topography image metadata + Ra/Rq per coupon",
		},
		{
			"instrument_id": "profilometer",
			"role": "film_thickness",
			"name": "Optical profilometry",
			"data_produced": "thickness map per coupon (10x10 grid)",
		},
		{
			"instrument_id": "simtra",
			"role": "simulation",
			"name": "SIMTRA sputter transport Monte-Carlo",
			"data_produced": "deposition atom-energy / composition profiles per deposition run",
		},
	],
	"deposition_matrix": [
		{
			"batch_code": "B721",
			"pt_gun_tilt_deg": 20,
			"au_gun_tilt_deg": 0,
			"pt_power_W": 150,
			"au_power_W": 50,
			"pt_discharge_V": 432,
			"au_discharge_V": 311,
		},
		{
			"batch_code": "B722",
			"pt_gun_tilt_deg": 20,
			"au_gun_tilt_deg": 20,
			"pt_power_W": 100,
			"au_power_W": 100,
			"pt_discharge_V": 399,
			"au_discharge_V": 352,
		},
		{
			"batch_code": "B723",
			"pt_gun_tilt_deg": 20,
			"au_gun_tilt_deg": 20,
			"pt_power_W": 150,
			"au_power_W": 50,
			"pt_discharge_V": 430,
			"au_discharge_V": 311,
		},
		{
			"batch_code": "B724",
			"pt_gun_tilt_deg": 0,
			"au_gun_tilt_deg": 20,
			"pt_power_W": 50,
			"au_power_W": 150,
			"pt_discharge_V": 376,
			"au_discharge_V": 340,
		},
	],
	"hierarchy": {
		"batches": 4,
		"wafers_per_batch": 3,
		"coupon_grid_per_wafer": [7, 7],
		"coupons_per_wafer": 49,
	},
	"friction_assignment": {
		"friction_coupons_total": 480,
		"reserve_coupons_total": 108,
		"runs": 4,
		"coupons_per_run": 120,
		"plates_per_run": 5,
		"probes_per_plate": 6,
		"coupons_per_probe_square": 4,
		"tracks_per_friction_coupon": 3,
		"counterfaces_per_holder": 3,
		"fresh_counterface_per_track": True,
		# QA witness pattern: rows 1/4/7 x cols 1/4/7 of the 7x7 grid stay
		# unworn (reserve), sampling low/mid/high Au columns on every wafer.
		"reserve_grid_positions": [1, 4, 7, 22, 25, 28, 43, 46, 49],
	},
	"volumes": {
		"coupons_total": 588,
		"tracks_total": 1440,
		"cycles_per_track": 1000,
		"cycle_rows_total": 1440000,
		"loop_every_n_cycles": 100,
		"loops_per_track": 10,
		"loop_points_per_loop": 200,
		"loop_points_total": 2880000,
		"xrf_grid": [20, 20],
		"xrf_points_per_coupon": 400,
		"xrf_points_total": 235200,
		"profilometry_grid": [10, 10],
		"profilometry_points_per_coupon": 100,
		"nanoindentation_indents_per_coupon": 25,
		"simtra_rows_per_batch": 1000,
	},
	"runs": [
		{"run_code": "R1", "batch_code": "B721", "environment": "lab_air", "date": "2026-04-06", "operator": "A. Reyes"},
		{"run_code": "R2", "batch_code": "B722", "environment": "lab_air", "date": "2026-04-13", "operator": "K. Patel"},
		{"run_code": "R3", "batch_code": "B723", "environment": "dry_n2", "date": "2026-04-20", "operator": "A. Reyes"},
		{"run_code": "R4", "batch_code": "B724", "environment": "dry_n2", "date": "2026-04-27", "operator": "M. Novak"},
	],
	"schedule": {
		"deposition_start_date": "2026-03-02",
		"batch_interval_days": 7,
		"characterization_lag_days": 3,
		"run_start_time": "09:00:00",
		"track_interval_s": 2100,
	},
	"test_conditions": {
		"normal_load_mN": 100,
		"stroke_mm": 1.0,
		"speed_mm_s": 1.0,
		"counterface": {"material": "ruby (Al2O3)", "diameter_mm": 3.175},
		"rh_pct": {"lab_air": 45, "dry_n2": 2},
		"temperature_C": 23,
		"temperature_tolerance_C": 1,
	},
	"physical_models": {
		"au_gradient": {
			"description": "linear Au wt% gradient along wafer x, batch-dependent center",
			"batch_center_wtpct": {"B721": 8, "B722": 25, "B723": 10, "B724": 60},
			"wafer_span_wtpct": 5,
			"xrf_noise_sd_wtpct": 0.3,
		},
		"cof_run_in": {
			"formula": "cof(c) = cof_ss + (cof_0 - cof_ss) * exp(-c / tau) + N(0, sigma)",
			"cof_0_range": [0.35, 0.5],
			"tau_cycles_range": [100, 400],
			"noise_sd": 0.01,
		},
		"cof_steady_state": {
			"formula": "cof_ss = intercept + slope_per_au_wtpct * au_wtpct_mean, clamped to >= floor",
			"lab_air": {"intercept": 0.32, "slope_per_au_wtpct": -0.0015},
			"dry_n2": {"intercept": 0.24, "slope_per_au_wtpct": -0.0012},
			"floor": 0.12,
		},
		"hardness_GPa": {
			"formula": "H = intercept + slope_per_au_wtpct * au_wtpct_mean + N(0, noise_sd)",
			"intercept": 8.5,
			"slope_per_au_wtpct": -0.05,
			"noise_sd": 0.25,
		},
		"reduced_modulus_GPa": {
			"formula": "Er = intercept + slope_per_au_wtpct * au_wtpct_mean + N(0, noise_sd)",
			"intercept": 190,
			"slope_per_au_wtpct": -0.6,
			"noise_sd": 4,
		},
		"roughness_ra_nm": {
			"distribution": "lognormal",
			"median_nm": 5,
			"sigma_log": 0.3,
			"rq_over_ra": 1.25,
		},
		"thickness_um": {
			"profile": "radial parabolic across wafer plus noise",
			"center_um": 1.05,
			"edge_um": 0.35,
			"noise_sd_um": 0.02,
			"range_um": [0.3, 1.1],
		},
		"wear_archard": {
			# Spec 01 fixes only the model shape (V = k*F*s, k depends on Au%
			# and environment); these coefficients are project defaults chosen
			# to yield Sandia-plausible low-wear volumes at 100 mN / 2 m.
			"formula": "V = k * F * s; k = k0 * (1 + slope_per_au_wtpct * au_wtpct_mean), lognormal noise, clamped to >= k_floor",
			"k0_mm3_per_N_m": {"lab_air": 1.5e-07, "dry_n2": 6.0e-08},
			"slope_per_au_wtpct": {"lab_air": -0.008, "dry_n2": -0.010},
			"k_floor_mm3_per_N_m": 1.0e-09,
			"noise_sigma_log": 0.15,
		},
		"friction_loop": {
			"force_noise_sd_mN": 0.3,
		},
		"nanoindentation": {
			"max_load_mN": 10,
			"indent_grid": [5, 5],
			"grid_pitch_um": 20,
		},
		"simtra": {
			"angle_deg_range": [0, 90],
			"energy_eV_max": 50,
			"surface_binding_energy_eV": {"pt": 5.84, "au": 3.81},
			"flux_per_W": 0.02,
		},
	},
}

YAML_HEADER = (
	"# Generated by make_lab_config.py (task 01) from docs/specs/01_lab_configuration.md.\n"
	"# Do not hand-edit: regenerate via `python make_lab_config.py`.\n"
)


def validate_config(cfg: dict) -> None:
	"""Assert the arithmetic identities between declared volumes (spec 01)."""
	h = cfg["hierarchy"]
	f = cfg["friction_assignment"]
	v = cfg["volumes"]
	checks = [
		("coupons_total", v["coupons_total"], h["batches"] * h["wafers_per_batch"] * h["coupons_per_wafer"]),
		("coupons_per_wafer", h["coupons_per_wafer"], h["coupon_grid_per_wafer"][0] * h["coupon_grid_per_wafer"][1]),
		("friction+reserve", v["coupons_total"], f["friction_coupons_total"] + f["reserve_coupons_total"]),
		("friction_coupons_total", f["friction_coupons_total"], f["runs"] * f["coupons_per_run"]),
		("coupons_per_run", f["coupons_per_run"], f["plates_per_run"] * f["probes_per_plate"] * f["coupons_per_probe_square"]),
		("tracks_total", v["tracks_total"], f["friction_coupons_total"] * f["tracks_per_friction_coupon"]),
		("cycle_rows_total", v["cycle_rows_total"], v["tracks_total"] * v["cycles_per_track"]),
		("loops_per_track", v["loops_per_track"], v["cycles_per_track"] // v["loop_every_n_cycles"]),
		("loop_points_total", v["loop_points_total"], v["tracks_total"] * v["loops_per_track"] * v["loop_points_per_loop"]),
		("xrf_points_per_coupon", v["xrf_points_per_coupon"], v["xrf_grid"][0] * v["xrf_grid"][1]),
		("xrf_points_total", v["xrf_points_total"], v["coupons_total"] * v["xrf_points_per_coupon"]),
		("profilometry_points_per_coupon", v["profilometry_points_per_coupon"], v["profilometry_grid"][0] * v["profilometry_grid"][1]),
		("runs_count", f["runs"], len(cfg["runs"])),
		("batches_count", h["batches"], len(cfg["deposition_matrix"])),
		("reserve_total", f["reserve_coupons_total"], len(f["reserve_grid_positions"]) * h["wafers_per_batch"] * h["batches"]),
		("friction_per_batch", f["coupons_per_run"], (h["coupons_per_wafer"] - len(f["reserve_grid_positions"])) * h["wafers_per_batch"]),
	]
	for name, declared, derived in checks:
		if declared != derived:
			raise ValueError(f"volume identity failed: {name}: declared {declared} != derived {derived}")
	run_batches = [r["batch_code"] for r in cfg["runs"]]
	matrix_batches = [b["batch_code"] for b in cfg["deposition_matrix"]]
	if sorted(run_batches) != sorted(matrix_batches):
		raise ValueError(f"run-batch mapping mismatch: runs cover {run_batches}, matrix has {matrix_batches}")
	logger.info("all %d volume identities hold", len(checks))


def write_yaml(cfg: dict, config_path: Path) -> int:
	config_path.parent.mkdir(parents=True, exist_ok=True)
	body = yaml.safe_dump(cfg, sort_keys=False, allow_unicode=False, width=100)
	text = YAML_HEADER + body
	with open(config_path, "w", encoding="utf-8", newline="\n") as fh:
		fh.write(text)
	reloaded = yaml.safe_load(text)
	if reloaded != cfg:
		raise RuntimeError(f"YAML round-trip mismatch for {config_path}")
	return len(text.encode("utf-8"))


def main() -> int:
	args = parse_task_args("Task 01: write out/config/lab_config.yaml")
	out_root: Path = args.out_root
	config_path = out_root / "config" / "lab_config.yaml"
	try:
		remove_stale_marker(out_root, TASK_ID)
		validate_config(LAB_CONFIG)
		config_bytes = write_yaml(LAB_CONFIG, config_path)
		v = LAB_CONFIG["volumes"]
		f = LAB_CONFIG["friction_assignment"]
		h = LAB_CONFIG["hierarchy"]
		marker_path = write_marker(out_root, TASK_ID, {
			"config_file": config_path.as_posix(),
			"config_bytes": config_bytes,
			"seed": LAB_CONFIG["seed"],
			"batches": h["batches"],
			"wafers": h["batches"] * h["wafers_per_batch"],
			"coupons": v["coupons_total"],
			"friction_coupons": f["friction_coupons_total"],
			"tracks": v["tracks_total"],
			"cycle_rows": v["cycle_rows_total"],
			"loop_points": v["loop_points_total"],
			"xrf_points": v["xrf_points_total"],
		})
	except Exception:
		logger.critical("task %s failed", TASK_ID, exc_info=True)
		return 1

	v = LAB_CONFIG["volumes"]
	print(f"task 01 ok: {config_path} ({config_bytes} bytes)")
	print(
		f"volumes: batches=4 wafers=12 coupons={v['coupons_total']} tracks={v['tracks_total']} "
		f"cycle_rows={v['cycle_rows_total']} loop_points={v['loop_points_total']} xrf_points={v['xrf_points_total']}"
	)
	print(f"marker: {marker_path}")
	return 0


if __name__ == "__main__":
	sys.exit(main())