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

Computes the lattice constant for:
1. Any element,
2. HCP crystal,
3. At 0 K, 0 GPa,
by simplex minimization.

Date: 2015/08/25
Author: Junhao Li <streaver91@gmail.com>

"""
# Import External Modules
# ASE Modules
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.optimize import FIRE
from ase.data import chemical_symbols
from ase.data import reference_states

# KIM Modules
from kimcalculator import *
from kimservice import KIM_API_get_data_double

# Other Python Modules
import sys
import re
import json
import math
import simplejson
import jinja2
import os
import fileinput
import copy
import numpy as np
from scipy.optimize import fmin
from collections import OrderedDict

# Constants for result.edn
SPACE_GROUP = 'P63/mmc'
WCYKOFF_CODE = '2a'
KEY_SOURCE_VALUE = 'source-value'
KEY_SOURCE_UNIT = 'source-unit'
KEY_SOURCE_UNCERT = 'source-std-uncert-value'
UNIT_LENGTH = 'angstrom'
UNIT_TEMPERATURE = 'K'
UNIT_PRESSURE = 'GPa'
UNIT_ENERGY = 'eV'
STRUCTURE_PROP_ID = 'tag:staff@noreply.openkim.org,2014-04-15:property/structure-hexagonal-crystal-npt'
ENERGY_PROP_ID = 'tag:staff@noreply.openkim.org,2014-04-15:property/cohesive-free-energy-hexagonal-crystal'

# Minimization Convergence Criteria
FMIN_FTOL = 1e-10
FMIN_XTOL = 1e-10

def getModelInfo(model, elem):
    # Obtain cutoff and nbcname from model
    calc = KIMCalculator(model)
    atoms = bulk(elem, 'hcp', a = 1)
    atoms.set_calculator(calc)
    cutoff = KIM_API_get_data_double(calc.pkim, 'cutoff')[0]
    nbcname = calc.get_NBC_method()
    return cutoff, nbcname

def getEnergy(cellVector, atoms):
    # Obtain Cohesive tmpCell Corresponds to cellVector
    # Construct tmpAtoms based on the length of cellVector
    if len(cellVector) == 2:
        tmpAtoms = bulk('X', a = cellVector[0], c = cellVector[1], crystalstructure = 'hcp')
    else:
        tmpAtoms = bulk('X', a = cellVector[0], crystalstructure = 'hcp')
    tmpCell = tmpAtoms.get_cell()
    # Scale positions of atoms and obtain energy
    atoms.set_cell(tmpCell, scale_atoms = True)
    energy = atoms.get_potential_energy()
    return energy

def searchLatticeConstants(latticeConstantsGuess, atoms):
    # Simplex searching lattice constants that minimize potential energy of atoms
    # Searching starts from latticeConstantsGuess
    # Parameters and 'atoms' are used when calling 'getEnergy()'
    print 'Simplex Searching Start From:', latticeConstantsGuess
    tmpLatticeConstants, tmpEnergy = fmin(
        getEnergy,
        latticeConstantsGuess,
        args = (atoms, ),
        full_output = True,
        ftol = FMIN_FTOL,
        xtol = FMIN_XTOL,
    )[:2]
    print 'Tmp Lattice Constants:', tmpLatticeConstants
    print 'Tmp Energy:', tmpEnergy
    return tmpLatticeConstants, tmpEnergy
    
def getLatticeConstant(elem, model):
    # Obtain cutoff and nbcname
    cutoff, nbcname = getModelInfo(model, elem)
    print 'Model Cutoff:', cutoff
    print 'NBC Name:', nbcname
    
    # Create Calculator
    calc = KIMCalculator(model)
    
    # Support NEIGH_RVEC_F Models
    if nbcname != 'NEIGH_RVEC_F':
        print 'Model not supported.'
        sys.exit(0)
        
    atoms = bulk(elem, 'hcp', a = 1)
    atoms.set_calculator(calc)
    
    # Relaxation With c/a Ratio Fixed
    print 'Relaxation With c/a Ratio Fixed'
    minEnergy = 0
    for i in [x * 0.02 for x in range(5, 25)]:
        print 'Simplex Searching start from: cutoff *', i
        tmpLatticeConstants, tmpEnergy = searchLatticeConstants(
            [cutoff * i],
            atoms,
        )
        if tmpEnergy < minEnergy and tmpLatticeConstants[0] < cutoff:
            minLatticeConstants = tmpLatticeConstants
            minEnergy = tmpEnergy
    
    if minEnergy == 0:
        # Exit if previous step is not successful
        print 'Failed to Obtain Lattice Constants'
        print 'minEnergy = 0 After Fix c/a Ratio Relaxation'
        sys.exit(0)
    
    # Relaxation With c/a Ratio Relaxed
    print 'Relaxation With c/a Ratio Relaxed'
    minLatticeConstantsCopy = copy.deepcopy(minLatticeConstants)
    minEnergy = 0 # To make sure tmpLatticeConstants got updated at least once
    for i in [x * 0.02 + 1 for x in range(-10, 10)]:
        print 'Simplex Searching start from c/a ratio: 1.633 *', i
        tmpLatticeConstants, tmpEnergy = searchLatticeConstants(
            [minLatticeConstantsCopy[0], minLatticeConstantsCopy[0] * i],
            atoms,
        )
        if tmpEnergy < minEnergy and tmpLatticeConstants[0] < cutoff:
            minLatticeConstants = tmpLatticeConstants
            minEnergy = tmpEnergy
    print 'Lattice Constants:', minLatticeConstants
    print 'Energy:', minEnergy
    
    # Return Results
    return minLatticeConstants[0], minLatticeConstants[1], minEnergy

def V(value, unit = '', uncert = ''):
    # Generate OrderedDict for JSON Dump
    res = OrderedDict([
        (KEY_SOURCE_VALUE, value),
    ])
    if unit != '':
        res.update(OrderedDict([
            (KEY_SOURCE_UNIT, unit),
        ]))
    if uncert != '':
        res.update(OrderedDict([
            (KEY_SOURCE_UNCERT, uncert)
        ]))
    return res

def main():
    # Input Parameters
    elem = raw_input('Element = ')
    lattice = raw_input('Lattice = ')
    model = raw_input('Model = ')
    
    # Parameters for Debugging
    # elem = 'Co'
    # lattice = 'hcp'
    # model = 'EAM_Dynamo_PurjaPun_Mishin_Co__MO_885079680379_001'
    
    # elem = 'Al'
    # lattice = 'hcp'
    # model = 'EAM_Dynamo_Mishin_Farkas_Al__MO_651801486679_001'
    
    # elem = 'C'
    # lattice = 'hcp'
    # model = 'EAM_Dynamo_Hepburn_Ackland_FeC__MO_143977152728_001'
    # model = 'MEAM_2NN_Fe_to_Ga__MO_145522277939_001'
    # model = 'model_ArCHHeXe_BOP_AIREBO__MO_154399806462_001'
    # model = 'Tersoff_LAMMPS_Erhart_Albe_CSi__MO_903987585848_001'
    
    # Print Inputs
    print 'Element:', elem
    print 'Lattice:', lattice # Not used here
    print 'Model:', model
    
    # Obtain Lattice Constants and Cohesive Energy
    a, c, energy = getLatticeConstant(elem, model)
    print 'Lattice Constants:', a, c
    
    # Output Results
    structureResults = OrderedDict([
        ('property-id', STRUCTURE_PROP_ID),
        ('instance-id', 1),
        ('cauchy-stress', V([0, 0, 0, 0, 0], UNIT_PRESSURE)),
    ])
    cohesiveEnergyResults = OrderedDict([
        ('property-id', ENERGY_PROP_ID),
        ('instance-id', 2),
        ('cohesive-free-energy', V(energy, UNIT_ENERGY)),
    ])
    commonResults = OrderedDict([
        ('short-name', V(['hcp'])),
        ('species', V([elem])),
        ('a', V(a, UNIT_LENGTH)),
        ('c', V(c, UNIT_LENGTH)),
        ('basis-atom-coordinates', V(bulk(elem, 'hcp', a = 1, c = 1).get_positions().tolist())),
        ('space-group', V(SPACE_GROUP)),
        ('temperature', V(0, UNIT_TEMPERATURE)),
    ])
    structureResults.update(commonResults)
    cohesiveEnergyResults.update(commonResults)
    results = [
        structureResults,
        cohesiveEnergyResults,
    ]
    resultsString = json.dumps(results, separators = (' ', ' '), indent = 4)
    print resultsString
    with open(os.path.abspath('output/results.edn'), 'w') as f:
        f.write(resultsString)
    

if __name__ == '__main__':
    main()