"""Task 03 - Simulated CSV data generation (~150 MB).

Generates the canonical raw corpus under ./out/csv/ exactly as the
instruments would write it, driven entirely by out/config/lab_config.yaml
(volumes, physical-model coefficients, seed). Also writes
./out/csv/MANIFEST.csv (path, rows, bytes, sha256) and the
./out/.done/03.ok marker. Spec: docs/specs/03_csv_data_generation.md.
"""

from __future__ import annotations

import hashlib
import logging
import math
import shutil
import sys
from datetime import datetime, timedelta
from pathlib import Path

import numpy as np

from common.pipeline import (
	check_dependencies,
	load_lab_config,
	parse_task_args,
	remove_stale_marker,
	write_marker,
)

logger = logging.getLogger(__name__)

TASK_ID = "03"
DEPENDS_ON = ["01"]

SIZE_GATE_MIB = (120, 200)  # abort bounds on the generated tree, spec 03

# One float format per column family (docs/rules/data-determinism.md):
# repr drift between runs must never change file bytes. Precision of the two
# bulk families also sets corpus size; chosen to land in the 140-170 MB
# target band of spec 03.
FMT_COF = "%.16g"
FMT_POSITION = "%.14g"
FMT_FORCE = "%.14g"
FMT_MISC = "%.10g"  # wt%, um, nm, GPa, simtra, wear

# RNG stream tags for per-entity derived generators.
S_XRF, S_PROF, S_NANO, S_AFM, S_COF, S_LOOP, S_WEAR, S_SIMTRA = range(1, 9)


class ValidationError(RuntimeError):
	pass


