"""Thin-plate thermal FEA with three coupled lumped heater nodes, SI units.

P1 triangles, consistent heat-capacity matrix, distributed plate/air losses,
edge convection, finite heater-to-plate contact, and implicit time integration.
Temperature is uniform through the plate thickness by model assumption.
"""
from pathlib import Path
import json
import numpy as np
from scipy.sparse import bmat,csr_matrix,diags
from scipy.sparse.linalg import spsolve,factorized
from skfem import MeshTri,Basis,FacetBasis,ElementTriP1,LinearForm,asm
from skfem.models.poisson import laplace,mass
ROOT=Path(__file__).resolve().parent

def load():return json.loads((ROOT/'inputs.json').read_text())
def axis(L,h,extra):return np.unique(np.round(np.r_[np.linspace(-L/2,L/2,round(L/h)+1),extra],12))
def model(c,layout='even',h=.005,poor=False,convection=None):
 if any(not np.isfinite(c[k]) or c[k]<=0 for k in ['conductivity_W_mK','thickness_m','length_m','width_m','density_kg_m3','specific_heat_J_kgK','heater_width_m','heater_length_m','heater_heat_capacity_J_K','nominal_contact_W_m2K','poor_contact_W_m2K']):raise ValueError('Positive finite material and geometry inputs required')
 if h<=0 or c['working_half_x_m']>c['length_m']/2 or c['working_half_y_m']>c['width_m']/2:raise ValueError('Invalid mesh size or working area')
 if layout not in c['layouts'] or len(c['layouts'][layout])!=3:raise ValueError('Exactly three heaters required')
 if c['heater_length_m']>c['width_m'] or any(abs(x)+c['heater_width_m']/2>c['length_m']/2+1e-12 for x in c['layouts'][layout]):raise ValueError('Heaters must fit on the plate')
 centers=sorted(c['layouts'][layout])
 if min(np.diff(centers))<c['heater_width_m']-1e-12:raise ValueError('Heater footprints overlap')
 if any(abs(x)>c['length_m']/2 or abs(y)>c['width_m']/2 for x,y in c['thermocouples_m']):raise ValueError('Thermocouple lies outside plate')
 if any(c[k]<0 for k in ['top_convection_W_m2K','bottom_U_W_m2K','edge_convection_W_m2K','heater_back_loss_W_K']) or convection is not None and convection<0:raise ValueError('Heat losses must be nonnegative')
 hx=c['layouts'][layout];hw=c['heater_width_m'];hl=c['heater_length_m'];t=c['thickness_m']
 xs=axis(c['length_m'],h,[-c['working_half_x_m'],c['working_half_x_m'],0]+[v+o for v in hx for o in [-hw/2,hw/2]])
 ys=axis(c['width_m'],h,[-c['working_half_y_m'],c['working_half_y_m'],0,-hl/2,hl/2]);mesh=MeshTri.init_tensor(xs,ys);basis=Basis(mesh,ElementTriP1(),intorder=3);n=basis.N
 A=asm(mass,basis);K=c['conductivity_W_mK']*t*asm(laplace,basis)
 htop=c['top_convection_W_m2K'] if convection is None else convection
 K+=(htop+c['bottom_U_W_m2K'])*A
 K+=c['edge_convection_W_m2K']*t*asm(mass,FacetBasis(mesh,basis.elem,intorder=3))
 M=c['density_kg_m3']*c['specific_heat_J_kgK']*t*A
 @LinearForm
 def one(v,w):return v
 vectors=[];diagonal=[];areas=[]
 for i,x in enumerate(hx):
  elems=mesh.elements_satisfying(lambda q:(q[0]>x-hw/2-1e-12)&(q[0]<x+hw/2+1e-12)&(q[1]>-hl/2-1e-12)&(q[1]<hl/2+1e-12))
  sub=basis.with_elements(elems);area=float(sub.dx.sum());hc=c['poor_contact_W_m2K'] if poor and i==0 else c['nominal_contact_W_m2K']
  K+=hc*asm(mass,sub);v=hc*asm(one,sub);vectors.append(csr_matrix(-v.reshape(-1,1)));diagonal.append(hc*area+c['heater_back_loss_W_K']);areas.append(area)
 B=bmat([vectors],format='csr');K=bmat([[K,B],[B.T,diags(diagonal)]],format='csr');M=bmat([[M,None],[None,diags([c['heater_heat_capacity_J_K']]*3)]],format='csr')
 working=np.where((abs(mesh.p[0])<=c['working_half_x_m']+1e-12)&(abs(mesh.p[1])<=c['working_half_y_m']+1e-12))[0]
 tc=[int(np.argmin((mesh.p[0]-x)**2+(mesh.p[1]-y)**2)) for x,y in c['thermocouples_m']]
 sensor=int(np.argmin(mesh.p[0]**2+mesh.p[1]**2))
 return {'mesh':mesh,'basis':basis,'K':K,'M':M,'n':n,'working':working,'thermocouples':tc,'sensor':sensor,'areas':areas,'config':c,'layout':layout,'poor':poor,'h':h,'htop':htop}

