"""Quadratic triangular plane-stress linear elasticity, SI units.

K = thickness * integral(B.T D B) + distributed support springs.
The explicit two-dimensional model cannot establish 3D contact or notch strength.
"""
from pathlib import Path
import json
import numpy as np
from scipy.sparse.linalg import spsolve
from skfem import MeshTri,Basis,FacetBasis,ElementTriP2,ElementVector,BilinearForm,LinearForm,asm,condense
from skfem.models.elasticity import linear_elasticity
ROOT=Path(__file__).resolve().parent

def load():return json.loads((ROOT/'inputs.json').read_text())
def grid(length,h,extra=()):return np.unique(np.round(np.r_[np.linspace(0,length,round(length/h)+1),extra],12))
def mesh_for(c,profile,h):
 win=c['profiles'][profile];extra_x=[.015,.025,.215,.225,.11,.13,.14,.16];extra_y=[]
 if win:
  extra_x+=[v for pair in win['x_ranges_m'] for v in pair];extra_y=[win['bottom_m'],win['top_m']]
 mesh=MeshTri.init_tensor(grid(c['length_m'],h,extra_x),grid(c['height_m'],h,extra_y))
 if win:
  points=mesh.p[:,mesh.t].mean(axis=1);mask=np.zeros(mesh.nelements,dtype=bool)
  for lo,hi in win['x_ranges_m']:mask|=(points[0]>lo)&(points[0]<hi)&(points[1]>win['bottom_m'])&(points[1]<win['top_m'])
  mesh=mesh.remove_elements(np.where(mask)[0])
 return mesh

def solve(c,profile='solid',h=.0025,center=.12,stiffness=2e7):
 if c['E_Pa']<=0 or not (0<c['poisson']<.49) or c['thickness_m']<=0:raise ValueError('Invalid elastic inputs')
 m=mesh_for(c,profile,h);elem=ElementVector(ElementTriP2());basis=Basis(m,elem,intorder=4)
 E=c['E_Pa'];nu=c['poisson'];mu=E/(2*(1+nu));lam=E*nu/(1-nu**2);t=c['thickness_m']
 elastic=t*asm(linear_elasticity(lam,mu),basis)
 @LinearForm
 def traction(v,w):return v[1]*(-c['load_N']/c['load_patch_width_m'])
 load_facets=m.facets_satisfying(lambda x:np.isclose(x[1],c['height_m'])&(x[0]>=center-c['load_patch_width_m']/2-1e-10)&(x[0]<=center+c['load_patch_width_m']/2+1e-10))
 fb=FacetBasis(m,elem,facets=load_facets,intorder=4);f=asm(traction,fb)
 spring=elastic.copy()*0
 @BilinearForm
 def support(u,v,w):return (stiffness/c['support_patch_width_m'])*u[1]*v[1]
 for x in c['support_centers_m']:
  facets=m.facets_satisfying(lambda q:np.isclose(q[1],0)&(q[0]>=x-c['support_patch_width_m']/2-1e-10)&(q[0]<=x+c['support_patch_width_m']/2+1e-10))
  spring+=asm(support,FacetBasis(m,elem,facets=facets,intorder=4))
 K=elastic+spring
 node=np.where(np.isclose(m.p[0],c['support_centers_m'][0])&np.isclose(m.p[1],0))[0]
 D=basis.nodal_dofs[0,node];free=np.setdiff1d(np.arange(basis.N),D);u=np.zeros(basis.N);u[free]=spsolve(K[free][:,free],f[free])
 field=basis.interpolate(u);grad=field.grad
 ex=grad[0,0];ey=grad[1,1];gxy=grad[0,1]+grad[1,0]
 sx=E/(1-nu**2)*(ex+nu*ey);sy=E/(1-nu**2)*(ey+nu*ex);tau=mu*gxy
 vm=np.sqrt(sx*sx-sx*sy+sy*sy+3*tau*tau).mean(axis=1)
 centroids=m.p[:,m.t].mean(axis=1);mask=(centroids[1]>.004)&(centroids[1]<c['height_m']-.004)
 win=c['profiles'][profile]
 if win:
  for xr in win['x_ranges_m']:
   for xx in xr:
    for yy in [win['bottom_m'],win['top_m']]:mask &= np.hypot(centroids[0]-xx,centroids[1]-yy)>.004
 reaction=-(spring@u);yy=basis.split_indices()[1]
 # Compliance / applied force is the force-weighted displacement of the loading patch.
 load_disp=float(-f@u/c['load_N'])
 area=float(np.sum(basis.dx));mass=area*t*c['density_kg_m3']
 rx=K@u-f
 return {'mesh':m,'basis':basis,'u':u,'vm':vm,'profile':profile,'h_m':h,'load_center_m':center,'support_N_m':stiffness,'dofs':basis.N,'triangles':m.nelements,'load_displacement_mm':-load_disp*1000,'max_vertical_mm':float(-u[yy].min()*1000),'stress_away_from_corners_MPa':float(vm[mask].max()/1e6),'mass_kg':mass,'reaction_N':float(reaction[yy].sum()),'reaction_moment_Nm':float(reaction[yy]@basis.doflocs[0,yy]),'force_error_relative':float(abs(reaction[yy].sum()-c['load_N'])/c['load_N']),'moment_error_relative':float(abs(reaction[yy]@basis.doflocs[0,yy]-c['load_N']*center)/(c['load_N']*center)),'free_residual_relative':float(np.linalg.norm(rx[free])/np.linalg.norm(f[free]))}

def benchmark(h=.0025):
 L=.2;H=.02;t=.02;E=69e9;nu=.33;P=500
 m=MeshTri.init_tensor(grid(L,h),grid(H,h));basis=Basis(m,ElementVector(ElementTriP2()),intorder=4)
 K=t*asm(linear_elasticity(E*nu/(1-nu**2),E/(2*(1+nu))),basis)
 @LinearForm
 def traction(v,w):return -P/H*v[1]
 f=asm(traction,FacetBasis(m,basis.elem,facets=m.facets_satisfying(lambda x:np.isclose(x[0],L))))
 D=basis.get_dofs(lambda x:np.isclose(x[0],0)).all();free=np.setdiff1d(np.arange(basis.N),D);u=np.zeros(basis.N);u[free]=spsolve(K[free][:,free],f[free]);actual=float(f@u/P)
 I=t*H**3/12;eb=P*L**3/(3*E*I);shear=P*L/((5/6)*(E/(2*(1+nu)))*t*H);reference=eb+shear
 return {'mesh_size_m':h,'FEA_tip_displacement_mm':actual*1000,'Euler_Bernoulli_mm':eb*1000,'Timoshenko_reference_mm':reference*1000,'relative_difference':abs(actual-reference)/reference}

def public(r):return {k:v for k,v in r.items() if k not in ['mesh','basis','u','vm']}