# GVB - a GTK+/GNOME vibrations simulator
#
# Copyright (C) 2008 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)
#		print astep, aqstep

		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]]))
#		print self.mat

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

	def update(self, points, speeds, time):
#		print self.mat, points, speeds
#		print self.mat.shape, concatenate([points,speeds]).shape
		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:
#			print "rebuilding"
			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):
#		print self.n
		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 = -4*identity(self.n[1])+diag([1]*(self.n[1]-1),1)+diag([1]*(self.n[1]-1),-1)

		#rowlenght:
		rl = self.n[1]

#		print ([
		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]]))
#		print self.mat

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

	def update(self, points_matrix, speeds_matrix, time):
#		print self.mat, points, speeds
#		print self.mat.shape, concatenate([points,speeds]).shape

#		print points_matrix

		points = concatenate( points_matrix)
		speeds = concatenate(speeds_matrix)

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

#		rowlenght = self.n[1]
#		points_lenght = self.n[0]*self.n[1]

#		new_points = array([new_coords[i*rowlenght:(i+1)*rowlenght] for i in range(self.n[0])])
#		new_speeds = array([new_coords[points_lenght+i*rowlenght:points_lenght+(i+1)*rowlenght] for i in range(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 = -4*identity(self.n[1])+diag([1]*(self.n[1]-1),1)+diag([1]*(self.n[1]-1),-1)

		#rowlenght:
		rl = self.n[1]

#		print ([
		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_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):

#		print mhstep, mqhstep, mvstep, mqvstep, pos

		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, 'naif matrix': Calculator_matrixnaif, 'quad': Calculator_quad}

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

calculators_dict={1: calculators_1d, 2: calculators_2d}

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