"""Finite-volume enthalpy conduction solver; SI internally, STL in mm.
No fluid dynamics, gas transport, shrinkage or chemical equilibrium model.
"""
from dataclasses import dataclass, asdict
from pathlib import Path
import json, csv, math
import numpy as np
import trimesh
from scipy import sparse, ndimage
from scipy.sparse.linalg import cg, LinearOperator

@dataclass
class Material:
    rho: float = 2600.0
    cp: float = 1050.0
    k: float = 130.0
    latent: float = 390000.0
    solidus: float = 580.0
    liquidus: float = 650.0

# Engineering scenario constants, NOT measured material certificates.
PRESETS = {
 'Al: plechovky + profily (nejisty scenar)': Material(),
 'Al: cisty, priblizne konstanty': Material(2700,1050,160,397000,659.8,660.8),
 'Pb: ciste, priblizne konstanty': Material(11340,145,30,23000,327.0,328.0),
 'Cu: cista, priblizne konstanty': Material(8960,450,330,205000,1084.1,1085.1),
}

@dataclass
class Settings:
    pitch_mm: float = 3.0
    margin_mm: float = 15.0
    metal_temp: float = 720.0
    mold_temp: float = 20.0
    ambient: float = 20.0
    mold_rho: float = 1600.0
    mold_cp: float = 900.0
    mold_k: float = 0.7
    contact_h: float = 500.0
    outside_h: float = 10.0
    emissivity: float = 0.8
    duration: float = 180.0
    dt: float = 3.0
    snapshot_dt: float = 0.25
    gate_mm: float = 8.0
    vent_mm: float = 3.0
    riser_mm: float = 16.0
    head_mm: float = 100.0
    discharge: float = 0.6
    furnace_kw: float = 15.0
    efficiency: float = 0.2
    charge_kg: float = 1.0
    cans_percent: float = 50.0
    yield_percent: float = 85.0
    rx: float = 0.0
    ry: float = 0.0
    rz: float = 0.0
    length_mm: float = 300.0

LIMITATIONS = [
 'Vyzkumny prototyp: bez experimentalni validace pro tuto tavbu.',
 'Pocatecni podminka: cela dutina je jiz naplnena kovem se stejnou teplotou.',
 'Pouze vedeni tepla a tuhnuti; NEPOCITA proudeni, zachyceni plynu, oxidy, porozitu ani smrsteninove dutiny.',
 'Konstantni rho, cp a k; latentni teplo pres linearni podil kapaliny. Neni CALPHAD ani skutecny fazovy diagram srotu.',
 'Kandidatni vtoky/odvzdusneni jsou geometricke navrhy, ne optimalizace CFD.',
 'Ubytek hmoty pri taveni neni obsah necistot. Vytaznost i ucinnost pece jsou zadane odhady.',
 'Rozpocet pece neurcuje jeji vykon, ucinnost ani teplotu kovu.',
 'Simulace plneho tela nezahrnuje vtoky/nalitky, dokud nenactete exportovanou sjednocenou sestavu.',
]

def validate(s, m):
    for obj in (s,m):
        if not all(np.isfinite(v) for v in asdict(obj).values()):
            raise ValueError('All values must be finite numbers.')
    for v in (m.rho,m.cp,m.k,s.pitch_mm,s.margin_mm,s.mold_rho,s.mold_cp,s.mold_k,s.contact_h,s.duration,s.dt,s.snapshot_dt,s.length_mm,s.gate_mm,s.vent_mm,s.riser_mm,s.head_mm,s.furnace_kw,s.charge_kg):
        if v<=0: raise ValueError('Dimensions, times and material constants must be positive.')
    if m.latent<0 or m.liquidus<=m.solidus: raise ValueError('Liquidus must exceed solidus; latent heat must be nonnegative.')
    if s.metal_temp<=m.liquidus: raise ValueError('Initial metal temperature must exceed liquidus.')
    if min(s.metal_temp,s.mold_temp,s.ambient,m.solidus)<=-273.15: raise ValueError('Temperature below absolute zero.')
    if not (0<=s.emissivity<=1 and 0<s.efficiency<=1 and 0<s.discharge<=1 and 0<s.yield_percent<=100 and 0<=s.cans_percent<=100 and s.outside_h>=0):
        raise ValueError('Invalid efficiency, emissivity, recovery, fraction or transfer coefficient.')
    if s.duration/s.dt>10000: raise ValueError('Maximum 10000 requested time steps.')

