#! /usr/bin/python

""" this file was written by Paul Brossier 
  it is released under the GNU/GPL license.
"""

import sys,time
from aubio.task import taskbeat,taskparams
from aubio.aubioclass import fvec, aubio_autocorr
from aubio.gnuplot import gnuplot_create, gnuplot_addargs
from aubio.aubiowrapper import *
from math import exp,log

usage = "usage: %s [options] -i soundfile" % sys.argv[0]

def parse_args():
        from optparse import OptionParser
        parser = OptionParser(usage=usage)
        parser.add_option("-i","--input",
                          action="store", dest="filename", 
                          help="input sound file")
        parser.add_option("-n","--printframe",
                          action="store", dest="printframe", default=-1, 
                          help="make a plot of the n_th frame")
        gnuplot_addargs(parser)
        (options, args) = parser.parse_args()
        if not options.filename: 
                 print "no file name given\n", usage
                 sys.exit(1)
        return options, args

def plotdata(x,y,plottitle="",**keyw):
	import Gnuplot
	return Gnuplot.Data(x, y, title="%s" % plottitle,**keyw)

options, args = parse_args()
filename = options.filename
xsize = float(options.xsize)
ysize = float(options.ysize)

printframe = int(options.printframe)

if options.outplot and printframe > 0: 
  extension = options.outplot.split('.')[-1] 
  outplot = '.'.join(options.outplot.split('.')[:-1])
else: 
  extension = ''
  outplot = None
f = gnuplot_create(outplot=outplot,extension=extension,options=options)

params = taskparams()
params.onsetmode = 'specdiff'
task = taskbeat(filename,params=params)

hopsize = params.hopsize
bufsize = params.bufsize
btstep = task.btstep
winlen = task.btwinlen
laglen = winlen/4
step = winlen/4

timesig = 0
maxnumelem = 4
gp = 0
counter = 0
flagconst = 0
constthresh = 3.901
g_var      = 3.901
rp = 0
rp1 = 0
rp2 = 0
g_mu = 0

rayparam = 48/512.*winlen

#t     = [i for i in range(hopsize)]
#tlong = [i for i in range(hopsize*(btstep-1))]
#tall  = [i for i in range(hopsize*btstep)]
#a     = [0 for i in range(hopsize*btstep)]
dfx = [i for i in range(winlen)]
dfframe = [0 for i in range(winlen)]
dfrev = [0 for i in range(winlen)]
acframe = [0 for i in range(winlen)]

localacf = [0 for i in range(winlen)]
inds = [0 for i in range(maxnumelem)]

acx = [i for i in range(laglen)]
acfout  = [0 for i in range(laglen)]

phwv  = [0 for i in range(2*laglen)]
phwvx = [i for i in range(2*laglen)]

dfwvnorm = exp(log(2.0)*(winlen+2.)/rayparam);
dfwv = [exp(log(2.)*(i+1.)/rayparam)/dfwvnorm for i in range(winlen)]

gwv = [exp(-.5*(j+1.-g_mu)**2/g_var**2) for j in range(laglen)]
rwv = [(i+1.)/rayparam**2 * exp(-(i+1.)**2 / (2.*rayparam)**2)
        for i in range(0,laglen)] 
acf = fvec(winlen,1)

