import numpy as np
from ase import Atoms
from gpaw import GPAW

def ptable(spacing=2.5):
    '''Generates the periodic table as an Atoms oobject to help with visualizing rendering and color palette settings.'''
    # generates column, row positions for each element
    # zmax = 118
    zmax = 86
    x, y = 1, 1  # column, row , initial coordinates for Hydrogen
    positions = np.zeros((zmax + 1, 3))
    for z in range(1, zmax + 1): # z is atomic number not, position
        if z == 2:
            x += 16
        if z == 5:
            x += 10
        if z == 13:
            x += 10
        if z == 57 or z == 89:
            y += 3
        if z == 72 or z == 104:
            y -= 3
            x -= 14
        positions[z] = (x, -y, 0)
        x += 1
        if x > 18:
            x = 1
            y += 1
    atoms = Atoms(np.arange(1, zmax + 1),
                  positions[1:] * spacing)

    return atoms

atoms = ptable()
atoms.center(vacuum=3.4)

delete_syms = {*'Tc La Ce Pr Nd Pm Sm Eu Gd Tb Dy Ho Er Tm Yb Lu Po At'.split()}
mask = [sym not in delete_syms for sym in atoms.symbols]
atoms = atoms[mask]


setups_txt = """
Ag.11.LDA.gz
Ir.9.LDA.gz
Mg.2.LDA.gz
Mn.7.LDA.gz
Mo.6.LDA.gz
Na.1.LDA.gz
Nb.5.LDA.gz
Ni.10.LDA.gz
Os.8.LDA.gz
Pd.10.LDA.gz
Pt.10.LDA.gz
Rh.9.LDA.gz
Ru.8.LDA.gz
Ta.5.LDA.gz
Te.16.LDA.gz
V.5.LDA.gz
W.6.LDA.gz
"""

setups = {}
for line in setups_txt.split():
    tokens = line.split('.')
    setups[tokens[0]] = tokens[1]

# atoms.set_initial_magnetic_moments([2.0] * len(atoms))
calc = GPAW(mode='lcao', basis='sz(dzp)',
            nbands=442,
            #nbands='200%',
            txt='gpaw.txt',
            setups=setups,
            h=0.2,
            mixer={'beta': 0.01, 'nmaxold': 8, 'weight': 200.0},
            occupations={'name': 'fermi-dirac', 'width': 0.1})


atoms.calc = calc
atoms.get_potential_energy()