#!/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

"""
#try:
#    from ase.lattice import bulk
#    print 'Imported bulk from ase.lattice' # For ASE version 3.9
#except ImportError:
#    from ase.structure import bulk
#    print 'Imported bulk from ase.structure' # For ASE version <= 3.8
from ase.build import bulk
import scipy.optimize as opt
from kimcalculator import KIMCalculator
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):

    calc = KIMCalculator(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 = 3.5
    aopt_arr, eopt, iterations, funccalls, warnflag = opt.fmin(energy, init_guess, args=(slab,cell,positions), full_output=True, xtol=1e-8)
    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))