# fmt: off

"""
IO support for the Gaussian cube format.

See the format specifications on:
http://local.wasp.uwa.edu.au/~pbourke/dataformats/cube/
"""

import time

import numpy as np

from ase.atoms import Atoms
from ase.io import read
from ase.units import Bohr

ATOMS = 'atoms'
CASTEP = 'castep'
DATA = 'data'


def write_cube(file_obj, atoms, data=None, origin=None, comment=None):
    """Function to write a cube file.

    file_obj: str or file object
        File to which output is written.
    atoms: Atoms
        The Atoms object specifying the atomic configuration.
    data : 3-dim numpy array, optional (default = None)
        Array containing volumetric data as e.g. electronic density
    origin : 3-tuple
        Origin of the volumetric data (units: Angstrom)
    comment : str, optional (default = None)
        Comment for the first line of the cube file.
    """

    if data is None:
        data = np.ones((2, 2, 2))
    data = np.asarray(data)

    if data.dtype == complex:
        data = np.abs(data)

    if comment is None:
        comment = "Cube file from ASE, written on " + time.strftime("%c")
    else:
        comment = comment.strip()
    file_obj.write(comment)

    file_obj.write("\nOUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n")

    if origin is None:
        origin = np.zeros(3)
    else:
        origin = np.asarray(origin) / Bohr

    file_obj.write(
        "{:5}{:12.6f}{:12.6f}{:12.6f}\n".format(
            len(atoms), *origin))

    for i in range(3):
        n = data.shape[i]
        d = atoms.cell[i] / n / Bohr
        file_obj.write("{:5}{:12.6f}{:12.6f}{:12.6f}\n".format(n, *d))

    positions = atoms.positions / Bohr
    numbers = atoms.numbers
    for Z, (x, y, z) in zip(numbers, positions):
        file_obj.write(
            "{:5}{:12.6f}{:12.6f}{:12.6f}{:12.6f}\n".format(
                Z, 0.0, x, y, z)
        )

    data.tofile(file_obj, sep="\n", format="%e")


def read_cube(file_obj, read_data=True, program=None, verbose=False):
    """Read atoms and data from CUBE file.

    file_obj : str or file
        Location to the cube file.
    read_data : boolean
        If set true, the actual cube file content, i.e. an array
        containing the electronic density (or something else )on a grid
        and the dimensions of the corresponding voxels are read.
    program: str
        Use program='castep' to follow the PBC convention that first and
        last voxel along a direction are mirror images, thus the last
        voxel is to be removed.  If program=None, the routine will try
        to catch castep files from the comment lines.
    verbose : bool
        Print some more information to stdout.

    Returns a dict with the following keys:

    * 'atoms': Atoms object
    * 'data' : (Nx, Ny, Nz) ndarray
    * 'origin': (3,) ndarray, specifying the cube_data origin.
    * 'spacing': (3, 3) ndarray, representing voxel size
    """

    readline = file_obj.readline
    line = readline()  # the first comment line
    line = readline()  # the second comment line

    # The second comment line *CAN* contain information on the axes
    # But this is by far not the case for all programs
    axes = []
    if "OUTER LOOP" in line.upper():
        axes = ["XYZ".index(s[0]) for s in line.upper().split()[2::3]]
    if not axes:
        axes = [0, 1, 2]

    # castep2cube files have a specific comment in the second line ...
    if "castep2cube" in line:
        program = CASTEP
        if verbose:
            print("read_cube identified program: castep")

    # Third line contains actual system information:
    line = readline().split()
    num_atoms = int(line[0])

    # num_atoms can be negative.
    # Negative num_atoms indicates we have extra data to parse after
    # the coordinate information.
    has_labels = num_atoms < 0
    num_atoms = abs(num_atoms)

    # There is an optional last field on this line which indicates
    # the number of values at each point. It is typically 1 (the default)
    # in which case it can be omitted, but it may also be > 1,
    # for example if there are multiple orbitals stored in the same cube.
    num_val = int(line[4]) if len(line) == 5 else 1

    # Origin around which the volumetric data is centered
    # (at least in FHI aims):
    origin = np.array([float(x) * Bohr for x in line[1:4:]])

    cell = np.empty((3, 3))
    shape = []
    spacing = np.empty((3, 3))

    # The upcoming three lines contain the cell information
    for i in range(3):
        n, x, y, z = (float(s) for s in readline().split())
        shape.append(int(n))

        # different PBC treatment in castep, basically the last voxel row is
        # identical to the first one
        if program == CASTEP:
            n -= 1
        cell[i] = n * Bohr * np.array([x, y, z])
        spacing[i] = np.array([x, y, z]) * Bohr
    pbc = [(v != 0).any() for v in cell]

    numbers = np.empty(num_atoms, int)
    positions = np.empty((num_atoms, 3))
    for i in range(num_atoms):
        line = readline().split()
        numbers[i] = int(line[0])
        positions[i] = [float(s) for s in line[2:]]

    positions *= Bohr

    atoms = Atoms(numbers=numbers, positions=positions, cell=cell, pbc=pbc)

    # CASTEP will always have PBC, although the cube format does not
    # contain this kind of information
    if program == CASTEP:
        atoms.pbc = True

    dct = {ATOMS: atoms}
    labels = []

    # If we originally had a negative num_atoms, parse the extra fields now.
    # The first field of the first line tells us how many other fields there
    # are to parse, but we have to guess how many rows this information is
    # split over.
    if has_labels:
        # Can't think of a more elegant way of doing this...
        fields = readline().split()
        nfields = int(fields[0])
        labels.extend(fields[1:])

        while len(labels) < nfields:
            fields = readline().split()
            labels.extend(fields)

    labels = [int(x) for x in labels]

    if read_data:
        # Cube files can contain more than one density,
        # so we need to be a little bit careful about where one ends
        # and the next begins.
        raw_volume = [float(s) for s in file_obj.read().split()]
        # Split each value at each point into a separate list.
        raw_volumes = [np.array(raw_volume[offset::num_val])
                       for offset in range(num_val)]

        datas = []

        # Adjust each volume in turn.
        for data in raw_volumes:
            data = data.reshape(shape)
            if axes != [0, 1, 2]:
                data = data.transpose(axes).copy()

            if program == CASTEP:
                # Due to the PBC applied in castep2cube, the last entry
                # along each dimension equals the very first one.
                data = data[:-1, :-1, :-1]

            datas.append(data)

        datas = np.array(datas)

        dct[DATA] = datas[0]
        dct["origin"] = origin
        dct["spacing"] = spacing
        dct["labels"] = labels
        dct["datas"] = datas

    return dct


def read_cube_data(filename):
    """Wrapper function to read not only the atoms information from a cube file
    but also the contained volumetric data.
    """
    dct = read(filename, format="cube", read_data=True, full_output=True)
    return dct[DATA], dct[ATOMS]