# drawing.py
# ----------
# By ____(you)___________
#
# (Briefly describe what this program does.)

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-10. What is your element?")
        elements = ["H", "He", "Li", "Be", "B", "C", "N","O","F","Ne"]
        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):
                     pencolor("Blue")
                     dot(ELECTRONSIZE)
                     right(36)
                     penup()
                     forward(sqrt(7.5))
                     pendown()
                     pencolor("Black")
       return_to_center()




'''This function indicates the number of rings that will form, upt ot 2 electrons will eventually
go in ring 1 and 8 in ring 2. Has capability to have up to 4 rings if more atoms are added later. '''

def number_of_rings(electron_number):                          
       if electron_number < 3:
              rings = 1
              return rings
       if electron_number < 11:
                rings = 2
                return rings
       if electron_number < 18:
                rings = 3
                return rings
       elif electron_number < 21:
                rings = 4
                return rings

'''This function draws the rings. Radius for smallest ring can be adjusted with 
variable above.'''

def draw_ring(ring):
        radius = ring*ring_1_size
        circle(radius,360)
        penup()
        right(90)
        forward(ring_1_size)
        left(90)
        pendown()

def electron(radius):
    """Assumes the turtle already has the proper heading. Draws the electron at a certain
    distance from the center and then returns the turtle to the center to draw the next 
    electron.
    """
    with restore_state_when_finished():
        pencolor("red")
        penup()
        forward(radius)
        pendown()
        dot(ELECTRONSIZE)

'''This function animates the electrons and draws the appropraite number based on the
stored electron_number'''

def draw_and_move_electrons(electron_number): 
                if electron_number == 1:
                        for frame in animate(360, loop=True):
                                 with frame.rotate(0, 359, cycles=cycles):
                                        right(90)
                                        electron(ring_1_size)
                if electron_number ==2:
                        for frame in animate(360, loop=True):
                                 with frame.rotate(0, 359, cycles=cycles):
                                        right(90)
                                        electron(ring_1_size)
                                        right(180)
                                        electron(ring_1_size)
                if electron_number ==3:
                        for frame in animate(360, loop=True):
                                 with frame.rotate(0, 359, cycles=cycles):
                                        right(90)
                                        electron(ring_1_size)
                                        right(180)
                                        electron(ring_1_size)
                                 with frame.rotate(0, -359, cycles=cycles):
                                        electron(ring_2_size)
                if electron_number ==4:
                        for frame in animate(360, loop=True):
                                 with frame.rotate(0, 359, cycles=cycles):
                                        right(90)
                                        electron(ring_1_size)
                                        right(180)
                                        electron(ring_1_size)
                                 with frame.rotate(0, -359, cycles=cycles):
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                if electron_number ==5:
                        for frame in animate(360, loop=True):
                                 with frame.rotate(0, 359, cycles=cycles):
                                        right(90)
                                        electron(ring_1_size)
                                        right(180)
                                        electron(ring_1_size)
                                 with frame.rotate(0, -359, cycles=cycles):
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(90)
                                        electron(ring_2_size)
                if electron_number == 6:
                        for frame in animate(360, loop=True):
                                 with frame.rotate(0, 359, cycles=cycles):
                                        right(90)
                                        electron(ring_1_size)
                                        right(180)
                                        electron(ring_1_size)
                                 with frame.rotate(0, -359, cycles=cycles):
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(90)
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                if electron_number == 7:
                        for frame in animate(360, loop=True):
                                 with frame.rotate(0, 359, cycles=cycles):
                                        right(90)
                                        electron(ring_1_size)
                                        right(180)
                                        electron(ring_1_size)
                                 with frame.rotate(0, -359, cycles=cycles):
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(90)
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(45)
                                        electron(ring_2_size)
                if electron_number == 8:
                        for frame in animate(360, loop=True):
                                 with frame.rotate(0, 359, cycles=cycles):
                                        right(90)
                                        electron(ring_1_size)
                                        right(180)
                                        electron(ring_1_size)
                                 with frame.rotate(0, -359, cycles=cycles):
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(90)
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(45)
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                if electron_number == 9:
                        for frame in animate(360, loop=True):
                                 with frame.rotate(0, 359, cycles=cycles):
                                        right(90)
                                        electron(ring_1_size)
                                        right(180)
                                        electron(ring_1_size)
                                 with frame.rotate(0, -359, cycles=cycles):
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(90)
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(45)
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(90)
                                        electron(ring_2_size)
                if electron_number == 10:
                        for frame in animate(360, loop=True):
                                 with frame.rotate(0, 359, cycles=cycles):
                                        right(90)
                                        electron(ring_1_size)
                                        right(180)
                                        electron(ring_1_size)
                                 with frame.rotate(0, -359, cycles=cycles):
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(90)
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(45)
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)
                                        right(90)
                                        electron(ring_2_size)
                                        right(180)
                                        electron(ring_2_size)

#Code being run starts here

electron_number = select_element()    #sets output from select_element and stores as variable

draw_nucleus(electron_number)
return_to_center()

rings= number_of_rings(electron_number) +1       #iterates to draw each ring
for ring in range(rings):                  
       draw_ring(ring)

draw_and_move_electrons(electron_number)         #draws electrons and moves them



                
input()