def load_model(path, s):
    mesh=trimesh.load(path,force='mesh',process=True)
    if not isinstance(mesh,trimesh.Trimesh) or not len(mesh.faces): raise ValueError('Empty mesh.')
    if not np.isfinite(mesh.vertices).all(): raise ValueError('Invalid coordinates.')
    if not mesh.is_watertight or not mesh.is_winding_consistent:
        raise ValueError('Mesh is not closed or consistently oriented. Repair the STL first.')
    if mesh.volume<0: mesh.invert()
    mesh.apply_scale(s.length_mm/max(mesh.extents))
    mesh.apply_transform(trimesh.transformations.euler_matrix(*np.radians([s.rx,s.ry,s.rz])))
    mesh.apply_translation(-mesh.bounds[0])
    return mesh

GRID_CELL_LIMIT=300000

def grid_plan(mesh,s,pitch=None):
    pitch=s.pitch_mm if pitch is None else pitch
    pad=max(2,int(math.ceil(s.margin_mm/pitch)))
    shape=tuple(math.ceil((float(e)+2*pad*pitch)/pitch) for e in mesh.extents)
    estimated=math.prod(math.ceil((float(e)+2*s.margin_mm)/pitch)+5 for e in mesh.extents)
    return {'pad':pad,'shape':shape,'cells':math.prod(shape),'estimated':estimated}

def recommended_pitch(mesh,s,limit=GRID_CELL_LIMIT):
    target=int(limit*.9);lo=s.pitch_mm;hi=lo
    while max(grid_plan(mesh,s,hi)['cells'],grid_plan(mesh,s,hi)['estimated'])>target:
        hi*=1.25
    for _ in range(48):
        mid=(lo+hi)/2;plan=grid_plan(mesh,s,mid)
        if max(plan['cells'],plan['estimated'])>target:lo=mid
        else:hi=mid
    pitch=math.ceil(hi*100)/100
    while max(grid_plan(mesh,s,pitch)['cells'],grid_plan(mesh,s,pitch)['estimated'])>target:
        pitch+=.01
    return pitch

class GridResolutionError(ValueError):
    def __init__(self,mesh,s,plan):
        pitch=recommended_pitch(mesh,s);suggested=grid_plan(mesh,s,pitch)
        self.details={'code':'grid-limit','estimated_cells':plan['estimated'],'actual_cells':plan['cells'],'limit':GRID_CELL_LIMIT,'pitch_mm':s.pitch_mm,'suggested_pitch_mm':pitch,'suggested_estimated_cells':suggested['estimated'],'suggested_cells':suggested['cells']}
        super().__init__(f'Grid exceeds {GRID_CELL_LIMIT:,} cells. Use a cell size of at least {pitch:g} mm for this geometry.')

