import numpy as np
from ase.build import bulk
from gpaw import GPAW, FermiDirac
from gpaw.response.df import DielectricFunction
from gpaw.lrtddft import LrTDDFT, photoabsorption_spectrum
from gpaw.lcaotddft import LCAOTDDFT
from gpaw.lcaotddft.dipolemomentwriter import DipoleMomentWriter


calc=GPAW('Al2x2x3.gpw')
str=calc.get_atoms()
calc=GPAW(mode='lcao', basis='dzp', xc='PBE', kpts=(2,2,1), occupations=FermiDirac(0.1), txt='Al.log')
str.set_calculator(calc)
str.get_potential_energy()
calc.write('Al_sup.gpw')
# df = DielectricFunction(calc='Al_sup.gpw', domega0=0.05) 
# df.get_dielectric_function()
td=LCAOTDDFT('Al_sup.gpw')

DipoleMomentWriter(td, 'dm.dat')
timestep=2.0
iterations=500
ks=[0.0, 0.0, 1e-3]
td.absorption_kick(ks)
td.propagate(timestep, iterations)

photoabsorption_spectrum('dm.dat', 'sp.dat')