#!/usr/bin/env python3
"""
Author:      Brandon Runnels
Date:        June 2019
Institution: Department of Mechanical and Aerospace Engineering
             University of Colorado Colorado Springs
Email:       brunnels@uccs.edu

Description: This file contains a test driver for generating GB
             energy data given an arbitrary tilt axis, tilt range,
             crystal structure, material, and potential.
             The purpose of this file is to interface between the
             OpenKIM framework and the compute_gb_tilt_energy python
             script that generates LAMMPS input and runs the LAMMPS
             tests.
             Please see compute_gb_tilt_energy.py for additional
             information.


Copyright:   2016-2019 University of Colorado Colorado Springs
"""

import sys
import json

from compute_isolated_atom_energy import compute_isolated_atom_energy
from compute_gb_tilt_energy import compute_gb_tilt_energy

#
# Import specific parameters from test
#
model_name = input("Model name: ")
print("model_name={}".format(model_name))
a1 = [int(x) for x in input("a1: ").split(" ")]
print("a1={}".format(a1))
a2 = [int(x) for x in input("a2: ").split(" ")]
print("a2={}".format(a2))
a3 = [int(x) for x in input("a3: ").split(" ")]
print("a3={}".format(a3))
species = input("Species: ")
print("species={}".format(species))
short_name = input("Lattice type: ")
print("short_name=" + short_name)
a_cohesive_energy = [
    float(x)
    for x in json.loads(input("Lattice constant (meters) and cohesive energy (J): "))
]
a = a_cohesive_energy[0] * 1e10  # convert meters to Angstroms
print("a={}".format(a))
cohesive_energy = (
    a_cohesive_energy[1] * 6.242e18
)  # convert from Joules to eV
print("cohesive_energy={}".format(cohesive_energy))
mass = float(sys.stdin.readline().replace("\n", ""))
print("mass={}".format(mass))
theta_min = float(input("Min. tilt angle (degrees): "))
print("theta_min={}".format(theta_min))
theta_max = float(input("Max. tilt angle (degrees): "))
print("theta_max={}".format(theta_max))
max_denominator = int(input("Max. denominator: "))
print("max_denominator={}".format(max_denominator))
cutoff_distance = float(
    input("Atom deletion cutoff distance (as fraction of lattice parameter): ")
)
print("cutoff_distance={}".format(cutoff_distance))
x_repeat = int(input("Min. number unit cell repetitions along x direction: "))
print("x_repeat={}".format(x_repeat))
z_repeat = int(input("Min. number unit cell repetitions along z direction: "))
print("z_repeat={}".format(z_repeat))
min_cell_height = float(input("Min. cell height (Angstroms): "))
print("min_cell_height={}".format(min_cell_height))
offset_grid_frac = float(
    input("Offset grid fraction (as fraction of lattice parameter): ")
)
print("offset_grid_frac={}".format(offset_grid_frac))

#
# Compute energy of an isolated atom of this species for this model in eV
#
isolated_atom_energy = compute_isolated_atom_energy(
    output_dir="output", model_name=model_name, species=species
)

print("")
print("")
print("Isolated atom energy: {} eV".format(isolated_atom_energy))

#
# Compute GB energy using comput_gb_tilt_energy python script
#
# NOTE: If you want to specify a specific grid of offset translation vectors between the
#       grains, you can do so using the offsets_top and offsets_bottom args.  These
#       consist of lists of lists, where each sublist contains two floats corresponding
#       to the x and z translations as fractions of the lattice constant.  The two lists
#       given in offsets_top and offsets_bottom and zipped together to form the final
#       grid of shifts considered for each tilt angle.

thetas, relaxed_energies, minimum_distances, sigmas = compute_gb_tilt_energy(
    output_dir="output",
    model_name=model_name,
    species=species,
    a1=a1,
    a2=a2,
    a3=a3,
    theta_min=theta_min,
    theta_max=theta_max,
    max_denominator=max_denominator,
    lattice_style=short_name,
    lattice_constant=a,
    cutoff_distance=cutoff_distance,
    x_repeat=x_repeat,
    z_repeat=z_repeat,
    offset_grid_frac=offset_grid_frac,
    min_cell_height=min_cell_height,
    #    offsets_top=offsets_top,
    #    offsets_bottom=offsets_bottom,
    num_mpi_processors=1,
    cohesive_energy=cohesive_energy,
    mass=mass,
    isolated_atom_energy=isolated_atom_energy,
    #    dump_format = 'cfg',
    #    return_positions = 'files',
    gzip=False,
    verbose=True,
)

