# This may take a bit...
import os
import urllib2
from lsl.reader import drx

if not os.path.exists('temp.drspec'):
    fh1 = urllib2.urlopen('http://lda10g.alliance.unm.edu/tutorial/B0329+54/056770_000044687')
    fh2 = open('temp.drspec', 'wb')
    fh2.write(fh1.read(drx.FrameSize*300))
    fh1.close()
    fh2.close()
    
print "DR Spectrometer Size: %.1f kB" % (os.path.getsize('temp.drspec')/1024.)

from lsl.reader import drspec

# Open the file
fh = open('temp.drspec', 'rb')

# What is the size in bytes of each frame?
print "Frame size:", drspec.getFrameSize(fh), "bytes"

# What is the sample rate?
print "Sample Rate:", drspec.getSampleRate(fh)/1e6, "MHz"

# How many frequency channels are in the data?
print "FFT length:", drspec.getTransformSize(fh), "channels"
print "-> %.3f kHz/channel" % (drspec.getSampleRate(fh)/1e3/drspec.getTransformSize(fh),)

# What type of data do we have?
print "Linear (XX, YY, etc.)?", drspec.containsLinearData(fh)
print "Stokes (I, V, etc.)?", drspec.containsStokesData(fh)

# What are the exact data products in the file?
print "Data products:", ", ".join(drspec.getDataProducts(fh))

# Close it out
fh.close()

# Open the file and read in a frame
fh = open('temp.drspec', 'rb')
frame = drspec.readFrame(fh)

# Load in some basic information
beam = frame.parseID()
srate = frame.getSampleRate()
LFFT = drspec.getTransformSize(fh)
tInt = frame.header.nInts*LFFT/srate

# Print out some basic information
print "Beam:", beam
print "Timestamp:", frame.getTime()
print "Tuning 1:", frame.getCentralFreq(1)/1e6, "MHz"
print "Tuning 2:", frame.getCentralFreq(2)/1e6, "MHz"
print " "
print "Transform Length:", LFFT
print "Integration Time:", tInt, "s"

# Done
fh.close()

from datetime import datetime

# Open the file and read in a frame
fh = open('temp.drspec', 'rb')
frame = drspec.readFrame(fh)

# Load in some basic information
beam = frame.parseID()
srate = frame.getSampleRate()
LFFT = drspec.getTransformSize(fh)
tInt = frame.header.nInts*LFFT/srate

# Print out some basic information
print "Beam:", beam
print "Timestamp:", frame.getTime(), "(", datetime.utcfromtimestamp(frame.getTime()), ")"
print "Tuning 1:", frame.getCentralFreq(1)/1e6, "MHz"
print "Tuning 2:", frame.getCentralFreq(2)/1e6, "MHz"
print " "
print "Transform Length:", LFFT
print "Integration Time:", tInt, "s"

# Done
fh.close()

# Open the file and read in a frame
fh = open('temp.drspec', 'rb')
frame = drspec.readFrame(fh)

print "Tuning 1 XX mean:", frame.data.XX0.mean()
print "Tuning 2 XX mean:", frame.data.XX1.mean()
print " "

for tuning in (1, 2):
    for prod in frame.getDataProducts():
        data = getattr(frame.data, "%s%i" % (prod, tuning-1), None)
        print "Tuning %i %s standard deviation:" % (tuning, prod), data.std()

# Done
fh.close()

%matplotlib inline
import numpy
from lsl.misc.mathutil import to_dB
from matplotlib import pyplot as plt

# Open the file
fh = open('temp.drspec', 'rb')

# Create the array that will hold the data
products = drspec.getDataProducts(fh)
nProd = len(products)*2
nChan = drspec.getTransformSize(fh)
data = numpy.zeros((nProd, 5, nChan), dtype=numpy.float32)

# Read in the first 5 frames
for i in xrange(5):
    try:
        frame = drspec.readFrame(fh)
    except errors.syncError:
        continue
    except errors.eofError:
        break
        
    j = 0
    for tuning in (1, 2):
        for prod in products:
            data[j,i,:] = getattr(frame.data, "%s%i" % (prod, tuning-1), None)
            j += 1
fh.close()

# Integrate across the frames
spec = data.mean(axis=1)

# Compute the frequency bins
freq = numpy.fft.fftfreq(spec.shape[1], d=1.0/frame.getSampleRate())
freq = numpy.fft.fftshift(freq)
freq1 = freq + frame.getCentralFreq(1)
freq2 = freq + frame.getCentralFreq(2)

# Plot
beam = frame.parseID()

fig = plt.figure()
ax = fig.gca()
for i in xrange(spec.shape[0]):
    if i / len(products) == 0:
        freq = freq1
    else:
        freq = freq2
        
    ax.plot(freq/1e6, to_dB(spec[i,:]), label='B%i,T1,%s' % (beam, products[i % len(products)]))
ax.set_xlabel('Frequency [MHz]')
ax.set_ylabel('PSD [arb. dB]')
ax.legend(loc=0)
plt.show()