import math
import scipy
import scipy.fftpack
import scipy.integrate
import fileManager
import evaluator
import sys

import os
os.environ['MPLCONFIGDIR']='/tmp'
import pylab

from mod_python import apache

def cablemagic(formula, startVal, stopVal, numBins, supplyOutput, attenDB, attenFreqMHz, cableLength, debug=False):

    if(startVal<stopVal):
        xrange = scipy.linspace(startVal, stopVal, num=numBins)
    else:
        xrange = scipy.linspace(stopVal, startVal, num=numBins)

    if(debug):
        print "Formula: %s<br>" % formula

    # get a file manager and a temporary directory
    fm = fileManager.fileManager('/var/www/html/cable-magic/data', fileTimeOut=3600)
    dir = fm.getDir()

    # -----------------------------------------------------------
    # user supplied function
    eval = evaluator.Evaluator(formula, "t", xrange)
    xdata, ydata = eval.getData()

    # get the length of the data and figure out the timestep
    N = len(xdata)
    timestep = math.fabs(max(xdata)-min(xdata))/N
    
    # graph the origional function
    # origional function
    pylab.suptitle('Input Function', fontsize=16)
    pylab.subplot(211)
    origRealFunc = pylab.plot(xdata, ydata.real)
    pylab.title("Real Part", fontsize=12)
    pylab.ylabel("Amplitude")
    pylab.subplot(212)
    origImagFunc = pylab.plot(xdata, ydata.imag)
    pylab.title("Imag Part", fontsize=12)
    pylab.xlabel("Time")
    pylab.ylabel("Amplitude")
    origFuncPath = os.path.join(dir, "orig.png")
    origFuncPDFPath = os.path.join(dir, "orig.pdf")
    pylab.savefig(origFuncPath)
    pylab.savefig(origFuncPDFPath)
    pylab.clf()

    # integrate abs(input)^2 over time
    inputIntegral = scipy.integrate.trapz(abs(ydata)**2.0, dx=timestep)

    # bob's scaling factor
    #scaleFactor = N/(math.fabs(max(xdata)-min(xdata)))
    scaleFactor = (math.fabs(max(xdata)-min(xdata)))/N
    # get the input sum after scaling
    inputSum = sum(abs(ydata)**2.0)
    
    ydata = ydata * scaleFactor


    


    # ---------------------------------------------------------
    # get the fft of the user supplied function
    fourier = scipy.fft(ydata)
    
    freqs = scipy.fftpack.fftfreq(N, d = timestep)

    inputFFTSum = 1.0/N * sum(abs(fourier/scaleFactor)**2.0)
    
    freqstep = math.fabs(max(freqs)-min(freqs))/N

    # plot the fft of the origional function
    x = scipy.fftpack.fftshift(freqs)
    y = scipy.fftpack.fftshift(fourier)

    origFFTy = y
    
    pylab.plot(x, abs(y)**2)
    pylab.title("Input Power Spectrum", fontsize=16)
    pylab.xlabel("Frequency")
    pylab.ylabel("Power")
    fftFuncPath = os.path.join(dir, "fft.png")
    fftFuncPDFPath = os.path.join(dir, "fft.pdf")
    pylab.savefig(fftFuncPath)
    pylab.savefig(fftFuncPDFPath)
    pylab.clf()


    # -----------------------------------------------------------------------
    atten0 = attenDB * 1.151 * math.pow(10.0, -3)
    attenFreqGHz = attenFreqMHz/1000.0

    phi = ((cableLength * atten0) / math.sqrt(attenFreqGHz))
    respFunc = "exp(-1.0*%f * sqrt(abs(f)))*cos(%f * sqrt(abs(f))) - j * ( exp(-1.0*%f*sqrt(abs(f))) * sin(%f * sqrt(abs(f))))" % ( phi, phi, phi, phi)

    if(debug):
        print "response function: %s<br>" % respFunc
        
    respEval = evaluator.Evaluator(respFunc, "f", freqs)
    respXData, respYData = respEval.getData()

    dataLen = len(respXData)
    i=0
    while(i<dataLen):
        if(respXData[i]<0):
            respYData[i] = respYData[i].conjugate()
        i = i+1


    # plot the response
    # plot the system response
    # this is generated in the same order as the output for the fft
    pylab.suptitle('Cable Transfer Function', fontsize=16)
    pylab.subplot(211)

    x = scipy.fftpack.fftshift(respXData)
    y = scipy.fftpack.fftshift(respYData)
    pylab.plot(x, y.real)
    pylab.title("Real Part", fontsize=12)
    pylab.ylabel("Amplitude")
    pylab.subplot(212)
    pylab.plot(x, y.imag)
    pylab.title("Imaginary Part", fontsize=12)
    pylab.xlabel("Frequency")
    pylab.ylabel("Amplitude")
    transPath = os.path.join(dir, "trans.png")
    transPDFPath = os.path.join(dir, "trans.pdf")
    pylab.savefig(transPath)
    pylab.savefig(transPDFPath)
    pylab.clf()

    # bob also wants the transfer function
    # plot the cable system response
    pylab.suptitle('Cable System Response', fontsize=16)
    pylab.subplot(211)

    sysRespIFFT = scipy.ifft(respYData)
    
    pylab.plot(xdata, sysRespIFFT.real)
    pylab.title("Real Part", fontsize=12)
    pylab.ylabel("Amplitude")
    pylab.subplot(212)
    pylab.plot(xdata, sysRespIFFT.imag)
    pylab.title("Imaginary Part", fontsize=12)
    pylab.xlabel("Time")
    pylab.ylabel("Amplitude")
    sysRespPath = os.path.join(dir, "sysResp.png")
    sysRespPDFPath = os.path.join(dir, "sysResp.pdf")
    pylab.savefig(sysRespPath)
    pylab.savefig(sysRespPDFPath)
    pylab.clf()    

    # -----------------------------------------------------------------
    # the respXData and freqs array should be identical
    powerOutput = None
    if(supplyOutput):
        # we have the output
        powerOutput = (fourier/respYData)
        outData = scipy.ifft(powerOutput)
    else:
        # we have the input - fourier
        powerOutput = (respYData * fourier)
        outData = scipy.ifft(powerOutput)

    # take out bob's scaling factor
    outData = outData / scaleFactor
    
    # do this after removing the scaling factor
    outputSum = sum(abs(outData)**2.0) 


    
    # plot the power spectrum of the output
    x = scipy.fftpack.fftshift(respXData)
    y = scipy.fftpack.fftshift(powerOutput)

    # calculate the power in the output power spectrum
    outputFFTSum = (1.0/N)*sum(abs(y/scaleFactor)**2.0) 
    
    pylab.plot(x, (abs(y)**2))
    pylab.title("Output Power Spectrum", fontsize=16)
    pylab.xlabel("Frequency")
    pylab.ylabel("Power")
    outPowerPath = os.path.join(dir, "outPower.png")
    outPowerPDFPath = os.path.join(dir, "outPower.pdf")
    pylab.savefig(outPowerPath)
    pylab.savefig(outPowerPDFPath)
    pylab.clf()

    # plot ratio of input to output power spectrum
    x = scipy.fftpack.fftshift(respXData)
    y = abs(y/origFFTy)**2.0
    pylab.plot(x,y)
    pylab.title('Ratio Of Output To Input Power Spection', fontsize=16)
    pylab.xlabel('Frequency')
    pylab.ylabel('Power')
    ratioPath = os.path.join(dir, 'ratio.png')
    ratioPDFPath = os.path.join(dir, 'ratio.pdf')
    pylab.savefig(ratioPath)
    pylab.savefig(ratioPDFPath)
    pylab.clf()

    # plot the output
    pylab.suptitle('Output', fontsize=16)
    pylab.subplot(211)
    outReal = pylab.plot(xdata, outData.real)
    pylab.title("Real Part", fontsize=12)
    pylab.ylabel("Amplitude")
    pylab.subplot(212)
    outImag = pylab.plot(xdata, outData.imag)
    pylab.title("Imaginary Part", fontsize=12)
    pylab.xlabel("Time")
    pylab.ylabel("Amplitude")
    outPath = os.path.join(dir, "out.png")
    outPDFPath = os.path.join(dir, "out.pdf")
    pylab.savefig(outPath)
    pylab.savefig(outPDFPath)
    pylab.clf()

    # calculate the power in the output time spectrum
    outputIntegral = scipy.integrate.trapz(abs(outData)**2.0, dx=timestep)
    
    # bob wants a plot of abs(outData)**2, why I'm not sure but who am I to question why?
    pylab.title("$|output|^2$", fontsize=16)
    pylab.xlabel("Time")
    pylab.ylabel("Power")
    pylab.plot(xdata, abs(outData)**2.0)
    absOutPath = os.path.join(dir, "absOut.png")
    absOutPDFPath = os.path.join(dir, "absOut.pdf")
    pylab.savefig(absOutPath)
    pylab.savefig(absOutPDFPath)
    pylab.clf()

    
    baseDir = os.path.basename(dir)
    baseURL = os.path.join("http://cerulean.icecube.wisc.edu/cable-magic/data", baseDir)

    origURL = os.path.join(baseURL, "orig.png")
    fftURL = os.path.join(baseURL, "fft.png")
    transURL = os.path.join(baseURL, "trans.png")
    sysRespURL = os.path.join(baseURL, "sysResp.png")
    outPowerURL = os.path.join(baseURL, "outPower.png")
    outURL = os.path.join(baseURL, "out.png")
    absOutURL = os.path.join(baseURL, "absOut.png")
    ratioURL = os.path.join(baseURL, "ratio.png")


    origPDFURL = os.path.join(baseURL, "orig.pdf")
    fftPDFURL = os.path.join(baseURL, "fft.pdf")
    transPDFURL = os.path.join(baseURL, "trans.pdf")
    sysRespPDFURL = os.path.join(baseURL, "sysResp.pdf")
    outPowerPDFURL = os.path.join(baseURL, "outPower.pdf")
    outPDFURL = os.path.join(baseURL, "out.pdf")
    absOutPDFURL = os.path.join(baseURL, "absOut.pdf")
    ratioPDFURL = os.path.join(baseURL, 'ratio.pdf')

    
    returnVal = ";".join([origURL, origPDFURL, fftURL, fftPDFURL, transURL, transPDFURL, outPowerURL, outPowerPDFURL, outURL, outPDFURL, absOutURL, absOutPDFURL, ratioURL, ratioPDFURL, sysRespURL, sysRespPDFURL, "%5.4f" % inputSum, "%5.4f" % inputFFTSum, "%5.4f" % outputSum, "%5.4f" % outputFFTSum, "%5.4f" % inputIntegral, "%5.4f" % outputIntegral ])

    return returnVal
    
    
    
