"""Reduced hydraulic/geometry screening. NOT Navier-Stokes/VOF CFD.
A single effective stream path and an incompressible vented cavity are assumed.
Arrival ordering is a graph approximation, not a resolved free surface.
"""
import heapq, math
from dataclasses import dataclass,asdict
import numpy as np
from scipy import ndimage
from scipy.optimize import brentq

@dataclass
class FillSettings:
    viscosity: float = .0013       # Pa s; scenario for liquid aluminum
    liquid_density: float = 2380.0 # kg/m3; scenario, not solid density
    minor_loss: float = 2.0
    roughness_mm: float = .05
    backpressure_pa: float = 0.0
    vented: bool = True

FILL_LIMITS=[
 'REDUKOVANY HYDRAULICKY SCREENING, nikoli CFD: neresi Navier-Stokes ani VOF.',
 'Poradi fronty je odvozeno z drah v mrizce; skutecne hladiny, vetveni proudu a hybnost nejsou reseny.',
 'Predpoklad staly sloupec kovu a volny unik vzduchu. Neoveruje kapacitu navrzenych odvzdusneni.',
 'Kov se v hydraulickem modelu nezastavuje tuhnutim. Tepelny ukazatel je pouze screening.',
 'Odpar, oxidace, vodik, plyn z pojiva a povrchove napeti nejsou vypocteny. Nejsou nahrazeny procentem necistot.',
 'Cas plneni ani tepelny pomer nejsou potvrzenim uplneho zaliti detailu.',
]

def pipe_velocity(head,rho,mu,diameter,length,K,roughness=0):
    """Darcy-Weisbach head balance; laminar f=64/Re, turbulent Haaland.
    Smooth blending at 2300-4000 is a numerical engineering assumption.
    """
    if head<=0:return 0.
    if min(rho,mu,diameter)<=0 or min(length,K,roughness)<0:raise ValueError('Invalid hydraulic parameters.')
    dp=rho*9.80665*head
    def loss(v):
        if v<=0:return -dp
        re=rho*v*diameter/mu
        fl=64/re
        ft=(-1.8*math.log10((roughness/diameter/3.7)**1.11+6.9/max(re,2300)))**-2
        blend=np.clip((re-2300)/1700,0,1)
        f=fl*(1-blend)+ft*blend
        return .5*rho*v*v*(K+f*length/diameter)-dp
    high=max(1.,math.sqrt(2*9.80665*head))
    while loss(high)<0:
        high*=2
        if high>1e5:raise ValueError('Unbounded velocity: hydraulic resistance is missing.')
    return brentq(loss,0,high,xtol=1e-10)

