# -*- coding: utf-8 -*-
"""
Created on Sun Mar 27 08:16:24 2016

@author: gtruch
"""

# -*- coding: utf-8 -*-
"""
Animation of a mars transfer orbit.
"""

import math
import numpy as np
import copy as cp
import matplotlib.pyplot as plt
from Body import GravBody
from Vector import Vector
from Physics import UniformGravity
from Solver import RK4
from Simulation import TrajectorySim
from Search import Bisect, NewtonsMethod

import pygame as pg


# Define some colors
BLACK    = (   0,   0,   0)
WHITE    = ( 255, 255, 255)
GREEN    = (   0, 255,   0)
RED      = ( 255,   0,   0)



class Laser(pg.sprite.Sprite):
    def __init__(self,screen_size):
        pg.sprite.Sprite.__init__(self)
        self.screen_size = screen_size
        self.org_image = pg.image.load('laser.png').convert_alpha()
        self.image = pg.transform.rotate(self.org_image,0)
        self.rect = self.image.get_rect()
        

    def update(self,lander_rect):
        mouse = pg.mouse.get_pos()
        dx = lander_rect.centerx - mouse[0]
        dy = lander_rect.centery - mouse[1]
        rot = math.atan2(dy,dx)*180.0/math.pi
        self.image = pg.transform.rotate(self.org_image, -(rot+180))
        self.rect.x = lander_rect.centerx
        self.rect.y = lander_rect.centery
        return

class Lander(pg.sprite.Sprite):
    
    def __init__(self,screen_size,pixels_per_meter):
        pg.sprite.Sprite.__init__(self)
        
        self.screen_size=screen_size
        self.pixels_per_meter = pixels_per_meter
        self.on = pg.image.load('Lander_on.png').convert_alpha()
        self.off = pg.image.load('lander_off.png').convert_alpha()

        self.image = self.off
        self.rect = self.image.get_rect()
        self.body = GravBody(Vector(100.0,500.0),Vector(0.0,0.0))        
        self.update()

    def rocket_on(self):
        self.image = self.on
        
    def rocket_off(self):
        self.image = self.off

    def update(self):
        
        if self.body.position.y < 100:
            self.body.position.y = 100
            self.body.velocity.y = -self.body.velocity.y*0.7
        if self.body.position.y > self.screen_size[1]:
            self.body.position.y = self.screen_size[1]
            self.body.velocity.y = -self.body.velocity.y*0.7
            
        if self.body.position.x < 0:
            self.body.position.x = 0
            self.body.velocity.x = -self.body.velocity.x*0.8
        if self.body.position.x > self.screen_size[0]-100:
            self.body.position.x = self.screen_size[0]-100
            self.body.velocity.x = -self.body.velocity.x*0.8
            
        self.rect.y = ((self.screen_size[1]-self.body.position.y)*
                        self.pixels_per_meter)
        self.rect.x = (self.body.position.x*self.pixels_per_meter)

        

class App(object):
    
    def __init__(self):
           
        # Animation Parameters
        self.seconds_per_tick = 0.01
        self.pixels_per_meter = 1.0

        pg.init()       # Initialize pygame
        self.screen_size = (800,800)

        self.screen = pg.display.set_mode((int(self.screen_size[0]),
                                           int(self.screen_size[1])))
                                           
        self.background = pg.Surface(self.screen_size)
        self.background.fill(BLACK)

        self.lander = Lander(self.screen_size,self.pixels_per_meter)
        self.laser = Laser(self.screen_size)
        self.sgroup = pg.sprite.RenderUpdates([self.lander,self.laser])
        
        self.sim = TrajectorySim(UniformGravity(RK4(0.01),a = Vector(0.0,-4.0))
                                , self.lander.body,save_history=False)

        # Set up the timer.
        self.clock = pg.time.Clock()        
        self.ticks = pg.time.get_ticks()
        
#        self.screen.blit(self.background,(0,0))
        self.sgroup.clear(self.screen,self.background)
        dirty = self.sgroup.draw(self.screen)
        pg.display.update(dirty)
                        
                                                                                    
    def run(self):
    
#        # Run the simulation
        done = False
                                           
        while not done:
            
            self.clock.tick(60)
            
            # Process events
            for event in pg.event.get():
                if event.type == pg.QUIT:
                    done = True
                elif event.type == pg.KEYDOWN:
                    if event.key == pg.K_DOWN:
                        self.lander.rocket_on()
                        self.sim.physics.a.y = 9.8
                    if event.key == pg.K_RIGHT:
                        print "Key Right"
                        self.sim.physics.a.x = -4.0
                    if event.key == pg.K_LEFT:
                        print "Key Left"
                        self.sim.physics.a.x = 4.0                    
                    if event.key == pg.K_ESCAPE:
                        done = True
                elif event.type == pg.KEYUP:
                    if event.key == pg.K_DOWN:
                        self.lander.rocket_off()      
                        self.sim.physics.a.y = -2.0
                    if event.key == pg.K_RIGHT or event.key == pg.K_LEFT:
                        self.sim.physics.a.x = 0.0
                        

            self.advance()
                                    
        pg.quit()

    def advance(self):
            
        # Calculate the time step
        ticks = pg.time.get_ticks()
        elapsed_time = (ticks - self.ticks)*self.seconds_per_tick
        self.ticks = ticks

        # Set the drag
#        if sim.body.position.y <= 100:
#            sim.physics.            

        # Advance the simulation
        self.sim.advance(time=elapsed_time)
        
        # Update the display
        self.screen.blit(self.background,(0,0))
#        self.sgroup.update()
        self.lander.update()
        self.laser.update(self.lander.rect)
        dirty = self.sgroup.draw(self.screen)
        
        pg.display.update(dirty)
    
if __name__=='__main__':
    app = App()
    app.run()