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(message='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)