#
# Export results to EDN file
#

edn_out = open("output/results.edn", "w")

edn_out.write("{\n")
edn_out.write(
    '  "property-id" "tag:brunnels@noreply.openkim.org,2016-02-18:property/grain-boundary-symmetric-tilt-energy-relaxed-relation-cubic-crystal" \n'
)
edn_out.write('  "instance-id" 1 \n')
# a
edn_out.write('  "a" { \n')
edn_out.write('    "source-value" ' + str(a) + "\n")
edn_out.write('    "source-unit" "angstrom"\n')
edn_out.write("  }\n")
# basis-atom-coordinates TODO
edn_out.write('  "basis-atom-coordinates" { \n')
if short_name == "fcc":
    edn_out.write(
        '    "source-value" [ [0 0 0] [0 0.5 0.5] [0.5 0 0.5] [0.5 0.5 0] ] \n'
    )
elif short_name == "bcc":
    edn_out.write('    "source-value" [ [0 0 0] [0.5 0.5 0.5] ] \n')
elif short_name == "sc":
    edn_out.write('    "source-value" [ [0 0 0] ] \n')
else:
    raise ("short_name is not of type fcc, bcc, or sc")
edn_out.write("  }\n")
# interface-offset (ignored)
# minimum-atom-separation
edn_out.write('  "minimum-atom-separation" {\n')
edn_out.write('    "source-value" [ ')
for minimum_distance in minimum_distances:
    edn_out.write(str(minimum_distance) + " ")
edn_out.write("] \n")
edn_out.write('    "source-unit" "angstrom"\n')
edn_out.write("  }\n")
# relaxed-grain-boundary-energy
edn_out.write('  "relaxed-grain-boundary-energy" {\n')
edn_out.write('    "source-value" [ ')
for energy in relaxed_energies:
    edn_out.write(str(energy) + " ")
edn_out.write("] \n")
edn_out.write('    "source-unit" "J/m^2"\n')
edn_out.write("  }\n")
# relaxed-interface-positions
# edn_out.write('  "relaxed-interface-positions" { \n')
# edn_out.write('    "source-value" [\n')
# for f in files:
#    edn_out.write('           "'+f.split('output/')[1]+'"\n' )
# edn_out.write('                   ]\n')
# edn_out.write('  }\n')
# short-name
edn_out.write('  "short-name" {\n')
edn_out.write('    "source-value" [ "' + short_name + '" ] \n')
edn_out.write("  }\n")
# sigma
edn_out.write('  "sigma" {\n')
edn_out.write('    "source-value" [ ')
for sigma in sigmas:
    edn_out.write(str(sigma) + " ")
edn_out.write("] \n")
edn_out.write("  }\n")
# space-group (ignored)
# species
edn_out.write('  "species" {\n')
edn_out.write(
    '    "source-value" ["'
    + species
    + '" "'
    + species
    + '" "'
    + species
    + '" "'
    + species
    + '" ]\n'
)
edn_out.write("  }\n")
# tilt-angle
edn_out.write('  "tilt-angle" {\n')
edn_out.write('    "source-value" [ ')
for theta in thetas:
    edn_out.write(str(theta) + " ")
edn_out.write("] \n")
edn_out.write('    "source-unit" "degrees"\n')
edn_out.write("  }\n")
# tilt-axis
edn_out.write('  "tilt-axis" {\n')
edn_out.write(
    '    "source-value" [' + str(a3[0]) + " " + str(a3[1]) + " " + str(a3[2]) + "]\n"
)
edn_out.write("  }\n")
edn_out.write("}\n")
edn_out.close()
# wyckoff-coordinates (ignored)
# wyckoff-multiplicity-and-letter (ignored)
# wyckoff-species

exit(0)