"""Strict adapters for prepared POC schemas; no guessed arrays or invented topology."""
from __future__ import annotations
from pathlib import Path
import os
import re
from typing import Any
import numpy as np

ROOT = Path(os.environ.get("MOLTEX_EXPERIMENT_ROOT", Path(__file__).resolve().parents[1])).resolve()
class InputContractError(ValueError): pass

def split_embedded_units(names: np.ndarray | list[str]) -> tuple[np.ndarray,np.ndarray]:
    """Split prepared labels such as ``상부온도01 [°C]`` without guessing units."""
    base=[]; units=[]
    for value in names:
        match=re.fullmatch(r"(.*?)\s*\[([^]]+)\]",str(value))
        base.append(match.group(1).strip() if match else str(value))
        units.append(match.group(2).strip() if match else "")
    return np.asarray(base),np.asarray(units)

def _load(path: Path, required: set[str]) -> Any:
    if not path.exists(): raise InputContractError(f"missing prepared input: {path.relative_to(ROOT)}")
    d=np.load(path,allow_pickle=False); missing=required-set(d.files)
    if missing: raise InputContractError(f"{path.relative_to(ROOT)} missing required schema keys {sorted(missing)}")
    return d

def _masked_values(x: np.ndarray, valid: np.ndarray, label: str) -> np.ndarray:
    x=np.asarray(x,np.float32); valid=np.asarray(valid,bool)
    if x.shape != valid.shape: raise InputContractError(f"{label}: x and feature_valid must have exactly equal shapes")
    if np.any(~np.isfinite(x[valid])): raise InputContractError(f"{label}: a known feature is non-finite")
    return np.where(valid,x,0.).astype(np.float32)

def fixed_physical_normalize(x: np.ndarray, names: np.ndarray, units: np.ndarray | None=None) -> np.ndarray:
    """Unit-aware, non-fitted normalization used only for unlabeled smoke inputs.

    Coordinates/depth/edge lengths in mm share a 1000-mm reference, areas use
    1e6 mm², and curvature remains in its inverse-length physical scale under a
    signed-log transform. ``*_native`` scan coordinates use one common observed
    coordinate magnitude across x/y/z, preserving scanner-frame shape rather
    than independently distorting axes. Embedded process units use fixed
    references: 500 °C, 100 seconds and 1000 m/s. This is explicitly not a
    train scaler.
    """
    y=x.copy(); names=[str(n).lower() for n in names]; units=[str(u).lower() for u in (units if units is not None else [""]*len(names))]
    native_xyz=[j for j,n in enumerate(names) if n in {"x_native","y_native","z_native"}]
    native_scale=max(1.,float(np.max(np.abs(y[...,native_xyz])))) if native_xyz else 1.
    for j,(name,unit) in enumerate(zip(names,units)):
        if name in {"pos_norm_x","pos_norm_y","pos_norm_z"} or "normal" in name or "direction" in name or "valid" in name or "silhouette" in name or "visibility" in name or "boundary" in name or "feature_" in name: pass
        elif j in native_xyz: y[...,j]/=native_scale
        elif "1/mm" in name or "per_mm" in name or unit.startswith("1/mm"): y[...,j]=np.sign(y[...,j])*np.log1p(np.abs(y[...,j]))
        elif "mm2" in name or unit == "mm2": y[...,j]/=1e6
        elif "_mm" in name or unit == "mm": y[...,j]/=1000.
        elif unit in {"°c","degc"}: y[...,j]/=500.
        elif unit in {"sec","s"}: y[...,j]/=100.
        elif unit in {"m/sec","m/s"}: y[...,j]/=1000.
        elif "deg" in name or unit == "deg": y[...,j]/=180.
        elif "rad" in name or unit == "rad": y[...,j]/=np.pi
        else: y[...,j]=np.sign(y[...,j])*np.log1p(np.abs(y[...,j]))
    return y.astype(np.float32)