def screen(mask,origin,s,m,c,fs,progress=None,cancel=None):
    if not fs.vented:raise ValueError('This screening model requires a freely vented cavity. Use the extended front model for gas compression.')
    if not all(np.isfinite(v) for v in asdict(fs).values()):raise ValueError('Invalid filling inputs.')
    if min(fs.viscosity,fs.liquid_density)<=0 or min(fs.minor_loss,fs.roughness_mm,fs.backpressure_pa)<0:raise ValueError('Invalid filling inputs.')
    dx=s.pitch_mm/1000; coords=np.argwhere(mask);world=origin+coords*s.pitch_mm
    gate=np.asarray(c['gate']); start=tuple(coords[np.argmin(np.linalg.norm(world-gate,axis=1))])
    # Constant inlet free surface used by channel generator.
    surface=max(float(origin[2]+np.max(coords[:,2])*s.pitch_mm+s.pitch_mm/2+s.head_mm),c.get('feeding_geometry',{}).get('top_z_mm',0))/1000
    dist=ndimage.distance_transform_edt(mask)*dx
    # Distances estimated to cell faces, not voxel centres.
    diameter=np.maximum(dx,2*dist-dx)
    costs=np.full(mask.shape,np.inf);lengths=np.zeros(mask.shape);bottle=np.zeros(mask.shape)
    costs[start]=0.;bottle[start]=min(s.gate_mm/1000,float(diameter[start]))
    heap=[(0.,start)]; visited=np.zeros(mask.shape,bool); order=[]
    offsets=[(1,0,0),(-1,0,0),(0,1,0),(0,-1,0),(0,0,1),(0,0,-1)]
    while heap:
        cost,u=heapq.heappop(heap)
        if visited[u]:continue
        visited[u]=True;order.append(u)
        if cancel and len(order)%500==0 and cancel():raise InterruptedError('Filling calculation stopped.')
        for off in offsets:
            v=tuple(u[i]+off[i] for i in range(3))
            if any(v[i]<0 or v[i]>=mask.shape[i] for i in range(3)) or not mask[v] or visited[v]:continue
            head=surface-(origin[2]+v[2]*s.pitch_mm)/1000-fs.backpressure_pa/(fs.liquid_density*9.80665)
            if head<=0:continue
            d=min(bottle[u],diameter[v]);length=lengths[u]+dx
            # Proxy path cost: laminar drag plus gravity travel scale, no CFD.
            dp=fs.liquid_density*9.80665*head
            a=.5*fs.liquid_density*max(fs.minor_loss,1e-6); b=32*fs.viscosity*length/(d*d)
            speed=2*dp/(b+math.sqrt(b*b+4*a*dp))
            new=cost+dx/max(speed,1e-9)
            if new<costs[v]:costs[v]=new;lengths[v]=length;bottle[v]=d;heapq.heappush(heap,(new,v))
    arrival=np.full(mask.shape,np.nan); residence=np.full(mask.shape,np.nan);risk=np.full(mask.shape,np.nan)
    elapsed=0.;gatearea=math.pi*(s.gate_mm/1000)**2/4
    if progress:progress(.1,'Hydraulicky screening: vyhodnoceni drah')
    for i,u in enumerate(order):
        z=(origin[2]+u[2]*s.pitch_mm)/1000
        head=surface-z-fs.backpressure_pa/(fs.liquid_density*9.80665)
        if head<=0:continue
        d=bottle[u];length=max(lengths[u],dx)+s.head_mm/1000+.02
        velocity=pipe_velocity(head,fs.liquid_density,fs.viscosity,d,length,fs.minor_loss,fs.roughness_mm/1000)
        q=velocity*min(gatearea,math.pi*d*d/4)
        elapsed+=dx**3/q;arrival[u]=elapsed;residence[u]=length/velocity
        # Local lumped thermal-time proxy to solidus (not coupled advection).
        deltaH=m.cp*(s.metal_temp-m.solidus)+m.latent
        driving=max((s.metal_temp+m.solidus)/2-s.mold_temp,1e-6)
        # V/A of a sphere with radius equal to local EDT distance.
        tau=fs.liquid_density*deltaH*max(dist[u]-dx/2,dx/2)/(3*s.contact_h*driving)
        risk[u]=residence[u]/max(tau,1e-12)
        if cancel and i%500==0 and cancel():raise InterruptedError('Filling calculation stopped.')
    finite=np.isfinite(arrival)&mask
    return {'arrival_s':arrival,'residence_s':residence,'thermal_ratio':risk,'local_diameter_mm':diameter*1000,
      'settings':asdict(fs),'thermal_inputs':asdict(s),'material':asdict(m),'channels':c,
      'geometry':{'shape':list(mask.shape),'origin_mm':origin.tolist(),'pitch_mm':s.pitch_mm},'limitations':FILL_LIMITS,
      'summary':{'estimated_fill_s':float(np.nanmax(arrival)) if finite.any() else None,
       'geometrically_reached_fraction':float(finite.sum()/mask.sum()),
       'max_thermal_screen_ratio':float(np.nanmax(risk)) if finite.any() else None,
       'interpretation':'Residence/local thermal time ratio, not a defect probability. Free venting is assumed.'}}
