from PDEs import *


SCREEN_SIZE = (800,600)

from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLE import *

import math
from random import random
import pygame
from pygame.locals import *
from render_functions import *

def resize(width, height):
    glViewport(0, 0, width, height)
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    gluPerspective(45.0, float(width)/height, 0.1, 1000.)
    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()
    gluLookAt(15,0,12,0,0,0,0,0,1)
    
       
def init():
    glEnable(GL_DEPTH_TEST)
    glEnable(GL_NORMALIZE);
    glShadeModel(GL_SMOOTH)
    glClearColor(1, 1, 1, 0.0)
    glEnable(GL_LINE_SMOOTH)
    glEnable(GL_BLEND)
    glHint( GL_LINE_SMOOTH_HINT, GL_NICEST );
    glHint( GL_POLYGON_SMOOTH_HINT, GL_NICEST );
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
    glEnable(GL_COLOR_MATERIAL)
    glEnable(GL_LIGHTING)
    glEnable(GL_LIGHT0)        
    glLightfv(GL_LIGHT0, GL_POSITION,  (1.5, 1, 1,0))
    #glEnable(GL_TEXTURE_1D)


def noise(data, amp):
    for x in range(len(data)):
        for y in range(len(data[0])):
            data[x][y] += (random()-.5)*amp

    
def run():
    CORNERS = 0
    EDGES = 1
    mode = EDGES
    deltaT = .5 #setting this to  1 creates noise in the edges
    deltaX = 1.0
    D = 1.0
    alpha = .5*D*deltaT/deltaX**2 # needed for aig
    #alpha = 2*D*deltaT/deltaX**2 # created smooth sor transition
    #alpha = 3.9*D*deltaT/deltaX**2 # created fast sor solution
    #alpha = 4.0*D*deltaT/deltaX**2 #contained virbations in sor
    rotation = [0,0]
    pygame.init()
    screen = pygame.display.set_mode(SCREEN_SIZE,HWSURFACE|OPENGL|DOUBLEBUF)
    resize(*SCREEN_SIZE)
    init()
    clock = pygame.time.Clock()      
    max_fps = 40
    t = 0
    data = [ [ 0 for i in range(80) ] for j in range(80) ]
    rendermode = 1
    dAngle = 1
    alg = 1
    data[0][0] = 5.0
    data[-1][0] = 5.0
    data[-1][-1] = 5.0
    data[0][-1] = -5.0    

    while True:  
        rotation[0] += dAngle
        for event in pygame.event.get():
            if event.type == QUIT:
                return
            if event.type == KEYUP and event.key == K_ESCAPE:
                return
            if event.type == KEYUP and event.key == K_SPACE:
                rendermode = rendermode +1
                if rendermode == 5:
                    rendermode = 0
            if event.type == KEYUP and event.key == K_c:  
                print 'corner mode'  
                mode = CORNERS
            if event.type == KEYUP and event.key == K_e:    
                print 'edge mode'
                mode = EDGES
            
            ###adding noise
            if event.type == KEYUP and event.key == K_n:    
                noise(data,10.0)
            if event.type == KEYUP and event.key == K_m:    
                for x in range(1,len(data)-1):
                    for y in range(1,len(data[0])-1):
                        data[x][y] += (random()-.5)*4
            ##turntable
            if event.type == KEYUP and event.key == K_t:    
                dAngle = 1-dAngle
                
            ##select SOR / ADI solver
            if event.type == KEYUP and event.key == K_a:    
                alg = 1-alg
                if alg==0: print 'sor'
                else: print 'adi'
            
            ## plane presets
            if event.type == KEYUP and event.key == K_1:    
                n = len(data)
                m = len(data[0])
                for j in range(n):
                    for k in range(m):
                        data[j][k] = 0
                        if j==n-1 : 
                            data[j][k] = 0
                        if k == 0 : 
                            data[j][k] = 5
                        if k == m-1 : 
                            data[j][k] = -5
            
            if event.type == KEYUP and event.key == K_2:    
                n = len(data)
                m = len(data[0])
                jc = int(n/2)
                kc = int(m/2)
                for j in range(n):
                    for k in range(m):
                        data[j][k] = 0
                        if j in range(jc-10,jc+10) and k in range(kc-10,kc+10): 
                            data[j][k] = 5
                            
        ### mouse functions
        if pygame.mouse.get_pressed()[0]:
            rotation=map(sum, zip(rotation,pygame.mouse.get_rel()))
        else: pygame.mouse.get_rel()

        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
        
        #RENDERFUNCTIONS  - use as you please
        if rendermode in (0,1) : shade_grid(data,rotation)
        if rendermode in (0,2):  wireframe_grid(data,rotation)
        if rendermode in (3,):   x_grid(data,rotation)
        if rendermode in (4,):   y_grid(data,rotation)
        
        if alg:
            alpha = .5*D*deltaT/deltaX**2 # created save adi solution
            data = adi_diffusion_step(data, alpha, mode)
            data = transpose(adi_diffusion_step(transpose(data), alpha, mode))
        else:
            # set to 3.999 for borderline behaviour
            alpha = 3.9*D*deltaT/deltaX**2 # created fast sor solution
            sor_diffusion_step(data, alpha, deltaX)
        
        # Show the screen
        pygame.display.flip()
        
        #constant framerate
        framerate = clock.get_fps()
        clock.tick_busy_loop(max_fps)
        
        # variable framereate for adaptive stepping 
        #passed = clock.tick_busy_loop()
        #pygame.time.wait(max(30-passed, 0))
        #clock.tick_busy_loop()
        
        pygame.image.save(screen, "png/"+str(t)+'.png')
        t += 1

        
run()