# GVB - a GTK+/GNOME vibrations simulator
#
# Copyright © 2008-2013 Pietro Battiston
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA


from scipy import array, dot, diag, identity, bmat, concatenate, exp, transpose, zeros, real, sum, ones, split, sqrt, sin, arccos, sign, repeat, matrix
from scipy.linalg import eig, inv, solve

import time as timer


class Calculator():
    def __init__(self, n, gr, pos, speeds):
        self.n=n
        self.gr=gr
        self.start=pos
        self.startspeeds=speeds

    def reconfigure(self, **args):
        if 'shape' in args:
            self.n=args['shape']
        if 'gr' in args:
            self.gr=args['gr']
        if 'pos' in args:
            self.start=args['pos']

class Calculator1d(Calculator):
    '''
    Overriding two methods just because I want "self.n" to be an int.
    '''
    def __init__(self, *args):
        Calculator.__init__(self, *args)
        self.n=self.n[0]

    def reconfigure(self, **args):
        Calculator.reconfigure(self, **args)
        if 'shape' in args and args['shape']:
            self.n=self.n[0]

class Calculator_naif(Calculator1d):
    discrete=True
    def update(self, points, speeds, time):
        accels=array([2*points[0]-points[1]]+[2*points[i]-points[i-1]-points[i+1] for i in range(1,self.n-1)]+[2*points[self.n-1]-points[self.n-2]])
        new_speeds=speeds-dot(self.gr, accels)
        new_points=points+dot(self.gr, new_speeds)

        return new_points, new_speeds

class Calculator_quad(Calculator1d):
    discrete=True
    def update(self, points, speeds, time):
        step=1/(self.n+1)
        qstep=step**2
        astep = repeat(step,self.n+1)
        aqstep = repeat(qstep,self.n+1)

        factor = (lambda ip : ip * sin(arccos(astep/ip)) )

        strenght = (lambda ydiff : factor(sqrt(ydiff*ydiff+aqstep)) * sign(ydiff) )

        strenghts = strenght(concatenate([points, zeros(1)]) - concatenate([zeros(1), points]))

        new_speeds=speeds+self.gr*(strenghts[1:]-strenghts[:-1])
        new_points=points+dot(self.gr, new_speeds)

        return new_points, new_speeds


class Calculator_matrixnaif(Calculator1d):
    discrete=True
    def __init__(self, *args):
        Calculator1d.__init__(self, *args)

        self.build_mat()

    def build_mat(self):
        n_id=identity(self.n)
        self.op=diag([-2]*self.n)+diag([1]*(self.n-1),1)+diag([1]*(self.n-1),-1)
        self.mat=array(bmat([[n_id, self.gr*n_id], [self.gr*self.op, n_id]]))

    def reconfigure(self, **args):
        Calculator1d.reconfigure(self, **args)
        self.build_mat()

    def update(self, points, speeds, time):
        new_coords=dot(self.mat, concatenate([points, speeds]))

        return new_coords[:self.n], new_coords[self.n:]

class Calculator_eig(Calculator1d):
    discrete=False
    def __init__(self, *args):
        Calculator1d.__init__(self, *args)

        self.build_mat()
        self.build_operator()
        self.time=0
        

    def build_mat(self):#FIXME:merge 
        n_id=identity(self.n)
        n_zero=zeros([self.n, self.n])
        self.op=diag([-2]*self.n)+diag([1]*(self.n-1),1)+diag([1]*(self.n-1),-1)
        self.mat=array(bmat([[n_zero, n_id], [self.op, n_zero]]))
        self.eigval, self.eigvec = eig(self.mat)

    def build_operator(self):
        self.startvec = solve(self.eigvec, concatenate([self.start, self.startspeeds]))
        self.operator=dot(self.eigvec, diag(self.startvec)).transpose()

    def reconfigure(self, **args): 
        Calculator1d.reconfigure(self, **args)
        if 'shape' in args:
            self.build_mat()
            self.build_operator()
            return
        if 'pos' in args:
            self.build_operator()

    def update(self, points, speeds, time):
        expmultiplier=exp(self.eigval*time)
        result=real(dot(expmultiplier, self.operator))    #This dot product is perfectly equivalent to mutliplying by a diagonal and then adding rows

        return result[:self.n], result[self.n:]



