"""Solver A: P1 finite elements on a criss-cross mesh aligned with the half-square tiling."""
import numpy as np, scipy.sparse as sp, scipy.sparse.linalg as spla
from geom import inside, on_boundary
def mesh(P,N):
    h=1.0/N; xs=np.arange(0,3*N+1)*h
    X,Y=np.meshgrid(xs,xs,indexing='ij'); corners=np.c_[X.ravel(),Y.ravel()]
    cx=(np.arange(3*N)+0.5)*h; CX,CY=np.meshgrid(cx,cx,indexing='ij'); cent=np.c_[CX.ravel(),CY.ravel()]
    pts=np.vstack([corners,cent]); nc=len(corners); M=3*N+1
    I,J=np.meshgrid(np.arange(3*N),np.arange(3*N),indexing='ij'); I=I.ravel();J=J.ravel()
    a=I*M+J; b=(I+1)*M+J; c=(I+1)*M+J+1; d=I*M+J+1; e=nc+I*(3*N)+J
    T=np.vstack([np.c_[a,b,e],np.c_[b,c,e],np.c_[c,d,e],np.c_[d,a,e]])
    T=T[inside(P,pts[T].mean(1))]
    used=np.unique(T); remap=-np.ones(len(pts),int); remap[used]=np.arange(len(used))
    return pts[used],remap[T]
def assemble(p,T):
    v=p[T]; d1=v[:,1]-v[:,0]; d2=v[:,2]-v[:,0]; A=0.5*np.abs(d1[:,0]*d2[:,1]-d1[:,1]*d2[:,0])
    # gradients of barycentric functions
    G=np.empty((len(T),3,2))
    for k in range(3):
        i,j=(k+1)%3,(k+2)%3; e=v[:,j]-v[:,i]; G[:,k,0]=-e[:,1]; G[:,k,1]=e[:,0]
    G/= (2*A)[:,None,None]
    # orientation-free: sign of G doesn't matter in G G^T only if consistent; fix by sign of cross
    s=np.sign(d1[:,0]*d2[:,1]-d1[:,1]*d2[:,0]); G*=s[:,None,None]
    Ke=np.einsum('tik,tjk->tij',G,G)*A[:,None,None]
    Me=(np.ones((3,3))+np.eye(3))[None]*(A/12)[:,None,None]
    r=np.repeat(T,3,1).ravel(); c=np.tile(T,(1,3)).ravel(); n=len(p)
    return sp.csr_matrix((Ke.ravel(),(r,c)),(n,n)),sp.csr_matrix((Me.ravel(),(r,c)),(n,n))
def solve(P,N,k=50,bc='D',vecs=False):
    p,T=mesh(P,N); K,M=assemble(p,T)
    if bc=='D':
        free=np.where(~on_boundary(P,p))[0]; K=K[free][:,free]; M=M[free][:,free]
    else: free=np.arange(len(p))
    w,V=spla.eigsh(K,k,M,sigma=-0.1 if bc=='N' else 0.0,which='LM')
    o=np.argsort(w); w=w[o]; V=V[:,o]
    return (w,V,p,T,free) if vecs else w