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

if not os.path.exists('temp.tbw'):
    fh1 = urllib2.urlopen('http://fornax.phys.unm.edu/lwa/data/firstlight/firstLightTBW.dat')
    fh2 = open('temp.tbw', 'wb')
    fh2.write(fh1.read(tbw.FrameSize*1000))
    fh1.close()
    fh2.close()
print "TBW Size: %.1f kB" % (os.path.getsize('temp.tbw')/1024.)

if not os.path.exists('temp.tbn'):
    fh1 = urllib2.urlopen('http://fornax.phys.unm.edu/lwa/data/TBN/055987_000496599_DEFAULT_TBN_00000_30.dat')
    fh2 = open('temp.tbn', 'wb')
    fh2.write(fh1.read(tbn.FrameSize*1000))
    fh1.close()
    fh2.close()
print "TBN Size: %.1f kB" % (os.path.getsize('temp.tbn')/1024.)

if not os.path.exists('temp.drx'):
    fh1 = urllib2.urlopen('http://lda10g.alliance.unm.edu/data1/055917_000007389_Jupiter.dat')
    fh2 = open('temp.drx', 'wb')
    fh2.write(fh1.read(drx.FrameSize*300))
    fh1.close()
    fh2.close()
print "DRX Size: %.1f kB" % (os.path.getsize('temp.drx')/1024.)

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.ldp import LWA1DataFile

idfD = LWA1DataFile('temp.drx')
print "'temp.drx' is of type:", type(idfD)

idfS = LWA1DataFile('temp.drspec')
print "'temp.drspec' is of type:", type(idfS)

fh = open('temp.tbw', 'rb')
idfW = LWA1DataFile(fh=fh)
print "'temp.tbw' is of type:", type(idfW)

idfN = LWA1DataFile('temp.tbn')
print "'temp.tbn' is of type:", type(idfN)

# Poll the TBN file for its specifics
print "TBN Metadata:"
for key,value in idfN.getInfo().iteritems():
    print "  %s:" % key, value
print " "

# Poll the DRX file for the sample rate
print "DRX Sample Rate %.3f MS/s" % (idfD.getInfo('sampleRate')/1e6,)
print " "

# Poll the DR spectrometer file for its specifics
print "DR Spectrometer Metadata:"
for key,value in idfS.getInfo().iteritems():
    print "  %s:" % key, value
print " "

duration, tStart, data = idfD.read(0.010)
print "Duration read: %.3f ms" % (duration*1e3,)
print "Time of first sample: %.6f" % tStart
print "Data specifics:", data.shape, data.dtype

from datetime import datetime
tStartDT = datetime.utcfromtimestamp(tStart)
print tStartDT

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

# Compute the spectra
freq, spec = fxc.SpecMaster(data, LFFT=512, SampleRate=idfD.getInfo('sampleRate'), CentralFreq=0.0)
freq1 = freq + idfD.getInfo('freq1')
freq2 = freq + idfD.getInfo('freq2')

# Plot
fig = plt.figure()
ax = fig.gca()
ax.plot(freq1/1e6, to_dB(spec[0,:]), label='B%i,T1,P0' % idfD.getInfo('beam'))
ax.plot(freq1/1e6, to_dB(spec[1,:]), label='B%i,T1,P1' % idfD.getInfo('beam'))
ax.set_xlabel('Frequency [MHz]')
ax.set_ylabel('PSD [arb. dB]')
ax.legend(loc=0)
plt.show()

durationS, tStartS, dataS = idfS.read(0.100)
print "Duration read: %.3f ms" % (durationS*1e3,)
print "Time of first sample: %.6f" % tStartS
print "Data specifics:", dataS.shape, dataS.dtype

%matplotlib inline
# Get the metadata
b = idfS.getInfo('beam')
dp = idfS.getInfo('dataProducts')

# Compute the frequency bins
freq = numpy.fft.fftfreq(dataS.shape[2], d=1.0/idfS.getInfo('sampleRate'))
freq = numpy.fft.fftshift(freq)
freq1 = freq + idfS.getInfo('freq1')
freq2 = freq + idfS.getInfo('freq2')

# Intgrate
spec = dataS.mean(axis=1)

# Plot
fig = plt.figure()
ax = fig.gca()
ax.plot(freq1/1e6, to_dB(spec[0,:]), label='B%i,T1,%s' % (b, dp[0 % len(dp)]))
ax.plot(freq1/1e6, to_dB(spec[1,:]), label='B%i,T1,%s' % (b, dp[1 % len(dp)]))
ax.plot(freq2/1e6, to_dB(spec[2,:]), label='B%i,T1,%s' % (b, dp[2 % len(dp)]))
ax.plot(freq2/1e6, to_dB(spec[3,:]), label='B%i,T1,%s' % (b, dp[3 % len(dp)]))
ax.set_xlabel('Frequency [MHz]')
ax.set_ylabel('PSD [arb. dB]')
ax.legend(loc=0)
plt.show()

# Offset 10 ms into the file from the current location
skip = idfD.offset(0.010)
print "Skipped: %.3f ms" % (skip*1e3,)

# Read in another 50 ms of data
duration2, tStart2, data2 = idfD.read(0.050)
print "Duration read: %.3f ms" % (duration*1e3,)
print "Time of first sample: %.6f" % tStart
print "Data specifics:", data.shape, data.dtype
print " "
print "Time skip: %.3f ms" % ((tStart2-tStart)*1e3,)

# Go back to the beginning
idfD.reset()

# Read in the first 10 ms again
duaration3, tStart3, data3 = idfD.read(0.010)

# Are the start times the same?
print "Time good?", True if tStart3-tStart == 0 else False

# Are the data the same?
print "Data good?", True if (numpy.abs(data-data3)**2).sum() == 0 else False

idfD.close()
idfS.close()
idfN.close()
idfW.close()

try:
    idfD.read(0.010)
except Exception, e:
    print "ERROR: %s" % str(e)