"""Independent real-valued reference solver and diagnostics."""
import numpy as np
from scipy.integrate import solve_ivp


def spectral_laplacian(fields,length):
    n=fields.shape[-1]
    k=2*np.pi*np.fft.fftfreq(n,d=length/n)
    return np.fft.ifft(-k*k*np.fft.fft(fields,axis=-1),axis=-1).real


def spectral_derivative(fields,length):
    n=fields.shape[-1]
    k=2*np.pi*np.fft.fftfreq(n,d=length/n)
    return np.fft.ifft(1j*k*np.fft.fft(fields,axis=-1),axis=-1).real


def reference_solve(model,displacement,velocity,times,rtol=2e-12,atol=2e-13,max_step_factor=.5):
    displacement=np.asarray(displacement,float)
    velocity=np.asarray(velocity,float)
    species,n=displacement.shape
    k2=(2*np.pi*np.fft.fftfreq(n,d=float(model.length)/n))**2
    def rhs(t,y):
        u=y[:species*n].reshape(species,n)
        p=y[species*n:].reshape(species,n)
        lap=np.fft.ifft(-k2*np.fft.fft(u,axis=-1),axis=-1).real
        return np.concatenate([p.ravel(),(lap-model.continuum_force(u)).ravel()])
    sol=solve_ivp(rhs,(float(times[0]),float(times[-1])),
                  np.concatenate([displacement.ravel(),velocity.ravel()]),
                  method='DOP853',t_eval=times,rtol=rtol,atol=atol,
                  max_step=max_step_factor*float(model.length)/n)
    if not sol.success:
        raise RuntimeError(sol.message)
    return sol.y[:species*n].T.reshape(len(times),species,n), sol.y[species*n:].T.reshape(len(times),species,n),dict(
        solver='scipy.integrate.solve_ivp DOP853 + Fourier collocation',
        spatial_sites=n,rtol=rtol,atol=atol,max_step=max_step_factor*float(model.length)/n,nfev=sol.nfev)


def source_reference_solve(displacement,velocity,length,eta,coupling,g,times,rtol=2e-12,atol=2e-13,max_step_factor=.5):
    """Solve the symmetric u background and its first epsilon derivative v1.

    v1_tt-v1_xx+[2 eta+(3 lambda-g)u^2/2] v1=-u^3/2.
    The initial v1 and its physical time derivative both vanish.
    """
    displacement=np.asarray(displacement,float);velocity=np.asarray(velocity,float)
    n=len(displacement)
    k2=(2*np.pi*np.fft.fftfreq(n,d=length/n))**2
    def rhs(t,y):
        u,v,w,z=y.reshape(4,n)
        lap=np.fft.ifft(-k2*np.fft.fft(np.stack([u,v]),axis=-1),axis=-1).real
        fu=.5*(coupling+g)*u**3
        fv=(2*eta+.5*(3*coupling-g)*u*u)*v+.5*u**3
        return np.stack([w,z,lap[0]-fu,lap[1]-fv]).ravel()
    y0=np.stack([displacement,np.zeros(n),velocity,np.zeros(n)]).ravel()
    sol=solve_ivp(rhs,(float(times[0]),float(times[-1])),y0,method='DOP853',
                  t_eval=times,rtol=rtol,atol=atol,max_step=max_step_factor*length/n)
    if not sol.success:raise RuntimeError(sol.message)
    history=sol.y.T.reshape(len(times),4,n)
    return history,dict(solver='DOP853 + Fourier collocation of background and first variation',
        spatial_sites=n,rtol=rtol,atol=atol,max_step=max_step_factor*length/n,nfev=sol.nfev)


def unrounded_evolve(model,displacement,velocity,laplacian,steps):
    """Independent real leapfrog baseline with the same Taylor initialization."""
    u=np.asarray(displacement,float).copy()
    tau,h=float(model.tau),float(model.h)
    prev=u-tau*np.asarray(velocity)+.5*tau*tau*(laplacian-model.continuum_force(u))
    history=[u.copy()]
    for _ in range(steps):
        lap=(np.roll(u,-1,axis=-1)-2*u+np.roll(u,1,axis=-1))/(h*h)
        nxt=2*u-prev+tau*tau*(lap-model.continuum_force(u))
        prev,u=u,nxt
        history.append(u.copy())
    return np.asarray(history)


def branch_fields(fields):
    if fields.shape[1]!=2:
        raise ValueError('two fields required')
    return (fields[:,0]+fields[:,1])/np.sqrt(2),(fields[:,0]-fields[:,1])/np.sqrt(2)


def field_energy(model,fields):
    """Continuum-type energy using centered time differences, interior times only.

    This is a diagnostic, not an exactly conserved energy of the automaton.
    """
    tau,h=float(model.tau),float(model.h)
    q=fields[1:-1]
    vel=(fields[2:]-fields[:-2])/(2*tau)
    grad=(np.roll(q,-1,axis=-1)-q)/h
    density=.5*np.sum(vel*vel+grad*grad,axis=1)
    phi=q[:,0]
    density+=.5*float(model.mass_phi_sq)*phi**2+.25*float(model.lambda_phi)*phi**4
    if model.species==2:
        chi=q[:,1]
        density+=(.5*float(model.mass_chi_sq)*chi**2+.5*float(model.eta)*(phi-chi)**2
                  +.25*float(model.lambda_chi)*chi**4+.5*float(model.g)*phi**2*chi**2)
    return h*np.sum(density,axis=-1)


def uv_fraction(fields,cutoff_fraction=.5):
    """Power at |dimensionless lattice k| > cutoff_fraction*pi."""
    power=np.abs(np.fft.fft(fields,axis=-1))**2
    k=2*np.pi*np.fft.fftfreq(fields.shape[-1])
    numerator=power[...,abs(k)>cutoff_fraction*np.pi+1e-14].sum(axis=-1)
    denominator=power.sum(axis=-1)
    return np.divide(numerator,denominator,out=np.zeros_like(numerator),where=denominator>1e-30)


def discrete_h1(error,h):
    grad=(np.roll(error,-1,axis=-1)-error)/h
    return np.sqrt(h*np.sum(error*error+grad*grad))