from util.colorschemes import hex_to_rgb_float

def set_atom_color(symbol, hexcolor):
    from ase.data import atomic_numbers
    from ase.data.colors import jmol_colors, cpk_colors
    float_color = hex_to_rgb_float(hexcolor)
    n = atomic_numbers[symbol]
    jmol_colors[n] = float_color
    cpk_colors[n] = float_color


from scipy.spatial.distance import pdist, squareform
import numpy as np
def get_bondatoms(atoms):
    site_positions = [site.position for site in atoms]
    pair_distances = squareform(pdist(np.stack(site_positions)))
    vs = pair_distances
    bondatoms = []
    for i in range(vs.shape[0]):
        for j in range(i):
            if vs[i, j] < 3: # up to 3 angstrom distance show a bond TODO
                bondatoms.append((i, j))
    return bondatoms
# returns to many
# from ase.io.pov import get_bondpairs
# bondatoms=get_bondpairs(lat, 5)


TARGET_DPI = 300
# doc: https://github.com/WMD-group/ASE-Tutorials/blob/master/povray-tools/ase_povray.py
def get_pov_settings(size, COLORSCHEME, bondatoms=None):
    white = hex_to_rgb_float(COLORSCHEME["bg0"])
    other = hex_to_rgb_float(COLORSCHEME["fg-yellow"])
    pixels = TARGET_DPI * size[0]
    pov_settings=dict(
        transparent=True,
        display=False,
        # camera_type='orthographic',
        camera_type='perspective',
        canvas_width=pixels,
        # point_lights : [], #[(18,20,40), 'White'],[(60,20,40),'White'],             # [[loc1, color1], [loc2, color2],...]
        point_lights=[[(18,20,40), white],[(60,20,40),other]],             # [[loc1, color1], [loc2, color2],...]
        background=(0, 0, 0, 1.,),
        bondlinewidth=0.07,
        bondatoms=bondatoms
    )
    return pov_settings