import numpy,os,wave
import matplotlib.pyplot as p
import waveforms

# This file contains support for reading and playing sound files
# and for computing a low-pass filter.

#read_sound(filename) returns a numpy array of samples
#
#play_sound(sound) takes a numpy array of samples, assumes a 8000 samples/second
#                  rate and plays the sound
#
#plot_freq_response(omega,mhejw,title) takes numpy arrays of frequencies
#                                      and frequency response magnitudes
#                                      and plots them.
#plot_mehjw_usr(omega,mhejw,usr) takes numpy arrays of frequencies,
#                                frequency response magnitudes, and
#                                a unit sample response and plots them 
#                                together.
#
#lpf(omega_pass, omega_stop, Nt=5) Returns the unit sample response of a low 
#                                  pass filter with magnitude one in pass band
#                                  and magnitude zero in stop band.


# set up for audio playback
def read_sound(filename):
    w = waveforms.read_wavfile(filename)
    samples = w[0].samples
    sample_rate = w[0].sample_rate
    return samples, sample_rate

def play_sound(sound,sample_rate):
    waveforms.play(sound,sample_rate)

def plot_freq_response(omega,mhejw,title):
    p.figure()
    p.plot(omega,mhejw)
    p.axis((-numpy.pi,numpy.pi,0.0,float(int(max(mhejw)))+1))
    p.title("Magnitude of $|H(e^{j\Omega})|$ for "+title)
    p.xlabel("$\Omega$")
    p.grid()

def plot_freq_res_usr(omega,mhejw,usr,title):
    p.figure()
    p.subplots_adjust(hspace=0.6)
    omegaf = list(omega)
    mhejwf = list(mhejw)
    p.subplot(211)
    p.stem(range(len(usr)),usr)
    p.title("Unit-sample response for "+title)
    p.subplot(212)
    p.plot(omegaf,mhejwf)
    p.title("Magnitude of $|H(e^{j\Omega})|$ for "+title)
    p.xlabel("$\Omega$")
    p.axis((-numpy.pi,numpy.pi,0.0,float(int(max(mhejw)))+1))
    p.grid()


# Compute performance ratio for filter
def filter_quality(stop_band,pass_band,omega,mag_hejw):
    maxstopband = 0.0
    minpassband = 1.0
    for i in xrange(len(omega)):
        w = omega[i]
        if w >= stop_band[0] and w <= stop_band[1]:
            maxstopband = max(maxstopband,mag_hejw[i])
        if w >= pass_band[0] and w <= pass_band[1]:
            minpassband = min(minpassband,mag_hejw[i])

    return minpassband/maxstopband

def hum_it(filename):
    w = waveforms.read_wavfile(filename)
    A = max(w[0].samples)
    Am = min(w[0].samples)
    print "amplitude", A, Am
    sample_rate = w[0].sample_rate
    # 400 Hz hum
    omega = numpy.pi*(800.0/sample_rate)
    n_list = numpy.array(range(len(w[0].samples)))
    hum = numpy.sin(omega*n_list)
    w[0].samples += hum
    w[0].samples /= max(abs(w[0].samples))
    waveforms.write_wavfile(w[0],filename="hum_" + filename)
    
# Solve a possibly complex linear system of equations and return real
# part of the result
def solve_sys(matrix, rhs):
    n = len(rhs)
    assert len(matrix[0][:]) == len(matrix), "matrix is not square."
    assert len(matrix[0][:]) == len(rhs), "mismatch between rows in matrix and length of the right hand side"

    # Who knows what data type they used for matrix and rhs
    A = numpy.zeros((n,n),dtype = complex)
    rhscol = numpy.zeros((n,1),dtype = complex)
    for i in range(n):
        rhscol[i][0] = rhs[i]
        for j in range(n):
            A[i][j] = matrix[i][j]
            
    x = numpy.real(((numpy.linalg.solve(A,rhscol)).T)[0])
    return x

