from GA import *
from Shane import *
from Integrators import *
from numpy import *
import random as R

import unittest

def gaussA(xv, sigmav):
    result = []
    for i in range(0, len(xv)):
        result.append(gauss(xv[i], sigmav[i]))
    return array(result)

class RandomVectorGene(VectorGene):
    def __init__(self):
        self.generation = 0
        self.sigma = 1.0
        #self.v = []
        #for i in range(0, 12):
        #    self.v.append(random())
        #self.v = array(self.v)
        self.v = random.standard_normal(19) 
        self.v[12] = R.random()
        self.v[13] = R.random()
        self.v[14] = R.random()
        #v[15] vx + v[16] vy 
        #v[17] vx + v[18] vy

class Particle:
    tailCount = 10
    sigma_rv = 0.2
    sigma_vv = 0.1
    haveDrag = False
    defaultRadius = 0.05
    def __init__(self, gene = None):
        self.gene = gene
        self.nTerms = 6
        self.rv = random.standard_normal(2)
        self.vv = random.standard_normal(2)
        self.radius = Particle.defaultRadius
        self.sigma_rv = 0.2
        self.sigma_vv = 0.1
        self.steps = 0
        self.t = 0.0
        self.reproTime = 0.0
        self.babies = 0
        #self.tailCount = 10
        self.tail = []
        self.selected = False
        self.arrows = None
        self.geneArrow = None
        if gene == None:
            self.gene = RandomVectorGene()

    def color(self):
        r = self.gene.v[12]
        g = self.gene.v[13]
        b = self.gene.v[14]
        return [r,g,b]

    def speed(self):
        return mag(self.vv)

    def mutate(self, rate):
        p = Particle(self.gene.mutate(rate))
        #p.rv = gaussA(self.rv, a(self.sigma_rv * self.radius, self.sigma_rv * self.radius))
        p.rv = gaussA(self.rv, a(3 * self.radius, 3 * self.radius))
        p.vv = gaussA(self.vv, a(self.sigma_vv, self.sigma_vv))
        self.babies += 1
        return p

    def xdotcoefs(self):
        return array(self.gene.v.tolist()[0:self.nTerms])

    def ydotcoefs(self):
        return array(self.gene.v.tolist()[self.nTerms:self.nTerms * 2])

    def powerSeries(self, coefficients, x):
        sum = 0
        power = 0
        for c in coefficients:
            sum += c * x**power            
            power += 1
        return sum

    def ddot(self, coefs, xv):
        xcoefs = coefs[0:self.nTerms/2]
        ycoefs = coefs[self.nTerms/2:self.nTerms]
        return self.powerSeries(xcoefs, xv[0]) + self.powerSeries(ycoefs, xv[1])

    # Specifies the behavior of the particle
    def behavior(self, xv, vv = a(0,0)):
        g = self.gene.v
        return array([self.ddot(self.xdotcoefs(), xv) + g[15] * vv[0] + g[16] * vv[1],
                      self.ddot(self.ydotcoefs(), xv) + g[17] * vv[0] + g[18] * vv[1]])

    def _displayEquation(self, coefs):
        xcoefs = coefs[0:self.nTerms/2]
        ycoefs = coefs[self.nTerms/2:self.nTerms]
        eq = "%f + %f x + %f x^2 + %f y + %f y^2\n" % (xcoefs[0] + ycoefs[0], xcoefs[1], xcoefs[2], ycoefs[1], ycoefs[2])
        return eq

    def displayEquation(self):
        return "d^2/dt^2 x = %s\nd^2/dt^2 y = %s" % (self._displayEquation(self.xdotcoefs()),
                                                     self._displayEquation(self.ydotcoefs()))

    # Returns an array of vectors arranged on a grid
    def gridBehavior(self):
        if not self.arrows:
            self.arrows = []
            for x in frange(-2, 2, 0.2):
                for y in frange(-2, 2, 0.2):
                    start = a(x,y)
                    end = self.behavior(a(x,y))
                    arrow = end - start 
                    arrow = normalize(arrow)
                    arrow *= 0.2
                    self.arrows.append([start, start + arrow])
        return self.arrows

    # So we take our second order differential equation and split it into two
    # d^2 x/dt^2 = f(x,y)
    # d^2 y/dt^2 = g(x,y)

    # d vx/dt = f(x,y)
    # d vy/dt = g(x,y)
    # d x/dt = vx
    # d y/dt = vy
    
    def timeStep(self, dt):
        if Particle.haveDrag:
            self.vv += RK3DIntegrator(rcurry(self.behavior, self.vv), self.rv, dt)
        else:
            self.vv += RK3DIntegrator(self.behavior, self.rv, dt)
        self.rv += self.vv * dt
        self.steps += 1
        self.t += dt
        self.reproTime += dt
        self.tail.append(self.rv.copy())
        tooBig = len(self.tail) - Particle.tailCount
        if tooBig > 0:
            del self.tail[0:tooBig]