class CorpusGenerator:
	def __init__(self, cfg: dict, csv_root: Path) -> None:
		self.cfg = cfg
		self.csv_root = csv_root
		self.seed = cfg["seed"]
		self.manifest: list[tuple[str, int, int, str]] = []
		self.total_rows = 0
		self.total_bytes = 0
		self.counts = {"coupons": 0, "tracks": 0, "cycle_rows": 0, "loop_rows": 0, "xrf_rows": 0}
		self.runs_by_batch = {r["batch_code"]: r for r in cfg["runs"]}

	def rng(self, *tags: int) -> np.random.Generator:
		return np.random.default_rng([self.seed, *tags])

	def write_csv(self, relpath: str, header: str, lines: list[str]) -> None:
		text = header + "\n" + "\n".join(lines) + "\n"
		data = text.encode("ascii")
		path = self.csv_root / relpath
		path.parent.mkdir(parents=True, exist_ok=True)
		path.write_bytes(data)
		sha = hashlib.sha256(data).hexdigest()
		self.manifest.append((relpath, len(lines), len(data), sha))
		self.total_rows += len(lines)
		self.total_bytes += len(data)

	# --- top-level flat files ---

	def gen_batches(self) -> None:
		sched = self.cfg["schedule"]
		start = datetime.fromisoformat(sched["deposition_start_date"])
		lines = []
		for i, b in enumerate(self.cfg["deposition_matrix"]):
			date = (start + timedelta(days=i * sched["batch_interval_days"])).date().isoformat()
			lines.append(
				f"{b['batch_code']},{b['pt_gun_tilt_deg']},{b['au_gun_tilt_deg']},"
				f"{b['pt_power_W']},{b['au_power_W']},{b['pt_discharge_V']},{b['au_discharge_V']},{date}"
			)
		self.write_csv(
			"batches.csv",
			"batch_code,pt_gun_tilt_deg,au_gun_tilt_deg,pt_power_W,au_power_W,pt_discharge_V,au_discharge_V,deposition_date",
			lines,
		)

	def gen_runs(self) -> None:
		plates = self.cfg["friction_assignment"]["plates_per_run"]
		lines = [
			f"{r['run_code']},{r['environment']},{r['date']},{plates},{r['operator']}"
			for r in self.cfg["runs"]
		]
		self.write_csv("runs.csv", "run_code,environment,date,plates,operator", lines)

	def gen_simtra(self) -> None:
		pm = self.cfg["physical_models"]["simtra"]
		n = self.cfg["volumes"]["simtra_rows_per_batch"]
		us_pt = pm["surface_binding_energy_eV"]["pt"]
		us_au = pm["surface_binding_energy_eV"]["au"]
		for b_idx, b in enumerate(self.cfg["deposition_matrix"], 1):
			rng = self.rng(b_idx, S_SIMTRA)
			angle = np.linspace(pm["angle_deg_range"][0], pm["angle_deg_range"][1], n)
			pt_flux = pm["flux_per_W"] * b["pt_power_W"] * np.cos(np.radians(angle - b["pt_gun_tilt_deg"])) ** 2
			au_flux = pm["flux_per_W"] * b["au_power_W"] * np.cos(np.radians(angle - b["au_gun_tilt_deg"])) ** 2
			pt_flux = np.maximum(pt_flux * (1 + rng.normal(0, 0.02, n)), 0)
			au_flux = np.maximum(au_flux * (1 + rng.normal(0, 0.02, n)), 0)
			# Thompson mean emission energy ~ 2 * surface binding energy,
			# flux-weighted across the two species.
			energy = (pt_flux * 2 * us_pt + au_flux * 2 * us_au) / (pt_flux + au_flux + 1e-12)
			energy = energy + rng.normal(0, 0.2, n)
			lines = [
				f"{FMT_MISC % a},{FMT_MISC % e},{FMT_MISC % p},{FMT_MISC % g}"
				for a, e, p, g in zip(angle, energy, pt_flux, au_flux)
			]
			self.write_csv(f"simtra/simtra_profile_{b['batch_code']}.csv", "angle_deg,energy_eV,pt_flux,au_flux", lines)

	# --- per-coupon measurements ---

	def gen_xrf(self, rel_dir: str, b_idx: int, w: int, c: int, au_nominal: float) -> float:
		gx_n, gy_n = self.cfg["volumes"]["xrf_grid"]
		noise_sd = self.cfg["physical_models"]["au_gradient"]["xrf_noise_sd_wtpct"]
		rng = self.rng(b_idx, w, c, S_XRF)
		au = au_nominal + rng.normal(0, noise_sd, gx_n * gy_n)
		lines = []
		i = 0
		for gy in range(1, gy_n + 1):
			for gx in range(1, gx_n + 1):
				lines.append(f"{gx},{gy},{FMT_MISC % (100.0 - au[i])},{FMT_MISC % au[i]}")
				i += 1
		self.write_csv(f"{rel_dir}/xrf_map.csv", "grid_x,grid_y,pt_wtpct,au_wtpct", lines)
		self.counts["xrf_rows"] += len(lines)
		return float(au.mean())

	def gen_profilometry(self, rel_dir: str, b_idx: int, w: int, c: int, row: int, col: int) -> float:
		pm = self.cfg["physical_models"]["thickness_um"]
		gx_n, gy_n = self.cfg["volumes"]["profilometry_grid"]
		# Normalized squared wafer radius: corner coupon (r2=1) is thinnest.
		r2 = (((col - 4) / 3.0) ** 2 + ((row - 4) / 3.0) ** 2) / 2.0
		nominal = pm["center_um"] - (pm["center_um"] - pm["edge_um"]) * r2
		rng = self.rng(b_idx, w, c, S_PROF)
		t = nominal + rng.normal(0, pm["noise_sd_um"], gx_n * gy_n)
		lines = []
		i = 0
		for gy in range(1, gy_n + 1):
			for gx in range(1, gx_n + 1):
				lines.append(f"{gx},{gy},{FMT_MISC % t[i]}")
				i += 1
		self.write_csv(f"{rel_dir}/profilometry.csv", "grid_x,grid_y,thickness_um", lines)
		return float(t.mean())

	def gen_nanoindentation(self, rel_dir: str, b_idx: int, w: int, c: int, au_mean: float) -> None:
		pm = self.cfg["physical_models"]
		n = self.cfg["volumes"]["nanoindentation_indents_per_coupon"]
		grid_w = pm["nanoindentation"]["indent_grid"][0]
		pitch = pm["nanoindentation"]["grid_pitch_um"]
		max_load = pm["nanoindentation"]["max_load_mN"]
		rng = self.rng(b_idx, w, c, S_NANO)
		h = pm["hardness_GPa"]["intercept"] + pm["hardness_GPa"]["slope_per_au_wtpct"] * au_mean \
			+ rng.normal(0, pm["hardness_GPa"]["noise_sd"], n)
		er = pm["reduced_modulus_GPa"]["intercept"] + pm["reduced_modulus_GPa"]["slope_per_au_wtpct"] * au_mean \
			+ rng.normal(0, pm["reduced_modulus_GPa"]["noise_sd"], n)
		lines = [
			f"{i + 1},{(i % grid_w) * pitch},{(i // grid_w) * pitch},{FMT_MISC % h[i]},{FMT_MISC % er[i]},{max_load}"
			for i in range(n)
		]
		self.write_csv(
			f"{rel_dir}/nanoindentation.csv",
			"indent_id,x_um,y_um,hardness_GPa,reduced_modulus_GPa,max_load_mN",
			lines,
		)

	def gen_afm(self, rel_dir: str, b_idx: int, w: int, c: int, coupon_code: str) -> float:
		pm = self.cfg["physical_models"]["roughness_ra_nm"]
		rng = self.rng(b_idx, w, c, S_AFM)
		ra = float(np.exp(rng.normal(math.log(pm["median_nm"]), pm["sigma_log"])))
		rq = ra * pm["rq_over_ra"]
		self.write_csv(
			f"{rel_dir}/afm.csv",
			"ra_nm,rq_nm,image_file",
			[f"{FMT_MISC % ra},{FMT_MISC % rq},images/afm/{coupon_code}.tiff"],
		)
		return ra

	# --- per-track friction data ---

	def cof_ss(self, au_mean: float, environment: str) -> float:
		pm = self.cfg["physical_models"]["cof_steady_state"]
		value = pm[environment]["intercept"] + pm[environment]["slope_per_au_wtpct"] * au_mean
		return max(value, pm["floor"])  # deliberate model floor, spec 01

	def gen_track(
		self, rel_dir: str, b_idx: int, w: int, c: int, t: int,
		track_code: str, run: dict, au_mean: float, plate: int, probe: int, square: int,
	) -> None:
		tc = self.cfg["test_conditions"]
		pm = self.cfg["physical_models"]
		v = self.cfg["volumes"]
		sched = self.cfg["schedule"]
		env = run["environment"]
		load_mn = tc["normal_load_mN"]

		started = datetime.fromisoformat(f"{run['date']}T{sched['run_start_time']}+00:00") + timedelta(
			seconds=((plate - 1) * 12 + (square - 1) * 3 + (t - 1)) * sched["track_interval_s"]
		)
		counterface_id = f"{run['run_code']}-P{plate}-PR{probe}-B{t}"
		self.write_csv(
			f"{rel_dir}/track_info.csv",
			"track_code,run_code,environment,load_mN,stroke_mm,speed_mm_s,counterface_id,started_at",
			[
				f"{track_code},{run['run_code']},{env},{load_mn},{FMT_MISC % tc['stroke_mm']},"
				f"{FMT_MISC % tc['speed_mm_s']},{counterface_id},{started.strftime('%Y-%m-%dT%H:%M:%SZ')}"
			],
		)

		cof_ss = self.cof_ss(au_mean, env)
		cof_rng = self.rng(b_idx, w, c, t, S_COF)
		ri = pm["cof_run_in"]
		cof_0 = cof_rng.uniform(*ri["cof_0_range"])
		tau = cof_rng.uniform(*ri["tau_cycles_range"])
		n_cycles = v["cycles_per_track"]
		cycles = np.arange(1, n_cycles + 1)
		decay = (cof_0 - cof_ss) * np.exp(-cycles / tau)
		cof = cof_ss + decay + cof_rng.normal(0, ri["noise_sd"], n_cycles)
		self.write_csv(
			f"{rel_dir}/cof_vs_cycle.csv",
			"cycle,cof",
			[f"{cyc},{FMT_COF % val}" for cyc, val in zip(cycles, cof)],
		)
		self.counts["cycle_rows"] += n_cycles

		loop_rng = self.rng(b_idx, w, c, t, S_LOOP)
		half_um = tc["stroke_mm"] * 1000.0 / 2.0
		pts_half = v["loop_points_per_loop"] // 2
		fwd = np.linspace(-half_um, half_um, pts_half)
		rev = np.linspace(half_um, -half_um, pts_half)
		force_sd = pm["friction_loop"]["force_noise_sd_mN"]
		lines: list[str] = []
		for k in range(1, v["loops_per_track"] + 1):
			cyc = k * v["loop_every_n_cycles"]
			mu = cof_ss + (cof_0 - cof_ss) * math.exp(-cyc / tau)
			f_fwd = mu * load_mn + loop_rng.normal(0, force_sd, pts_half)
			f_rev = -mu * load_mn + loop_rng.normal(0, force_sd, pts_half)
			lines.extend(f"{cyc},{FMT_POSITION % p},{FMT_FORCE % f}" for p, f in zip(fwd, f_fwd))
			lines.extend(f"{cyc},{FMT_POSITION % p},{FMT_FORCE % f}" for p, f in zip(rev, f_rev))
		self.write_csv(f"{rel_dir}/friction_loops.csv", "cycle,position_um,friction_force_mN", lines)
		self.counts["loop_rows"] += len(lines)

		wm = pm["wear_archard"]
		wear_rng = self.rng(b_idx, w, c, t, S_WEAR)
		k_arch = wm["k0_mm3_per_N_m"][env] * (1 + wm["slope_per_au_wtpct"][env] * au_mean)
		k_arch *= math.exp(wear_rng.normal(0, wm["noise_sigma_log"]))
		k_arch = max(k_arch, wm["k_floor_mm3_per_N_m"])  # deliberate model floor, spec 01
		sliding_m = n_cycles * 2 * tc["stroke_mm"] / 1000.0
		volume_um3 = k_arch * (load_mn / 1000.0) * sliding_m * 1e9
		self.write_csv(
			f"{rel_dir}/wear.csv",
			"wear_volume_um3,k_archard,sliding_distance_m",
			[f"{FMT_MISC % volume_um3},{FMT_MISC % k_arch},{FMT_MISC % sliding_m}"],
		)
		self.counts["tracks"] += 1

	# --- hierarchy walk ---

	def generate(self) -> None:
		self.gen_batches()
		self.gen_runs()
		self.gen_simtra()

		h = self.cfg["hierarchy"]
		fa = self.cfg["friction_assignment"]
		ag = self.cfg["physical_models"]["au_gradient"]
		sched = self.cfg["schedule"]
		dep_start = datetime.fromisoformat(sched["deposition_start_date"])
		reserve = set(fa["reserve_grid_positions"])
		grid_w = h["coupon_grid_per_wafer"][0]

		for b_idx, batch in enumerate(self.cfg["deposition_matrix"], 1):
			b_code = batch["batch_code"]
			run = self.runs_by_batch[b_code]
			au_center = ag["batch_center_wtpct"][b_code]
			dep_date = (dep_start + timedelta(days=(b_idx - 1) * sched["batch_interval_days"])).date().isoformat()
			friction_idx = 0
			for w in range(1, h["wafers_per_batch"] + 1):
				wafer_code = f"{b_code}-W{w}"
				wafer_dir = f"batch_{b_code}/wafer_W{w}"
				self.write_csv(
					f"{wafer_dir}/wafer_info.csv",
					"wafer_code,batch_code,wafer_index,deposition_date,coupons,friction_coupons,reserve_coupons",
					[
						f"{wafer_code},{b_code},{w},{dep_date},{h['coupons_per_wafer']},"
						f"{h['coupons_per_wafer'] - len(reserve)},{len(reserve)}"
					],
				)
				for c in range(1, h["coupons_per_wafer"] + 1):
					row = (c - 1) // grid_w + 1
					col = (c - 1) % grid_w + 1
					coupon_code = f"{wafer_code}-C{c:02d}"
					rel_dir = f"{wafer_dir}/coupon_C{c:02d}"
					au_nominal = au_center + ag["wafer_span_wtpct"] * (col - 4) / 3.0

					au_mean = self.gen_xrf(rel_dir, b_idx, w, c, au_nominal)
					thickness = self.gen_profilometry(rel_dir, b_idx, w, c, row, col)
					self.gen_nanoindentation(rel_dir, b_idx, w, c, au_mean)
					ra = self.gen_afm(rel_dir, b_idx, w, c, coupon_code)

					if c in reserve:
						assignment = "RESERVE,,,"
					else:
						plate = friction_idx // 24 + 1
						probe = (friction_idx % 24) // 4 + 1
						square = friction_idx % 4 + 1
						friction_idx += 1
						assignment = f"{run['run_code']},{plate},{probe},{square}"
					self.write_csv(
						f"{rel_dir}/coupon_info.csv",
						"coupon_code,batch_code,wafer_code,grid_row,grid_col,thickness_um,ra_nm,au_wtpct_mean,run_code,plate,probe,square",
						[
							f"{coupon_code},{b_code},{wafer_code},{row},{col},{FMT_MISC % thickness},"
							f"{FMT_MISC % ra},{FMT_MISC % au_mean},{assignment}"
						],
					)
					self.counts["coupons"] += 1

					if c not in reserve:
						plate, probe, square = (int(x) for x in assignment.split(",")[1:])
						for t in range(1, fa["tracks_per_friction_coupon"] + 1):
							self.gen_track(
								f"{rel_dir}/track_T{t}", b_idx, w, c, t,
								f"{coupon_code}-T{t}", run, au_mean, plate, probe, square,
							)
				logger.debug("wafer %s done", wafer_code)
			logger.info(
				"batch %s done: %d files, %.1f MiB cumulative",
				b_code, len(self.manifest), self.total_bytes / (1024 * 1024),
			)

	def validate_counts(self) -> None:
		v = self.cfg["volumes"]
		expected = {
			"coupons": v["coupons_total"],
			"tracks": v["tracks_total"],
			"cycle_rows": v["cycle_rows_total"],
			"loop_rows": v["loop_points_total"],
			"xrf_rows": v["xrf_points_total"],
		}
		for key, exp in expected.items():
			if self.counts[key] != exp:
				raise ValidationError(f"count mismatch: {key}: generated {self.counts[key]} != expected {exp}")
		logger.info("all entity counts match lab_config volumes: %s", self.counts)

	def write_manifest(self) -> int:
		lines = [f"{p},{r},{b},{s}" for p, r, b, s in self.manifest]
		text = "path,rows,bytes,sha256\n" + "\n".join(lines) + "\n"
		data = text.encode("ascii")
		(self.csv_root / "MANIFEST.csv").write_bytes(data)
		return len(data)