def magic(req):

    form_values = req.form
    req.content_type="text/html"

    sys.stdout=req

    attenFound = False
    attenFreqFound = False
    cableLenFound = False

    startValFound = False
    stopValFound = False
    numBinsFound = False
    supplyOutputFound = False
    formulaFound = False
    debug = False
    
    objs = form_values.list
    for obj in objs:
        if(obj.name=="atten"):
            attenFound = True
            atten = float(obj.value)
        elif(obj.name=="attenFreq"):
            attenFreqFound = True
            attenFreq = float(obj.value)
        elif(obj.name=="cablelen"):
            cableLenFound = True
            cableLen = float(obj.value)
        elif(obj.name=="startVal"):
            startValFound = True
            startVal = float(obj.value)
        elif(obj.name=="stopVal"):
            stopValFound = True
            stopVal = float(obj.value)
        elif(obj.name=="numBins"):
            numBinsFound = True
            numBins = int(obj.value)
        elif(obj.name=="formula"):
            formulaFound = True
            formula = obj.value
        elif(obj.name=="supplyoutput"):
            supplyOutputFound = True
            supplyOutput = int(obj.value)
        elif(obj.name=="debug"):
            debug=True
            
    # check to make sure all required variables have been found
    if(not attenFound or not attenFreqFound or not cableLenFound or not startValFound or not stopValFound or not numBinsFound or not supplyOutputFound or not formulaFound):
        print "Required form data not present!<br>"
        print "attenFound: ", attenFound, "<br>"
        print "attenFreqFound: ", attenFreqFound, "<br>"
        print "cableLenFound: ", cableLenFound, "<br>"
        print "startValFound: ", startValFound, "<br>"
        print "stopValFound: ", stopValFound, "<br>"
        print "numBinsFound: ", numBinsFound, "<br>"
        print "supplyOutputFound: ", supplyOutputFound, "<br>"
        print "formulaFound: ", formulaFound, "<br>"
        return 

    if(debug):
        print "Formula: %s<br>" % formula
        
    try:
        data = cablemagic(formula, startVal, stopVal, numBins, supplyOutput, atten, attenFreq, cableLen, debug)
        print data
    except Exception, e:
        print "Error: ", e