"""Independent SciPy sparse direct/Newton solution of the specified fixed-boundary problem.
No JavaScript results are inputs. Generated values are checked by the Node test suite.
"""
import json, pathlib
import numpy as np
import scipy
from scipy.sparse import lil_matrix, diags
from scipy.sparse.linalg import spsolve
MU0=1.25663706127e-6

def solve(n,mode):
    R=np.linspace(1.5,2.5,n); Z=np.linspace(-.5,.5,n); h=1/(n-1)
    m=(n-2)**2; A=lil_matrix((m,m)); rr=np.zeros(m); zz=np.zeros(m)
    for j in range(1,n-1):
        for i in range(1,n-1):
            k=(i-1)+(n-2)*(j-1); r=R[i]; z=Z[j]; rr[k]=r; zz[k]=z
            A[k,k]=-4/h**2
            for ii,jj,c in [(i+1,j,1/h**2-1/(2*r*h)),(i-1,j,1/h**2+1/(2*r*h)),(i,j+1,1/h**2),(i,j-1,1/h**2)]:
                if 0<ii<n-1 and 0<jj<n-1: A[k,(ii-1)+(n-2)*(jj-1)]=c
    A=A.tocsr(); exact=.1*np.sin(np.pi*(rr-1.5))*np.sin(np.pi*(zz+.5))
    if mode=='manufactured':
        b=-2*np.pi**2*exact-.1*np.pi*np.cos(np.pi*(rr-1.5))*np.sin(np.pi*(zz+.5))/rr
        p=spsolve(A,b); iterations=1
    else:
        p=np.zeros(m)
        for iterations in range(1,20):
            source=-MU0*rr**2*(-2000+20*p-6*p**2)-.005*p
            deriv=-MU0*rr**2*(20-12*p)-.005
            res=A@p-source
            if np.max(np.abs(res))<1e-13: break
            p+=spsolve(A-diags(deriv),-res)
        b=-MU0*rr**2*(-2000+20*p-6*p**2)-.005*p
    full=np.zeros((n,n)); full[1:-1,1:-1]=p.reshape((n-2,n-2)); axis=np.unravel_index(np.argmax(np.abs(full)),full.shape)
    return dict(mode=mode,n=n,maximumAbsResidual_T=float(np.max(np.abs(A@p-b))),iterations=iterations,psi_WbPerRad=full.ravel().tolist(),axis_m=[float(R[axis[1]]),float(Z[axis[0]])],maximumAbsManufacturedError_WbPerRad=float(np.max(np.abs(p-exact))) if mode=='manufactured' else None)

if __name__=='__main__':
    result={'schema':'hbf-gs-independent-reference@1','generator':'Python SciPy sparse direct solve / Newton with analytic Jacobian; independent implementation','numpy':np.__version__,'scipy':scipy.__version__,'coordinates':'R-fastest, Z up, SI, psi in Wb/rad','domain':{'R_m':[1.5,2.5],'Z_m':[-.5,.5]},'boundaryPsi_WbPerRad':0,'pressurePolynomial_Pa':[0,-2000,10,-2],'F2Polynomial_T2m2':[16,0,.005],'cases':[solve(n,mode) for mode in ['manufactured','profile-polynomial'] for n in [17,33]]}
    path=pathlib.Path(__file__).with_name('reference.json'); path.write_text(json.dumps(result,indent=2)+'\n'); print(path)