# Takes the constraints and imaginary frequencies, formats and shifts 
def shift_cons(rhs, omegas):
    jscale = 0.0 + 1.0j

    # Some checking on the omegas
    omegas = numpy.array(omegas)

    # Make the omegas imaginary if they are not
    if min(numpy.isreal(omegas)) == True:
        omegas *= jscale
        assert max(abs(numpy.real(omegas))) == 0, "omegas are not all reall or all imaginary!"

    # Check for correct number of equations
    omegass = numpy.array(omegas)
    omegass.sort()
    
    n = 2*len(omegass)

    if omegass[0] == 0.0:
        n -= 1
    if numpy.abs(omegass[-1] - jscale*numpy.pi) < 1.0e-6:
        n -= 1
    rhsL = len(rhs)

    assert rhsL == n, "Wrong number of equations, should be %d, is %d." % (n, rhsL)

    # Put in the shift by half the length of the unit sample response
    nshift = len(rhs)/2
    if nshift < 2: #Don't shift for examples with just a few terms
        nshift = 0

    #print nshift
    # Create shifted right hand side (multiply by e^(j*omega*nshift)
    rhss = numpy.zeros(len(rhs),dtype = complex)

    i = 0
    for l in range(len(omegas)):
        ispi = numpy.abs(omegas[l] - jscale*numpy.pi) < 1.0e-6
        if omegas[l] == 0 or ispi:
            rhss[i,] = rhs[i]*numpy.exp(omegas[l]*nshift)
            i += 1
        else:
            rhss[i] = rhs[i]*numpy.exp(omegas[l]*nshift)
            rhss[i+1] = rhs[i]*numpy.exp(-omegas[l]*nshift)
            i += 2
    return(rhss)

def solve_for_usr(omegas, mhejw):

    # Make frequencies imaginary numbers
    jscale = 0.0 + 1.0j
    omegas = jscale*numpy.array(omegas)

    # Fill in matrix A (list of lists) and put the appropriate
    # constraints from mhejw in the right hand side rhs (list).

    k = len(omegas)
    L = 2*k

    # Eliminate equations based on if there is a zero freq or pi freq
    omegass = numpy.array(omegas)  
    omegass.sort()
    
    if omegass[0] == 0.0:
        L -= 1
    if numpy.abs(omegass[-1] - jscale*numpy.pi) < 1.0e-6:
        L -= 1

    # Allocations for L equations
    klist = numpy.array(range(L))
    A = numpy.zeros((L,L),dtype = complex)
    rhs = numpy.zeros(L,dtype = complex)

    # Fill in the matrix and rhs
    i = 0 
    for l in range(len(omegas)):
        ispi = numpy.abs(omegas[l] - jscale*numpy.pi) < 1.0e-6
        if omegas[l] == 0 or ispi:
            A[i,:] = numpy.exp(omegas[l]*klist)
            A[i,:] = numpy.exp(omegas[l]*klist)
            rhs[i] = mhejw[l]
            i += 1
        else:
            A[i,:] = numpy.exp(omegas[l]*klist)
            rhs[i] = mhejw[l]
            A[i+1,:] = numpy.exp(-omegas[l]*klist)
            rhs[i+1] = mhejw[l]
            i += 2

    # Adjust the angle, but not the magnitude, of constraints
    rhsshift = shift_cons(rhs, omegas)

    # Solve the linear system of equations for a unit sample response
    h = solve_sys(A, rhsshift)

    return h


# Designs the unit sample response of a low pass filter with parameters
# omega_pass and omega_stop, 0 <  omega_pass <  omega_stop < pi, 
# where the magnitude of the filter's frequency response is one for
# 0 < omega < omega_pass and is zero for omega_step < omega <= pi.
# Note, the closer omega_pass is to omega_stop, the longer the unit
# sample response. Note, Nt is the number of points in the transition
# region between pass and stop band.

def lpf(omega_pass, omega_stop, Nt=5):

    pi = numpy.pi

    omega_delta = omega_stop - omega_pass

    assert omega_delta > 0, "The stop band begins before pass band ends!"
    
    assert omega_delta > 0.01, "The stop band-pass band seperation, %g, is too small!" % omega_delta

    
    # Determine the number of frequency points to use in pass and stop bands.
    domega = omega_delta/(Nt+1)
    Np = round(omega_pass/domega) + 1
    Ns = round((pi-omega_stop)/domega) + 1

    #print Np, Nt, Ns

    # Allocate space for the frequency constraints
    mags = numpy.zeros(Np+Nt+Ns)
    omegas = numpy.zeros(Np+Nt+Ns)

    # Compute passband omegas
    passmag = numpy.ones(Np)
    passomega = numpy.linspace(0,omega_pass,Np)

    # Compute transition mags using a cosine
    transmag = 0.5*(numpy.cos(numpy.linspace(pi/Nt,(pi*(Nt-1)/Nt),Nt)) + 1.0)
    transomega = numpy.linspace(omega_pass, omega_stop, Nt+2)
    transomega = transomega[1:-1]

    # Compute passband omegas
    stopmag = numpy.zeros(Ns)
    stopomega = numpy.linspace(omega_stop,pi,Ns)

    # Put pass, transition and stop magnitudes and frequencies together
    mags[0:Np] = passmag
    mags[Np:Np+Nt] = transmag
    mags[Np+Nt:Np+Nt+Ns] = stopmag

    omegas[0:Np] = passomega
    omegas[Np:Np+Nt] = transomega
    omegas[Np+Nt:Np+Nt+Ns] = stopomega

    # Solve for the usr and return
    return solve_for_usr(omegas, mags)

