import magbox
import numpy as np
import os
import pytest

testdata = [ # N, K, J, K_hard, fix_step, dh
    (256,1.0,0.5,1, True, 5e-2),
    (512,1.0,0.5,2, True, 5e-2),
    (1024,1.0,0.5,3, True, 5e-2),
    (256,1.0,1.0,4, True, 5e-2),
    (256,1.0,2.0,5, True, 5e-2),
    (256,1.0,3.0,5, True, 5e-2),

    (256,1.0,0.5,1, False, 5e-2),
    (512,1.0,0.5,2, False, 5e-2),
    (1024,1.0,0.5,3, False, 5e-2),
    (256,1.0,1.0,4, False, 5e-2),
    (256,1.0,2.0,5, False, 5e-2),
    (256,1.0,3.0,5, False, 5e-2),
]

def plot_fun(t,x,y,ft_abs,q,w,dispersion,dispersion_theory,err,mean_err,max_err):
    if "PYTEST_CURRENT_TEST" in os.environ: 
        print("测试模式不作图")
        return
    import matplotlib.pyplot as plt
    w_max=dispersion_theory.max()
    plt.figure(figsize=(10,10))

    plt.subplot(2,2,1)
    plt.imshow(ft_abs.T,
            origin='lower',
            extent=(q[0],q[-1],w[0],w[-1]),
    )
    plt.ylim(0,2*w_max)

    plt.subplot(2,2,2)
    plt.scatter(q,dispersion,label="Simulation")
    plt.plot(q,dispersion_theory,label="Theory")
    plt.legend()

    plt.subplot(2,2,3)
    plt.scatter(q,err,label="Error")
    plt.legend()
    plt.title(f"mean error: {mean_err:.2e}, max err: {max_err:.2e}")

    plt.subplot(2,2,4)
    plt.plot(t,x[:,0],label="x")
    plt.plot(t,y[:,0],label="y")
    plt.legend()
    # plt.title(f"mean error: {mean_err:.2e}, max err: {max_err:.2e}")
    plt.show()

@pytest.mark.parametrize("N,K,J,K_hard, fix_step, dh",testdata)
def test_spin_chain(N,K,J,K_hard,fix_step,dh):
    # N=256
    # K=1.0
    # J=0.5
    # dt=1
    # T=50
    rng=np.random.default_rng()
    theta0=np.ones(N)*0.01
    phi0=rng.random(N)*2*np.pi
    z=np.cos(theta0)
    x=np.sin(theta0)*np.cos(phi0)
    y=np.sin(theta0)*np.sin(phi0)
    # phi0=10*np.arange(N)/N *2*np.pi
    hard_kernel = K_hard * np.array([[0,0,0],[0,1,0],[0,0,0]])
    def hard_axis_heff(kernel,t, cartS):
        return -cartS @ kernel[0]
    LT = magbox.Lattice(type="square", size=[N], periodic=True)
    vars = magbox.Vars(K1=K, J=J, custom_kernel=(hard_kernel,), custom_heff=hard_axis_heff)

    def dispersion_fun(qf):
        return np.sqrt((K+J*(1-np.cos(qf))*np.cos(np.mean(theta0))+K_hard)*(K+J*(1-np.cos(qf))*np.cos(np.mean(theta0))))
    q=np.fft.fftfreq(N,1)*2*np.pi
    q=np.fft.fftshift(q)
    dispersion_theory=dispersion_fun(q)
    w_max=dispersion_theory.max()
    W_diff=np.min(np.abs(np.diff(dispersion_theory)))
    print(f"freq max: {w_max:.3e}, freq diff: {W_diff:.3e}")
    dt=np.max([2*np.pi/(4*w_max),0.05])
    T=np.min([int(3*2*np.pi/W_diff),2e4])
    print(f"use dt: {dt:.3e}, Total Time: {T:.3e}")

    sf=magbox.llg3(x,y,z,LT,vars,
                   device='cpu', dtype='f32', alpha=0,T=T,dt=dt, fix_step=fix_step, dh=dh)

    t_tc,S,stats,err_info=sf.run()
    print(S.shape)

    t=t_tc.cpu().detach().numpy()
    S=S.reshape(len(t),N,3)
    

    x=S[...,0].detach().cpu().numpy().squeeze()
    y=S[...,1].detach().cpu().numpy().squeeze()
    z=S[...,2].detach().cpu().numpy().squeeze()

    # x=S[:,::3,:].detach().cpu().numpy().squeeze()
    # y=S[:,1::3,:].detach().cpu().numpy().squeeze()
    # z=S[:,2::3,:].detach().cpu().numpy().squeeze()

    u=x+1j*y
    ft=np.fft.fft2(u.T)
    ft_abs=np.abs(ft)
    w=np.fft.fftfreq(len(t), dt)*2*np.pi
    

    ft_abs=np.fft.fftshift(ft_abs)
    w=np.fft.fftshift(w)
    

    dispersion=np.zeros(len(q))
    for idx in range(len(q)):
        arg_max=np.argmax(ft_abs[idx,:])
        dispersion[idx]=w[arg_max]
    dispersion=np.abs(dispersion)

    err=dispersion/dispersion_theory-1
    max_err=np.max(np.abs(err))
    mean_err=np.mean(np.abs(err))

    print(f"mean error: {mean_err:.2e}, max err: {max_err:.2e}")

    plot_fun(t,x,y,ft_abs,q,w,dispersion,dispersion_theory,err,mean_err,max_err)

    assert mean_err<1e-2

if __name__=="__main__":
    test_spin_chain(32,1,1,0, True, 1e-1)


# spin_chain_test(256,1,0.5,1,50)
