#!/usr/bin/env python
#
# Simulation of a guitar string
#
# Copyright (C) 2007 Maximilian Hoegner
#
# 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 2 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., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
#

import pylab, numpy
import Tkinter, tkMessageBox, tkFileDialog
import wave, struct

# settings
settings = [
("Length of fretboard", "fretboard_length", Tkinter.DoubleVar, 0.65),
("Length of string", "string_length", Tkinter.DoubleVar, 0.617),
("Plucking position (%)", "plucking_position", Tkinter.DoubleVar, 75.0),
("String bulge", "string_bulge", Tkinter.DoubleVar, 0.01),
("Number of points", "points", Tkinter.IntVar, 100),
("Simulation time", "time", Tkinter.DoubleVar, 5.0),
("Spring constant of string", "spring_constant", Tkinter.DoubleVar, 145.76),
("String mass", "string_mass", Tkinter.DoubleVar, 0.00027243),
("Frame rate", "framerate", Tkinter.IntVar, 44100),
("Attenuation", "attenuation", Tkinter.DoubleVar, 0.01),
("Steps per frame", "steps_per_frame", Tkinter.IntVar, 10),
("Frames per plot", "frames_per_plot", Tkinter.IntVar, 500),
("Wave filename", "wave_file", Tkinter.StringVar, 'test.wav', 'openfile'),
("Wave amplification", "amplification", Tkinter.DoubleVar, 17.0)
]

class TKsettings:
	def __init__(self, settings, title="Settings"):
		self.tkroot = Tkinter.Tk()
		self.vars = {}
		self.types = {}
		self.tkroot.title(title)
		Tkinter.Label(self.tkroot,text="Settings").grid(columnspan=3)

		for i in xrange(len(settings)):
			Tkinter.Label(self.tkroot,text=settings[i][0]).grid(row=i+1,column=0)

			self.vars[settings[i][1]] = settings[i][2]()
			self.vars[settings[i][1]].set(settings[i][3])

			if len(settings[i])==5 and settings[i][4]=='openfile':
				Tkinter.Entry(self.tkroot,textvariable=self.vars[settings[i][1]]).grid(row=i+1,column=1)
				Tkinter.Button(self.tkroot,text="Select",command=lambda x=settings[i][1]: self.file(x)).grid(row=i+1,column=2)
				self.types[settings[i][1]] = settings[i][4]
			elif len(settings[i])<5 or settings[i][4]=='getvalue':
				Tkinter.Entry(self.tkroot,textvariable=self.vars[settings[i][1]]).grid(row=i+1,column=1)
				self.types[settings[i][1]] = 'getvalue'
			else:
				raise Exception("Unknown setting type "+str(settings[i][4]))

		Tkinter.Button(self.tkroot,text="OK",command=self.ok).grid(row=len(settings)+1,columnspan=2)

		self.tkroot.mainloop()

	def file(self, var):
		self.vars[var].set(tkFileDialog.asksaveasfilename())

	def ok(self):
		try:
			for i in self.vars.iteritems():
				if self.types[i[0]] == 'openfile':
					try:
						globals()[i[0]].close()
					except:
						pass

					globals()[i[0]] = open(i[1].get(), 'w')
				else:
					globals()[i[0]] = i[1].get()

			self.tkroot.destroy()
		except Exception,e:
			tkMessageBox.showerror('Error',e)

TKsettings(settings, "Simulation of a guitar string")

# set variables from settings

steps_per_plot = steps_per_frame*frames_per_plot

time_interval = 1.0/steps_per_frame/framerate

steps = int(time/time_interval)

scm = 1.0*(points-1)*spring_constant/string_mass * points
						# spring constant / mass for string part

part_length = string_length/(points-1)		# length of a part of a string (not stretched)

N = 0.0 + 0.0j					# position of nut
B = fretboard_length + 0.0j			# position of bridge

P = fretboard_length*plucking_position/100.0 + string_bulge*1j
						# plucking position

NP = numpy.abs(N-P)				# distance of nut and plucking position
BP = numpy.abs(B-P)				# distance of bridge and plucking position

point_dist = (NP+BP) / (points-1)		# distance between points (stretched)

Ppos = numpy.zeros(points, numpy.complex)	# point positions
Pvel = numpy.zeros(points, numpy.complex)	# point velocities

attenuation_factor = 1-attenuation

# create initial point positions
pointsright = numpy.floor(NP/point_dist+1.0)
Ppos[:pointsright] = (P-N)/NP*point_dist*numpy.arange(pointsright) + N

pointsleft = numpy.floor(BP/point_dist+1.0)
Ppos[-pointsleft:] = (P-B)/BP*point_dist*numpy.arange(pointsleft-1,-1,-1) + B

# open wave
sample_width = 4

w = wave.open(wave_file)
w.setnchannels(1)
w.setsampwidth(sample_width)
w.setframerate(framerate)
pcodes = {1:'B',2:'h',4:'i'}
pcode = pcodes[sample_width]
loudest = numpy.power(256,sample_width)/2
wavmul = loudest*amplification/string_bulge

# prepare plotting
pylab.ion()

line, = pylab.plot(numpy.real(Ppos), numpy.imag(Ppos),'ro')
pylab.axis(( numpy.real(N), numpy.real(B), -numpy.imag(P)*1.5, numpy.imag(P)*1.5 ))

pylab.draw()

# simulation
for i in xrange(steps):
	# calculate diff of positions
	d = numpy.diff(Ppos)
	d = d*(1-part_length/numpy.abs(d))
	dm = d[1:]
	d = d[:-1]

	# calculate elongation for Hooke's law
	s = dm - d

	# calculate new velocities
	Pvel[1:-1] += scm*time_interval*s

	# steal energy
	Pvel[-2] *= attenuation_factor

	# calculate new positions
	Ppos += Pvel*time_interval

	# record sound
	if i%steps_per_frame == 0:
		w.writeframes(struct.pack(pcode, numpy.imag(Ppos[-2]) * wavmul ))

	# plot image
	if i%steps_per_plot == 0:
		print  "%d / %d\t%f / %f " % (i, steps, time_interval*i, time)
		line.set_xdata(numpy.real(Ppos))
		line.set_ydata(numpy.imag(Ppos))
		pylab.draw()