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, 4, 19, 66), 'cgbind.log.logger.info', 'logger.info', ({(19, 16, 19, 65): '"""Calculating the ESP and generating a .cube file"""'}, {}), "('Calculating the ESP and generating a .cube file')", False, 'from cgbind.log import logger\n'), ((20, 17, 20, 23), 'time.time', 'time', ({}, {}), '()', False, 'from time import time\n'), ((33, 13, 33, 53), 'numpy.array', 'np.array', ({(33, 22, 33, 52): '[atom.coord for atom in atoms]'}, {}), '([atom.coord for atom in atoms])', True, 'import numpy as np\n'), ((34, 14, 34, 31), 'numpy.array', 'np.array', ({(34, 23, 34, 30): 'charges'}, {}), '(charges)', True, 'import numpy as np\n'), ((37, 22, 37, 44), 'numpy.max', 'np.max', (), '', True, 'import numpy as np\n'), ((38, 21, 38, 43), 'numpy.min', 'np.min', (), '', True, 'import numpy as np\n'), ((71, 39, 71, 103), 'esp_gen.get_cube_lines', 'get_cube_lines', ({(71, 54, 71, 56): 'nx', (71, 58, 71, 60): 'ny', (71, 62, 71, 64): 'nz', (71, 66, 71, 72): 'coords', (71, 74, 71, 83): 'min_carts', (71, 85, 71, 92): 'charges', (71, 94, 71, 102): 'vox_size'}, {}), '(nx, ny, nz, coords, min_carts, charges, vox_size)', False, 'from esp_gen import get_cube_lines\n'), ((30, 8, 30, 75), 'cgbind.log.logger.error', 'logger.error', ({(30, 21, 30, 74): '"""Could not generate an .cube file, charges were None"""'}, {}), "('Could not generate an .cube file, charges were None')", False, 'from cgbind.log import logger\n'), ((26, 14, 27, 65), 'cgbind.exceptions.CgbindCritical', 'CgbindCritical', ({(26, 29, 27, 64): '"""esp_gen not available. cgbind must be installed with the --esp_gen flag"""'}, {}), "(\n    'esp_gen not available. cgbind must be installed with the --esp_gen flag')", False, 'from cgbind.exceptions import CgbindCritical\n'), ((43, 57, 43, 67), 'numpy.ones', 'np.ones', ({(43, 65, 43, 66): '(3)'}, {}), '(3)', True, 'import numpy as np\n'), ((44, 58, 44, 68), 'numpy.ones', 'np.ones', ({(44, 66, 44, 67): '(3)'}, {}), '(3)', True, 'import numpy as np\n'), ((46, 42, 46, 57), 'numpy.array', 'np.array', ({(46, 51, 46, 56): 'coord'}, {}), '(coord)', True, 'import numpy as np\n'), ((67, 32, 67, 55), 'cgbind.atoms.get_atomic_number', 'get_atomic_number', ({(67, 50, 67, 54): 'atom'}, {}), '(atom)', False, 'from cgbind.atoms import get_atomic_number\n'), ((74, 17, 74, 23), 'time.time', 'time', ({}, {}), '()', False, 'from time import time\n')]