#!/usr/bin/env python
"""
Lattice constant Test Driver

Computes the lattice constant for any material and any cubic crystal structure
by minimizing the energy using a simplex minimization

Date: 2012/09/03
Author: Alex Alemi & Matt Bierbaum

Last Update: 2016/02/26 Junhao Li

             2017/08/21 Ellad Tadmor

             2018/02/17 Dan Karls

"""
from ase.build import bulk
import scipy.optimize as opt
from ase.calculators.kim.kim import KIM
import numpy as np
import jinja2
import os
import json
import sys

def energy(a, slab, cell, positions):
    """ Compute the energy of a lattice given the lattice constant.
     Multiply the positions for every lattice type since in this
     test only diamond has multiple sites, all other positions are zero.
    """
    slab.set_positions( positions * a )
    slab.set_cell( cell * a )
    energy = slab.get_potential_energy()
    return energy

def get_lattice_constant(model, symbol, lattice):
    """ Here, we start with an initial guess of 2.5 Angstroms for the lattice
    constant. We then attempt to minimize the potential energy using this
    initial value and, if the optimization fails for any reason, we increment
    the initial guess by 0.5 Angstroms and repeat the calculation. This
    process repeats until either the calculation succeeds or until a maximum
    initial guess is exceeded. This is the same approach used in the
    ForcesNumerDeriv Verification Check.
    """
    calc = KIM(model)
    repeat = 0

    slab = bulk(symbol, lattice, a=1)
    cell = slab.get_cell()
    positions = slab.get_positions()
    slab.set_calculator(calc)

    particles = len(slab)

    init_guess = 2.5
    max_lattice_constant = 10.0
    done = False
    while not done:
        print("\nAttempting to perform relaxation using initial lattice constant guess of %r Angstroms\n" % init_guess)
        try:
            aopt_arr, eopt, iterations, funccalls, warnflag = opt.fmin(energy, init_guess, args=(slab,cell,positions), full_output=True, xtol=1e-8)
            done = True
        except:
            print("\nFailed to perform relaxation using initial lattice constant guess of %r Angstroms\n" % init_guess)
            init_guess += 0.5
            if init_guess > max_lattice_constant:
                done = True
                raise Exception("Exceeded maximum of %r for initial lattice constant guess for relaxation." % max_lattice_constant)

    info = {"iterations": iterations,
            "func_calls": funccalls,
            "warnflag": warnflag,
            "repeat": repeat}
    return aopt_arr[0], -eopt/particles, info

if __name__ == "__main__":
    #grab from stdin (or a file)
    symbol  = raw_input("element=")
    lattice = raw_input("lattice type=")
    model   = raw_input("modelname=")
    print symbol, lattice, model

    aopt, eopt, info = get_lattice_constant(model=model, symbol=symbol, lattice=lattice)

    space_groups = {"fcc": "Fm-3m", "bcc": "Im-3m", "sc": "Pm-3m", "diamond": "Fd-3m"}
    wyckoff_codes = {"fcc": "4a", "bcc": "2a", "sc": "1a", "diamond": "8a"}
    normed_basis = {
          lattice: json.dumps(bulk(symbol, lattice, a=1, cubic=True).positions.tolist(), separators=(' ', ' '))
          for lattice in space_groups.keys()
        }
    count = bulk(symbol, lattice, a=1, cubic=True).positions.shape[0]  # count atoms in a unit cell
    # pack the results in a dictionary
    results = {"lattice_constant": aopt,
                "cohesive_energy": eopt,
                "element": symbol,
                "species": '" "'.join([symbol] * count), # repeat symbols to match normed_basis
                "crystal_structure": lattice,
                "space_group": space_groups[lattice],
                "wyckoff_code": wyckoff_codes[lattice],
                "basis_atoms": normed_basis[lattice]}
    results.update(info)
    print results
    template_environment = jinja2.Environment(
        loader=jinja2.FileSystemLoader('/'),
        block_start_string='@[',
        block_end_string=']@',
        variable_start_string='@<',
        variable_end_string='>@',
        comment_start_string='@#',
        comment_end_string='#@',
        undefined=jinja2.StrictUndefined,
        )

    #template the EDN output
    with open(os.path.abspath("output/results.edn"), "w") as f:
        template = template_environment.get_template(os.path.abspath("results.edn.tpl"))
        f.write(template.render(**results))