"""Tube shell FEA — does the focuser aperture soften the tube unacceptably?
Units: mm, N, MPa, tonne.  CalculiX 2.23 + gmsh 4.15.
"""
import subprocess, os, sys, math, re

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__))

R, L      = 125.0, 720.0     # tube outer radius, length
Z_FOC     = 560.0            # focuser station (CELL_STANDOFF + PRI_TO_SEC = 30+530)
R_HOLE    = 30.5             # 2in drawtube clearance
Z_RING1, Z_RING2 = 210.0, 450.0   # tube ring stations
E, NU, RHO = 69000.0, 0.33, 2.70e-9
G          = 9810.0          # mm/s^2
M_TRAIN    = 1.9             # kg imaging train hung off the focuser
MESH_SZ    = 9.0

def build(tag, thickness, hole=True):
    geo = f"""SetFactory("OpenCASCADE");
Circle(1) = {{0,0,0, {R}}};
Extrude {{0,0,{L}}} {{ Curve{{1}}; }}
"""
    if hole:
        geo += f"""Cylinder(1) = {{-300,0,{Z_FOC}, 600,0,0, {R_HOLE}}};
BooleanDifference{{ Surface{{1}}; Delete; }}{{ Volume{{1}}; Delete; }}
"""
    geo += f"Mesh.CharacteristicLengthMax = {MESH_SZ};\nMesh.ElementOrder = 1;\n"
    gp = f"{HERE}/{tag}.geo"; open(gp,"w").write(geo)
    inp = f"{HERE}/{tag}.inp"
    r = subprocess.run([GMSH, gp, "-2", "-format", "inp", "-o", inp],
                       capture_output=True, text=True)
    if not os.path.exists(inp):
        print(r.stdout[-1500:]); print(r.stderr[-1500:]); sys.exit(1)
    return inp

def parse_inp(path):
    nodes, elems = {}, []
    mode = None
    for line in open(path):
        s = line.strip()
        if s.startswith("*"):
            u = s.upper()
            mode = "N" if u.startswith("*NODE") else ("E" if u.startswith("*ELEMENT") and "CPS3" in u.replace(" ","") or (u.startswith("*ELEMENT") and "TYPE=CPS3" in u.replace(" ","")) else None)
            if u.startswith("*ELEMENT"):
                mode = "E" if "CPS3" in u.replace(" ","") else 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])] = (float(p[1]), float(p[2]), float(p[3]))
        elif mode == "E" and len(p) >= 4:
            elems.append((int(p[0]), int(p[1]), int(p[2]), int(p[3])))
    return nodes, elems

def run(tag, thickness, hole=True):
    inp = build(tag, thickness, hole)
    nodes, elems = parse_inp(inp)
    if not elems:
        print(f"  {tag}: no CPS3 elements parsed"); return None

    fixed = [n for n,(x,y,z) in nodes.items()
             if abs(z-Z_RING1) < 12 or abs(z-Z_RING2) < 12]
    # load applied around the focuser aperture rim (x>0 side)
    rim = [n for n,(x,y,z) in nodes.items()
           if x > 0 and abs(math.hypot(y, z-Z_FOC) - R_HOLE) < 7]
    if not rim:
        rim = [n for n,(x,y,z) in nodes.items() if x > R*0.9 and abs(z-Z_FOC) < 20]
    F_node = -M_TRAIN * 9.81 / max(len(rim),1)   # -Y, tube horizontal

    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("*ELEMENT, TYPE=S3, ELSET=Eshell\n")
        for e in elems: f.write(f"{e[0]},{e[1]},{e[2]},{e[3]}\n")
        f.write("*NSET, NSET=Nfix\n")
        for i in range(0,len(fixed),10): f.write(",".join(map(str,fixed[i:i+10]))+"\n")
        f.write("*NSET, NSET=Nrim\n")
        for i in range(0,len(rim),10): f.write(",".join(map(str,rim[i:i+10]))+"\n")
        f.write(f"""*MATERIAL, NAME=ALU
*ELASTIC
{E},{NU}
*DENSITY
{RHO}
*SHELL SECTION, ELSET=Eshell, MATERIAL=ALU
{thickness}
*BOUNDARY
Nfix,1,6
*STEP
*STATIC
*DLOAD
Eshell,GRAV,{G},0.,-1.,0.
*CLOAD
Nrim,2,{F_node}
*NODE PRINT, NSET=Nall
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
    disp, started = {}, False
    for line in open(dat):
        if "displacements" in line.lower(): started = True; continue
        p = line.split()
        if started and len(p) == 4:
            try: disp[int(p[0])] = tuple(float(v) for v in p[1:])
            except ValueError: pass
    if not disp: return None
    uy_foc = min((disp[n][1] for n in rim if n in disp), default=0.0)
    zc = [n for n,(x,y,z) in nodes.items() if z < 40]          # primary end
    uy_pri = min((disp[n][1] for n in zc if n in disp), default=0.0)
    umax = max(math.sqrt(sum(v*v for v in u)) for u in disp.values())
    return dict(nodes=len(nodes), elems=len(elems), umax=umax,
                uy_foc=uy_foc, uy_pri=uy_pri, diff=abs(uy_foc-uy_pri))

print("="*74)
print("  TUBE SHELL FEA — CalculiX 2.23, S3 shell, tube horizontal (worst case)")
print("  Load: self-weight + 1.9 kg imaging train on the focuser rim")
print("="*74)
CFZ_HALF, COMA_R = 21.5/1000.0, 1.41
results = {}
for tag, thk, hole in [("t20_hole",2.0,True), ("t15_hole",1.5,True), ("t20_solid",2.0,False)]:
    r = run(tag, thk, hole)
    results[tag] = r
    if r is None: print(f"  {tag:<12} FAILED"); continue
    print(f"\n  {tag}  ({r['nodes']} nodes, {r['elems']} S3)")
    print(f"    max total displacement        {r['umax']:.4f} mm")
    print(f"    focuser rim   dY              {r['uy_foc']:+.4f} mm")
    print(f"    primary end   dY              {r['uy_pri']:+.4f} mm")
    print(f"    DIFFERENTIAL (focuser-primary){r['diff']:.4f} mm   <- collimation shift")
    v = "PASS" if r['diff'] < COMA_R else "FAIL"
    print(f"    vs coma-free radius {COMA_R} mm  -> {v} ({r['diff']/COMA_R*100:.1f}% of budget)")

a, b, c = results.get("t20_hole"), results.get("t15_hole"), results.get("t20_solid")
if a and c:
    print(f"\n  APERTURE PENALTY: hole softens the tube {a['diff']/c['diff']:.2f}x "
          f"({c['diff']:.4f} -> {a['diff']:.4f} mm)")
if a and b:
    print(f"  WALL 2.0 -> 1.5 mm: {b['diff']/a['diff']:.2f}x more deflection "
          f"({a['diff']:.4f} -> {b['diff']:.4f} mm), saves 0.75 kg")
    print(f"  1.5 mm verdict: {'ACCEPT' if b['diff'] < COMA_R*0.5 else 'REJECT - keep 2.0 mm'}"
          f"  (target: stay under 50% of the {COMA_R} mm budget)")