from __future__ import print_function

import os
import numpy as np

from ase.units import Ry, eV
from ase.calculators.siesta.siesta import Siesta
from ase.calculators.siesta.parameters import Specie, PAOBasisBlock
from ase.calculators.calculator import FileIOCalculator
from ase import Atoms

# Create temporary directory for running tests.
test_path = 'tmp_siesta'
if not os.path.exists(test_path):
    os.makedirs(test_path)
os.chdir(test_path)
run_path = 'run_directory'
pseudo_path = 'pseudos'
if not os.path.exists(pseudo_path):
    os.makedirs(pseudo_path)
if not os.path.exists(run_path):
    os.makedirs(run_path)

# Make dummy pseudopotentials.
for symbol in 'HCO':
    with open('{0}/{1}.lda.psf'.format(pseudo_path, symbol), 'w') as fd:
        fd.close()

# Change to test directory.
os.chdir(run_path)

# Setup test structures.
h = Atoms('H', [(0.0, 0.0, 0.0)])
co2 = Atoms('CO2', [(0.0, 0.0, 0.0), (-1.178, 0.0, 0.0), (1.178, 0.0, 0.0)])
ch4 = Atoms('CH4', np.array([
    [0.000000, 0.000000, 0.000000],
    [0.682793, 0.682793, 0.682793],
    [-0.682793, -0.682793, 0.682790],
    [-0.682793, 0.682793, -0.682793],
    [0.682793, -0.682793, -0.682793]]))

# Setup required environment variables.
os.environ['SIESTA_PP_PATH'] = '../' + pseudo_path

# Test the initialization.
siesta = Siesta()
assert isinstance(siesta, FileIOCalculator)
assert isinstance(siesta.implemented_properties, tuple)
assert isinstance(siesta.default_parameters, dict)
assert isinstance(siesta.name, str)
assert isinstance(siesta.default_parameters, dict)

# Test simple fdf-argument case.
atoms = h.copy()
siesta = Siesta(
    label='test_label',
    fdf_arguments={'DM.Tolerance': 1e-3})
atoms.set_calculator(siesta)
siesta.write_input(atoms, properties=['energy'])
atoms = h.copy()
atoms.set_calculator(siesta)
siesta.write_input(atoms, properties=['energy'])
with open('test_label.fdf', 'r') as f:
    lines = f.readlines()
assert 'DM.Tolerance  0.001\n' == lines[0]

# Test (slightly) more complex case of setting fdf-arguments.
siesta = Siesta(
    label='test_label',
    mesh_cutoff=3000 * eV,
    fdf_arguments={
        'DM.Tolerance': 1e-3,
        'ON.eta': 5 * Ry})
atoms.set_calculator(siesta)
siesta.write_input(atoms, properties=['energy'])
atoms = h.copy()
atoms.set_calculator(siesta)
siesta.write_input(atoms, properties=['energy'])
with open('test_label.fdf', 'r') as f:
    lines = f.readlines()
assert 'MeshCutoff  3000.0000 eV\n' in lines
assert 'DM.Tolerance  0.001\n' in lines
assert 'ON.eta  68.02846506 eV\n' in lines

# Test setting fdf-arguments after initiation.
siesta.set_fdf_arguments(
        {'DM.Tolerance': 1e-2,
         'ON.eta': 2 * Ry})
siesta.write_input(atoms, properties=['energy'])
with open('test_label.fdf', 'r') as f:
    lines = f.readlines()
assert 'MeshCutoff  3000.0000 eV\n' in lines
assert 'DM.Tolerance  0.01\n' in lines
assert 'ON.eta  27.21138602 eV\n' in lines

# Test initiation using Species.
atoms = ch4.copy()
species, numbers = siesta.species(atoms)
assert all(numbers == np.array([1, 2, 2, 2, 2]))
siesta = Siesta(species=[Specie(symbol='C', tag=1)])
species, numbers = siesta.species(atoms)
assert all(numbers == np.array([1, 2, 2, 2, 2]))
atoms.set_tags([0, 0, 0, 1, 0])
species, numbers = siesta.species(atoms)
assert all(numbers == np.array([1, 2, 2, 2, 2]))
siesta = Siesta(species=[Specie(symbol='H', tag=1, basis_set='SZ')])
species, numbers = siesta.species(atoms)
assert all(numbers == np.array([1, 2, 2, 3, 2]))
siesta = Siesta(label='test_label', species=species)
siesta.write_input(atoms, properties=['energy'])
with open('test_label.fdf', 'r') as f:
    lines = f.readlines()
lines = [line.split() for line in lines]
assert ['1', '6', 'C.lda.1'] in lines
assert ['2', '1', 'H.lda.2'] in lines
assert ['3', '1', 'H.lda.3'] in lines
assert ['C.lda.1', 'DZP'] in lines
assert ['H.lda.2', 'DZP'] in lines
assert ['H.lda.3', 'SZ'] in lines

# Test if PAO block can be given as species.
c_basis = """2 nodes 1.00
0 1 S 0.20 P 1 0.20 6.00
5.00
1.00
1 2 S 0.20 P 1 E 0.20 6.00
6.00 5.00
1.00 0.95"""
basis_set = PAOBasisBlock(c_basis)
specie = Specie(symbol='C', basis_set=basis_set)
siesta = Siesta(label='test_label', species=[specie])
siesta.write_input(atoms, properties=['energy'])
with open('test_label.fdf', 'r') as f:
    lines = f.readlines()
lines = [line.split() for line in lines]
assert ['%block', 'PAO.Basis'] in lines
assert ['%endblock', 'PAO.Basis'] in lines

# Remove the test directory.
os.chdir('../..')
os.system('rm -rf %s' % test_path)
