#!/usr/bin/env python3
'''
Test 2D Ewald summation consistency for basis VSWFs, using varying values of eta.
'''
import sys
# Test parameters, rather arbitrary

npoints = 10

etafactors = [0.9, 1.1]

zpointmax = 0. if len(sys.argv) == 1 else float(sys.argv[1])
lattice_basis=[(580e-9,0),(0,580e-9)]
positions = [(0,0)] # particle positions
wavevector = [0.,0.,0.]
bgparam = 1.52
omega_eh = 2.1
lMax = 2


import numpy as np
from qpms import Particle, CTMatrix, lorentz_drude, EpsMuGenerator, TMatrixGenerator, EwaldPart, BaseSpec, FinitePointGroup, ScatteringSystem, TMatrixInterpolator, EpsMu, dbgmsg_enable, dbgmsg_disable, dbgmsg_active, BesselType,eV, hbar, c
from qpms.symmetries import point_group_info
eh = eV/hbar
np.random.seed(666)

omega = omega_eh * eh

bspec = BaseSpec(lMax=lMax)
medium = EpsMuGenerator(bgparam)
emg_dummy = EpsMuGenerator(10)
tmg_dummy = TMatrixGenerator.sphere(medium, emg_dummy, 50e-9)

particles = [Particle(pos, tmg_dummy, bspec) for pos in positions]

ss, ssw = ScatteringSystem.create(particles, medium, omega, latticebasis=lattice_basis)
wavevector = np.array(wavevector)

sswk = ssw._sswk(wavevector)
eta_orig = sswk.eta

print("default eta: %g" % (eta_orig,))
print("lattice basis:\n", ss.lattice_basis)

points =  (np.random.rand(npoints) - .5)[:,None] * ss.lattice_basis[0] + (np.random.rand(npoints) - .5)[:,None] * ss.lattice_basis[1]
points += (np.random.rand(npoints) - .5)[:,None] * np.array([0.,0.,zpointmax])


fields_reference = sswk.scattered_field_basis(points)

fields = np.empty((len(etafactors),) + fields_reference.shape, dtype=complex)

fails = 0

for n, etafactor in enumerate(etafactors):
    sswk.eta = eta_orig * etafactor
    fields[n] = sswk.scattered_field_basis(points)
    print(fields[n])
    print(fields[n]-fields_reference)
    if np.sum(1-np.isclose(fields[n], fields_reference)) > 0: 
        fails += 1

print(fails)

sys.exit(int(fails))