def main() -> int:
	args = parse_task_args("Task 03: generate the CSV corpus")
	out_root: Path = args.out_root
	csv_root = out_root / "csv"
	try:
		check_dependencies(out_root, DEPENDS_ON)
		remove_stale_marker(out_root, TASK_ID)
		if csv_root.exists():
			logger.info("re-run: removing previous corpus %s", csv_root)
			shutil.rmtree(csv_root)
		cfg = load_lab_config(out_root)

		gen = CorpusGenerator(cfg, csv_root)
		gen.generate()
		gen.validate_counts()
		manifest_bytes = gen.write_manifest()

		tree_bytes = gen.total_bytes + manifest_bytes
		tree_mib = tree_bytes / (1024 * 1024)
		print(f"generated tree size: {tree_mib:.1f} MiB ({tree_bytes} bytes, {len(gen.manifest) + 1} files)")
		if not SIZE_GATE_MIB[0] <= tree_mib <= SIZE_GATE_MIB[1]:
			raise ValidationError(
				f"tree size {tree_mib:.1f} MiB outside gate {SIZE_GATE_MIB[0]}-{SIZE_GATE_MIB[1]} MiB"
			)
		marker_path = write_marker(out_root, TASK_ID, {
			"csv_root": gen.csv_root.as_posix(),
			"files": len(gen.manifest),
			"total_rows": gen.total_rows,
			"total_bytes": tree_bytes,
			"coupons": gen.counts["coupons"],
			"tracks": gen.counts["tracks"],
			"cycle_rows": gen.counts["cycle_rows"],
			"loop_points": gen.counts["loop_rows"],
			"xrf_points": gen.counts["xrf_rows"],
		})
	except Exception:
		logger.critical("task %s failed", TASK_ID, exc_info=True)
		return 1

	c = gen.counts
	print(
		f"rows: total={gen.total_rows} cycles={c['cycle_rows']} loops={c['loop_rows']} "
		f"xrf={c['xrf_rows']} coupons={c['coupons']} tracks={c['tracks']}"
	)
	print(f"manifest: {csv_root / 'MANIFEST.csv'} ({len(gen.manifest)} entries)")
	print(f"marker: {marker_path}")
	return 0


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