"""Phase A — mirror cell + spider solid FEA from the STEP exports.
Units mm/N/MPa/tonne. gmsh meshes the real CAD; CalculiX solves."""
import subprocess, os, math, sys
GMSH = os.path.expanduser("~/.local/bin/gmsh")
CCX  = os.path.expanduser("~/opt/fea/envs/solver/bin/ccx")
HERE = os.path.dirname(os.path.abspath(__file__))
STEP = os.path.expanduser("~/astro-lab/ota-cad/exports")
E, NU, RHO = 69000.0, 0.33, 2.70e-9

def mesh(step, tag, size):
    inp = f"{HERE}/{tag}.inp"
    subprocess.run([GMSH, step, "-3", "-format", "inp", "-o", inp,
                    "-clmax", str(size), "-order", "2"],
                   capture_output=True, text=True)
    return inp if os.path.exists(inp) else None

def parse(path):
    nodes, elems, etype = {}, [], None
    mode = None
    for line in open(path):
        s = line.strip()
        if s.startswith("*"):
            u = s.upper().replace(" ", "")
            if u.startswith("*NODE"): mode = "N"
            elif u.startswith("*ELEMENT"):
                for t in ("C3D10", "C3D4"):
                    if f"TYPE={t}" in u: etype, mode = t, "E"; break
                else: mode = None
            else: mode = None
            continue
        if not s or mode is None: continue
        p = [x for x in s.replace(",", " ").split() if x]
        if mode == "N" and len(p) >= 4:
            nodes[int(p[0])] = tuple(float(v) for v in p[1:4])
        elif mode == "E":
            elems.append([int(x) for x in p])
    # gmsh wraps long element lines; stitch by expected node count
    want = 10 if etype == "C3D10" else 4
    fixed, buf = [], []
    for row in elems:
        if not buf: buf = row
        else: buf += row
        if len(buf) >= want + 1:
            fixed.append(buf[:want+1]); buf = []
    return nodes, fixed, etype

def solve(tag, nodes, elems, etype, fix_fn, load_fn, total_N, gravdir):
    fixed = [n for n, c in nodes.items() if fix_fn(*c)]
    loaded = [n for n, c in nodes.items() if load_fn(*c)]
    if not fixed or not loaded:
        return None, f"fixed={len(fixed)} loaded={len(loaded)}"
    fN = -total_N / len(loaded)
    d = f"{HERE}/{tag}_job.inp"
    with open(d, "w") as f:
        f.write("*NODE, NSET=Nall\n")
        for n, (x, y, z) in nodes.items(): f.write(f"{n},{x:.4f},{y:.4f},{z:.4f}\n")
        f.write(f"*ELEMENT, TYPE={etype}, ELSET=Esol\n")
        for e in elems: f.write(",".join(map(str, e)) + "\n")
        for nm, lst in (("Nfix", fixed), ("Nload", loaded)):
            f.write(f"*NSET, NSET={nm}\n")
            for i in range(0, len(lst), 10): f.write(",".join(map(str, lst[i:i+10])) + "\n")
        f.write(f"""*MATERIAL, NAME=ALU
*ELASTIC
{E},{NU}
*DENSITY
{RHO}
*SOLID SECTION, ELSET=Esol, MATERIAL=ALU
*BOUNDARY
Nfix,1,3
*STEP
*STATIC
*DLOAD
Esol,GRAV,9810.,{gravdir}
*CLOAD
Nload,{2 if gravdir.startswith('0.,-1') else 3},{fN}
*NODE PRINT, NSET=Nload
U
*END STEP
""")
    subprocess.run([CCX, d[:-4]], capture_output=True, text=True, cwd=HERE)
    dat = d[:-4] + ".dat"
    if not os.path.exists(dat): return None, "no .dat"
    disp, on = [], False
    for line in open(dat):
        if "displacements" in line.lower(): on = True; continue
        p = line.split()
        if on and len(p) == 4:
            try: disp.append(tuple(float(v) for v in p[1:]))
            except ValueError: pass
    if not disp: return None, "no displacements"
    mags = [math.sqrt(sum(v*v for v in u)) for u in disp]
    return dict(n=len(nodes), e=len(elems), et=etype,
                umax=max(mags), umin=min(mags), diff=max(mags)-min(mags)), None

print("="*74); print("  PHASE A — MIRROR CELL + SPIDER, solid FEA from the STEP CAD"); print("="*74)

# ---- MIRROR CELL: 1.79 kg primary on 3 pads, cell held at 3 collimation bosses
inp = mesh(f"{STEP}/mirror_cell.step", "cell", 7.0)
if inp:
    nodes, elems, et = parse(inp)
    zmax = max(c[2] for c in nodes.values())
    r = lambda x, y: math.hypot(x, y)
    res, err = solve("cell", nodes, elems, et,
        fix_fn=lambda x, y, z: r(x, y) > 105 and z < 3,          # outer rim seats
        load_fn=lambda x, y, z: 65 < r(x, y) < 82 and z > zmax-3, # 0.707R pads
        total_N=1.79*9.81, gravdir="0.,0.,-1.")
    if res:
        print(f"\n  MIRROR CELL  ({res['n']} nodes, {res['e']} {res['et']})")
        print(f"    max deflection at pads      {res['umax']*1000:.3f} um")
        print(f"    DIFFERENTIAL across pads    {res['diff']*1000:.3f} um  <- mirror tilt")
        tilt = math.degrees(math.atan(res['diff']/141.4))*3600
        print(f"    equivalent mirror tilt      {tilt:.3f} arcsec")
        print(f"    Nelson glass deflection     0.84 nm (analytic, separate)")
        print(f"    -> {'PASS' if res['diff']*1e6 < 27.5 else 'CHECK'}: cell compliance "
              f"{res['diff']*1e6:.1f} nm vs lambda/20 = 27.5 nm")
    else: print("  cell:", err)

# ---- SPIDER: 0.48 kg secondary assembly, tube horizontal (worst)
inp = mesh(f"{STEP}/spider.step", "spider", 3.5)
if inp:
    nodes, elems, et = parse(inp)
    res, err = solve("spider", nodes, elems, et,
        fix_fn=lambda x, y, z: math.hypot(x, y) > 108,      # vane tips at the tube
        load_fn=lambda x, y, z: math.hypot(x, y-5.0) < 24,  # hub (offset 5 mm)
        total_N=0.48*9.81, gravdir="0.,-1.,0.")
    if res:
        print(f"\n  SPIDER  ({res['n']} nodes, {res['e']} {res['et']})")
        print(f"    hub displacement            {res['umax']*1000:.2f} um")
        print(f"    -> optical axis shift       {res['umax']:.4f} mm")
        print(f"    vs coma-free radius 1.41 mm -> "
              f"{'PASS' if res['umax'] < 1.41*0.1 else 'CHECK'} "
              f"({res['umax']/1.41*100:.2f}% of budget)")
    else: print("  spider:", err)
print()