import numpy as np
import netket as nk
import sys
import scipy.optimize as spo
import netket.custom.utils as utls

from quspin.operators import hamiltonian # operators


from quspin.basis import spin_basis_general # spin basis constructor


import scipy as sp

from netket.operator import local_values as _local_values


from netket.utils import (
    MPI_comm as _MPI_comm,
    n_nodes as _n_nodes,
    node_number as _rank
)

from netket.stats import sum_inplace as _sum_inplace

# 1D Lattice
Lx = 4

Ly=6


J1=1.
J2 = 0.5

ha = nk.custom.J1J2(Lx=Lx,Ly=Ly, J2=J2, msr=False)

hi = ha.hilbert

transl = nk.custom.get_symms_square_lattice(Lx, Ly, translations=True, point_group=True)

ma = nk.machine.QGPSProdSym(hi, n_bond=3, automorphisms=transl, spin_flip_sym=True, cluster_ids=None, dtype=complex)
ma._exp_kern_representation = False
ma._fast_update = False


ma.init_random_parameters(sigma=0.5, start_from_uniform=True, small_arg=False)

N_2d = Lx*Ly
s = np.arange(N_2d)
x = s%Lx
y = s//Lx
T_x = (x+1)%Lx + Lx*y
T_x_minus = (x-1)%Lx + Lx*y
T_y = x +Lx*((y+1)%Ly)
P_x = x + Lx*(Ly-y-1)
P_y = (Lx-x-1) + Lx*y
Z   = -(s+1)
basis_2d = spin_basis_general(N_2d, N_2d//2, pauli=0, kxblock=(T_x,0),kyblock=(T_y,0),pxblock=(P_x,0),pyblock=(P_y,0),zblock=(Z,0))


J1term = [[J1,i,T_x[i]] for i in range(N_2d)]
J1term += [[J1,i,T_y[i]] for i in range(N_2d)]
J2term = [[J2,i,T_y[T_x[i]]] for i in range(N_2d)]
J2term += [[J2,i,T_y[T_x_minus[i]]] for i in range(N_2d)]
static=[["xx",J1term],["yy",J1term],["zz",J1term],["xx",J2term],["yy",J2term],["zz",J2term]]
H=hamiltonian(static,[],basis=basis_2d,dtype=np.float64)


groundstate = H.eigsh(k=1,which="SA",return_eigenvectors=True)

print(groundstate[0])

basis = basis_2d.states
amplitudes = groundstate[1][:,0]

trn_basis = []
trn_amplitudes = []

for (i,k) in enumerate(basis):
    trn_basis.append(2 * np.fromstring(basis_2d.int_to_state(k)[1:-1], sep=" ") - 1)
    trn_amplitudes.append(amplitudes[i] * basis_2d.get_amp(k))

probs = np.array([1/(abs(basis_2d.get_amp(k))**2) for k in basis])
probs /= probs.sum()

trn_basis = np.array(trn_basis)
trn_amplitudes = np.array(trn_amplitudes)
trn_weightings = probs



alpha_guess = 1.0
learning = nk.custom.QGPSLearningExp(ma, init_noise_tilde=1.0, init_alpha=alpha_guess, complex_expand=True)


scale_factor = np.exp(learning.get_bias(trn_amplitudes, weightings=trn_weightings))
trn_amplitudes_scaled = trn_amplitudes/scale_factor

total_no_samples = 5000
rng = np.random.default_rng(1234)
reduced_inds = rng.choice(len(trn_amplitudes_scaled), size=total_no_samples, p = probs)

start_id = _rank * (total_no_samples//_n_nodes + 1)
end_id = min(start_id + total_no_samples//_n_nodes + 1, total_no_samples)
local_inds = reduced_inds[start_id:end_id]


start_id_complete = _rank * (len(trn_amplitudes)//_n_nodes + 1)
end_id_complete = min(start_id + len(trn_amplitudes)//_n_nodes + 1, len(trn_amplitudes))

trn_basis_small = trn_basis[local_inds]
trn_amplitudes_small = trn_amplitudes_scaled[local_inds]

np.save("trn_set_{}.npy".format(_rank), trn_basis_small)

j = 0
log_ml = 1.0
log_ml_old = 1.0
mean_squared_error_full = []
mean_squared_error_train = []
energies= []
while(j < 5000):
    exp_reset = 0
    print("<", flush=True)
    full_pred = np.zeros(len(trn_basis))
    full_pred[start_id_complete:end_id_complete] = np.exp(ma.log_val(trn_basis[start_id_complete:end_id_complete, :]))
    full_pred = _sum_inplace(full_pred)
    full_amps = full_pred*np.sqrt(probs)
    full_amps /= (abs(full_amps)).sum()
    full_amps /= np.sqrt((abs(full_amps)**2).sum())
    energies.append(H.expt_value(full_amps))
    mean_squared_error_full.append(np.sum(abs(full_pred - trn_amplitudes_scaled)**2*trn_weightings)/np.sum(trn_weightings))
    print(">", flush=True)
    mean_squared_error_train.append(learning.mean_squared_error(trn_basis_small, trn_amplitudes_small, np.ones(len(trn_amplitudes_small)))/(_n_nodes * len(trn_amplitudes_small)))
    for i in range(Lx*Ly):
        if learning.ref_site is not None:
            for n in range(ma._epsilon.shape[1]):
                if np.sum(np.abs(ma._epsilon[learning.ref_site, n, :])) < 1.e-15:
                    exp_reset += 1
                    ma._epsilon[learning.ref_site, n, :] = np.random.normal(loc=1.0, scale=0.2, size=2) + 1.j * np.random.normal(scale=0.2, size=2)
                    learning.alpha_mat[learning.ref_site, 2 * n] = alpha_guess
                    learning.alpha_mat[learning.ref_site, 2 * n + 1] = alpha_guess
        learning.fit_step(trn_basis_small, trn_amplitudes_small, i, opt_alpha=True, opt_noise=True, rvm=True, alpha_iterations=10, noise_iterations=1)
        ma._opt_params = ma._epsilon[ma._der_ids >= 0].copy()
        learning.alpha_mat[i, ~learning.active_elements] = 1.0
    j += 1
    print(mean_squared_error_full[-1], mean_squared_error_train[-1], energies[-1], learning.noise_tilde, np.sum(learning.alpha_mat), exp_reset)

