import numpy as np
import matplotlib.pyplot as plt
from gpaw.lcao.overlap import (FourierTransformer, TwoSiteOverlapCalculator,
                               ManySiteOverlapCalculator,
                               DerivativeAtomicDisplacement)
from gpaw.lcao.tci import get_lvalues, TCICalculator, TCIExpansions
from gpaw.spline import Spline

rcmax = 5.0
r_g = np.linspace(0, rcmax, 100)
f_g = r_g * np.exp(-r_g**2)

l = 0
fspline = Spline(l, rmax=rcmax, f_g=f_g)
assert np.allclose(f_g, fspline.map(r_g))

fspline_Ij = [[fspline]]
#l_Ij = [[l]]

atom_indices_a = np.zeros(2, int)


# Pass R(r) * dV(r)/dr as phit_Ij and R(r) as pt_Ij.
# All these things must be splines.
# "atom indices" or "I" is the ID of each atom: The lists fspline_Ij
# have length I, meaning we must have an element for each /kind/ of atom.
# Then atom_indices_a must specify, for each /actual/ atom, the I value
# identifying its /kind/.
#
# For example if all atoms are Au, then I_a=[0, 0, 0, 0].
# If there is also Ag, and spline_Ij[0] are the Au splines, then
# spline_Ij[1] would be for Ag and we would have
#   = I_a[0 or 1 for each atom, as appropriate for Au/Ag]
exps = TCIExpansions(phit_Ij=fspline_Ij, pt_Ij=fspline_Ij, I_a=atom_indices_a)


# Cell and positions can come from Atoms (remember units are Bohr
# here).  ibzk_qc can come from the kpoint descriptor (I think).
tcicalc = exps.get_tci_calculator(
    cell_cv=np.diag([3., 4., 5.]),
    spos_ac=[[0., 0., 0.], [.1, .1, .1]],
    pbc_c=[1, 1, 1],
    ibzk_qc=np.zeros((1, 3)),
    dtype=float)

# This is the matrix of overlap derivatives between functions on atoms
# 0 and 1:
x = tcicalc.dPdR(0, 1)

print(x)  # .flat[0])