def process_conditions() -> tuple[np.ndarray,np.ndarray,list[str],str]:
    p=ROOT/"data_processed/tabular/process_and_geometry.npz"; d=_load(p,{"x","feature_valid","process_x","process_feature_valid","run_id","training_valid"})
    x,m=_masked_values(d["x"],d["feature_valid"],"tabular process contract"),d["feature_valid"].astype(bool)
    if x.ndim!=2 or not len(x): raise InputContractError("tabular x must be [runs,features]")
    process_names,process_units=split_embedded_units(d["process_feature_names"].astype(str))
    names=process_names.tolist()+d["geometry_feature_names"].astype(str).tolist()
    units=process_units.tolist()+[""]*len(d["geometry_feature_names"])
    if len(names)!=x.shape[1]: raise InputContractError("tabular process and geometry names must align to x")
    return fixed_physical_normalize(x,np.asarray(names),np.asarray(units)),m,[str(r) for r in d["run_id"]],str(p.relative_to(ROOT))

def image_input() -> dict[str,Any]:
    p=ROOT/"data_processed/images/cad_multiview_128.npz"; d=_load(p,{"x","feature_valid","feature_names","training_valid"})
    x,m=_masked_values(d["x"],d["feature_valid"],"image contract"),d["feature_valid"].astype(bool)
    if x.ndim!=4 or x.shape[0]!=6: raise InputContractError("image x must be [6,C,H,W] (six prepared views)")
    names=d["feature_names"]
    if names.shape != (x.shape[1],): raise InputContractError("image feature_names must align to channels")
    return {"x":fixed_physical_normalize(x.transpose(0,2,3,1),names,d["feature_units"]).transpose(0,3,1,2),"mask":m,"source":str(p.relative_to(ROOT))}

def voxel_input() -> dict[str,Any]:
    p=ROOT/"data_processed/voxels/cad_surface_dense_32.npz"; d=_load(p,{"x","feature_valid","feature_names"})
    x,m=_masked_values(d["x"],d["feature_valid"],"dense voxel contract"),d["feature_valid"].astype(bool)
    if x.ndim!=4 or d["feature_names"].shape!=(x.shape[0],): raise InputContractError("dense voxel x [C,X,Y,Z] and names [C] required")
    return {"x":fixed_physical_normalize(x.transpose(1,2,3,0),d["feature_names"]).transpose(3,0,1,2),"mask":m,"source":str(p.relative_to(ROOT))}

def sparse_voxel_input() -> dict[str,Any]:
    """Load the explicit sparse-surface schema, never deriving it from dense voxels."""
    p=ROOT/"data_processed/voxels/cad_surface_sparse_32.npz"
    d=_load(p,{"voxel_coords","x","feature_valid","feature_names","training_valid"})
    coord=np.asarray(d["voxel_coords"],np.int32); x,m=_masked_values(d["x"],d["feature_valid"],"sparse voxel contract"),d["feature_valid"].astype(bool)
    if coord.ndim!=2 or coord.shape[1]!=3 or len(coord)!=len(x) or x.ndim!=2 or d["feature_names"].shape!=(x.shape[1],):
        raise InputContractError("sparse voxel contract requires voxel_coords [N,3], x/feature_valid [N,C], feature_names [C]")
    if np.any(coord<0) or len(np.unique(coord,axis=0))!=len(coord): raise InputContractError("sparse voxel coordinates must be non-negative and unique")
    values=fixed_physical_normalize(x,d["feature_names"])
    return {"coords":coord,"features":np.concatenate([values,m.astype(np.float32)],1),"source":str(p.relative_to(ROOT)),"grid_dims":np.asarray(d["grid_dims"],np.int32)}

def sparse_batch_collate(coords: np.ndarray, features: np.ndarray, batch_index: int=0) -> tuple[np.ndarray,np.ndarray]:
    """Return MinkowskiEngine convention coordinates [batch,x,y,z] and features."""
    coords=np.asarray(coords); features=np.asarray(features,np.float32)
    if coords.ndim!=2 or coords.shape[1]!=3 or features.ndim!=2 or len(coords)!=len(features): raise InputContractError("sparse collation requires coords [N,3] and features [N,C]")
    if not np.issubdtype(coords.dtype,np.integer) or np.any(coords<0) or len(np.unique(coords,axis=0))!=len(coords): raise InputContractError("sparse coordinates must be unique non-negative integers")
    if not np.isfinite(features).all() or batch_index<0: raise InputContractError("sparse features must be finite and batch index non-negative")
    batch=np.full((len(coords),1),batch_index,dtype=np.int32)
    return np.concatenate([batch,coords.astype(np.int32,copy=False)],1),features

