from __future__ import division
from math import sqrt, ceil, atan, sin, cos
import random, pygame, sys
import os
import json
from pygame.locals import *

pygame.init()

GRAY = ( 182, 182, 182)
VIOLET = (150, 100, 190)
RED = (255,  0,  0)
GREEN = (0, 150, 0)
BLUE = (30, 30, 180)
VERYLIGHT = (210, 210, 210)
BLACK = (5,5,5)
WHITE = (255, 255, 255)

PINK = (255, 62, 150)
MAGENTA = (255, 0, 255)
ORANGE = (255, 69, 0)
GREENYYELLOW = (173, 255, 47)
CRIMSON = (220, 20,  60)
GRAYISH = (210, 210, 255)
ROYALBLUE = (65, 105, 225)

TIME_STEP = 0.01

K = 120000

X_SCREEN_BORDER = 1366
Y_SCREEN_BORDER = 768

INITIAL_SPEED_X = 300
INITIAL_SPEED_Y = 0
INITIAL_Y = 400
DEFAULT_OLD_X = 2 - INITIAL_SPEED_X*TIME_STEP
DEFAULT_OLD_Y = INITIAL_Y - INITIAL_SPEED_Y*TIME_STEP
BEAM_WIDHT = 40

def main():
    #create the screen
    window = pygame.display.set_mode((1366, 768), FULLSCREEN)
    colors = [GRAY, VIOLET, RED, GREEN, BLUE, ROYALBLUE, BLACK, PINK, ORANGE, GREENYYELLOW, CRIMSON]
    
    settings = SimulationSettings(K, TIME_STEP, INITIAL_Y, INITIAL_SPEED_X, BEAM_WIDHT, colors)
    #Initialize universe with some atoms
    universe = Universe(settings)

    screen = Screen(window, colors, universe)
    screen.draw_surface()
    keep_running = True
    simulation_running = True
    while keep_running:
        if(simulation_running):
            screen.draw_surface()
            universe.update_positions()
        else:
            screen.draw_static_surface()
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                keep_running = False
            elif event.type == pygame.KEYUP:
                if event.key == K_ESCAPE:
                    keep_running = False
                elif event.key == K_UP :
                    universe.settings.initial_y -= 10
                elif event.key == K_DOWN:
                    universe.settings.initial_y += 10
                elif event.key == K_RIGHT:
                    universe.settings.speed_x += 20
                elif event.key == K_LEFT:
                    universe.settings.speed_x -= 20
                    if(universe.settings.speed_x < 0):
                        universe.settings.speed_x = 0
                elif event.key == K_SPACE:
                    universe.add_particle()
                elif event.key == K_h:
                    if(universe.settings.show_instructions):
                        universe.settings.show_instructions = False
                    else:
                        universe.settings.show_instructions = True
                elif event.key == K_RETURN:
                    if(simulation_running):
                        simulation_running = False
                    else:
                        simulation_running = True 
                elif event.key == K_COMMA:
                    universe.settings.beam_width -= 5
                    if(universe.settings.beam_width < 5):
                        universe.settings.beam_width = 5
                elif event.key == K_PERIOD:
                    universe.settings.beam_width += 5

    print "Symulation is over."

class Screen:
    def __init__(self, window, colors, universe):
        self.communicates = []
        self.communicates.append("Rutherford scattering")
        self.window = window
        self.universe = universe
        self.colors = colors
        self.color = WHITE

    def change_color(self):
        self.color = random.choice(self.colors)

    def draw_universe(self):
        self.draw_guide()
        for particle in self.universe.particles:
            self.draw_particle(particle)

    def draw_static_universe(self):
        self.draw_guide()
        for particle in self.universe.particles:
            self.draw_particle(particle)
        for x in range(1, len(self.universe.particles)):
            self.draw_particle_vector(self.universe.particles[x])

    def draw_particle(self, particle):
        position = (int(ceil(particle.position.x)), int(ceil(particle.position.y)))
        pygame.draw.circle(self.window, particle.color, position, 6+ int(particle.mass/100) )

    def draw_particle_vector(self, particle):
        factor = 20
        arrow_head_x = 0.50
        arrow_head_y = 0.30
        start_position = (particle.position.x, particle.position.y)
        end_position = (int((particle.position.x - particle.old_position.x)*factor +particle.position.x), 
            int((particle.position.y -particle.old_position.y)*factor + particle.position.y))
        
        pygame.draw.aaline(self.window, BLUE, start_position, end_position, 4)
        pygame.draw.circle(self.window, RED, end_position, 2)

    def draw_guide(self):
        y = self.universe.settings.initial_y
        offset = self.universe.settings.beam_width
        for x in range(0, 39, 2):
            pygame.draw.aaline(self.window, BLACK, (X_SCREEN_BORDER*x/40, y), (X_SCREEN_BORDER*(x+1)/40, y), 4)
        for x in range(0, 79, 2):
            pygame.draw.aaline(self.window, RED, (X_SCREEN_BORDER*x/80, y + offset), (X_SCREEN_BORDER*(x+1)/80, y + offset), 2)
        for x in range(0, 79, 2):
            pygame.draw.aaline(self.window, RED, (X_SCREEN_BORDER*x/80, y - offset), (X_SCREEN_BORDER*(x+1)/80, y - offset), 2)

    def draw_static_surface(self):
        self.window.fill(GRAYISH)
        self.draw_static_universe()
        description_of_simulation_settings = self.prepare_settings_description(self.universe.settings)
        for i, text in enumerate(self.communicates + description_of_simulation_settings):
            if(i == 0):
                self.print_text(text, 20, 14 + i*20, CRIMSON, 28, self.window)
            else:
                self.print_text(text, 20, 20 + i*20, BLACK, 24, self.window)
        if(self.universe.settings.show_instructions):
            for i, text in enumerate(self.universe.settings.instructions):
                if(i == 0):
                    self.print_text(text, 350, 14 + i*20, GREEN, 28, self.window)
                else:
                    self.print_text(text, 350, 20 + i*20, BLACK, 24, self.window)
        pygame.display.flip()

    def prepare_settings_description(self, settings):
        description = []
        line = "K = {0:d}".format(settings.K)
        description.append(line)
        line = "Position of the beam is {0:d}".format(settings.initial_y)
        description.append(line)
        line = "Initial speed of the beam is {0:d}".format(settings.speed_x)
        description.append(line)
        line = "Width of the beam is {0:d}".format(settings.beam_width)
        description.append(line)
        line = "Time step is set to {0:f}".format(settings.time_step)
        description.append(line)
        line = "Press H to show/hide instructions"
        description.append(line)
        return description

    def draw_surface(self):
        self.window.fill(self.color)
        self.draw_universe()
        description_of_simulation_settings = self.prepare_settings_description(self.universe.settings)
        for i, text in enumerate(self.communicates + description_of_simulation_settings):
            if(i == 0):
                self.print_text(text, 20, 14 + i*20, CRIMSON, 28, self.window)
            else:
                self.print_text(text, 20, 20 + i*20, BLACK, 24, self.window)
        if(self.universe.settings.show_instructions):
            for i, text in enumerate(self.universe.settings.instructions):
                if(i == 0):
                    self.print_text(text, 350, 14 + i*20, GREEN, 28, self.window)
                else:
                    self.print_text(text, 350, 20 + i*20, BLACK, 24, self.window)
        pygame.display.flip()

    def print_text(self, text,xx,yy,color,text_size, screen):
        font = pygame.font.SysFont(None,text_size)
        ren = font.render(text,1,color)
        screen.blit(ren, (xx,yy))

