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

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

import mpi4py.MPI as mpi

from shutil import move

from netket.stats import (
    statistics as _statistics,
    mean as _mean,
    sum_inplace as _sum_inplace,
)

import netket.custom.fermionic_hilbert as fermhi

import gpw.utils as utls

import numpy_ml as ml

import netket.custom.hubbard as hub

from netket.custom.fermionic_hopping_sampler import FermionicHoppingSampler,FermionicHoppingKernel

from netket.custom.ab_initio_ham import AbInitio

from pyscf import scf, gto, ao2mo, fci

import scipy.optimize as spo

from netket.custom.exact_opt import ExactOpt

import netket.custom.utils as symmetry_pruning


N = int(sys.argv[1])
U = float(sys.argv[2])
basis_rep = int(sys.argv[3]) # 0: local basis, 1: canonical basis
symmetric = bool(int(sys.argv[4]))

Lx = 3
Ly = 4


ha = hub.Hubbard(Lx=Lx, Ly=Ly, n_elec=((Lx*Ly)//2,(Lx*Ly)//2), U=U, pbx=-1, pby=-1)
hi = ha.hilbert

if basis_rep == 1:
    hopping_prefactors = ha._operators[0]._prefactors
    hopping_sites = ha._operators[0]._sites
    hopping_spins = ha._operators[0]._spins

    h1 = np.zeros((Lx*Ly, Lx*Ly))

    for i in range(len(hopping_prefactors)):
        if hopping_spins[i][0] == 0:
            h1[hopping_sites[i][0][0], hopping_sites[i][0][1]] = hopping_prefactors[i].real


    h2 = np.zeros((Lx*Ly, Lx*Ly, Lx*Ly, Lx*Ly))

    np.fill_diagonal(h2, U)

    transformation_mat = np.linalg.eigh(h1)[1]

    h1_canonical = transformation_mat.T.dot(h1.dot(transformation_mat))
    h2_canonical = np.einsum("ijkl,ia->ajkl",h2, transformation_mat)
    h2_canonical = np.einsum("ajkl,jb->abkl",h2_canonical, transformation_mat)
    h2_canonical = np.einsum("abkl,kc->abcl",h2_canonical, transformation_mat)
    h2_canonical = np.einsum("abcl,ld->abcd",h2_canonical, transformation_mat)
    h1 = h1_canonical
    h2 = h2_canonical

    norb=Lx*Ly
    nelec=Lx*Ly
    hi = fermhi.Fermions(norb, n_elec=(nelec//2,nelec//2))
    ha = AbInitio(hi, h1, h2)

basis = hi.all_states()

if symmetric:
    transl = nk.custom.get_symms_square_lattice(Lx, Ly, translations=True, point_group=True, pbx=-1, pby=-1)
    ma = nk.machine.qgps.QGPSFermionicSumSym(hi, n_bond=N, automorphisms=transl, spin_flip_sym=True, dtype=complex)
    basis, _, weightings = symmetry_pruning.get_symmetric_inequivalent_set(ma, basis, np.ones(basis.shape[0]))
else:
    ma = nk.machine.qgps.QGPSFermionicProdSym(hi, n_bond=N, automorphisms=None, spin_flip_sym=False, dtype=complex)
    weightings = None


ma._exp_kern_representation = False
ma._fast_update = True
ma.init_random_parameters(sigma=0.1, start_from_uniform=False)

op = nk.optimizer.Sgd(ma, learning_rate=0.02)

sr = nk.optimizer.SR(ma, lsq_solver="cg")

basis = np.array_split(basis, _n_nodes)[_rank]

if weightings is not None:
    weightings = np.array_split(weightings, _n_nodes)[_rank]

gs = ExactOpt(ha, basis, ma, optimizer=op, sr=sr, weightings=weightings)

if mpi.COMM_WORLD.Get_rank() == 0:
    with open("out.txt", "w") as fl:
        fl.write("")
    np.save("parameters.npy", ma.parameters)

for it in gs.iter(2000,1):
    bias = _MPI_comm.allreduce(ma.log_val(basis).real.max(), op=mpi.MAX)
    ma.bias -= bias
    if mpi.COMM_WORLD.Get_rank() == 0:
        move("parameters.npy", "parameters_old.npy")
        np.save("parameters.npy", ma.parameters)
        print(it, gs.energy.mean, flush=True)
        with open("out.txt", "a") as fl:
            fl.write("{}  {}  {}\n".format(np.real(gs.energy.mean), np.imag(gs.energy.mean), gs.energy.error_of_mean))
import numpy as np
import netket as nk
import sys
import scipy.optimize as spo
import netket.custom.utils as utls

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

import mpi4py.MPI as mpi

from shutil import move

from netket.stats import (
    statistics as _statistics,
    mean as _mean,
    sum_inplace as _sum_inplace,
)

import netket.custom.fermionic_hilbert as fermhi

import gpw.utils as utls

import numpy_ml as ml

import netket.custom.hubbard as hub

from netket.custom.fermionic_hopping_sampler import FermionicHoppingSampler,FermionicHoppingKernel

from netket.custom.ab_initio_ham import AbInitio

from pyscf import scf, gto, ao2mo, fci

import scipy.optimize as spo

from netket.custom.exact_opt import ExactOpt

import netket.custom.utils as symmetry_pruning


N = int(sys.argv[1])
U = float(sys.argv[2])
basis_rep = int(sys.argv[3]) # 0: local basis, 1: canonical basis
symmetric = bool(int(sys.argv[4]))

Lx = 3
Ly = 4


ha = hub.Hubbard(Lx=Lx, Ly=Ly, n_elec=((Lx*Ly)//2,(Lx*Ly)//2), U=U, pbx=-1, pby=-1)
hi = ha.hilbert

if basis_rep == 1:
    hopping_prefactors = ha._operators[0]._prefactors
    hopping_sites = ha._operators[0]._sites
    hopping_spins = ha._operators[0]._spins

    h1 = np.zeros((Lx*Ly, Lx*Ly))

    for i in range(len(hopping_prefactors)):
        if hopping_spins[i][0] == 0:
            h1[hopping_sites[i][0][0], hopping_sites[i][0][1]] = hopping_prefactors[i].real


    h2 = np.zeros((Lx*Ly, Lx*Ly, Lx*Ly, Lx*Ly))

    np.fill_diagonal(h2, U)

    transformation_mat = np.linalg.eigh(h1)[1]

    h1_canonical = transformation_mat.T.dot(h1.dot(transformation_mat))
    h2_canonical = np.einsum("ijkl,ia->ajkl",h2, transformation_mat)
    h2_canonical = np.einsum("ajkl,jb->abkl",h2_canonical, transformation_mat)
    h2_canonical = np.einsum("abkl,kc->abcl",h2_canonical, transformation_mat)
    h2_canonical = np.einsum("abcl,ld->abcd",h2_canonical, transformation_mat)
    h1 = h1_canonical
    h2 = h2_canonical

    norb=Lx*Ly
    nelec=Lx*Ly
    hi = fermhi.Fermions(norb, n_elec=(nelec//2,nelec//2))
    ha = AbInitio(hi, h1, h2)

basis = hi.all_states()

if symmetric:
    transl = nk.custom.get_symms_square_lattice(Lx, Ly, translations=True, point_group=True, pbx=-1, pby=-1)
    ma = nk.machine.qgps.QGPSFermionicSumSym(hi, n_bond=N, automorphisms=transl, spin_flip_sym=True, dtype=complex)
    basis, _, weightings = symmetry_pruning.get_symmetric_inequivalent_set(ma, basis, np.ones(basis.shape[0]))
else:
    ma = nk.machine.qgps.QGPSFermionicProdSym(hi, n_bond=N, automorphisms=None, spin_flip_sym=False, dtype=complex)
    weightings = None


ma._exp_kern_representation = False
ma._fast_update = True
ma.init_random_parameters(sigma=0.1, start_from_uniform=False)

op = nk.optimizer.Sgd(ma, learning_rate=0.02)

sr = nk.optimizer.SR(ma, lsq_solver="cg")

basis = np.array_split(basis, _n_nodes)[_rank]

if weightings is not None:
    weightings = np.array_split(weightings, _n_nodes)[_rank]

gs = ExactOpt(ha, basis, ma, optimizer=op, sr=sr, weightings=weightings)

if mpi.COMM_WORLD.Get_rank() == 0:
    with open("out.txt", "w") as fl:
        fl.write("")
    np.save("parameters.npy", ma.parameters)

for it in gs.iter(2000,1):
    bias = _MPI_comm.allreduce(ma.log_val(basis).real.max(), op=mpi.MAX)
    ma.bias -= bias
    if mpi.COMM_WORLD.Get_rank() == 0:
        move("parameters.npy", "parameters_old.npy")
        np.save("parameters.npy", ma.parameters)
        print(it, gs.energy.mean, flush=True)
        with open("out.txt", "a") as fl:
            fl.write("{}  {}  {}\n".format(np.real(gs.energy.mean), np.imag(gs.energy.mean), gs.energy.error_of_mean))
# 0: local basis, 1: canonical basis