import sys
import os
import io

import numpy as np
import ase
import ase.io
import ase.io.cube
from ase.calculators.vasp import VaspChargeDensity

def main():
    if len(sys.argv) > 1:
        cubefilepath = sys.argv[1]
    else:
        print("Usage: python cube_to_iri.py cube_file_path")
        sys.exit(1)

    if cubefilepath.endswith(".cube"):
        density, atoms, origin = read_cube(cubefilepath)
    else:
        density, atoms, origin = read_vasp(cubefilepath)

    cell = atoms.get_cell()

    if not cell.orthorhombic:
        print("Only works with orthorombic cell")
        sys.exit(1)

    step_sizes = cell.lengths()/density.shape[0:3]
    print(step_sizes)

    grad = np.gradient(density, *step_sizes)

    ##
    ## IRI(r) = ||grad(p(r))||/(p(r)**a)
    ##
    norm_grad = np.linalg.norm(grad, axis=0)
    a = 1.1 # Default value for a in IRI
    iri = norm_grad / (np.abs(density) ** 1.1)

    with open("iri.cube", "w") as f:
        ase.io.cube.write_cube(f, atoms, iri, origin, "IRI computed by finite difference")

def read_vasp(filepath):
    vasp_charge = VaspChargeDensity(filename=filepath)
    density = vasp_charge.chg[-1]  # separate density
    atoms = vasp_charge.atoms[-1]  # separate atom positions

    return density, atoms, np.zeros(3)  # TODO: Can we always assume origin at 0,0,0?


def read_cube(filepath):
    with open(filepath, "r") as f:
        cube = ase.io.cube.read_cube(f)
    # sometimes there is an entry at index 3
    # denoting the number of values for each grid position
    origin = cube["origin"][0:3]
    # by convention the cube electron density is given in electrons/Bohr^3,
    # and ase read_cube does not convert to electrons/Å^3, so we do the conversion here
    cube["data"] *= 1.0 / ase.units.Bohr ** 3
    return cube["data"], cube["atoms"], origin

if __name__ == "__main__":
    main()