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

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.)

from lsl.reader import drx, errors

# Read in the first 5 frames
fh = open('temp.drx', 'rb')
for i in xrange(5):
    try:
        frame = drx.readFrame(fh)
    except errors.syncError:
        continue
    except errors.eofError:
        break
        
    print "Beam: %i, Tuning: %i, Pol.: %i" % frame.parseID()
    print "-> Time tag: %.6f s" % frame.getTime()
    print "-> Mean value: %.3f + %.3f j" % (frame.data.iq.mean().real, frame.data.iq.mean().imag)
fh.close()

from datetime import datetime
from lsl.reader import drx, errors

# Read in the first 5 frames
fh = open('temp.drx', 'rb')
for i in xrange(5):
    try:
        frame = drx.readFrame(fh)
    except errors.syncError:
        continue
    except errors.eofError:
        break
        
    print "Beam: %i, Tuning: %i, Pol.: %i" % frame.parseID()
    print "-> Time tag: %.6f s (%s)" % (frame.getTime(), datetime.utcfromtimestamp(frame.getTime()))
    print "-> Mean value: %.3f + %.3f j" % (frame.data.iq.mean().real, frame.data.iq.mean().imag)
fh.close()

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

fh = open('temp.drx', 'rb')
frame1 = drx.readFrame(fh)
frame2 = drx.readFrame(fh)
srate = frame1.getSampleRate()      # Data sample rate in Hz
cFreq = frame1.getCentralFreq()     # Data tuning in Hz

# SpecMaster expects 2-D data
data = numpy.zeros((2,frame1.data.iq.size), dtype=frame1.data.iq.dtype)
data[0,:] = frame1.data.iq
data[1,:] = frame2.data.iq

freq, spec = fxc.SpecMaster(data, LFFT=512, SampleRate=srate, CentralFreq=cFreq) 

fig = plt.figure()
ax = fig.gca()
ax.plot(freq/1e6, to_dB(spec[0,:]), label='B%i,T%i,P%i' % frame1.parseID())
ax.plot(freq/1e6, to_dB(spec[1,:]), label='B%i,T%i,P%i' % frame2.parseID())
ax.set_xlabel('Frequency [MHz]')
ax.set_ylabel('PSD [arb. dB]')
ax.legend(loc=0)
plt.show()

data = numpy.zeros((4,50*4096), dtype=numpy.complex64)
count = [0,0,0,0]

for i in xrange(50):
    for j in xrange(2):
        frame = drx.readFrame(fh)
        srate = frame.getSampleRate()      # Data sample rate in Hz
        cFreq = frame.getCentralFreq()     # Data tuning in Hz
        
        beam,tune,pol = frame.parseID()
        k = 2*(tune-1) + pol
        data[k, count[k]*4096:(count[k]+1)*4096] = frame.data.iq
        count[k] += 1
        
freq, spec = fxc.SpecMaster(data, LFFT=512, SampleRate=srate, CentralFreq=cFreq) 

fig = plt.figure()
ax = fig.gca()
ax.plot(freq/1e6, to_dB(spec[0,:]), label='B%i,T%i,P%i' % (beam,tune,0))
ax.plot(freq/1e6, to_dB(spec[1,:]), label='B%i,T%i,P%i' % (beam,tune,1))
ax.set_xlabel('Frequency [MHz]')
ax.set_ylabel('PSD [arb. dB]')
ax.legend(loc=0)
plt.show()

from lsl.common import stations

antennas = []
for i in xrange(2):
    if i / 2 == 0:
        newAnt = stations.Antenna(1)
    else:
        newAnt = stations.Antenna(2)
    if i % 2 == 0:
        newAnt.pol = 0
    else:
        newAnt.pol = 1
    newAnt.digitizer = i+1
    antennas.append(newAnt)
    
for ant in antennas:
    print str(ant)

freq, spec = fxc.StokesMaster(data, antennas, LFFT=512, SampleRate=srate, CentralFreq=cFreq)
print freq.shape, spec.shape

fig = plt.figure()
ax = fig.gca()
ax.plot(freq/1e6, spec[0,0,:], label='I')
ax.plot(freq/1e6, spec[1,0,:], label='Q')
ax.plot(freq/1e6, spec[2,0,:], label='U')
ax.plot(freq/1e6, spec[3,0,:], label='V')
ax.set_xlabel('Frequency [MHz]')
ax.set_ylabel('PSD [arb. dB]')
ax.legend(loc=0)
plt.show()