def point_inputs() -> list[dict[str,Any]]:
    out=[]
    point_paths = [ROOT/"data_processed/points/cad_surface_points.npz"]
    point_paths.extend(sorted((ROOT/"data_processed/points").glob("scan_*_diagnostic.npz")))
    for p in point_paths:
        d=_load(p,{"x","feature_valid","feature_names","training_valid"}); x,m=_masked_values(d["x"],d["feature_valid"],f"point {p.name}"),d["feature_valid"].astype(bool)
        if x.ndim!=2 or d["feature_names"].shape!=(x.shape[1],): raise InputContractError(f"{p.name}: x [N,C], feature_names [C] required")
        timing="post_scan_diagnostic" if "diagnostic_post_process" in d and bool(d["diagnostic_post_process"]) else "cad_pre_process"
        out.append({"x":fixed_physical_normalize(x,d["feature_names"],d["feature_units"]),"mask":m,"source":str(p.relative_to(ROOT)),"timing":timing})
    return out

def graph_inputs() -> list[dict[str,Any]]:
    paths=sorted((ROOT/"data_processed/graphs").glob("graph_*.npz"))
    if not paths: raise InputContractError("no graph_<budget>.npz; graph preparation has not completed")
    out=[]; schema=None
    for p in paths:
        d=_load(p,{"x","feature_valid","feature_names","edge_index","edge_attr","edge_feature_names","training_valid"})
        x,m=_masked_values(d["x"],d["feature_valid"],f"graph {p.name}"),d["feature_valid"].astype(bool); e=np.asarray(d["edge_index"],np.int64); a=np.asarray(d["edge_attr"],np.float32)
        if x.ndim!=2 or e.ndim!=2 or e.shape[0]!=2 or a.ndim!=2 or a.shape[0]!=e.shape[1] or (e.size and (e.min()<0 or e.max()>=len(x))): raise InputContractError(f"{p.name}: full topology schema invalid")
        sig=(tuple(d["feature_names"].astype(str)),tuple(d["edge_feature_names"].astype(str)),x.shape[1],a.shape[1])
        if schema is None: schema=sig
        elif schema!=sig: raise InputContractError("graph resolutions must share feature and edge schema")
        out.append({"x":fixed_physical_normalize(x,d["feature_names"]),"mask":m,"edge_index":e,"edge_attr":a,"source":str(p.relative_to(ROOT)),"nodes":len(x)})
    return out

def brep_input() -> dict[str,Any]:
    p=ROOT/"data_processed/brep/faces_uv.npz"; d=_load(p,{"pos_mm","normal","curvature","trim_valid","normal_valid","curvature_valid","edge_index","edge_attr"})
    pos,nor,curv=np.asarray(d["pos_mm"],np.float32),np.asarray(d["normal"],np.float32),np.asarray(d["curvature"],np.float32); trim,nv,cv=np.asarray(d["trim_valid"],bool),np.asarray(d["normal_valid"],bool),np.asarray(d["curvature_valid"],bool)
    if pos.ndim!=4 or pos.shape[-1]!=3 or nor.shape!=pos.shape or curv.shape!=pos.shape: raise InputContractError("B-Rep UV grid shapes invalid")
    raw=np.concatenate([pos,nor,curv,trim[...,None],nv[...,None],cv[...,None]],-1); valid=np.concatenate([np.repeat(trim[...,None],3,-1),np.repeat(nv[...,None],3,-1),np.repeat(cv[...,None],3,-1),np.ones((*trim.shape,3),bool)],-1)
    names=np.asarray(["x_mm","y_mm","z_mm","normal_x","normal_y","normal_z","kmax_1_per_mm","kmin_1_per_mm","gaussian_1_per_mm2","trim_valid","normal_valid","curvature_valid"]); grids=fixed_physical_normalize(_masked_values(raw,valid,"B-Rep grids"),names).transpose(0,3,1,2); e=np.asarray(d["edge_index"],np.int64); a=np.asarray(d["edge_attr"],np.float32)
    if e.shape[0]!=2 or a.shape[0]!=e.shape[1] or (e.size and (e.min()<0 or e.max()>=len(grids))): raise InputContractError("B-Rep adjacency invalid")
    return {"grids":grids,"adjacency":e,"source":str(p.relative_to(ROOT)),"faces":len(grids),"adjacency_edges":int(e.shape[1])}

def tabular_input() -> dict[str,Any]:
    x,m,runs,source=process_conditions(); return {"x":x,"mask":m,"runs":runs,"source":source}