class Calculator_naif2d(Calculator):
    discrete=True
    def update(self, points, speeds, time):
        n=self.n[0]
        m=self.n[1]
        accels=array(
                    [
                    [ 4*points[0][0]-points[1][0]-points[0][1] ]
                    +[ 4*points[0][i]-points[0][i-1]-points[0][i+1]-points[1][i] for i in range(1,n-1) ]
                    +[ 4*points[0][n-1]-points[0][n-2]-points[1][n-1] ]
                    ]
                +
                    [
                    [ 4*points[j][0]-points[j+1][0]-points[j-1][0]-points[j][1] ]
                    +[ 4*points[j][i]-points[j][i-1]-points[j][i+1]-points[j+1][i]-points[j-1][i] for i in range(1,n-1) ]
                    +[ 4*points[j][n-1]-points[j][n-2]-points[j+1][n-1]-points[j-1][n-1] ]
                    for j in range(1, m-1)]
                +
                    [
                    [ 4*points[m-1][0]-points[n-2][0]-points[n-1][1] ]
                    +[ 4*points[m-1][i]-points[m-1][i-1]-points[m-1][i+1]-points[m-2][i] for i in range(1,n-1)]
                    +[ 4*points[m-1][n-1]-points[m-1][n-2]-points[m-2][n-1] ]
                    ]
                    )

        new_speeds=speeds-dot(self.gr, accels)
        new_points=points+dot(self.gr, new_speeds)

        return new_points, new_speeds

class Calculator_matrixnaif2d(Calculator):
    discrete=True
    def __init__(self, *args):
        Calculator.__init__(self, *args)

        self.build_mat()

    def build_mat(self):
        self.edge = self.n[0]*self.n[1]
        small_id = identity(self.n[1])
        small_zeros = zeros((self.n[1], self.n[1]))
        n_id = identity(self.edge)
        small_op = diag(ones(self.n[1]-1),1) - 4*identity(self.n[1]) + diag(ones(self.n[1]-1),-1)
#        self.op=diag([-2]*self.n)+diag([1]*(self.n-1),1)+diag([1]*(self.n-1),-1)
        #rowlenght:
        rl = self.n[1]

        self.op = bmat([
                        [small_op, small_id] + [small_zeros for i in range(2, self.n[0])]
                        ]+[
                        [small_zeros for i in range(j)] + [small_id, small_op, small_id] + [small_zeros for i in range(j+1, rl-2)]
                        for j in range(rl-2)] + [
                        [small_zeros for i in range(2, self.n[0])] + [small_id, small_op]
                        ])

        self.mat=array(bmat([[n_id, self.gr*n_id], [self.gr*self.op, n_id]]))

    def reconfigure(self, **args):
        Calculator.reconfigure(self, **args)
        self.build_mat()

    def update(self, points_matrix, speeds_matrix, time):
        points = concatenate( points_matrix)
        speeds = concatenate(speeds_matrix)

        new_coords=array(split(dot(self.mat, concatenate([points, speeds])), 2*self.n[0]))

        return new_coords[:self.n[0]], new_coords[self.n[0]:]

