"""Exploratory finite-band Hamiltonian with a quadratic dark-subspace constraint.

No identification with a completed microscopic theory of gravity is asserted.
"""
from __future__ import annotations
import numpy as np
from scipy.special import roots_legendre


def matrices(k,x,m,omega,wilson=0.):
    k,x,wilson=np.broadcast_arrays(k,x,wilson)
    shape=k.shape+(4,4)
    h=np.zeros(shape,dtype=complex)
    r=1/np.sqrt(1+x*x)
    f=x*r
    h[...,0,0]=k
    h[...,1,1]=k+omega
    h[...,2,2]=-k
    h[...,3,3]=-k+omega*(1+x*x)
    for i,j,value in [(0,2,m*f-1j*wilson*r),(0,3,m*r+1j*wilson*f),
                      (1,2,m*r-1j*wilson*f),(1,3,-m*f-1j*wilson*r)]:
        h[...,i,j]=value
        h[...,j,i]=np.conj(value)
    derivative=np.zeros(shape)
    derivative[...,2,3]=-omega
    derivative[...,3,2]=-omega
    derivative[...,3,3]=2*omega*x
    second=np.zeros(shape)
    second[...,2,2]=2*omega*r*r
    second[...,2,3]=-2*omega*f*r
    second[...,3,2]=-2*omega*f*r
    second[...,3,3]=2*omega*f*f
    return h,derivative,second


def spectral(k,x,m,omega,wilson=0.):
    h,first,second=matrices(k,x,m,omega,wilson)
    values,vectors=np.linalg.eigh(h)
    ground=vectors[...,0]
    derivative=np.einsum("...i,...ij,...j->...",ground.conj(),first,ground).real
    contact=np.einsum("...i,...ij,...j->...",ground.conj(),second,ground).real
    transition=np.einsum("...ij,...jk,...k->...i",vectors.conj().swapaxes(-1,-2),first,ground)[...,1:]
    gaps=values[...,1:]-values[...,:1]
    response=2*np.sum(np.abs(transition)**2/gaps,axis=-1)
    fidelity=np.sum(np.abs(transition)**2/gaps**2,axis=-1)
    return values,derivative,contact-response,contact,response,fidelity


def nodes(order=64):
    z,w=roots_legendre(order)
    return (z+1)/2,w/2


def band(x,m=1.,omega=10.,order=64):
    t,w=nodes(order)
    k=np.concatenate((t*t,-t*t))
    weight=np.tile(w*t**3,2)
    values,d1,d2,contact,response,fidelity=spectral(k,np.asarray(x)[...,None],m,omega)
    quantities=[np.sum(a*weight,axis=-1) for a in (values[...,0],d1,d2,contact,response,fidelity)]
    return quantities


def mond(x,m=1.,omega=10.,order=64):
    a2=-band(0.,m,omega,order)[2]/2
    u,w=nodes(order)
    energy,first,second,*_=band(np.asarray(x)[...,None]*u,m,omega,order)
    mu=1+3*np.sum(first*u*w,axis=-1)/(2*a2*np.asarray(x))
    radial=1+3*np.sum(second*u*u*w,axis=-1)/(2*a2)
    return mu,radial


def lattice(x,m=1.,omega=20.,wilson=.5,order=64):
    t,w=nodes(order)
    # q=pi*t² resolves the massless Dirac cone near the origin.
    q=np.pi*t*t
    weight=(2*t*w)[:,None]*(2*t*w)[None,:]
    k=np.sqrt(np.sin(q[:,None])**2+np.sin(q[None,:])**2)
    wterm=wilson*(2-np.cos(q[:,None])-np.cos(q[None,:]))
    ks=np.concatenate((k.ravel(),-k.ravel()))
    ws=np.tile(wterm.ravel(),2)
    weights=np.tile(weight.ravel()/2,2)
    result=spectral(ks,np.asarray(x)[...,None],m,omega,ws)
    return [np.sum(a*weights,axis=-1) for a in (result[0][...,0],*result[1:])]


def lattice_polar(x,m=1.,omega=20.,wilson=.5,order=64,angle_order=24):
    t,wt=nodes(order)
    angle,wa=nodes(angle_order)
    angle=angle*np.pi/4
    wa=wa*np.pi/4
    rmax=np.pi/np.cos(angle)
    radius=rmax[:,None]*t[None,:]**2
    q1=radius*np.cos(angle[:,None])
    q2=radius*np.sin(angle[:,None])
    k=np.sqrt(np.sin(q1)**2+np.sin(q2)**2)
    wterm=wilson*(2-np.cos(q1)-np.cos(q2))
    weight=2/np.pi**2*(rmax*rmax*wa)[:,None]*(t**3*wt)[None,:]
    ks=np.r_[k.ravel(),-k.ravel()]
    ws=np.tile(wterm.ravel(),2)
    weights=np.tile(weight.ravel(),2)
    result=spectral(ks,np.asarray(x)[...,None],m,omega,ws)
    return [np.sum(a*weights,axis=-1) for a in (result[0][...,0],*result[1:])]


if __name__=="__main__":
    for omega in [5.,10.,30.,100.]:
        m=1.
        e0,_,second,*_=band(0.,m,omega,128)
        a2=-second/2
        v0=omega/np.sqrt(omega*omega+4*m*m)
        coefficient=3*v0*m**3/(8*a2)
        print("parameters",omega,"a2",a2,"v0",v0,"deep_mu_slope",coefficient,flush=True)
        for x in [.001,.003,.01]:
            value=band(x,m,omega,128)[0]-e0
            print("cube",x,(value+a2*x*x)/x**3,"target",v0*m**3/3,flush=True)
        xs=np.array([.01,.1,.3,1.,3.,10.,30.])
        mus,radials=mond(xs,m,omega,96)
        print("response",list(zip(xs,mus,radials)),flush=True)
    values,*_=spectral(np.array([.2,-.2]),np.array([.1,1.,3.,10.])[:,None],0.,10.)
    print("zero_M_ground",values[...,0],flush=True)
