import numpy as np
from time import time
from cgbind.atoms import get_atomic_number
from cgbind.log import logger
from cgbind.constants import Constants
from cgbind.exceptions import CgbindCritical


def get_esp_cube_lines(charges, atoms):
    """
    From a list of charges and a set of xyzs create the electrostatic potential
    map grid-ed uniformly between the most negative x, y, z values -5 Å
    and the largest x, y, z +5 Å

    :param charges: (list(float))
    :param atoms: (list(autode.atoms.Atom))
    :return: (list(str)), (min ESP value, max ESP value)
    """
    logger.info('Calculating the ESP and generating a .cube file')
    start_time = time()

    try:
        from esp_gen import get_cube_lines

    except ModuleNotFoundError:
        raise CgbindCritical('esp_gen not available. cgbind must be '
                             'installed with the --esp_gen flag')

    if charges is None:
        logger.error('Could not generate an .cube file, charges were None')
        return [], (None, None)

    coords = np.array([atom.coord for atom in atoms])
    charges = np.array(charges)

    # Get the max and min points from the coordinates
    max_cart_values = np.max(coords, axis=0)
    min_cat_values = np.min(coords, axis=0)

    # The grid needs to be slightly larger than the smallest/largest Cartesian
    # coordinate
    # NOTE: All distances from here are in Bohr (a0) i.e. atomic units
    min_carts = Constants.ang2a0 * (min_cat_values - 5 * np.ones(3))
    max_carts = Constants.ang2a0 * (max_cart_values + 5 * np.ones(3))

    coords = np.array([Constants.ang2a0 * np.array(coord) for coord in coords])

    # Number of voxels will be nx * ny * nz
    nx, ny, nz = 50, 50, 50

    vox_size = max_carts - min_carts
    rx, ry, rz = vox_size[0] / nx, vox_size[1] / ny, vox_size[2] / nz

    # Write the .cube file lines
    cube_file_lines = ['Generated by cgbind\n', 'ESP\n']

    n_atoms = len(coords)
    min_x, min_y, min_z = min_carts
    cube_file_lines.append(f'{n_atoms:>5d}{min_x:>12f}{min_y:>12f}{min_z:>12f}\n')  # n_atoms origin(x y z)

    cube_file_lines.append(f'{nx:>5d}{rx:>12f}{0.0:>12f}{0.0:>12f}\n')  # Number of voxels and their size
    cube_file_lines.append(f'{ny:>5d}{0.0:>12f}{ry:>12f}{0.0:>12f}\n')
    cube_file_lines.append(f'{nz:>5d}{0.0:>12f}{0.0:>12f}{rz:>12f}\n')

    for atom in atoms:
        x, y, z = atom.coord
        cube_file_lines.append(f'{get_atomic_number(atom):>5d}{0.0:>12f}'
                               f'{Constants.ang2a0*x:>12f}{Constants.ang2a0*y:>12f}{Constants.ang2a0*z:>12f}\n')

    # Looping over x, y, z is slow in python so use Cython extension
    cube_val_lines, min_val, max_val = get_cube_lines(nx, ny, nz, coords, min_carts, charges, vox_size)
    cube_file_lines += cube_val_lines

    logger.info(f'ESP generated in {time()-start_time:.3f} s')
    return cube_file_lines, (min_val, max_val)

[(19, 'cgbind.log.logger.info', 'logger.info', (['"""Calculating the ESP and generating a .cube file"""'], {}), False, 'from cgbind.log import logger\n'), (20, 'time.time', 'time', ([], {}), False, 'from time import time\n'), (33, 'numpy.array', 'np.array', (['[atom.coord for atom in atoms]'], {}), True, 'import numpy as np\n'), (34, 'numpy.array', 'np.array', (['charges'], {}), True, 'import numpy as np\n'), (37, 'numpy.max', 'np.max', (['coords'], {'axis': '(0)'}), True, 'import numpy as np\n'), (38, 'numpy.min', 'np.min', (['coords'], {'axis': '(0)'}), True, 'import numpy as np\n'), (71, 'esp_gen.get_cube_lines', 'get_cube_lines', (['nx', 'ny', 'nz', 'coords', 'min_carts', 'charges', 'vox_size'], {}), False, 'from esp_gen import get_cube_lines\n'), (30, 'cgbind.log.logger.error', 'logger.error', (['"""Could not generate an .cube file, charges were None"""'], {}), False, 'from cgbind.log import logger\n'), (26, 'cgbind.exceptions.CgbindCritical', 'CgbindCritical', (['"""esp_gen not available. cgbind must be installed with the --esp_gen flag"""'], {}), False, 'from cgbind.exceptions import CgbindCritical\n'), (43, 'numpy.ones', 'np.ones', (['(3)'], {}), True, 'import numpy as np\n'), (44, 'numpy.ones', 'np.ones', (['(3)'], {}), True, 'import numpy as np\n'), (46, 'numpy.array', 'np.array', (['coord'], {}), True, 'import numpy as np\n'), (67, 'cgbind.atoms.get_atomic_number', 'get_atomic_number', (['atom'], {}), False, 'from cgbind.atoms import get_atomic_number\n'), (74, 'time.time', 'time', ([], {}), False, 'from time import time\n')]