class TestParticle(unittest.TestCase):

    def assertEqualA(self, a, b):
        # Assert Equal Array
        self.assertEqual(len(a), len(b), "%s != %s" % (str(a), str(b)))
        for i in range(0, len(a)):
            self.assertEqual(a[i], b[i], "%s != %s" % (str(a), str(b)))

    def testTimeStep(self):
        g = VectorGene(array([1,1,0,0,0,0, 1,1,0,0,0,0]), 1.0)
        p = Particle(g)
        self.assertEqualA(a(0,0), p.rv)
        self.assertEqualA(a(0,0), p.vv)
        self.assertEqualA(a(1,1), p.behavior(a(0,0)))
        p.timeStep(1)
        self.assertEqualA(a(1,1), p.rv)
        self.assertEqualA(a(1,1), p.vv)
        p.timeStep(0.2)
        self.assertEqualA(a(5,1), p.rv)
        self.assertEqualA(a(1,1), p.vv)

    def testCoefs(self):
        g = VectorGene(array(range(0, 12)), 1.0)
        p = Particle(g)
        self.assertEqualA(a(0,1,2,3,4,5), p.xdotcoefs())
        self.assertEqualA(a(6,7,8,9,10,11), p.ydotcoefs())
        coefs = p.xdotcoefs()
        xcoefs = coefs[0:p.nTerms/2]
        ycoefs = coefs[p.nTerms/2:p.nTerms]
        self.assertEqualA(a(0,1,2), xcoefs)
        self.assertEqualA(a(3,4,5), ycoefs)

    def testStuff(self):
        g = VectorGene(array([1,1,0,0,0,0, 1,1,0,0,0,0]), 1.0)
        p = Particle(g)
        z = a(0,0)
        o = a(1,1)
        #self.assertEqualA(array([0,1]), array([1,0]))
        #self.assertEqualA(z, p.xdotcoefs())
        self.assertEqualA(o, p.behavior(z))
        self.assertEqualA(a(2,2), p.behavior(o))
        self.assertEqualA(a(3,3), p.behavior(a(2,2)))
        g = VectorGene(a(1,1,0,0,0,0, 0,0,0,1,1,0), 1.0)
        p = Particle(g)
        self.assertEqualA(a(2,1), p.behavior(a(1,0)))
        self.assertEqualA(a(1,2), p.behavior(a(0,1)))

        self.assertEqualA(a(3,1), p.behavior(a(2,0)))
        self.assertEqualA(a(1,3), p.behavior(a(0,2)))

        g = VectorGene(a(1,1,1,0,0,0, 0,0,0,1,1,1), 1.0)
        p = Particle(g)
        self.assertEqualA(a(3,1), p.behavior(a(1,0)))
        self.assertEqualA(a(1,3), p.behavior(a(0,1)))
        self.assertEqualA(a(7,1), p.behavior(a(2,0)))
        self.assertEqualA(a(1,7), p.behavior(a(0,2)))


if __name__ == '__main__':
    unittest.main()