# Jerome, 5th November 2012:
# This is a more recent version of large.py that Martin sent us by email on 31th July 2012
# He called it large2.py but I've renamed it large.py to fit the code.
import os
import traceback
from _collections import defaultdict
from numpy import *
import numpy as np
from scipy.misc import imsave
import matplotlib.pyplot as plt
MAXHARMONICS=3
FREQRANGE=375
def largeplot(matrixview, x=None, y=None, figname='largeplot'):
    rgbmatrix=zeros(matrixview.shape + (3,))
    rgbmatrix[:, :, 0:3] = matrixview[:, :, newaxis]
    if y !=None and x!=None:
        rgbmatrix[y, x, 1] = 40 #green is call
    flippedmatrix=np.flipud(rgbmatrix)
    imsave(figname+'.png', flippedmatrix)
    return rgbmatrix


def rgbplot(rgbmatrix, x=None, y=None, figname='rgbplot'):
    if y !=None and x!=None:
            rgbmatrix[y, x, 1] = 40 #green is call
    flippedmatrix=np.flipud(rgbmatrix)
    imsave(figname+'.png', flippedmatrix)
    return rgbmatrix



def trajectory(spectrogram, harmonics, w, z, cutoff=0.01, dirname='.'):
    harm = array(np.round(np.average(harmonics, axis = 1, weights=w)), dtype=int)  # consensus number of harmonics
    freqs = {}
    def average(z, axis, weights ):
        #find rows with zero weights and replace with a constant
        if np.sum(weights)==0:
            return 0
        zerorows = np.sum(weights, axis=1)==0
        numcols=weights.shape[1]
        weights[zerorows,:] = 1.0/numcols
        try:
            return np.average(z, axis, weights)
        except Exception, e:
            print traceback.format_exc(e)
            print weights
            return 0
    for numharmo in range(1, MAXHARMONICS+1):
        freqs[numharmo] = zeros(z.shape[0])
        #check there is at least one particle matching the consensus harmonic
        passed = ((harm == numharmo) & (sum((harmonics==numharmo), axis=1)!=0))
        for thisharmo in range(1, numharmo+1):
            # fill out upper harmonics
            # get frequency conditional on number of harmonic
            freqs[thisharmo][passed, :] = average(z[passed,:] - 1, axis=1, weights=w[passed,:] * (harmonics[passed,:]==numharmo)) *thisharmo
            freqs[thisharmo] = where(freqs[thisharmo] >= FREQRANGE, FREQRANGE-1, freqs[thisharmo])

    infavour = np.sum(ma.array(w, mask = z!=0), axis=1)
    notzero = freqs[1][infavour<cutoff]!=0
    x1, y1=arange(freqs[1].shape[0])[infavour<cutoff][notzero], array(freqs[1], dtype=int)[infavour<cutoff][notzero]
    rgb1=largeplot(spectrogram, x1,y1)
    notzero = freqs[2][infavour<cutoff]!=0
    x2, y2=arange(freqs[2].shape[0])[infavour<cutoff][notzero], array(freqs[2], dtype=int)[infavour<cutoff][notzero]
    rgb2=rgbplot(rgb1, x2, y2)
    notzero = freqs[3][infavour<cutoff]!=0
    x3, y3=arange(freqs[3].shape[0])[infavour<cutoff][notzero], array(freqs[3], dtype=int)[infavour<cutoff][notzero],
    rgb3=rgbplot(rgb2, x3, y3, os.path.join(dirname, 'finaltrajectory'))
    return ((x1,y1), (x2,y2), (x3,y3))

def matplot(themat, figname='rgbplot.test.'):
    x=themat[:,0]
    y=themat[:,1]
    rgbmatrix = zeros((max(y)+1, max(x)+1, 3))
    if y !=None and x!=None:
        rgbmatrix[y, x, 1] = 40 #green is call
    #test plot of a matrix
    flippedmatrix=np.flipud(rgbmatrix)
    imsave(figname+'.png', flippedmatrix)
    return rgbmatrix