class Calculator_eig2d(Calculator):
    discrete=False
    def __init__(self, *args):
        Calculator.__init__(self, *args)

        self.build_mat()
        self.build_operator()
        self.time=0
        

    def build_mat(self):#FIXME:merge 
        self.edge = self.n[0]*self.n[1]
        small_id = identity(self.n[1])
        small_zeros = zeros((self.n[1], self.n[1]))
        n_id = identity(self.edge)
        n_zero = zeros([self.edge, self.edge])
        small_op = diag(ones(self.n[1]-1),1) - 4*identity(self.n[1]) + diag(ones(self.n[1]-1),-1)
        # Number of blocks (per dimension):
        rl = self.n[0]
        self.op = bmat([
                        [small_op, small_id] + [small_zeros for i in range(2, rl)]
                        ]+[
                        [small_zeros for i in range(j)] + [small_id, small_op, small_id] + [small_zeros for i in range(j+1, rl-2)]
                        for j in range(rl-2)] + [
                        [small_zeros for i in range(rl-2)] + [small_id, small_op]
                        ])
        self.mat=array(bmat([[n_zero, n_id], [self.op, n_zero]]))
        self.eigval, self.eigvec = eig(self.mat)

    def build_operator(self):
        self.startvec = solve(self.eigvec, concatenate([concatenate(self.start), concatenate(self.startspeeds)]))
        self.operator=dot(self.eigvec, diag(self.startvec)).transpose()

    def reconfigure(self, **args): 
        Calculator.reconfigure(self, **args)
        if 'shape' in args:
            self.build_mat()
            self.build_operator()
            return
        if 'pos' in args:
#            print "rebuilding"
            self.build_operator()

    def update(self, points, speeds, time):
        expmultiplier=exp(self.eigval*time)
        result=array(split(real(dot(expmultiplier, self.operator)), self.n[0]*2))    #This dot product is perfectly equivalent to mutliplying by a diagonal and then adding rows

        return result[:self.n[0]], result[self.n[0]:]

class Calculator_quad_2D(Calculator):
    discrete=True
    def __init__(self, *args):
        Calculator.__init__(self, *args)

        self.build_lambdas()
        self.time=0

    def build_lambdas(self):
        hstep=1/(self.n[1]+1)
        qhstep=hstep**2
        mhstep = array(repeat(matrix(repeat(hstep,self.n[1]+1)), self.n[0], 0))
        mqhstep = array(repeat(matrix(repeat(qhstep,self.n[1]+1)), self.n[0], 0))
        vstep=1/(self.n[0]+1)
        qvstep=vstep**2
        mvstep = array(repeat(matrix(repeat(vstep,self.n[1])), self.n[0]+1, 0))
        mqvstep = array(repeat(matrix(repeat(qvstep,self.n[1])), self.n[0]+1, 0))

        hfactor = (lambda ip : ip * sin(arccos(mhstep/ip)) )
        self.hstrenght = (lambda zdiff : hfactor(sqrt(zdiff*zdiff+mqhstep)) * sign(zdiff) )

        vfactor = (lambda ip : ip * sin(arccos(mvstep/ip)) )
        self.vstrenght = (lambda zdiff : vfactor(sqrt(zdiff*zdiff+mqvstep)) * sign(zdiff) )

    def reconfigure(self, **args): 
        Calculator.reconfigure(self, **args)
        if 'shape' in args:
            self.build_lambdas()

    def update(self, pos, speeds, time):

        hstrenghts = self.hstrenght(concatenate([pos, zeros([self.n[0],1])], 1 ) - concatenate([zeros([self.n[0],1]), pos],1 ))
        vstrenghts = self.vstrenght(concatenate([pos, [zeros(self.n[1])]]) - concatenate([[zeros(self.n[1])], pos]))

        new_speeds=speeds+self.gr*(hstrenghts[:,1:]-hstrenghts[:,:-1] + vstrenghts[1:]-vstrenghts[:-1])
        new_points=pos+self.gr*new_speeds

        return new_points, new_speeds



calculators_1d={'eig': Calculator_eig, 'naif': Calculator_naif, 'matrix': Calculator_matrixnaif, 'quad': Calculator_quad}

calculators_2d={'eig': Calculator_eig2d, 'naif': Calculator_naif2d, 'matrix': Calculator_matrixnaif2d,  'quad': Calculator_quad_2D}#, 'eig': Calculator_eig2}

calculators_dict={1: calculators_1d, 2: calculators_2d}

calculators={
            1: ['eig', 'quad', 'matrix', 'naif'],
            2: ['eig', 'quad', 'matrix', 'naif']
            }
