#!/usr/bin/python
# mp3_sync by Nick Fankhauser in 2009
# Synchronizes movie frame progression to beats detected in mp3 audio
import sys,os,shutil,progressbar,numpy.fft,Image
fps=25.0
num_bands=32 # number of frequency bands
bps=40 # beats per second
threshold=25 # C value default = 25
bgdir='/home/nyk/unm/' # directory containing the images
default_widgets=[progressbar.Percentage(), ' ', progressbar.Bar(), progressbar.ETA()]

def empty_directory(d):
 for f in os.listdir(d):
  os.unlink('%s/%s' % (d.rstrip('/'),f))

def wav_vect(fn):
 """Reads 16bit stereo PCM Wave File"""
 fromLittleEndianDWORD=lambda s : ord(s[0]) + (ord(s[1])<<8) + (ord(s[2])<<16) + (ord(s[3])<<24)
 fromLittleEndianWORD=lambda s: ord(s[0]) + (ord(s[1])<<8)
 wav=open(fn).read()
 format=fromLittleEndianWORD(wav[20:22])
 assert format==1 # PCM
 channels=fromLittleEndianWORD(wav[22:24])
 assert channels==2 # Stereo
 data_pos=wav.find('data') + 8 # start of wave data
 samp2wav=lambda x : int(x) * 2 + data_pos
 samplerate=fromLittleEndianDWORD(wav[24:28]) # samples per second
 bits=fromLittleEndianWORD(wav[34:36]) # samples per second
 assert bits==16 # 16 bit
 block_align=channels*bits/8
 data_len=len(wav)-data_pos
 total=float(data_len)/float(block_align)/float(samplerate)
 print 'Samplerate = %d, Duration [s] = %2.1f, Channels = %d, Bits = %d' % (samplerate,total,channels,bits)
 values=[]
 pbar = progressbar.ProgressBar(widgets=['Loading samples: '] + default_widgets).start()
 data_range=data_len/block_align
 for n in range(data_range):
  p=data_pos + n * block_align
  v1=fromLittleEndianWORD(wav[p:(p+2)])
  v2=fromLittleEndianWORD(wav[(p+2):(p+4)])
  values.append(complex(v1,v2))
  pbar.update(float(n)/float(data_range)*100)
 pbar.finish()
 return values,samplerate,total

fn=sys.argv[1]
# load and analyze MP3 file
assert fn[-3:]=='mp3'
wfn='/tmp/'+fn.lower().replace('mp3','wav')
if not os.path.isfile(wfn): os.system('madplay %s -o %s' % (fn,wfn))
samples,samplerate,duration=wav_vect(wfn)
beat=samplerate/bps
band_len=beat/num_bands
bandhist=[[]]*(num_bands+1)
beats=[]
beat_range=len(samples)/beat
average=lambda x: sum(x)/len(x)
complex_module=lambda x : (x.real**2) + (x.imag**2)
pbar = progressbar.ProgressBar(widgets=['Detecting beats: '] + default_widgets).start()
for i in range(beat_range):
 bt=samples[(i*beat):(i*beat+beat)]
 fftm=map(complex_module,numpy.fft.fft(bt))
 bands=[fftm[x:x+band_len] for x in range(0,len(fftm),band_len)]
 bandsm=map(average,bands)
 for b in range(len(bandsm)):
  if len(bandhist[b])>bps:
   if bandsm[b]>average(bandhist[b])*threshold: beats.append([float(i)/float(bps),b,bandsm[b]])
   bandhist[b].pop(0)
  bandhist[b].append(bandsm[b])
 pbar.update(float(i)/float(beat_range)*100)
pbar.finish()
subbands=map(lambda x : x[1],beats)
subband_occ=list(reversed(sorted(map(lambda b : (subbands.count(b),b),set(subbands)))))
print '\n' + '\n'.join(map(lambda x : '%d peaks in subband %d' % x,subband_occ))
top_subband=subband_occ[0][1]
print '-> Using subband %d' % top_subband
top_beats=filter(lambda x : x[1]==top_subband,beats)
bpm=int(float(len(top_beats))/float(duration)*60)
print 'Beats: %d -> %d BPM' % (len(top_beats),bpm)
energies=map(lambda x : x[2],beats)
mne,mxe=min(energies),max(energies)
enorm=lambda x : (x-mne)/(mxe-mne)

# empty directory for frames
odir='/tmp/' + os.path.basename(fn).split('.')[0]
if os.path.isdir(odir): empty_directory(odir)
else: os.mkdir(odir)

# create frames
bgfiles=sorted(os.listdir(bgdir))
print 'Images:',len(bgfiles)
frames=[0.1]*(int(duration*fps)+1)
for t,b,e in top_beats: frames[int(t*fps)]=enorm(e)
step=float(len(bgfiles))/sum(frames)
print 'Step:',step
pbar = progressbar.ProgressBar(widgets=['Creating frames: ']+default_widgets).start()
pos1,pos2=0,0
for n,f in enumerate(frames):
 pos1+=step*f
 if int(pos1)>int(pos2) and pos2<len(bgfiles)-1: pos2+=1
 shutil.copy(bgdir+bgfiles[pos2],'%s/image%05d.jpg' % (odir,n))
 pbar.update(float(n)/float(len(frames))*100)
pbar.finish()
afn='/tmp/'+fn.lower().replace('mp3','avi')
afn2=fn.lower().replace('mp3','avi')
cmd='mencoder mf://%s/*.jpg -ovc x264 -x264encopts crf=20 -o %s -mf fps=25 >/dev/null 2>&1' % (odir,afn) 
print cmd
os.system(cmd)
if os.path.isfile(afn2): os.unlink(afn2)
cmd='ffmpeg -i %s -i %s -vcodec copy %s -acodec copy -newaudio >/dev/null 2>&1' % (afn,fn,afn2)
print cmd
os.system(cmd)