"""Axisymmetric creeping-flow FEM in a body-fitted meridional mesh.

Coordinates and velocities use mm and mm/s; viscosity is Pa s, pressure Pa.
Taylor-Hood P2 velocity / P1 pressure, full cylindrical symmetric strain.
"""
from dataclasses import dataclass
import numpy as np
from scipy.sparse import bmat, csr_matrix
from scipy.sparse.linalg import spsolve
from skfem import MeshTri, Basis, ElementVector, ElementTriP2, ElementTriP1, BilinearForm, asm, condense
from skfem.helpers import sym_grad, ddot

@dataclass
class Flow:
    mesh: object
    velocity_basis: object
    pressure_basis: object
    velocity: np.ndarray
    pressure: np.ndarray
    z_knots: np.ndarray
    radii: np.ndarray
    inner_radius: float
    flow_rate: float
    viscosity: float
    residual: float

    def sample(self, z, r):
        z,r=np.broadcast_arrays(np.asarray(z),np.asarray(r))
        # Locate boundary samples a negligible distance inside the mesh to avoid
        # floating-point element-finder misses on sloping exterior facets.
        tol=1e-12*max(1.,self.z_knots[-1])
        outer=np.interp(z,self.z_knots,self.radii)
        if np.any(z < -tol) or np.any(z > self.z_knots[-1]+tol) or np.any(r < self.inner_radius-tol) or np.any(r > outer+tol):
            raise ValueError('Sample lies outside the fluid domain')
        zz=np.clip(z,tol,self.z_knots[-1]-tol)
        rr=np.clip(r,self.inner_radius+tol,np.interp(zz,self.z_knots,self.radii)-tol)
        xy=np.vstack([zz.ravel(),rr.ravel()])
        v=self.velocity_basis.interpolator(self.velocity)(xy)
        p=self.pressure_basis.interpolator(self.pressure)(xy)
        return v.reshape((2,)+z.shape),p.reshape(z.shape)

    def section(self,z,n=81):
        r=np.linspace(self.inner_radius,np.interp(z,self.z_knots,self.radii),n)
        v,p=self.sample(np.full_like(r,z),r)
        q=np.trapezoid(2*np.pi*r*v[0],r)
        pmean=np.trapezoid(2*r*p,r)/(r[-1]**2-r[0]**2)
        return q,pmean

    def diagnostics(self):
        qs=[];ps=[]
        for z in np.linspace(0,self.z_knots[-1],21):
            q,p=self.section(z);qs.append(float(q));ps.append(float(p))
        return {'pressure_drop_Pa':ps[0]-ps[-1], 'section_flow_min_mm3_s':min(qs),
                'section_flow_max_mm3_s':max(qs),'section_flow_deviation_fraction':max(abs(np.array(qs)/self.flow_rate-1)),
                'inlet_outlet_imbalance_fraction':abs(qs[0]-qs[-1])/self.flow_rate,
                'free_equation_relative_residual':self.residual,'elements':int(self.mesh.nelements),
                'degrees_of_freedom':int(self.velocity_basis.N+self.pressure_basis.N)}

def annular_profile(r,a,b,q):
    r=np.asarray(r)
    if a==0:
        return 2*q/(np.pi*b*b)*(1-(r/b)**2)
    shape=b*b-r*r-(b*b-a*a)*np.log(b/r)/np.log(b/a)
    denom=np.pi/2*(b**4-a**4-(b*b-a*a)**2/np.log(b/a))
    return q*shape/denom

def solve_flow(z_knots,radii,q=50.,mu=.25,inner=0.,nz=80,nr=16):
    z_knots=np.asarray(z_knots,float);radii=np.asarray(radii,float)
    if not np.all(np.isfinite(z_knots)) or not np.all(np.isfinite(radii)) or not np.all(np.isfinite([q,mu,inner])): raise ValueError('Finite geometry and fluid inputs required')
    if len(z_knots)<2 or len(z_knots)!=len(radii) or z_knots[0]!=0 or np.any(np.diff(z_knots)<=0): raise ValueError('Increasing z stations must start at zero')
    if np.min(radii)<=inner or inner<0 or q<=0 or mu<=0: raise ValueError('Positive clearance, flow and viscosity required')
    z=np.unique(np.r_[np.linspace(0,z_knots[-1],nz+1),z_knots])
    eta=np.linspace(0,1,nr+1)
    m=MeshTri.init_tensor(z,eta)
    coords=m.p.copy();coords[1]=inner+coords[1]*(np.interp(coords[0],z_knots,radii)-inner)
    m=MeshTri(coords,m.t).with_boundaries({
        'inlet':lambda x:np.isclose(x[0],0),
        'outlet':lambda x:np.isclose(x[0],z_knots[-1]),
        'outer':lambda x:np.isclose(x[1],np.interp(x[0],z_knots,radii)),
        'inner':lambda x:np.isclose(x[1],inner)})
    bu=Basis(m,ElementVector(ElementTriP2()),intorder=5)
    bp=Basis(m,ElementTriP1(),intorder=5)
    @BilinearForm
    def visc(u,v,w):
        return 2*mu*(ddot(sym_grad(u),sym_grad(v))+u[1]*v[1]/w.x[1]**2)*w.x[1]
    @BilinearForm
    def continuity(u,p,w):
        return -(u.grad[0,0]+u.grad[1,1]+u[1]/w.x[1])*p*w.x[1]
    A=asm(visc,bu);B=asm(continuity,bu,bp)
    K=bmat([[A,B.T],[B,csr_matrix((bp.N,bp.N))]],format='csr')
    x=np.zeros(K.shape[0]);rhs=np.zeros_like(x)
    fixed=[bu.get_dofs('outer').all(),bu.get_dofs('inlet').all(),bu.get_dofs('outlet').all()]
    fixed.append(bu.get_dofs('inner').all() if inner>0 else bu.get_dofs('inner').all(['u^2']))
    for name,b in [('inlet',radii[0]),('outlet',radii[-1])]:
        ids=bu.get_dofs(name).all(['u^1']);r=bu.doflocs[1,ids]
        x[ids]=annular_profile(r,inner,b,q)
    pin=bu.N+int(np.argmin((bp.doflocs[0]-z_knots[-1])**2+(bp.doflocs[1]-inner)**2))
    D=np.unique(np.r_[*fixed,pin]);Kc,fc,xx,I=condense(K,rhs,x=x,D=D)
    xx[I]=spsolve(Kc,fc)
    residual=np.linalg.norm(Kc@xx[I]-fc)/max(np.linalg.norm(fc),1e-30)
    return Flow(m,bu,bp,xx[:bu.N],xx[bu.N:],z_knots,radii,inner,q,mu,float(residual))

if __name__=='__main__':
    import json
    f=solve_flow([0,20],[1,1],nz=40,nr=12)
    d=f.diagnostics();d['analytical_Pa']=8*.25*20*50/np.pi
    print(json.dumps(d,indent=2))