def getmatrix(spectrogram, harmonics, w, z, cutoff=0.01):
    harm = array(np.round(np.average(harmonics, axis = 1, weights=w)), dtype=int)  # consensus number of harmonics
    freqs = {}
    def average(z, axis, weights ):
        #find rows with zero weights and replace with a constant
        if np.sum(weights)==0:
            return 0
        zerorows = np.sum(weights, axis=1)==0
        numcols=weights.shape[1]
        weights[zerorows,:] = 1.0/numcols
        try:
            return np.average(z, axis, weights)
        except Exception, e:
            print traceback.format_exc(e)
            print weights
            return 0
    for numharmo in range(1, MAXHARMONICS+1):
        freqs[numharmo] = zeros(z.shape[0])
        #check there is at least one particle matching the consensus harmonic
        passed = ((harm == numharmo) & (sum((harmonics==numharmo), axis=1)!=0))
        for thisharmo in range(1, numharmo+1):
            # fill out upper harmonics
            # get frequency conditional on number of harmonic
            freqs[thisharmo][passed, :] = average(z[passed,:] - 1, axis=1, weights=w[passed,:] * (harmonics[passed,:]==numharmo)) *thisharmo
            freqs[thisharmo] = where(freqs[thisharmo] >= FREQRANGE, FREQRANGE-1, freqs[thisharmo])
    resx, resy = [], []
    infavour = np.sum(ma.array(w, mask = z!=0), axis=1)
    notzero = freqs[1][infavour<cutoff]!=0
    x1, y1=arange(freqs[1].shape[0])[infavour<cutoff][notzero], array(freqs[1], dtype=int)[infavour<cutoff][notzero]
    resx.extend(x1)
    resy.extend(y1)
    notzero = freqs[2][infavour<cutoff]!=0
    x2, y2=arange(freqs[2].shape[0])[infavour<cutoff][notzero], array(freqs[2], dtype=int)[infavour<cutoff][notzero]
    resx.extend(x2)
    resy.extend(y2)
    notzero = freqs[3][infavour<cutoff]!=0
    x3, y3=arange(freqs[3].shape[0])[infavour<cutoff][notzero], array(freqs[3], dtype=int)[infavour<cutoff][notzero],
    resx.extend(x3)
    resx.extend(y3)
    return  array(zip(resx, resy))






def summarize(self, i, wind=None, windharmonics=None, ws=None, z=None, h=None, ss=None, thisz=None, thiswind=None, thiswindharmonic=None, thiswindstrength=None, thish=None, thissignalstrength=None, windamplitude=None, plotgraph=False):
    """summarizes likelihood for testing"""
    windamplitude = windamplitude if windamplitude else self.windamplitude
    wind = wind if wind else self.wind
    windharmonics = windharmonics if windharmonics else self.windharmonics
    ws = ws if ws else self.windstrength
    z = z if z else self.z
    h = h if h else self.harmonics
    ss = ss if ss else self.signalstrength
    zs = z[i, :]
    winds = wind[i, :]
    thiswind = thiswind if thiswind else int(median(wind[i, winds!=0]))
    thiswindstrength = thiswindstrength if thiswindstrength else int(median(ws[i, winds!=0]))
    thisz =  thisz if thisz else int(median(z[i, zs!=0]))
    thish = thish if thish else int(median(h[i,:]))
    thiswindharmonic = thiswindharmonic if thiswindharmonic  else int(median(windharmonics[i,:]))
    thissignalstrength = thissignalstrength if thissignalstrength else median(ss[i, zs!=0])
    print "median, this wind at %s = %s, %s" % (i, median(wind[i,:]), thiswind)
    print "median, mean windstrength at %s = %s, %s" % (i, median(ws[i,:]), mean(ws[i,:]))
    print "z, thisz at %s = %s, %s" % (i, median(zs), thisz)
    print "h at %s = %s, %s" % (i, median(h[i,:]), mean(h[i,:]))
    print "signal at %s = %s, %s" % (i, median(ss[i,:]), mean(ss[i,:]))
    thislike = self.checklikelihood(i, thisz, thish, thiswind, thiswindharmonic, thiswindstrength,thissignalstrength)
    print "this likelihood = ",thislike[1]
    nowindlike = self.checklikelihood(i,thisz, thish, 0, 0, thiswindstrength, thissignalstrength)
    print "no wind likelihood = ", nowindlike[1]
    nosignallike = self.checklikelihood(i, 0, thish, thiswind, thiswindharmonic, thiswindstrength, thissignalstrength)
    print "no signal likelihood = ", nosignallike[1]
    if plotgraph:
        fig = plt.figure();ax = fig.add_subplot(111)
        ax.scatter(arange(0,375), self.spectrogram[:,i+self.analysisstart])
        ylabels=arange(0, ((250000/2)/375*375), 10000)/1000
        yticks=ylabels*1000/((250000/2)/375)

        ax.set_xticklabels(ylabels); ax.set_xticks(yticks); ax.set_xlabel("Frequency (kHz)");ax.set_ylabel("Amplitude")
        ax.plot(arange(375), thislike[0][0])
        ax.plot(arange(375), nowindlike[0][0])

    return (thislike[1])