# drawing.py
# ----------
# By Chris Mekelburg
#
# This program draws a Bohr model for any element from Hydrogen- Calcium (atomic #s 1-20)

from turtle import *
from math import sqrt
from superturtle.animation import animate
from superturtle.movement import restore_state_when_finished

ELECTRONSIZE = 15                  #electron size
ring_1_size = 75                   #radius of first shell
ring_2_size= ring_1_size*2         #radii of shells 2-4
ring_3_size = ring_1_size*3
ring_4_size = ring_1_size*4

cycles =1                          #relates to how fast the electrons move

'''This function enables the user to pick and element and then retruns the number of
electrons for that element'''

def select_element():
        choice = input("Enter the element symbol for any element with atomic number 1-20. What is your element?")
        elements = ["H", "He", "Li", "Be", "B", "C", "N","O","F","Ne", "Na", "Mg", "Al", "Si", "P", "S","Cl", "Ar","K","Ca"]
        while not (choice in elements):
              print("Try again. Enter the element symbol for any element with atomic number 1-20. What is your element symbol?")
        if choice in elements:
              electron_number = elements.index(choice) + 1
              return electron_number
        else:
                print("Error")

'''Returns turtle to center'''
def return_to_center():
       penup()
       goto(0,0)
       pendown()

'''Draws the nucleus, one dot per proton. As wrriten keeps the nucleus centered but 
protons are not clearly outlined, however size of nucleus is proprotional to atomic number'''

def draw_nucleus(electron_number):
       protons=electron_number
       penup()
       forward(sqrt(7.5))
       right(90)
       forward(sqrt(7.5))
       pendown()
       for proton in range(protons):
                     if proton <10:
                            pencolor("Blue")
                            dot(ELECTRONSIZE)
                            right(36)
                            penup()
                            forward(sqrt(7.5))
                            pendown()
                            pencolor("Black")
                     else:
                            pencolor("Blue")
                            dot(ELECTRONSIZE*1.75)
                            right(36)
                            penup()
                            forward(sqrt(7.5))
                            pendown()
                            pencolor("Black")  
       return_to_center()


'''This function indicates the number of rings that will form. Rings is called in a later function to
help determine where the electrons are located. '''

def number_of_rings(electron_number):                          
       if electron_number <= 2:
              rings = 1
       if 3 < electron_number <= 10:
                rings = 2 
       if 10 < electron_number <= 18:
                rings = 3
       if 18 < electron_number <= 20:
                rings = 4
       return rings

'''This function draws the rings. Radius for smallest ring can be adjusted with 
variable above. This function also outputs the radius of each ring which can be called
in a later function.'''

def draw_ring(ring):
       return_to_center()
       radius = ring*ring_1_size
       penup()
       right(90)
       forward(radius)
       left(90)
       pendown()
       pencolor("black")    
       circle(radius,360)
       return radius

'''This function determines how far the electrons will be from each other in each ring.'''
def get_electron_ring_offset(ring, electron_number):
       if ring ==1:
                offset = 180
       if ring ==2:
              offset = 45
       if ring ==3:
              offset = 45
       if ring ==4:                #using if instead of else here enables incorporating more elements later, after #20 ring 3 can hold more electrons
               offset = 45
       return offset

'''This function determines how many electrons belong in each ring (shell) so that it can be called in a later function.'''
def electrons_per_shell(electron_number, ring):
       if ring ==1:
              if electron_number >= 2:
                     electrons_per_shelly = 2
              else:
                     electrons_per_shelly = 1
       if ring ==2:
              if electron_number >=10:
                     electrons_per_shelly=8
              else:
                     electrons_per_shelly = electron_number - 2
       if ring == 3:
              if electron_number >=18:
                      electrons_per_shelly = 8
              else:
                     electrons_per_shelly = electron_number - 10
       if ring ==4:
              if electron_number >=36:           #Ar
                       electrons_per_shelly = 8
              else:
                      electrons_per_shelly = electron_number - 18
       return electrons_per_shelly

'''This function draws an electron based on the radius (detemrines ring) and offset (how far apart the electrons are).'''
def draw_electron(radius, offset):
                            return_to_center()
                            pencolor("red")
                            penup()
                            forward(radius)
                            pendown()
                            dot(ELECTRONSIZE)
                            penup()
                            return_to_center
                            right(offset)



electron_number = select_element()        #input for which element to draw
rings=number_of_rings(electron_number)    #saves rings variable for later use
#print(rings)                             helpful for troubleshooting
#print(electron_number)

for frame in animate(360, loop=True): 
       draw_nucleus(electron_number)
       return_to_center()
       rings_range = rings + 1       #adds 1 so that range function can be used later
       
      
       with frame.rotate(0, 360):
              with restore_state_when_finished():
                     for ring in range(1,rings_range):  #iterates to draw each ring 
                            radius2 = draw_ring(ring)   #saves radius of each ring for use, new value on each iteration
                            offset = get_electron_ring_offset(ring,electron_number) #saves offset for later use, also new value on each iteration
                            electrons_per_shelly2 = electrons_per_shell(electron_number, ring)    #saves electrons per shell so it can be called down below, also new on each iteration
                            #print(electrons_per_shelly2)             helpful for troubleshooting
                            for electron in range(electrons_per_shelly2):    #iterates to draw each electron
                                   draw_electron(radius2, offset)

input()              #pauses function so it does not close automatically