##################################################
# Verify the delay estimator.
##################################################
def verify_task2(f):
    points = 0.0
    h1 = numpy.array([1.0/11]*11)  # Should be five
    h2 = numpy.array([0.1,0.2,0.3,0.2,0.1]) # Should be two
    h3 = numpy.array([0.3,0.2,0.1,0.2,0.3]) # Should still be two
    h4 = numpy.cos((2.0*numpy.pi/51.0)*numpy.linspace(0.0,1.0,num=21)) #Ans=10

    if f(h1) == 5:
        points += 0.5
    else:
        print "failed with h1"

    if f(h2) == 2:
        points += 0.5
    else:
        print "failed with h2"

    if f(h3) == 2:
        points += 0.5
    else:
        print "failed with h3"

    if f(h4) == 10:
        points += 0.5
    else:
        print "failed with h4"

    print "Task 2 points = %f out of %f" % (points, 2.0)
    return points


##################################################
# Verify the notch filter.
##################################################
def verify_task6(f):
    omega = numpy.pi/5.0  # Notch filter frequency
    usr = f(omega)  # Unit sample response of the notch
    ten_periods = round(2.0*numpy.pi/omega) * 10
    n = max(len(usr), ten_periods*10) # At least 100 periods
    narray = numpy.array(range(n))
    last10 = range(n - ten_periods,n) # indices of last ten periods
    sine_in = numpy.sin(omega*narray)
    cos_in = numpy.cos(omega*narray)
    sine_low_in = numpy.sin(0.8*omega*narray)
    cos_low_in = numpy.cos(0.8*omega*narray)
    sine_high_in = numpy.sin(1.2*omega*narray)
    cos_high_in = numpy.cos(1.2*omega*narray)

    in_list = [sine_low_in,cos_low_in,sine_in,cos_in,sine_high_in,cos_high_in]
    maxs = [0]*len(in_list)
    for inseq in range(len(in_list)):
        notch_out = numpy.convolve(usr,in_list[inseq])
        maxs[inseq] = numpy.max(numpy.abs(notch_out[last10]))

    notch_ratio = max(maxs[2:4])/min(min(maxs[0:2]),min(maxs[4:6]))
    print "notch ratio", notch_ratio
    if(notch_ratio < 1.0e-2): # 100x smaller in notch
        points = 2
    else:
        points = 0

    print "Task 6 points = %f out of %f" % (points, 2.0)
    return points

##################################################
##
## Code to submit task to server.  Do not change.
## Task-specific code is in verify(), defined above.
##
##################################################

import Tkinter
class Dialog(Tkinter.Toplevel):
    def __init__(self, parent, title = None):
        Tkinter.Toplevel.__init__(self, parent)
        self.transient(parent)
        if title: self.title(title)
        self.parent = parent

        body = Tkinter.Frame(self)
        self.initial_focus = self.body(body)
        body.pack(padx=5, pady=5)

        self.buttonbox()
        self.grab_set()

        if not self.initial_focus:
            self.initial_focus = self

        self.protocol("WM_DELETE_WINDOW", self.cancel)
        self.geometry("+%d+%d" % (parent.winfo_rootx()+50,parent.winfo_rooty()+50))
        
        self.initial_focus.focus_set()
        self.wait_window(self)

    def body(self, master):
        return None

    # add standard button box
    def buttonbox(self):
        box = Tkinter.Frame(self)
        w = Tkinter.Button(box, text="Ok", width=10, command=self.ok, default=Tkinter.ACTIVE)
        w.pack(side=Tkinter.LEFT, padx=5, pady=5)
        box.pack()
        
    # standard button semantics
    def ok(self, event=None):
        if not self.validate():
            self.initial_focus.focus_set() # put focus back
            return
        self.withdraw()
        self.update_idletasks()
        self.apply()
        self.cancel()
        
    def cancel(self, event=None):
        # put focus back to the parent window
        self.parent.focus_set()
        self.destroy()
        
    # command hooks
    def validate(self):
        return 1 # override

    def apply(self):
        pass   # override

