from pathlib import Path
import re, math
import vtk

case=Path(r"C:\Users\tm\cae-ai\cases\cantilever")
inp=case/'cantilever_ccx.inp'
frd=case/'cantilever_ccx.frd'
vtu=case/'cantilever_ccx.vtu'
png=case/'cantilever_mises_deformed.png'

# --- read mesh from CalculiX INP ---
nodes={}; elems=[]; mode=None
for line in inp.read_text(errors='ignore').splitlines():
    s=line.strip()
    if not s or s.startswith('**'): continue
    if s.startswith('*'):
        u=s.upper()
        if u.startswith('*NODE'): mode='node'; continue
        if u.startswith('*ELEMENT'): mode='elem'; continue
        mode=None; continue
    if mode=='node':
        p=[x.strip() for x in s.split(',')]
        if len(p)>=4:
            nodes[int(p[0])] = tuple(float(x) for x in p[1:4])
    elif mode=='elem':
        p=[x.strip() for x in s.split(',')]
        elems.append((int(p[0]), [int(x) for x in p[1:] if x]))

# --- read FRD result blocks ---
float_re=re.compile(r'[+-]?\d\.\d{5}E[+-]\d{3}')

def read_block(name, ncomp):
    out={}; active=False; data_started=False
    for line in frd.read_text(errors='ignore').splitlines():
        s=line.strip()
        if s.startswith('-4'):
            active=(name in s)
            data_started=False
            continue
        if not active: continue
        if s=='-3': break
        if s.startswith('-1'):
            data_started=True
            nid=int(s[2:12].strip())
            vals=[float(x) for x in float_re.findall(s[12:])]
            if len(vals)>=ncomp: out[nid]=vals[:ncomp]
        elif data_started and s.startswith('-2'):
            pass
    return out

disp=read_block('DISP',3)
stress=read_block('STRESS',6)
if len(disp)!=len(nodes): raise RuntimeError(f'DISP count {len(disp)} != nodes {len(nodes)}')
if len(stress)!=len(nodes): raise RuntimeError(f'STRESS count {len(stress)} != nodes {len(nodes)}')

mises={}
for nid,(sxx,syy,szz,sxy,syz,szx) in stress.items():
    vm=math.sqrt(0.5*((sxx-syy)**2+(syy-szz)**2+(szz-sxx)**2)+3.0*(sxy*sxy+syz*syz+szx*szx))
    mises[nid]=vm

# --- build VTK grid ---
ids=sorted(nodes)
id_to_idx={nid:i for i,nid in enumerate(ids)}
pts=vtk.vtkPoints()
for nid in ids: pts.InsertNextPoint(*nodes[nid])

grid=vtk.vtkUnstructuredGrid(); grid.SetPoints(pts)
for eid,conn in elems:
    if len(conn)!=8: continue
    h=vtk.vtkHexahedron()
    for j,nid in enumerate(conn): h.GetPointIds().SetId(j,id_to_idx[nid])
    grid.InsertNextCell(h.GetCellType(),h.GetPointIds())

uarr=vtk.vtkDoubleArray(); uarr.SetName('U'); uarr.SetNumberOfComponents(3)
vmarr=vtk.vtkDoubleArray(); vmarr.SetName('S_Mises'); vmarr.SetNumberOfComponents(1)
for nid in ids:
    uarr.InsertNextTuple3(*disp[nid]); vmarr.InsertNextValue(mises[nid])
grid.GetPointData().AddArray(uarr); grid.GetPointData().SetVectors(uarr)
grid.GetPointData().AddArray(vmarr); grid.GetPointData().SetScalars(vmarr)

w=vtk.vtkXMLUnstructuredGridWriter(); w.SetFileName(str(vtu)); w.SetInputData(grid)
if not w.Write(): raise RuntimeError('VTU write failed')

# --- warp and render offscreen ---
warp=vtk.vtkWarpVector(); warp.SetInputData(grid); warp.SetScaleFactor(30.0); warp.Update()
mapper=vtk.vtkDataSetMapper(); mapper.SetInputConnection(warp.GetOutputPort()); mapper.SetScalarModeToUsePointData(); mapper.SelectColorArray('S_Mises'); mapper.SetScalarRange(min(mises.values()),max(mises.values()))
actor=vtk.vtkActor(); actor.SetMapper(mapper)

# original undeformed wireframe
orig_mapper=vtk.vtkDataSetMapper(); orig_mapper.SetInputData(grid); orig_mapper.ScalarVisibilityOff()
orig_actor=vtk.vtkActor(); orig_actor.SetMapper(orig_mapper); orig_actor.GetProperty().SetRepresentationToWireframe(); orig_actor.GetProperty().SetOpacity(0.25)

ren=vtk.vtkRenderer(); ren.AddActor(orig_actor); ren.AddActor(actor); ren.SetBackground(1,1,1)
win=vtk.vtkRenderWindow(); win.SetOffScreenRendering(1); win.SetSize(1400,700); win.AddRenderer(ren)
ren.ResetCamera(); cam=ren.GetActiveCamera(); cam.Azimuth(25); cam.Elevation(20); ren.ResetCameraClippingRange(); win.Render()

bar=vtk.vtkScalarBarActor(); bar.SetLookupTable(mapper.GetLookupTable()); bar.SetTitle('von Mises [MPa]'); bar.SetNumberOfLabels(5); ren.AddViewProp(bar); win.Render()

f=vtk.vtkWindowToImageFilter(); f.SetInput(win); f.Update()
pw=vtk.vtkPNGWriter(); pw.SetFileName(str(png)); pw.SetInputConnection(f.GetOutputPort()); pw.Write()

max_u=max((math.sqrt(sum(c*c for c in disp[n])) for n in ids))
print('nodes=',len(nodes),'elements=',len(elems))
print('max_displacement=',max_u)
print('max_mises=',max(mises.values()))
print('vtu=',vtu, vtu.exists(), vtu.stat().st_size if vtu.exists() else 0)
print('png=',png, png.exists(), png.stat().st_size if png.exists() else 0)