def power_vector(m,P):
 if P<0:raise ValueError('Heating power must be nonnegative')
 f=np.zeros(m['n']+3);f[-3:]=P/3;return f

def steady(m,P):return spsolve(m['K'],power_vector(m,P))
def summary(m,theta,P):
 c=m['config'];w=theta[m['working']]+c['ambient_C'];heater=theta[-3:]+c['ambient_C'];ones=np.ones(len(theta));loss=float(ones@(m['K']@theta))
 return {'layout':m['layout'],'mesh_m':m['h'],'poor_contact':m['poor'],'top_h_W_m2K':m['htop'],'power_W':P,'minimum_working_C':float(w.min()),'mean_working_C':float(w.mean()),'maximum_working_C':float(w.max()),'surface_range_C':float(w.max()-w.min()),'maximum_heater_C':float(heater.max()),'loss_W':loss,'energy_error_relative':abs(loss-P)/max(P,1),'in_band':bool(w.min()>=c['target_C']-c['band_C'] and w.max()<=c['target_C']+c['band_C'])}

def transient(m,P,dt=2,horizon=3000,controlled=False):
 if dt<=0 or horizon<0:raise ValueError('Positive time step and nonnegative duration required')
 c=m['config'];M=m['M'];K=m['K'];step=factorized((M/dt+K).tocsc());theta=np.zeros(m['n']+3);t=0;integ=0;rows=[];energy_input=0;energy_lost=0;max_residual=0;ones=np.ones(len(theta));target=c['target_C']-c['ambient_C']
 for j in range(round(horizon/dt)+1):
  w=theta[m['working']]+c['ambient_C']
  rows.append({'time_s':j*dt,'working_min_C':float(w.min()),'working_max_C':float(w.max()),'sensor_C':float(theta[m['sensor']]+c['ambient_C']),'heater_max_C':float(theta[-3:].max()+c['ambient_C']),'power_W':float(P),'stored_energy_J':float(ones@(M@theta))})
  if j==round(horizon/dt):break
  if controlled:
   error=target-theta[m['sensor']];control=c['controller'];raw=control['Kp_W_K']*error+integ;P=float(np.clip(raw,0,control['maximum_power_W']))
   # Conditional integration prevents the saturated heater from accumulating windup.
   if 0<raw<control['maximum_power_W'] or raw>=control['maximum_power_W'] and error<0 or raw<=0 and error>0:integ+=control['Ki_W_Ks']*error*dt
  f=power_vector(m,P);new=step(M@theta/dt+f);res=M@(new-theta)/dt+K@new-f;max_residual=max(max_residual,float(np.linalg.norm(res)/max(np.linalg.norm(f),1)))
  energy_input+=P*dt;energy_lost+=float(ones@(K@new))*dt;theta=new
 stored=float(ones@(M@theta));return rows,theta,{'input_J':energy_input,'loss_J':energy_lost,'stored_J':stored,'relative_energy_error':abs(energy_input-energy_lost-stored)/max(energy_input,1),'maximum_step_residual':max_residual}