def grid(mesh,s):
    plan=grid_plan(mesh,s)
    if max(plan['estimated'],plan['cells'])>GRID_CELL_LIMIT:
        raise GridResolutionError(mesh,s,plan)
    # Cell-centre occupancy by ray parity (avoids the volume inflation of
    # surface voxelization+fill). Ray intersects triangular surface along +X.
    pad=plan['pad']
    lo=mesh.bounds[0]-pad*s.pitch_mm
    shape=np.array(plan['shape'],dtype=int)
    origin=lo+s.pitch_mm/2
    axes=[origin[i]+np.arange(shape[i])*s.pitch_mm for i in range(3)]
    mask=np.zeros(shape,dtype=bool)
    tri=mesh.triangles; v0=tri[:,0]; e1=tri[:,1]-v0; e2=tri[:,2]-v0
    determinant=e1[:,1]*e2[:,2]-e1[:,2]*e2[:,1]
    valid=np.abs(determinant)>1e-12
    ymin=tri[:,:,1].min(1); ymax=tri[:,:,1].max(1)
    zmin=tri[:,:,2].min(1); zmax=tri[:,:,2].max(1)
    for j,y0 in enumerate(axes[1]):
        y=y0+s.pitch_mm*1.234567e-7
        for k,z0 in enumerate(axes[2]):
            z=z0+s.pitch_mm*2.345679e-7
            ix=np.flatnonzero(valid&(y>=ymin)&(y<=ymax)&(z>=zmin)&(z<=zmax))
            if not len(ix):continue
            dy=y-v0[ix,1]; dz=z-v0[ix,2]; det=determinant[ix]
            u=(dy*e2[ix,2]-dz*e2[ix,1])/det
            v=(e1[ix,1]*dz-e1[ix,2]*dy)/det
            good=(u>=0)&(v>=0)&(u+v<=1)
            xx=np.sort(v0[ix,0][good]+u[good]*e1[ix,0][good]+v[good]*e2[ix,0][good])
            if not len(xx):continue
            xx=xx[np.r_[True,np.diff(xx)>s.pitch_mm*1e-9]]
            if len(xx)%2:raise ValueError('Odd intersection count: self-intersecting or ambiguous surface.')
            mask[:,j,k]=np.searchsorted(xx,axes[0],side='right')%2==1
    if not mask.any():raise ValueError('No metal cells; decrease the cell size.')
    return mask,origin

def liquid_fraction(T,m): return np.clip((T-m.solidus)/(m.liquidus-m.solidus),0,1)
def specific_h(T,m): return m.cp*np.asarray(T)+m.latent*liquid_fraction(np.asarray(T),m)

def build_operator(mask,s,m):
    shape=mask.shape; ids=np.arange(mask.size).reshape(shape); dx=s.pitch_mm/1000
    k=np.where(mask,m.k,s.mold_k); rows=[]; cols=[]; data=[]
    for axis in range(3):
        a=[slice(None)]*3; b=a.copy(); a[axis]=slice(None,-1); b[axis]=slice(1,None); a=tuple(a); b=tuple(b)
        resistance=dx/(2*k[a])+dx/(2*k[b])+(mask[a]!=mask[b])/s.contact_h
        g=(1/(resistance*dx)).ravel(); i=ids[a].ravel(); j=ids[b].ravel()
        rows.extend([i,j,i,j]); cols.extend([i,j,j,i]); data.extend([g,g,-g,-g])
    A=sparse.coo_matrix((np.concatenate(data),(np.concatenate(rows),np.concatenate(cols))),shape=(mask.size,mask.size)).tocsr()
    boundary=np.zeros(shape)
    for axis in range(3):
        for end in (0,-1):
            sl=[slice(None)]*3; sl[axis]=end; boundary[tuple(sl)]+=1/dx
    return A,boundary.ravel()

