import periodictable
import numpy as np
import matplotlib.pyplot as plt
from mendeleev import element

import colorsys
import periodictable
import numpy as np
import colorsys
def atomic_number_to_color(atomic_number):
    # Define the total number of elements you want to support
    total_elements = 118  # Up to Oganesson, element 118

    # Normalize the atomic number to a value between 0 and 1
    normalized_value = atomic_number / total_elements

    # Convert normalized value to a hue in HSL color space
    # (Hue varies from 0 to 1 in colorsys)
    hue = normalized_value

    # Keep saturation and lightness constant
    saturation = 0.5  # Adjust as needed
    lightness = 0.5   # Adjust as needed

    # Convert HSL to RGB (values between 0 and 1)
    r, g, b = colorsys.hls_to_rgb(hue, lightness, saturation)

    # Convert RGB values to 0-255 scale
    r, g, b = int(r * 255), int(g * 255), int(b * 255)

    # Format as a hexadecimal color code
    color_code = f'#{r:02x}{g:02x}{b:02x}'

    return color_code

def plot_atom_list(atoms, marker="*"):
    """
    2D plot of atoms
    :param atoms:
    :param marker:
    :return:
    """
    color = []
    size = []

    x=[]
    y=[]
    for atom in atoms:
        xy=np.array(atom[1][:2])
        # xy=xy.dot(lattice)
        color.append(atomic_number_to_color(getattr(periodictable, atom[0]).number))
        x.append(xy[0])
        y.append(xy[1])


    plt.scatter(x,y, c=color, alpha=0.5, marker=marker)