"""Conservative axisymmetric cell-centered transport on a straight annulus.

First-order upwind axial advection, centered molecular diffusion, backward Euler.
C is the old-sample fraction. Inlet advective C=0; zero diffusive flux at ends;
no radial flux at either wall. Internal face fluxes cancel exactly in mass form.
"""
from dataclasses import dataclass
import numpy as np
from scipy.sparse import lil_matrix,diags
from scipy.sparse.linalg import factorized
from flow import annular_profile

@dataclass
class Transport:
 z:np.ndarray
 r:np.ndarray
 volumes:np.ndarray
 ring_flow:np.ndarray
 operator:object
 initial_mass:float
 profile_correction:float

def assemble(length,a,b,q,diffusion,nz,nr,flow_solution=None):
 if length<=0 or a<=0 or b<=a or q<=0 or diffusion<0 or nz<2 or nr<2:raise ValueError('Positive length, annulus and flow; nonnegative diffusivity; at least 2x2 cells required')
 dz=length/nz;dr=(b-a)/nr;rf=np.linspace(a,b,nr+1);rc=(rf[:-1]+rf[1:])/2;zc=(np.arange(nz)+.5)*dz
 ring_area=np.pi*np.diff(rf**2);vol=np.tile(ring_area*dz,nz)
 # Three-point Gauss integration of FEM velocity into annular face flows.
 gp,gw=np.polynomial.legendre.leggauss(3);rr=rc[:,None]+dr/2*gp
 if flow_solution is None: vel=annular_profile(rr,a,b,q)
 else:vel=flow_solution.sample(np.full(rr.shape,length/2),rr)[0][0]
 qj=(2*np.pi*rr*vel*gw*dr/2).sum(axis=1)
 if np.any(qj<=0):raise ValueError('Axial transport requires forward ring flows')
 correction=q/qj.sum();qj*=correction
 size=nz*nr;op=lil_matrix((size,size))
 for i in range(nz):
  for j in range(nr):
   k=i*nr+j;op[k,k]-=qj[j]
   if i>0:op[k,k-nr]+=qj[j]
   if i+1<nz:
    g=diffusion*ring_area[j]/dz;ell=k+nr
    op[k,k]-=g;op[ell,ell]-=g;op[k,ell]+=g;op[ell,k]+=g
   if j+1<nr:
    g=diffusion*2*np.pi*rf[j+1]*dz/dr;ell=k+1
    op[k,k]-=g;op[ell,ell]-=g;op[k,ell]+=g;op[ell,k]+=g
 return Transport(zc,rc,vol,qj,op.tocsc(),float(vol.sum()),float(correction))

def crossing(t,y,threshold=.01):
 ids=np.flatnonzero(np.asarray(y)<=threshold)
 if not len(ids):return None
 k=ids[0]
 if k==0:return float(t[0])
 return float(t[k-1]+(threshold-y[k-1])*(t[k]-t[k-1])/(y[k]-y[k-1]))

def run_transport(model,dt,probe_region,threshold=.01,stop_fraction=.002,max_turnovers=80,collect_frames=False):
 if dt<=0:raise ValueError('Positive time step required')
 nz,nr=len(model.z),len(model.r);v=model.volumes;c=np.ones(nz*nr);tau=model.initial_mass/model.ring_flow.sum()
 dz=model.z[1]-model.z[0];dr=model.r[1]-model.r[0]
 z_overlap=np.maximum(0,np.minimum(model.z+dz/2,probe_region['z_max_mm'])-np.maximum(model.z-dz/2,probe_region['z_min_mm']))
 radial_hi=np.minimum(model.r+dr/2,probe_region['inner_radius_mm']+probe_region['radial_depth_mm'])
 radial_lo=np.maximum(model.r-dr/2,probe_region['inner_radius_mm'])
 radial_overlap=np.pi*np.maximum(0,radial_hi**2-radial_lo**2)
 probe_weights=(z_overlap[:,None]*radial_overlap[None,:]).ravel()
 if probe_weights.sum()<=0:raise ValueError('Probe measurement region does not intersect the mesh')
 solve=factorized(diags(v,format='csc')-dt*model.operator)
 times=[0.];whole=[1.];local=[1.];outlet=[1.];emitted=[0.];mass_error=[0.];frames=[];frame_times=[]
 if collect_frames:frames.append(c.reshape(nz,nr).copy());frame_times.append(0.)
 next_frame=.5*tau;minimum=1.;maximum=1.
 for n in range(1,int(np.ceil(max_turnovers*tau/dt))+1):
  c=solve(v*c);t=n*dt;out=float(np.dot(model.ring_flow,c[-nr:]));mass=float(np.dot(v,c));probe_c=float(np.dot(probe_weights,c)/probe_weights.sum())
  times.append(t);whole.append(mass/model.initial_mass);local.append(probe_c);outlet.append(out/model.ring_flow.sum());emitted.append(emitted[-1]+dt*out);mass_error.append((mass+emitted[-1]-model.initial_mass)/model.initial_mass)
  minimum=min(minimum,float(c.min()));maximum=max(maximum,float(c.max()))
  if collect_frames and t>=next_frame and len(frames)<30:frames.append(c.reshape(nz,nr).copy());frame_times.append(t);next_frame+=.5*tau
  if max(whole[-1],local[-1])<stop_fraction:break
 result={'time_s':np.array(times),'whole_fraction':np.array(whole),'probe_fraction':np.array(local),'outlet_fraction':np.array(outlet),'emitted_mass_mm3':np.array(emitted),'mass_balance_fraction':np.array(mass_error),'t99_whole_s':crossing(times,whole,threshold),'t99_probe_s':crossing(times,local,threshold),'t99_outlet_s':crossing(times,outlet,threshold),'residence_time_s':tau,'min_fraction':minimum,'max_fraction':maximum,'maximum_mass_balance_fraction':max(abs(np.array(mass_error))),'frames':frames,'frame_times':frame_times,'nz':nz,'nr':nr,'dt_s':dt,'profile_correction':model.profile_correction}
 return result