def snapshot_interval(s,shape,duration,plane=True):
    if not np.isfinite(s.snapshot_dt) or s.snapshot_dt<=0:
        raise ValueError('Snapshot interval must be positive.')
    cells=shape[0]*shape[2] if plane else int(np.prod(shape))
    max_frames=max(2,min(1201,2000000//max(cells,1)))
    return max(s.snapshot_dt,duration/(max_frames-1))

def solve(mask,s,m,progress=None,cancel=None,target_temperature=None):
    validate(s,m)
    if target_temperature is not None and (not np.isfinite(target_temperature) or not s.ambient<target_temperature<s.metal_temp):
        raise ValueError('Final temperature must be above ambient and below initial metal temperature.')
    A, boundary=build_operator(mask,s,m)
    metal=mask.ravel(); rho=np.where(metal,m.rho,s.mold_rho); cp=np.where(metal,m.cp,s.mold_cp)
    T=np.where(metal,s.metal_temp,s.mold_temp).astype(float)
    latent=np.where(metal,m.rho*m.latent,0.)
    def H(t): return rho*cp*t+latent*liquid_fraction(t,m)
    def C(t): return rho*cp+latent*((t>m.solidus)&(t<m.liquidus))/(m.liquidus-m.solidus)
    sigma=5.670374419e-8; ta=s.ambient+273.15
    def flux(t): return boundary*(s.outside_h*(t-s.ambient)+s.emissivity*sigma*((t+273.15)**4-ta**4))
    def flux_deriv(t): return boundary*(s.outside_h+4*s.emissivity*sigma*(t+273.15)**3)
    dx=s.pitch_mm/1000; cellvol=dx**3; initial=float(H(T).sum()*cellvol)
    history=[]; time=0.; escaped=0.; rejected=0; min_dt=s.dt; freeze=np.full(mask.size,np.nan)
    snapshot_times=[]; temperature_slices=[]; phase_history=[]
    freeze[metal&(T<=m.solidus)]=0.
    interval=snapshot_interval(s,mask.shape,s.duration)
    next_snapshot=0.
    def record():
        nonlocal next_snapshot
        f=liquid_fraction(T[metal],m)
        energy=float(H(T).sum()*cellvol)
        err=energy-initial+escaped
        phase_history.append([time,float(np.mean(T[metal]<=m.solidus)),float(np.mean((T[metal]>m.solidus)&(T[metal]<m.liquidus))),float(np.mean(T[metal]>=m.liquidus))])
        history.append([time,float(T[metal].min()),float(T[metal].mean()),float(T[metal].max()),float(f.mean()),escaped,err])
        if time>=next_snapshot-1e-9 or time>=s.duration-1e-10 or (target_temperature is not None and T[metal].max()<=target_temperature):
            snapshot_times.append(time)
            temperature_slices.append(T.reshape(mask.shape)[:,mask.shape[1]//2,:].T.copy())
            next_snapshot=time+interval
    record()
    while time<s.duration-1e-10 and (target_temperature is None or T[metal].max()>target_temperature):
        if cancel and cancel(): raise InterruptedError('Computation stopped.')
        dt=min(s.dt,s.duration-time,next_snapshot-time); previous=H(T); old=T.copy(); success=False
        while not success:
            candidate=old.copy()
            for it in range(35):
                residual=H(candidate)-previous+dt*(A@candidate+flux(candidate))
                # convergence in volumetric energy: equivalent 1e-5 K max
                norm=np.max(np.abs(residual)/(rho*cp))
                if norm<1e-5: success=True; break
                diag=C(candidate)+dt*flux_deriv(candidate)
                J=dt*A+sparse.diags(diag)
                pre=LinearOperator(J.shape,matvec=lambda x:x/J.diagonal())
                step,info=cg(J,-residual,M=pre,rtol=1e-8,atol=0,maxiter=600)
                if info!=0: break
                base=np.linalg.norm(residual); alpha=1.
                for _ in range(20):
                    trial=candidate+alpha*step
                    r=H(trial)-previous+dt*(A@trial+flux(trial))
                    if np.linalg.norm(r)<base: candidate=trial; break
                    alpha*=.5
                else: break
            if not success:
                dt*=.5; rejected+=1
                if dt<max(1e-5,s.dt/65536): raise RuntimeError('Nonlinear solver did not converge. Review time step and material data.')
                if cancel and cancel(): raise InterruptedError('Computation stopped.')
        T=candidate; time+=dt; min_dt=min(min_dt,dt)
        escaped+=float(flux(T).sum()*cellvol*dt)
        newly=metal&np.isnan(freeze)&(T<=m.solidus)
        freeze[newly]=time
        record()
        if progress: progress(time/s.duration,f'{time:.1f}/{s.duration:.0f} s | kapalina {history[-1][4]*100:.1f} %')
    return {'temperature':T.reshape(mask.shape),'freeze_time':freeze.reshape(mask.shape),'history':np.array(history),
            'snapshot_times':np.array(snapshot_times),'temperature_slices':np.array(temperature_slices),
            'snapshot_interval_s':interval,'phase_history':np.array(phase_history),
            'target_temperature_C':target_temperature,'target_reached':bool(target_temperature is not None and T[metal].max()<=target_temperature),
            'initial_energy_J':initial,'min_dt_s':min_dt,'rejected_steps':rejected,
            'energy_error_J':history[-1][-1], 'energy_error_relative':abs(history[-1][-1])/max(abs(escaped),1)}

def candidates(mesh,mask,origin,s):
    # Thick interior proxy via distance transform; not a feeding/porosity prediction.
    dist=ndimage.distance_transform_edt(mask)*s.pitch_mm
    thick=np.array(np.unravel_index(np.argmax(dist),mask.shape)); p=origin+thick*s.pitch_mm
    vertices=np.asarray(mesh.vertices)
    # Gate at lower surface near x coordinate of thick region.
    score=((vertices[:,0]-p[0])/max(mesh.extents[0]*.12,1))**2 + ((vertices[:,2]-mesh.bounds[0,2])/max(mesh.extents[2]*.25,1))**2
    g=vertices[np.argmin(score)]
    # Upper envelope bins, local maxima and spatial nonmaximum suppression.
    cell=max(s.pitch_mm, max(mesh.extents)/20)
    bins={}
    for v in vertices:
        key=tuple(np.floor(v[:2]/cell).astype(int))
        if key not in bins or v[2]>bins[key][2]: bins[key]=v
    peaks=[]
    for key,v in bins.items():
        if all(v[2]>=bins.get((key[0]+i,key[1]+j),v)[2]-1e-6 for i in (-1,0,1) for j in (-1,0,1)):
            peaks.append(v)
    chosen=[]
    for v in sorted(peaks,key=lambda a:-a[2]):
        if all(np.linalg.norm(v[:2]-w[:2])>max(mesh.extents)*.15 for w in chosen): chosen.append(v)
        if len(chosen)==4: break
    # Riser point on upper envelope near maximum thickness location.
    near=np.linalg.norm(vertices[:,:2]-p[:2],axis=1)
    valid=vertices[near<=max(s.pitch_mm*2,float(near.min())+.01)]
    r=valid[np.argmax(valid[:,2])]
    return {'gate':g.tolist(),'vents':[v.tolist() for v in chosen],'riser':r.tolist(),
            'thick_point':p.tolist(),'thickness_proxy_mm':float(2*dist.max())}

def cylinder(a,b,d,sections=32):
    a=np.asarray(a,float); b=np.asarray(b,float)
    return trimesh.creation.cylinder(radius=d/2,sections=sections,segment=np.array([a,b]))

def channels(mesh,c,s):
    result={}; g=np.asarray(c['gate']); top=mesh.bounds[1,2]+s.head_mm
    feeding=c.get('feeding_geometry')
    if feeding:top=feeding['top_z_mm']
    elbow=np.array([g[0],mesh.bounds[0,1]-20,g[2]])
    # Penetrate body slightly to permit solid boolean union.
    center=mesh.center_mass; direction=center-g; direction/=max(np.linalg.norm(direction),1e-9)
    result['vtok_spojka']=cylinder(g+direction*s.gate_mm*.4,elbow,s.gate_mm)
    result['vtok_svisly']=cylinder(elbow,[elbow[0],elbow[1],top],s.gate_mm)
    for i,v in enumerate(c['vents']):
        v=np.asarray(v); result[f'odvzdusneni_{i+1}']=cylinder(v-[0,0,s.vent_mm*.4],[v[0],v[1],top],s.vent_mm)
    r=np.array(c['riser'])
    if feeding:
        base=np.array([r[0],r[1],feeding['base_z_mm']])
        result['hrdlo_nalitku']=cylinder(r-[0,0,feeding['neck_diameter_mm']*.4],base+[0,0,.05],feeding['neck_diameter_mm'],64)
        result['nalitek']=cylinder(base,[r[0],r[1],feeding['top_z_mm']],feeding['diameter_mm'],64)
    else:result['nalitek']=cylinder(r-[0,0,s.riser_mm*.3],r+[0,0,s.riser_mm*1.5],s.riser_mm)
    return result

def estimates(mesh,s,m):
    volume=abs(mesh.volume)*1e-9; mass=volume*m.rho
    area=math.pi*(s.gate_mm/1000)**2/4
    velocity=s.discharge*math.sqrt(2*9.80665*s.head_mm/1000)
    q=area*velocity
    energy=s.charge_kg*float(specific_h(s.metal_temp,m)-specific_h(s.ambient,m))
    f=s.cans_percent/100
    # Illustrative Mg balance: ONLY if cans 1.3-1.6% and profiles 6060 0.35-0.6%.
    return {'objem_STL_cm3':volume*1e6,'hmotnost_scenar_kg':mass,
      'idealni_prutok_cm3_s':q*1e6,'idealni_cas_plneni_s':volume/q,
      'teplo_pouze_vsazka_kJ':energy/1000,'cas_ohrevu_pri_zadane_ucinnosti_min':energy/(s.furnace_kw*1000*s.efficiency)/60,
      'vsazka_pro_samotny_odlitek_kg':mass/(s.yield_percent/100),
      'Mg_hmotnostni_pct_podmineny_interval':[f*1.3+(1-f)*.35,f*1.6+(1-f)*.6]}

def export(folder,mesh,mask,origin,s,m,c,result=None):
    folder=Path(folder); folder.mkdir(parents=True,exist_ok=True)
    mesh.export(folder/'model_orientovany.stl')
    parts=channels(mesh,c,s)
    for name,part in parts.items(): part.export(folder/(name+'.stl'))
    union_status='Nevytvoreno'
    try:
        union=trimesh.boolean.union([mesh,*parts.values()],engine='manifold')
        if not union.is_watertight: raise ValueError('Union is not watertight.')
        union.export(folder/'sestava_sjednocena.stl'); union_status='Uzavrene boolean sjednoceni; vyjmutelnost z pisku neoverena.'
    except Exception as exc: union_status='Sjednoceni selhalo: '+str(exc)
    report={'verze':'1.0','omezeni':LIMITATIONS,'nastaveni':asdict(s),'material':asdict(m),'odhady':estimates(mesh,s,m),
      'kandidati':c,'sjednoceni':union_status,'voxel_objem_cm3':float(mask.sum()*(s.pitch_mm/10)**3),
      'STL_jednotky':'mm; vstup skalovan na zadanou nejdelsi osu pred otocenim','pocet_bunek':int(mask.size)}
    if result is not None:
        np.savez_compressed(folder/'pole.npz',mask=mask,origin_mm=origin,pitch_mm=s.pitch_mm,temperature_C=result['temperature'],freeze_time_s=result['freeze_time'])
        with open(folder/'prubeh.csv','w',newline='',encoding='utf-8-sig') as f:
            w=csv.writer(f); w.writerow(['cas_s','T_min_C','T_prumer_C','T_max_C','podil_kapaliny','teplo_do_okoli_J','chyba_bilance_J']); w.writerows(result['history'])
        report['resic']={k:v for k,v in result.items() if not isinstance(v,np.ndarray)}
        report['konec']={'cas_s':float(result['history'][-1,0]),'podil_kapaliny':float(result['history'][-1,4])}
    (folder/'protokol.json').write_text(json.dumps(report,ensure_ascii=False,indent=2),encoding='utf-8')
    return report