nrframe = 0
while (task.readsize == params.hopsize):
  task()
  #print task.pos2
  #a[:-hopsize] = [i for i in a[-(btstep-1)*hopsize:]]
  #a[-hopsize:] = [task.myvec.get(i,0) for i in t]

  #g('set xrange [%f:%f]' % (t[0],t[-1]))
  #time.sleep(.2)
  if task.pos2==btstep-1:
    nrframe += 1
    dfframe = [task.dfframe.get(i,0) for i in range(winlen)]
    if printframe == nrframe or printframe == -1:
      d  = [[plotdata(range(-winlen,0),dfframe,plottitle="onset detection", with='lines')]]
    # start beattracking_do
    for i in range(winlen):
      dfrev[winlen-1-i] = 0.
      dfrev[winlen-1-i] = dfframe[i]*dfwv[i]
    aubio_autocorr(task.dfframe(),acf()); 
    acframe = [acf.get(i,0) for i in range(winlen)]
    if not timesig:
      numelem = 4
    else:
      numelem = timesig

    old = 0
    acfout = [0 for i in range(winlen/4)]
    for i in range(1,laglen-1):
      for a in range(1,numelem+1):
        for b in range (1-a,a):
          acfout[i] += acframe[a*(i+1)+b-1] * 1./(2.*a-1.)*rwv[i]
          if old < acfout[i]:
            old = acfout[i]
            maxi = i
    rp = max(maxi,1);

    if printframe == nrframe or printframe == -1:
      rwvs = [rwv[i]*max(acframe) for i in range(len(rwv))]
      d += [[plotdata(acx,acfout,plottitle="comb filterbank", with='lines', axes='x1y1'),
          plotdata([rp,rp],[1.2*old,min(acfout)],plottitle="period", with='impulses', axes='x1y1'),
          plotdata(acx,rwvs,plottitle="L_w", with='lines', axes='x1y1')]]

    # getperiod
    inds = [0 for i in range(maxnumelem)]
    localacf = [0 for i in range(winlen)]
    period = 0
    for a in range(1,4+1):
      for b in range(1-a,a):
        localacf[a*rp+b-1] = acframe[a*rp+b-1]
    for i in range(numelem):
      maxindex = 0
      maxval = 0.0
      for j in range(rp*(i+1)+i):
        if localacf[j] > maxval:
          maxval = localacf[j]
          maxind = j
        localacf[j] = 0
      inds[i] = maxind
    for i in range(numelem):
      period += inds[i]/(i+1.)
    period = period/numelem
    #print "period", period

    # checkstate 
    if gp:
      # context dependant model
      acfout = [0 for i in range(winlen/4)]
      old = 0
      for i in range(laglen-1):
        for a in range(timesig):
          for b in range(1-a,a):
            acfout[i] += acframe[a*(i+1)+b-1] * gwv[i]
        if old < acfout[i]:
          old = acfout[i]
          maxi = i
      gp = maxi
    else:
      # general model
      gp = 0
    #print "gp", gp
    if printframe == nrframe or printframe == -1:
      gwvs = [gwv[i]*max(acfout) for i in range(len(gwv))]
      d += [[plotdata(acx,acfout,plottitle="comb filterbank", with='lines', axes='x1y1'),
          plotdata(gp,old,plottitle="period", with='impulses', axes='x1y1'),
          plotdata(acx,gwvs,plottitle="L_{gw}", with='lines', axes='x1y1')]]

    if counter == 0:
      # initial step
      if abs(gp-rp) > 2.*constthresh:
        flagstep = 1
        counter  = 3
      else:
        flagstep = 0
    #print "flagstep", flagstep
    #print "rp2,rp1,rp", rp2,rp1,rp
    acfw = [dfframe[i]*dfwv[i] for i in range(winlen)]

    if counter == 1 and flagstep == 1:
      # "3rd frame after flagstep set"
      if abs(2.*rp-rp1- rp2) < constthresh:
        flagconst = 1
        counter = 0
      else:
        flagconst = 0
        counter = 2
    elif counter > 0:
      counter -= 1

    rp2 = rp1; rp1 = rp

    if flagconst:
      # "first run of new hypothesis"
      gp = rp
      g_mu = gp
      timesig = 4 #FIXME
      gwv = [exp(-.5*(j+1.-g_mu)**2/g_var**2) for j in range(laglen)]
      flagconst = 0
      bp = gp
      phwv = [1 for i in range(2*laglen)]
    elif timesig:
      # "contex dependant"
      bp = gp
      if step > lastbeat:
        phwv = [exp(-.5*(1.+j-step+lastbeat)**2/(bp/8.)) for j in range(2*laglen)]
      else:
        print "NOT using phase weighting"
        phwv = [1 for i in range(2*laglen)]
    else:
      # "initial state"
      bp = rp
      phwv = [1 for i in range(2*laglen)]

    while bp < 25:
      print "WARNING, doubling the beat period"
      bp *= 2

    # 
    phout = [0. for i in range(winlen)]

    kmax = int(winlen/float(bp));

    old = 0
    for i in range(bp):
      phout[i] = 0.
      for k in range(kmax):
        phout[i] += dfrev[i+bp*k] * phwv[i]
      if phout[i] > old:
        old = phout[i]
        maxi = i
    maxindex = maxi 
    if (maxindex == winlen - 1): maxindex = 0
    phase = 1 + maxindex
    i = 1
    beat = bp - phase
    beats= []
    if beat >= 0: beats.append(beat)
    while beat+bp < step: 
      beat += bp
      beats.append(beat)
    lastbeat = beat
    #print beats,
    #print "the lastbeat is", lastbeat

    # plot all this
    if printframe == nrframe or printframe == -1:
      phwvs = [phwv[i]*max(phout) for i in range(len(phwv))]
      d += [[plotdata(range(-laglen,0),phwvs[laglen:0:-1],plottitle="A_{gw}", with='lines',axes='x1y1'),
          plotdata(range(-laglen,0),phout[laglen:0:-1],plottitle="df", with='lines'),
          plotdata(-phase,old,plottitle="phase", with='impulses', axes='x1y1'),
          plotdata([i for i in beats],[old for i in beats],plottitle="predicted", with='impulses')
          ]]
    #d += [[plotdata(dfx,dfwv,plottitle="phase weighting", with='lines', axes='x1y2'),
    #    plotdata(dfx,dfrev,plottitle="df reverse", with='lines', axes='x1y1')]]
    #d += [[plotdata(dfx,phout,plottitle="phase", with='lines', axes='x1y2')]]
    #d += [[plotdata(dfx,dfwv,plottitle="phase weighting", with='lines', axes='x1y2'),
    #    plotdata(dfx,dfrev,plottitle="df reverse", with='lines', axes='x1y1')]]
    #d += [[plotdata(dfx,phout,plottitle="phase", with='lines', axes='x1y2')]]

      f('set lmargin 4')
      f('set rmargin 4')
      f('set size %f,%f' % (1.0*xsize,1.0*ysize) )
      f('set key spacing 1.3')
      f('set multiplot')

      f('set size %f,%f' % (1.0*xsize,0.33*ysize) )
      f('set orig %f,%f' % (0.0*xsize,0.66*ysize) )
      f('set xrange [%f:%f]' % (-winlen,0) )
      f.title('Onset detection function')
      f.xlabel('time (df samples)')
      f.plot(*d[0])
      f('set size %f,%f' % (0.5*xsize,0.33*ysize) )
      f('set orig %f,%f' % (0.0*xsize,0.33*ysize) )
      f('set xrange [%f:%f]' % (0,laglen) )
      f.title('Period detection: Rayleygh weighting')
      f.xlabel('lag (df samples)')
      f.plot(*d[1])
      f('set size %f,%f' % (0.5*xsize,0.33*ysize) )
      f('set orig %f,%f' % (0.5*xsize,0.33*ysize) )
      f('set xrange [%f:%f]' % (0,laglen) )
      f.title('Period detection: Gaussian weighting')
      f.xlabel('lag (df samples)')
      f.plot(*d[2])
      f('set size %f,%f' % (1.0*xsize,0.33*ysize) )
      f('set orig %f,%f' % (0.0*xsize,0.00*ysize) )
      f('set xrange [%f:%f]' % (-laglen,laglen) )
      f.title('Phase detection and predicted beats')
      f.xlabel('time (df samples)')
      f.plot(*d[3])
      f('set nomultiplot')
    if printframe == -1: a = sys.stdin.read()
    elif 0 < printframe and printframe < nrframe:
      break