class Vector2:
    def __init__(self, x, y):
        self.x = x
        self.y = y


class Particle:
    def __init__(self, mass, position, old_position, charge, color):
        self.mass = mass
        self.position = position
        self.old_position = old_position
        self.acc = Vector2(0,0)
        self.color = color
        self.charge = charge

class SimulationSettings:
    def __init__(self, K, time_step, initial_y, speed_x, beam_width, colors):
        self.K = K
        self.time_step = time_step
        self.initial_y = initial_y
        self.speed_x = speed_x
        self.beam_width = beam_width
        self.colors = colors
        self.show_instructions = True
        self.instructions = [ "Instructions:",
                                "To change speed, use right and left arrow keys",
                                "To change position of the beam, use up and down arrow keys",
                                "To change width of the beam, use , and . keys",
                                "To add particle, press space",
                                "To pause/continue simulation, press enter"]


class Universe:
    def __init__(self, settings):
        self.settings = settings
        atom1 = Particle(1000, Vector2(700,400), Vector2(700,400), 20, ORANGE)
        atom2 = Particle(1, Vector2(2,400), Vector2(2-TIME_STEP*INITIAL_SPEED_X,
            INITIAL_Y- TIME_STEP*INITIAL_SPEED_Y), 1, RED)

        self.particles = [atom1, atom2]
        print("Universe has just been created!")

    def update_accelerations(self):
        for particle in self.particles:
            particle.acc = Vector2(0,0)
        for x in range(1, len(self.particles)):
            accelerations = self.compute_acc(self.particles[x], self.particles[0])
            self.particles[x].acc = accelerations[0]

    def update_positions(self):
        self.update_accelerations()
        for particle in self.particles:
            temporary_x = particle.position.x 
            temporary_y = particle.position.y
            particle.position.x = 2 * particle.position.x - particle.old_position.x 
            particle.position.x += particle.acc.x * TIME_STEP * TIME_STEP
            particle.position.y = 2 * particle.position.y - particle.old_position.y 
            particle.position.y += particle.acc.y * TIME_STEP * TIME_STEP
            particle.old_position.x = temporary_x
            particle.old_position.y = temporary_y
            particle = self.limit_position(particle)

    def limit_position(self, particle):
        if(particle.position.x > X_SCREEN_BORDER or particle.position.y > Y_SCREEN_BORDER or
            particle.position.x < 0 or particle.position.y < 0):
            y = self.settings.initial_y
            speed = self.settings.speed_x
            random_offset = random.randrange(-self.settings.beam_width, self.settings.beam_width)
            particle.position.x = 2
            particle.old_position.x = 2 - speed*self.settings.time_step
            particle.position.y = y + random_offset
            particle.old_position.y = y + random_offset
        return particle

    def compute_acc(self, particle1, particle2):
        vertical_distance = particle1.position.y - particle2.position.y
        horizontal_distance = particle1.position.x - particle2.position.x
        distance = sqrt((vertical_distance)**2 + (horizontal_distance)**2)
        force = particle1.charge * particle2.charge *K/distance**2
        acc1x = force/particle1.mass*horizontal_distance/distance
        acc1y = force/particle1.mass*vertical_distance/distance
        acc2x = -force/particle2.mass*horizontal_distance/distance
        acc2y = -force/particle2.mass*vertical_distance/distance
        return [Vector2(acc1x, acc1y), Vector2(acc2x, acc2y)]

    def add_particle(self):
        y = self.settings.initial_y
        speed = self.settings.speed_x
        colors = self.settings.colors
        random_offset = random.randrange(-self.settings.beam_width, self.settings.beam_width)
        particle = Particle(1, Vector2(2, y + random_offset), 
                        Vector2(2-self.settings.time_step*speed, y + random_offset), 1, random.choice(colors))
        self.particles.append(particle)

if __name__ == '__main__':
    main()