#!/usr/bin/env python

from game import *

class WaterSim(Game):
    alpha = 100000
    u = lambda Cm1,C,C1,dx: WaterSim.alpha * (Cm1+C1-2*C) * 0.5 * dx
    v = lambda u,v,d,dt: ( v + (u - d * v) * dt )
    w = lambda w,v,dt: w + v * dt

    def __init__(self,N):
        self.N = N
        self.init_data(N//2)
        self.mousepos = (0,0)
        super(WaterSim, self).__init__( 1000, 600, 120)

    def init_data(self, middle):
        self.C = [50] * self.N
        for i in range(middle - self.N//4,middle + self.N//4):
            if i < 0 or i >= self.N: continue
            self.C[i] += (1-math.cos(2 * math.pi * (i - middle + self.N//4) / (self.N//2))) * 100
        self.C1 = [0] * self.N

    def on_event(self, event):
        if event.type == MOUSEBUTTONDOWN and event.button == 1:
            self.init_data( self.N*event.pos[0] // self._width )
        elif event.type == MOUSEMOTION:
            self.mousepos = event.pos[:]
        else:
            super(WaterSim, self).on_event(event)

    def on_update(self, dtime):
        C2 = [0] * self.N
        C2[0] =  WaterSim.u(self.C[1]  , self.C[0]  , self.C[1]  , 1.0/self.N)
        C2[-1] = WaterSim.u(self.C[-2] , self.C[-1] , self.C[-2] , 1.0/self.N)
        #C2[0] = 100 * (self.C[1] - self.C[0]) * 100/self.N
        #C2[-1] = 100 * (self.C[-2] - self.C[-1]) * 100/self.N

        for i in range(1,self.N-1):
            C2[i] = WaterSim.u(self.C[i-1],self.C[i], self.C[i+1], 1.0/self.N)
            #C2[i] =  1000 * (self.C[i+1] + self.C[i-1] - 2*self.C[i]) * 0.5 * 100 /(self.N)

        for i in range(0,self.N):
            if i == 0 or i == self.N-1:
                d = 10
                #d = 0.95
            else:
                d = 0.1
                #d = 0.9999
            #self.C1[i] = d * (self.C1[i] + C2[i] * dtime)
            self.C1[i] = WaterSim.v(C2[i], self.C1[i], d, dtime)
            #self.C[i] = ( self.C[i] + self.C1[i] * dtime )
            self.C[i] = WaterSim.w(self.C[i], self.C1[i], dtime)

            if self.C[i] < 0:
                self.C[i] = 0

        n = self.mousepos[0] * self.N // self._width
        alpha = 0.05
        v = self._height - self.mousepos[1]
        if self.C[n] > v:
            self.C[n] =  self.C[n]*(1-alpha) + alpha*v

    def on_render(self,srfc):
        BLACK = (0,0,0)
        BLUE = (100,100,255)

        srfc.fill(BLACK)
        prev_point = ( 0, self._height-int(self.C[0]) )
        for i in range(1,self.N):
            new_point = ( (i*self._width)//(self.N-1), self._height - int( self.C[i] ) )
            pygame.draw.line( srfc, BLUE, prev_point, new_point, 3 )
            prev_point = new_point

if __name__ == '__main__':
    WaterSim(50).on_execute()