# ask user for Athena username and MIT ID
class SubmitDialog(Dialog):
    def __init__(self,parent,error=None,title = None):
        self.error = error
        self.athena_name = None
        self.mit_id = None
        Dialog.__init__(self,parent,title=title)

    def body(self, master):
        row = 0
        if self.error:
            l = Tkinter.Label(master,text=self.error,
                              anchor=Tkinter.W,justify=Tkinter.LEFT,fg="red")
            l.grid(row=row,sticky=Tkinter.W,columnspan=2)
            row += 1
        Tkinter.Label(master, text="Athena username:").grid(row=row,sticky=Tkinter.E)
        self.e1 = Tkinter.Entry(master)
        self.e1.grid(row=row, column=1)

        row += 1
        Tkinter.Label(master, text="MIT ID:").grid(row=row,sticky=Tkinter.E)
        self.e2 = Tkinter.Entry(master)
        self.e2.grid(row=row, column=1)

        return self.e1 # initial focus

    # add standard button box
    def buttonbox(self):
        box = Tkinter.Frame(self)
        w = Tkinter.Button(box, text="Submit", width=10, command=self.ok,
                           default=Tkinter.ACTIVE)
        w.pack(side=Tkinter.LEFT, padx=5, pady=5)
        w = Tkinter.Button(box, text="Cancel", width=10, command=self.cancel)
        w.pack(side=Tkinter.LEFT, padx=5, pady=5)
        box.pack()
        
    def apply(self):
        self.athena_name = self.e1.get()
        self.mit_id = self.e2.get()

# Let user know what server said
class MessageDialog(Dialog):
    def __init__(self, parent,message = '',title = None):
        self.message = message
        Dialog.__init__(self,parent,title=title)

    def body(self, master):
        l = Tkinter.Label(master, text=self.message,anchor=Tkinter.W,justify=Tkinter.LEFT)
        l.grid(row=0)

# return contents of file as a string
def file_contents(fname):
    # use universal mode to ensure cross-platform consistency in hash
    f = open(fname,'U')
    result = f.read()
    f.close()
    return result

import hashlib
def digest(s):
    m = hashlib.md5()
    m.update(s)
    return m.hexdigest()

# if verify(f) indicates points have been earned, submit results
# to server if requested to do so
import inspect,os,urllib,urllib2
def checkoff(f,task='???',submit=True):
    if task == 'L6_1':
        points = 1
    elif task == 'L6_2':
        points = verify_task2(f)
    elif task == 'L6_3':
        points = 1
    elif task == 'L6_4':
        points = 1
    elif task == 'L6_5':
        points = 1
    elif task == 'L6_6':
        points = verify_task6(f)
    else:
        raise ValueError,"task must be one of L6_1, L6_2, L6_3, L6_4, L6_5, L_6"
    """
    tasks = ('L6_1','L6_2','L6_3','L6_4','L6_5','L6_6')
    if task in tasks:
        points = 'tba'
    else:
        raise ValueError,"task must be one of %s" % ', '.join(tasks)
    """
    
    if submit and points:
        root = Tkinter.Tk(); #root.withdraw()
        error = None
        while submit:
            sd = SubmitDialog(root,error=error,title="Submit Task %s?"%task)
            if sd.athena_name:
                if isinstance(f,str): fname = os.path.abspath(f)
                else: fname = os.path.abspath(inspect.getsourcefile(f))
                post = {
                    'user': sd.athena_name,
                    'id': sd.mit_id,
                    'task': task,
                    'digest': digest(file_contents(os.path.abspath(inspect.getsourcefile(checkoff)))),
                    'points': points,
                    'filename': fname,
                    'file': file_contents(fname)
                    }
                try:
                    response = urllib2.urlopen('http://scripts.mit.edu/~6.02/currentsemester/submit_task.cgi',
                                               urllib.urlencode(post)).read()
                except Exception,e:
                    response = 'Error\n'+str(e)
                if response.startswith('Error\n'):
                    error = response[6:]
                else:
                    MessageDialog(root,message=response,title='Submission response')
                    break
            else: break

        root.destroy()