#!/usr/bin/env python3

# Python 2-3 compatible code issues
from __future__ import print_function

import os
import json
from functools import partial
import hashlib

import jinja2
import matplotlib
import pylab as plt
import numpy as np
import scipy.optimize as opt
from ase.build import bulk
from ase.phonons import Phonons
from ase.calculators.kim.kim import KIM
import kim_python_utils.ase as kim_ase_utils

matplotlib.use("Agg")

try:
    input = raw_input
except NameError:
    pass

symbol = input("Enter an atomic species:\n")
print(symbol)
lattice = input("Enter a cubic lattice type:\n")
print(lattice)
model = input("Enter a model name:\n")
print(model)
latticeconstant = input("Enter a lattice constant (m):\n")
print(latticeconstant)
print("")
# Convert lattice constant from query to angstroms
latticeconstant = float(latticeconstant) * 1e10

if lattice.lower() != "fcc":
    raise ValueError("Only fcc lattices are currently supported")

# Check if atoms have an energy interaction for this model
atoms_interacting = kim_ase_utils.check_if_atoms_interacting(
    model, symbols=[symbol, symbol], check_force=False, etol=1e-6
)
if not atoms_interacting:
    raise RuntimeError(
        "The model provided, {}, does not possess a non-trivial "
        "energy interaction for species {} as required by this Test.  Aborting."
        "".format(model, symbol)
    )

# Set up crystal and calculator
calc = KIM(model)
atoms = bulk(symbol, lattice, a=latticeconstant)

# Phonon calculator
N = 7
ph = Phonons(atoms, calc, supercell=(N, N, N), delta=0.05)
ph.run()

# Read forces and assemble the dynamical matrix
ph.read(acoustic=True)
ph.clean()

# Band structure
path = atoms.cell.bandpath('GXULGK', npoints=100)
bs = ph.get_band_structure(path)
branches = ["TA", "TA", "LA"] # Hard-coded to FCC for now (1-atom basis)
omega_kn = 1000*bs.energies[0,:,:]  # Extract 100 x 3 array of frequencies and convert eV -> meV
wave_numbers, special_points, special_point_names = bs.get_labels()

for i, name in enumerate(special_point_names):
    if name == "G":
        special_point_names[i] = r"$\Gamma$"

# Calculate phonon DOS
dos_energies, dos_weights = ph.dos(kpts=(50, 50, 50), npts=1000, delta=5e-4)
dos_energies *= 1000 # Convert to meV

# Plot the band structure and DOS
emax=40 # meV
fig = plt.figure(1, figsize=(8, 6))
ax = fig.add_axes([.1, .07, .67, .85])
for n, branch in enumerate(branches):
    omega_n = omega_kn[:,n]
    plt.plot(wave_numbers, omega_n, 'r-', lw=2)

plt.xticks(special_points, special_point_names, fontsize=18)
plt.yticks(fontsize=18)
plt.xlim(wave_numbers[0], wave_numbers[-1])
plt.ylabel("Frequency ($\mathrm{meV}$)", fontsize=22)
plt.grid(True)

plt.axes([.8, .07, .17, .85])
plt.fill_between(dos_weights, dos_energies, y2=0, color='blue', edgecolor='k', lw=1)
plt.ylim(0, emax)
plt.xticks([], [])
plt.yticks([], [])
plt.xlabel("DOS", fontsize=18)

filename = "output/phonons-%s-%s-%s.png" % (model, symbol, lattice)
fig.savefig(filename)

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()
}

results = {
    "lattice_constant": latticeconstant,
    "crystal_structure": lattice,
    "element": symbol,
    "basis_atoms": normed_basis[lattice],
    "dos_energies": dos_energies.tolist(),
    "dos_densities": dos_weights.tolist(),
    "wavevectors": np.asarray(path.kpts).tolist(),
    "branch_labels": branches,
    "wavenumbers": wave_numbers.tolist(),
    "frequencies": omega_kn.tolist()
}

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,
)

jsondump = partial(json.dumps, separators=(" ", " "), indent=4)
jsonlinedump = partial(json.dumps, separators=(" ", " "))
template_environment.filters.update({"json": jsondump, "jsonl": jsonlinedump})

# 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))