#!/usr/bin/env python
"""
# filterbank.py

Python class and command line utility for reading and plotting filterbank files.

This provides a class, Filterbank(), which can be used to read a .fil file:

    fil = Filterbank('test_psr.fil')
    print(fil.header)
    print(fil.data.shape)
    print(fil.freqs)

    plt.figure()
    fil.plot_spectrum(t=0)
    plt.show()

TODO: check the file seek logic works correctly for multiple IFs

"""

import sys
import six

from astropy.time import Time
import scipy.stats
from matplotlib.ticker import NullFormatter

import logging as logger

try:
    from .utils import db, lin, rebin, closest, unpack_2to8
    from .sigproc import *
except:
    from utils import db, lin, rebin, closest, unpack_2to8
    from sigproc import *

try:
    import h5py
    HAS_HDF5 = True
except ImportError:
    HAS_HDF5 = False

try:
    from pyslalib import slalib as s
    HAS_SLALIB = True
except ImportError:
    HAS_SLALIB = False

#import pdb #pdb.set_trace()

# Check if $DISPLAY is set (for handling plotting on remote machines with no X-forwarding)
if 'DISPLAY' in os.environ.keys():
    import pylab as plt
else:
    import matplotlib
    matplotlib.use('Agg')
    import pylab as plt

plt.rcParams['axes.formatter.useoffset'] = False

###
# Config values
###

MAX_PLT_POINTS      = 65536                  # Max number of points in matplotlib plot
MAX_IMSHOW_POINTS   = (8192, 4096)           # Max number of points in imshow plot
MAX_DATA_ARRAY_SIZE = 1024 * 1024 * 1024 * 2 # Max size of data array to load into memory
MAX_HEADER_BLOCKS   = 100                    # Max size of header (in 512-byte blocks)

# Telescope coordinates (needed for LSR calc)
parkes_coords = (-32.998370, 148.263659,  324.00)
gbt_coords    = (38.4331294, 79.8398397, 824.36)

###
# Main blimpy class
###

class Filterbank(object):
    """ Class for loading and plotting blimpy data """

    def __repr__(self):
        return "Filterbank data: %s" % self.filename

    def __init__(self, filename=None, f_start=None, f_stop=None,
                 t_start=None, t_stop=None, load_data=True,
                 header_dict=None, data_array=None, blank_dc=False,cal_band_pass=False):
        """ Class for loading and plotting blimpy data.

        This class parses the blimpy file and stores the header and data
        as objects:
            fb = Filterbank('filename_here.fil')
            fb.header        # blimpy header, as a dictionary
            fb.data          # blimpy data, as a numpy array

        Args:
            filename (str): filename of blimpy file.
            f_start (float): start frequency in MHz
            f_stop (float): stop frequency in MHz
            t_start (int): start integration ID
            t_stop (int): stop integration ID
            load_data (bool): load data. If set to False, only header will be read.
            header_dict (dict): Create blimpy from header dictionary + data array
            data_array (np.array): Create blimpy from header dict + data array
            blank_dc (bool): Blanks the DC bin or not.
        """

        if filename:
            self.filename = filename
            if HAS_HDF5:
                if h5py.is_hdf5(filename):
                    # TODO: self.read_hdf5(filename, f_start, f_stop, t_start, t_stop, load_data)
                    self.read_hdf5(filename, f_start, f_stop, t_start, t_stop, load_data)
                else:
                    self.read_filterbank(filename, f_start, f_stop, t_start, t_stop, load_data)
            else:
                self.read_filterbank(filename, f_start, f_stop, t_start, t_stop, load_data)
        elif header_dict is not None and data_array is not None:
            self.filename = b''
            self.header = header_dict
            self.data = data_array
            self.n_ints_in_file = 0

            self._setup_freqs()

        else:
            pass

        if blank_dc:
            print("Blanking DC bin")
            n_coarse_chan = self.calc_n_coarse_chan()
            self.blank_dc(n_coarse_chan)

        if cal_band_pass:
            print("Calibrating the band pass.")
            self.calibrate_band_pass_N1()

    def read_hdf5(self, filename, f_start=None, f_stop=None,
                        t_start=None, t_stop=None, load_data=True):
        """ Populate Filterbank instance with data from HDF5 file

        Note:
            This is to be deprecated in future, please use Waterfall() to open files.
        """
        print("Warning: this function will be deprecated in the future. Please use Waterfall to open HDF5 files.")
#        raise DeprecationWarning('Please use Waterfall to open HDF5 files.')

        self.header = {}
        self.filename = filename
        self.h5 = h5py.File(filename)
        for key, val in self.h5[b'data'].attrs.items():
            if six.PY3:
                key = bytes(key, 'ascii')
            if key == b'src_raj':
                self.header[key] = Angle(val, unit='hr')
            elif key == b'src_dej':
                self.header[key] = Angle(val, unit='deg')
            else:
                self.header[key] = val

        self.n_ints_in_file = self.h5[b"data"].shape[0]
        i_start, i_stop, chan_start_idx, chan_stop_idx = self._setup_freqs(f_start=f_start, f_stop=f_stop)
        ii_start, ii_stop, n_ints = self._setup_time_axis(t_start=t_start, t_stop=t_stop)

        if load_data:
            self.data = self.h5[b"data"][ii_start:ii_stop, :, chan_start_idx:chan_stop_idx]

            self.file_size_bytes = os.path.getsize(self.filename)

#         if self.header[b'foff'] < 0:
#             self.data = self.data[..., ::-1] # Reverse data

        else:
            print("Skipping data load...")
            self.data = np.array([0])
            self.n_ints_in_file  = 0
            self.file_size_bytes = os.path.getsize(self.filename)

    def _setup_freqs(self, f_start=None, f_stop=None):
        """ Setup frequency axis """
        ## Setup frequency axis
        f0 = self.header[b'fch1']
        f_delt = self.header[b'foff']

        i_start, i_stop = 0, self.header[b'nchans']
        if f_start:
            i_start = int((f_start - f0) / f_delt)
        if f_stop:
            i_stop  = int((f_stop - f0)  / f_delt)

        #calculate closest true index value
        chan_start_idx = np.int(i_start)
        chan_stop_idx  = np.int(i_stop)

        #create freq array
        if i_start < i_stop:
            i_vals = np.arange(chan_start_idx, chan_stop_idx)
        else:
            i_vals = np.arange(chan_stop_idx, chan_start_idx)

        self.freqs = f_delt * i_vals + f0

#         if f_delt < 0:
#             self.freqs = self.freqs[::-1]

        if chan_stop_idx < chan_start_idx:
            chan_stop_idx, chan_start_idx = chan_start_idx,chan_stop_idx

        return i_start, i_stop, chan_start_idx, chan_stop_idx

    def _setup_time_axis(self, t_start=None, t_stop=None):
        """  Setup time axis. """

        # now check to see how many integrations requested
        ii_start, ii_stop = 0, self.n_ints_in_file
        if t_start:
            ii_start = t_start
        if t_stop:
            ii_stop = t_stop
        n_ints = ii_stop - ii_start

        ## Setup time axis
        t0 = self.header[b'tstart']
        t_delt = self.header[b'tsamp']

        self.timestamps = np.arange(0, n_ints) * t_delt / 24./60./60 + t0

        return ii_start, ii_stop, n_ints

    def read_filterbank(self, filename=None, f_start=None, f_stop=None,
                        t_start=None, t_stop=None, load_data=True):
        """ Populate Filterbank instance with data from Filterbank file

        Note:
            This is to be deprecated in future, please use Waterfall() to open files.
        """
        if filename is None:
            filename = self.filename
        else:
            self.filename = filename

        self.header = read_header(filename)

        #convert input frequencies into what their corresponding index would be
        i_start, i_stop, chan_start_idx, chan_stop_idx = self._setup_freqs(f_start=f_start, f_stop=f_stop)

        n_bits  = self.header[b'nbits']
        n_bytes  = int(self.header[b'nbits'] / 8)
        n_chans = self.header[b'nchans']
        n_chans_selected = self.freqs.shape[0]
        n_ifs   = self.header[b'nifs']

        # Load binary data
        self.idx_data = len_header(filename)
        f = open(filename, 'rb')
        f.seek(self.idx_data)
        filesize = os.path.getsize(self.filename)
        n_bytes_data = filesize - self.idx_data

        # Finally add some other info to the class as objects
        self.n_ints_in_file  = calc_n_ints_in_file(self.filename)
        self.file_size_bytes = filesize

        ## Setup time axis
        ii_start, ii_stop, n_ints = self._setup_time_axis(t_start=t_start, t_stop=t_stop)

        # Seek to first integration
        f.seek(int(ii_start * n_bits * n_ifs * n_chans / 8), 1)

        # Set up indexes used in file read (taken out of loop for speed)
        i0 = np.min((chan_start_idx, chan_stop_idx))
        i1 = np.max((chan_start_idx, chan_stop_idx))

        #Set up the data type (taken out of loop for speed)
        if n_bits == 2:
            dd_type = b'uint8'
            n_chans_selected = int(n_chans_selected/4)
        elif n_bytes == 4:
            dd_type = b'float32'
        elif n_bytes == 2:
            dd_type = b'uint16'
        elif n_bytes == 1:
            dd_type = b'uint8'

        if load_data:

            if n_ints * n_ifs * n_chans_selected > MAX_DATA_ARRAY_SIZE:
                print("[Filterbank]  Error: data array is too large to load. Either select fewer points or manually increase MAX_DATA_ARRAY_SIZE. Large files are now handle with Waterfall .")
                sys.exit()

            if n_bits == 2:
                self.data = np.zeros((n_ints, n_ifs, n_chans_selected*4), dtype=dd_type)
            else:
                self.data = np.zeros((n_ints, n_ifs, n_chans_selected), dtype=dd_type)

            for ii in range(n_ints):
                """d = f.read(n_bytes * n_chans * n_ifs)
                """

                for jj in range(n_ifs):

                    f.seek(n_bytes * i0, 1) # 1 = from current location
                    #d = f.read(n_bytes * n_chans_selected)
                    #bytes_to_read = n_bytes * n_chans_selected

                    dd = np.fromfile(f, count=n_chans_selected, dtype=dd_type)

                    # Reverse array if frequency axis is flipped
#                     if f_delt < 0:
#                         dd = dd[::-1]

                    if n_bits == 2:
                        dd = unpack_2to8(dd)
                    self.data[ii, jj] = dd

                    f.seek(n_bytes * (n_chans - i1), 1)  # Seek to start of next block
        else:
            print("Skipping data load...")
            self.data = np.array([0], dtype=dd_type)

    def compute_lst(self):
        """ Compute LST for observation """
        if self.header[b'telescope_id'] == 6:
            self.coords = gbt_coords
        elif self.header[b'telescope_id'] == 4:
            self.coords = parkes_coords
        else:
            raise RuntimeError("Currently only Parkes and GBT supported")
        if HAS_SLALIB:
            # dut1 = (0.2 /3600.0) * np.pi/12.0
            dut1 = 0.0
            mjd = self.header[b'tstart']
            tellong = np.deg2rad(self.coords[1])
            last = s.sla_gmst(mjd) - tellong + s.sla_eqeqx(mjd) + dut1
            # lmst = s.sla_gmst(mjd) - tellong
            if last < 0.0 : last = last + 2.0*np.pi
            return last
        else:
            raise RuntimeError("This method requires pySLALIB")

    def compute_lsrk(self):
        """ Computes the LSR in km/s

        uses the MJD, RA and DEC of observation to compute
        along with the telescope location. Requires pyslalib
        """
        ra = Angle(self.header[b'src_raj'], unit='hourangle')
        dec = Angle(self.header[b'src_dej'], unit='degree')
        mjdd = self.header[b'tstart']
        rarad = ra.to('radian').value
        dcrad = dec.to('radian').value
        last = self.compute_lst()
        tellat  = np.deg2rad(self.coords[0])
        tellong = np.deg2rad(self.coords[1])

        # convert star position to vector
        starvect = s.sla_dcs2c(rarad, dcrad)

        # velocity component in ra,dec due to Earth rotation
        Rgeo = s.sla_rverot( tellat, rarad, dcrad, last)

        # get Barycentric and heliocentric velocity and position of the Earth.
        evp = s.sla_evp(mjdd, 2000.0)
        dvb = evp[0]   # barycentric velocity vector, in AU/sec
        dpb = evp[1]   # barycentric position vector, in AU
        dvh = evp[2]   # heliocentric velocity vector, in AU/sec
        dph = evp[3]   # heliocentric position vector, in AU

        # dot product of vector to object and heliocentric velocity
        # convert AU/sec to km/sec
        vcorhelio = -s.sla_dvdv( starvect, dvh) *149.597870e6
        vcorbary  = -s.sla_dvdv( starvect, dvb) *149.597870e6

        # rvlsrd is velocity component in ra,dec direction due to the Sun's
        # motion with respect to the "dynamical" local standard of rest
        rvlsrd = s.sla_rvlsrd( rarad,dcrad)

        # rvlsrk is velocity component in ra,dec direction due to i
        # the Sun's motion w.r.t the "kinematic" local standard of rest
        rvlsrk = s.sla_rvlsrk( rarad,dcrad)

        # rvgalc is velocity component in ra,dec direction due to
        #the rotation of the Galaxy.
        rvgalc = s.sla_rvgalc( rarad,dcrad)
        totalhelio = Rgeo + vcorhelio
        totalbary  = Rgeo + vcorbary
        totallsrk = totalhelio + rvlsrk
        totalgal  = totalbary  + rvlsrd + rvgalc

        return totallsrk

    def blank_dc(self, n_coarse_chan):
        """ Blank DC bins in coarse channels.

        Note: currently only works if entire file is read
        """

        if n_coarse_chan < 1:
            logger.warning('Coarse channel number < 1, unable to blank DC bin.')
            return None

        if not n_coarse_chan % int(n_coarse_chan) == 0:
            logger.warning('Selection does not contain an interger number of coarse channels, unable to blank DC bin.')
            return None

        n_coarse_chan = int(n_coarse_chan)

        n_chan = self.data.shape[-1]
        n_chan_per_coarse = int(n_chan / n_coarse_chan)

        mid_chan = int(n_chan_per_coarse / 2)

        for ii in range(n_coarse_chan):
            ss = ii*n_chan_per_coarse
            self.data[..., ss+mid_chan] = np.median(self.data[..., ss+mid_chan+5:ss+mid_chan+10])

    def info(self):
        """ Print header information """

        for key, val in self.header.items():
            if key == b'src_raj':
                val = val.to_string(unit=u.hour, sep=':')
            if key == b'src_dej':
                val = val.to_string(unit=u.deg, sep=':')
            if key == b'tsamp':
                val *= u.second
            if key in ('foff', 'fch1'):
                val *= u.MHz
            if key == b'tstart':
                print("%16s : %32s" % ("tstart (ISOT)", Time(val, format='mjd').isot))
                key = "tstart (MJD)"
            print("%16s : %32s" % (key, val))

        print("\n%16s : %32s" % ("Num ints in file", self.n_ints_in_file))
        print("%16s : %32s" % ("Data shape", self.data.shape))
        print("%16s : %32s" % ("Start freq (MHz)", self.freqs[0]))
        print("%16s : %32s" % ("Stop freq (MHz)", self.freqs[-1]))

    def generate_freqs(self, f_start, f_stop):
        """
        returns frequency array [f_start...f_stop]
        """

        fch1 = self.header[b'fch1']
        foff = self.header[b'foff']

        #convert input frequencies into what their corresponding index would be
        i_start = int((f_start - fch1) / foff)
        i_stop  = int((f_stop - fch1)  / foff)

        #calculate closest true index value
        chan_start_idx = np.int(i_start)
        chan_stop_idx  = np.int(i_stop)

        #create freq array
        i_vals = np.arange(chan_stop_idx, chan_start_idx, 1)

        freqs = foff * i_vals + fch1

        return freqs

    def grab_data(self, f_start=None, f_stop=None, t_start=None, t_stop=None, if_id=0):
        """ Extract a portion of data by frequency range.

        Args:
            f_start (float): start frequency in MHz
            f_stop (float): stop frequency in MHz
            if_id (int): IF input identification (req. when multiple IFs in file)

        Returns:
            (freqs, data) (np.arrays): frequency axis in MHz and data subset
        """

        if f_start is None:
            f_start = self.freqs[0]
        if f_stop is None:
            f_stop = self.freqs[-1]

        i0 = np.argmin(np.abs(self.freqs - f_start))
        i1 = np.argmin(np.abs(self.freqs - f_stop))

        if i0 < i1:
            plot_f    = self.freqs[i0:i1 + 1]
            plot_data = np.squeeze(self.data[t_start:t_stop, ..., i0:i1 + 1])
        else:
            plot_f    = self.freqs[i1:i0 + 1]
            plot_data = np.squeeze(self.data[t_start:t_stop, ..., i1:i0 + 1])

        return plot_f, plot_data

    def _calc_extent(self,plot_f=None,plot_t=None,MJD_time=False):
        """ Setup ploting edges.
        """

        plot_f_begin = plot_f[0]
        plot_f_end = plot_f[-1] + (plot_f[1]-plot_f[0])

        plot_t_begin = self.timestamps[0]
        plot_t_end  = self.timestamps[-1] + (self.timestamps[1] - self.timestamps[0])

        if MJD_time:
            extent=(plot_f_begin, plot_f_begin_end, plot_t_begin, plot_t_end)
        else:
            extent=(plot_f_begin, plot_f_end, 0.0,(plot_t_end-plot_t_begin)*24.*60.*60)

        return extent

    def plot_spectrum(self, t=0, f_start=None, f_stop=None, logged=False, if_id=0, c=None, **kwargs):
        """ Plot frequency spectrum of a given file

        Args:
            t (int): integration number to plot (0 -> len(data))
            logged (bool): Plot in linear (False) or dB units (True)
            if_id (int): IF identification (if multiple IF signals in file)
            c: color for line
            kwargs: keyword args to be passed to matplotlib plot()
        """
        if self.header[b'nbits'] <=2:
            logged = False
            t='all'
        ax = plt.gca()

        plot_f, plot_data = self.grab_data(f_start, f_stop, if_id)

        #Using accending frequency for all plots.
        if self.header[b'foff'] < 0:
            plot_data = plot_data[..., ::-1] # Reverse data
            plot_f = plot_f[::-1]

        if isinstance(t, int):
            print("extracting integration %i..." % t)
            plot_data = plot_data[t]
        elif t == b'all':
            print("averaging along time axis...")
            #Since the data has been squeezed, the axis for time goes away if only one bin, causing a bug with axis=1
            if len(plot_data.shape) > 1:
                plot_data = plot_data.mean(axis=0)
            else:
                plot_data = plot_data.mean()
        else:
            raise RuntimeError("Unknown integration %s" % t)

        # Rebin to max number of points
        dec_fac_x = 1
        if plot_data.shape[0] > MAX_PLT_POINTS:
            dec_fac_x = int(plot_data.shape[0] / MAX_PLT_POINTS)

        plot_data = rebin(plot_data, dec_fac_x, 1)
        plot_f    = rebin(plot_f, dec_fac_x, 1)

        if not c:
            kwargs['c'] = '#333333'

        if logged:
            plt.plot(plot_f, db(plot_data),label='Stokes I', **kwargs)
            plt.ylabel("Power [dB]")
        else:

            plt.plot(plot_f, plot_data,label='Stokes I', **kwargs)
            plt.ylabel("Power [counts]")
        plt.xlabel("Frequency [MHz]")
        plt.legend()

        try:
            plt.title(self.header[b'source_name'])
        except KeyError:
            plt.title(self.filename)

        plt.xlim(plot_f[0], plot_f[-1])

    def plot_spectrum_min_max(self, t=0, f_start=None, f_stop=None, logged=False, if_id=0, c=None, **kwargs):
        """ Plot frequency spectrum of a given file

        Args:
            logged (bool): Plot in linear (False) or dB units (True)
            if_id (int): IF identification (if multiple IF signals in file)
            c: color for line
            kwargs: keyword args to be passed to matplotlib plot()
        """
        ax = plt.gca()

        plot_f, plot_data = self.grab_data(f_start, f_stop, if_id)

        #Using accending frequency for all plots.
        if self.header[b'foff'] < 0:
            plot_data = plot_data[..., ::-1] # Reverse data
            plot_f = plot_f[::-1]

        fig_max = plot_data[0].max()
        fig_min = plot_data[0].min()

        print("averaging along time axis...")

        #Since the data has been squeezed, the axis for time goes away if only one bin, causing a bug with axis=1
        if len(plot_data.shape) > 1:
            plot_max = plot_data.max(axis=0)
            plot_min = plot_data.min(axis=0)
            plot_data = plot_data.mean(axis=0)
        else:
            plot_max = plot_data.max()
            plot_min = plot_data.min()
            plot_data = plot_data.mean()

        # Rebin to max number of points
        dec_fac_x = 1
        MAX_PLT_POINTS = 8*64  # Low resoluition to see the difference.
        if plot_data.shape[0] > MAX_PLT_POINTS:
            dec_fac_x = int(plot_data.shape[0] / MAX_PLT_POINTS)

        plot_data = rebin(plot_data, dec_fac_x, 1)
        plot_min = rebin(plot_min, dec_fac_x, 1)
        plot_max = rebin(plot_max, dec_fac_x, 1)
        plot_f    = rebin(plot_f, dec_fac_x, 1)

        if logged:
            plt.plot(plot_f, db(plot_data), "#333333", label='mean', **kwargs)
            plt.plot(plot_f, db(plot_max),  "#e74c3c", label='max', **kwargs)
            plt.plot(plot_f, db(plot_min),  '#3b5b92', label='min', **kwargs)
            plt.ylabel("Power [dB]")
        else:
            plt.plot(plot_f, plot_data,  "#333333", label='mean', **kwargs)
            plt.plot(plot_f, plot_max,   "#e74c3c", label='max', **kwargs)
            plt.plot(plot_f, plot_min,   '#3b5b92', label='min', **kwargs)
            plt.ylabel("Power [counts]")
        plt.xlabel("Frequency [MHz]")
        plt.legend()

        try:
            plt.title(self.header[b'source_name'])
        except KeyError:
            plt.title(self.filename)

        plt.xlim(plot_f[0], plot_f[-1])
        if logged:
            plt.ylim(db(fig_min),db(fig_max))

    def plot_waterfall(self, f_start=None, f_stop=None, if_id=0, logged=True, cb=True, MJD_time=False, **kwargs):
        """ Plot waterfall of data

        Args:
            f_start (float): start frequency, in MHz
            f_stop (float): stop frequency, in MHz
            logged (bool): Plot in linear (False) or dB units (True),
            cb (bool): for plotting the colorbar
            kwargs: keyword args to be passed to matplotlib imshow()
        """

        plot_f, plot_data = self.grab_data(f_start, f_stop, if_id)

        #Using accending frequency for all plots.
        if self.header[b'foff'] < 0:
            plot_data = plot_data[..., ::-1] # Reverse data
            plot_f = plot_f[::-1]

        if logged:
            plot_data = db(plot_data)

        # Make sure waterfall plot is under 4k*4k
        dec_fac_x, dec_fac_y = 1, 1
        if plot_data.shape[0] > MAX_IMSHOW_POINTS[0]:
            dec_fac_x = int(plot_data.shape[0] / MAX_IMSHOW_POINTS[0])

        if plot_data.shape[1] > MAX_IMSHOW_POINTS[1]:
            dec_fac_y =  int(plot_data.shape[1] /  MAX_IMSHOW_POINTS[1])

        plot_data = rebin(plot_data, dec_fac_x, dec_fac_y)

        try:
            plt.title(self.header[b'source_name'])
        except KeyError:
            plt.title(self.filename)

        extent = self._calc_extent(plot_f=plot_f,plot_t=self.timestamps,MJD_time=MJD_time)

        plt.imshow(plot_data,
            aspect='auto',
            origin='lower',
            rasterized=True,
            interpolation='nearest',
            extent=extent,
            cmap='viridis',
            **kwargs
        )
        if cb:
            plt.colorbar()
        plt.xlabel("Frequency [MHz]")
        if MJD_time:
            plt.ylabel("Time [MJD]")
        else:
            plt.ylabel("Time [s]")

    def plot_time_series(self, f_start=None, f_stop=None, if_id=0, logged=True, orientation='h', MJD_time=False, **kwargs):
        """ Plot the time series.

         Args:
            f_start (float): start frequency, in MHz
            f_stop (float): stop frequency, in MHz
            logged (bool): Plot in linear (False) or dB units (True),
            kwargs: keyword args to be passed to matplotlib imshow()
        """

        ax = plt.gca()
        plot_f, plot_data = self.grab_data(f_start, f_stop, if_id)

        if logged and self.header[b'nbits'] >= 8:
            plot_data = db(plot_data)

        #Since the data has been squeezed, the axis for time goes away if only one bin, causing a bug with axis=1
        if len(plot_data.shape) > 1:
            plot_data = plot_data.mean(axis=1)
        else:
            plot_data = plot_data.mean()

        #Make proper time axis for plotting (but only for plotting!). Note that this makes the values inclusive.
        extent = self._calc_extent(plot_f=plot_f,plot_t=self.timestamps,MJD_time=MJD_time)
        plot_t = np.linspace(extent[2],extent[3],len(self.timestamps))

        if MJD_time:
            tlabel = "Time [MJD]"
        else:
            tlabel = "Time [s]"

        if logged:
            plabel = "Power [dB]"
        else:
            plabel = "Power [counts]"

        # Reverse oder if vertical orientation.
        if 'v' in orientation:
            plt.plot(plot_data, plot_t, **kwargs)
            plt.xlabel(plabel)

        else:
            plt.plot(plot_t, plot_data, **kwargs)
            plt.xlabel(tlabel)
            plt.ylabel(plabel)

        ax.autoscale(axis='both',tight=True)

    def plot_kurtosis(self, f_start=None, f_stop=None, if_id=0, **kwargs):
        """ Plot kurtosis

         Args:
            f_start (float): start frequency, in MHz
            f_stop (float): stop frequency, in MHz
            kwargs: keyword args to be passed to matplotlib imshow()
        """
        ax = plt.gca()

        plot_f, plot_data = self.grab_data(f_start, f_stop, if_id)

        #Using accending frequency for all plots.
        if self.header[b'foff'] < 0:
            plot_data = plot_data[..., ::-1] # Reverse data
            plot_f = plot_f[::-1]

        try:
            plot_kurtosis = scipy.stats.kurtosis(plot_data, axis=0, nan_policy='omit')
        except:
            plot_kurtosis = plot_data*0.0

        plt.plot(plot_f, plot_kurtosis, **kwargs)
        plt.ylabel("Kurtosis")
        plt.xlabel("Frequency [MHz]")

        plt.xlim(plot_f[0], plot_f[-1])

    def plot_all(self, t=0, f_start=None, f_stop=None, logged=False, if_id=0, kurtosis=True, **kwargs):
        """ Plot waterfall of data as well as spectrum; also, placeholder to make even more complicated plots in the future.

        Args:
            f_start (float): start frequency, in MHz
            f_stop (float): stop frequency, in MHz
            logged (bool): Plot in linear (False) or dB units (True),
            t (int): integration number to plot (0 -> len(data))
            logged (bool): Plot in linear (False) or dB units (True)
            if_id (int): IF identification (if multiple IF signals in file)
            kwargs: keyword args to be passed to matplotlib plot() and imshow()
        """
        if self.header[b'nbits'] <=2:
            logged = False

        nullfmt = NullFormatter()  # no labels

        # definitions for the axes
        left, width = 0.35, 0.5
        bottom, height = 0.45, 0.5
        width2, height2 = 0.1125, 0.15
        bottom2, left2 = bottom - height2 - .025, left - width2 - .02
        bottom3, left3 = bottom2 - height2 - .025, 0.075

        rect_waterfall = [left, bottom, width, height]
        rect_colorbar = [left + width, bottom, .025, height]
        rect_spectrum = [left, bottom2, width, height2]
        rect_min_max = [left, bottom3, width, height2]
        rect_timeseries = [left + width, bottom, width2, height]
        rect_kurtosis = [left3, bottom3, 0.25, height2]
        rect_header = [left3 - .05, bottom, 0.2, height]

        # --------
        #         axColorbar = plt.axes(rect_colorbar)
        #         print 'Ploting Colorbar'
        #         print plot_data.max()
        #         print plot_data.min()
        #
        #         plot_colorbar = range(plot_data.min(),plot_data.max(),int((plot_data.max()-plot_data.min())/plot_data.shape[0]))
        #         plot_colorbar = np.array([[plot_colorbar],[plot_colorbar]])
        #
        #         plt.imshow(plot_colorbar,aspect='auto', rasterized=True, interpolation='nearest',)

        #         axColorbar.xaxis.set_major_formatter(nullfmt)
        #         axColorbar.yaxis.set_major_formatter(nullfmt)

        #         heatmap = axColorbar.pcolor(plot_data, edgecolors = 'none', picker=True)
        #         plt.colorbar(heatmap, cax = axColorbar)

        # --------
        axMinMax = plt.axes(rect_min_max)
        print('Plotting Min Max')
        self.plot_spectrum_min_max(logged=logged, f_start=f_start, f_stop=f_stop, t=t)
        plt.title('')
        axMinMax.yaxis.tick_right()
        axMinMax.yaxis.set_label_position("right")

        # --------
        axSpectrum = plt.axes(rect_spectrum,sharex=axMinMax)
        print('Plotting Spectrum')
        self.plot_spectrum(logged=logged, f_start=f_start, f_stop=f_stop, t=t)
        plt.title('')
        axSpectrum.yaxis.tick_right()
        axSpectrum.yaxis.set_label_position("right")
        plt.xlabel('')
#        axSpectrum.xaxis.set_major_formatter(nullfmt)
        plt.setp(axSpectrum.get_xticklabels(), visible=False)

        # --------
        axWaterfall = plt.axes(rect_waterfall,sharex=axMinMax)
        print('Plotting Waterfall')
        self.plot_waterfall(f_start=f_start, f_stop=f_stop, logged=logged, cb=False)
        plt.xlabel('')

        # no labels
#        axWaterfall.xaxis.set_major_formatter(nullfmt)
        plt.setp(axWaterfall.get_xticklabels(), visible=False)

        # --------
        axTimeseries = plt.axes(rect_timeseries)
        print('Plotting Timeseries')
        self.plot_time_series(f_start=f_start, f_stop=f_stop, orientation='v')
        axTimeseries.yaxis.set_major_formatter(nullfmt)
#        axTimeseries.xaxis.set_major_formatter(nullfmt)

        # --------
        # Could exclude since it takes much longer to run than the other plots.
        if kurtosis:
            axKurtosis = plt.axes(rect_kurtosis)
            print('Plotting Kurtosis')
            self.plot_kurtosis(f_start=f_start, f_stop=f_stop)

        # --------
        axHeader = plt.axes(rect_header)
        print('Plotting Header')
        # Generate nicer header
        telescopes = {0: 'Fake data',
                      1: 'Arecibo',
                      2: 'Ooty',
                      3: 'Nancay',
                      4: 'Parkes',
                      5: 'Jodrell',
                      6: 'GBT',
                      8: 'Effelsberg',
                      10: 'SRT',
                      64: 'MeerKAT',
                      65: 'KAT7'
                      }

        telescope = telescopes.get(self.header[b"telescope_id"], self.header[b"telescope_id"])

        plot_header = "%14s: %s\n" % ("TELESCOPE_ID", telescope)
        for key in (b'SRC_RAJ', b'SRC_DEJ', b'TSTART', b'NCHANS', b'NBEAMS', b'NIFS', b'NBITS'):
            try:
                plot_header += "%14s: %s\n" % (key, self.header[key.lower()])
            except KeyError:
                pass
        fch1 = "%6.6f MHz" % self.header[b'fch1']

        foff = (self.header[b'foff'] * 1e6 * u.Hz)
        if np.abs(foff) > 1e6 * u.Hz:
            foff = str(foff.to('MHz'))
        elif np.abs(foff) > 1e3 * u.Hz:
            foff = str(foff.to('kHz'))
        else:
            foff = str(foff.to('Hz'))

        plot_header += "%14s: %s\n" % ("FCH1", fch1)
        plot_header += "%14s: %s\n" % ("FOFF", foff)

        plt.text(0.05, .95, plot_header, ha='left', va='top', wrap=True)

        axHeader.set_facecolor('white')
        axHeader.xaxis.set_major_formatter(nullfmt)
        axHeader.yaxis.set_major_formatter(nullfmt)

    def write_to_filterbank(self, filename_out):
        """ Write data to blimpy file.

        Args:
            filename_out (str): Name of output file
        """

        print("[Filterbank] Warning: Non-standard function to write in filterbank (.fil) format. Please use Waterfall.")

        n_bytes  = int(self.header[b'nbits'] / 8)
        with open(filename_out, "wb") as fileh:
            fileh.write(generate_sigproc_header(self))
            j = self.data
            if n_bytes == 4:
                np.float32(j.ravel()).tofile(fileh)
            elif n_bytes == 2:
                np.int16(j.ravel()).tofile(fileh)
            elif n_bytes == 1:
                np.int8(j.ravel()).tofile(fileh)

    def write_to_hdf5(self, filename_out, *args, **kwargs):
        """ Write data to HDF5 file.

        Args:
            filename_out (str): Name of output file
        """

        print("[Filterbank] Warning: Non-standard function to write in HDF5 (.h5) format. Please use Waterfall.")

        if not HAS_HDF5:
            raise RuntimeError("h5py package required for HDF5 output.")

        with h5py.File(filename_out, 'w') as h5:

            dset = h5.create_dataset(b'data',
                              data=self.data,
                              compression='lzf')

            dset_mask = h5.create_dataset(b'mask',
                                     shape=self.data.shape,
                                     compression='lzf',
                                     dtype='uint8')

            dset.dims[0].label = b"frequency"
            dset.dims[1].label = b"feed_id"
            dset.dims[2].label = b"time"

            dset_mask.dims[0].label = b"frequency"
            dset_mask.dims[1].label = b"feed_id"
            dset_mask.dims[2].label = b"time"

            # Copy over header information as attributes
            for key, value in self.header.items():
                dset.attrs[key] = value

    def calibrate_band_pass_N1(self):
        """ One way to calibrate the band pass is to take the median value
            for every frequency fine channel, and divide by it.
        """

        band_pass = np.median(self.data.squeeze(),axis=0)
        self.data = self.data/band_pass

def cmd_tool(args=None):
    """ Command line tool for plotting and viewing info on filterbank files """

    from argparse import ArgumentParser

    parser = ArgumentParser(description="Command line utility for reading and plotting filterbank files.")

    parser.add_argument('-p', action='store',  default='ank', dest='what_to_plot', type=str,
                        help='Show: "w" waterfall (freq vs. time) plot; "s" integrated spectrum plot; \
                        "t" for time series; "mm" for spectrum including min max; "k" for kurtosis; \
                        "a" for all available plots and information; and "ank" for all but kurtosis.')
    parser.add_argument('filename', type=str,
                        help='Name of file to read')
    parser.add_argument('-b', action='store', default=None, dest='f_start', type=float,
                        help='Start frequency (begin), in MHz')
    parser.add_argument('-e', action='store', default=None, dest='f_stop', type=float,
                        help='Stop frequency (end), in MHz')
    parser.add_argument('-B', action='store', default=None, dest='t_start', type=int,
                        help='Start integration (begin) ID')
    parser.add_argument('-E', action='store', default=None, dest='t_stop', type=int,
                        help='Stop integration (end) ID')
    parser.add_argument('-i', action='store_true', default=False, dest='info_only',
                        help='Show info only')
    parser.add_argument('-a', action='store_true', default=False, dest='average',
                       help='average along time axis (plot spectrum only)')
    parser.add_argument('-s', action='store', default='', dest='plt_filename', type=str,
                       help='save plot graphic to file (give filename as argument)')
    parser.add_argument('-S', action='store_true', default=False, dest='save_only',
                       help='Turn off plotting of data and only save to file.')
    parser.add_argument('-D', action='store_false', default=True, dest='blank_dc',
                       help='Use to not blank DC bin.')
    parser.add_argument('-c', action='store_true', default=False, dest='calibrate_band_pass',
                        help='Calibrate band pass.')
    args = parser.parse_args()

    # Open blimpy data
    filename = args.filename
    load_data = not args.info_only

    # only load one integration if looking at spectrum
    wtp = args.what_to_plot
    if not wtp or 's' in wtp:
        if args.t_start == None:
            t_start = 0
        else:
            t_start = args.t_start
        t_stop  = t_start + 1

        if args.average:
            t_start = None
            t_stop  = None
    else:
        t_start = args.t_start
        t_stop  = args.t_stop

    if args.info_only:
        args.blank_dc = False
        args.calibrate_band_pass = False

    fil = Filterbank(filename, f_start=args.f_start, f_stop=args.f_stop,
                     t_start=t_start, t_stop=t_stop,
                     load_data=load_data,blank_dc=args.blank_dc,
                     cal_band_pass=args.calibrate_band_pass)
    fil.info()

    # And if we want to plot data, then plot data.
    if not args.info_only:
        # check start & stop frequencies make sense
        #try:
        #    if args.f_start:
        #        print "Start freq: %2.2f" % args.f_start
        #        assert args.f_start >= fil.freqs[0] or np.isclose(args.f_start, fil.freqs[0])
        #
        #    if args.f_stop:
        #        print "Stop freq: %2.2f" % args.f_stop
        #        assert args.f_stop <= fil.freqs[-1] or np.isclose(args.f_stop, fil.freqs[-1])
        #except AssertionError:
        #    print "Error: Start and stop frequencies must lie inside file's frequency range."
        #    print "i.e. between %2.2f-%2.2f MHz." % (fil.freqs[0], fil.freqs[-1])
        #    exit()

        if args.what_to_plot == "w":
            plt.figure("waterfall", figsize=(8, 6))
            fil.plot_waterfall(f_start=args.f_start, f_stop=args.f_stop)
        elif args.what_to_plot == "s":
            plt.figure("Spectrum", figsize=(8, 6))
            fil.plot_spectrum(logged=True, f_start=args.f_start, f_stop=args.f_stop, t='all')
        elif args.what_to_plot == "mm":
            plt.figure("min max", figsize=(8, 6))
            fil.plot_spectrum_min_max(logged=True, f_start=args.f_start, f_stop=args.f_stop, t='all')
        elif args.what_to_plot == "k":
            plt.figure("kurtosis", figsize=(8, 6))
            fil.plot_kurtosis(f_start=args.f_start, f_stop=args.f_stop)
        elif args.what_to_plot == "t":
            plt.figure("Time Series", figsize=(8, 6))
            fil.plot_time_series(f_start=args.f_start, f_stop=args.f_stop,orientation='h')
        elif args.what_to_plot == "a":
            plt.figure("Multiple diagnostic plots", figsize=(12, 9),facecolor='white')
            fil.plot_all(logged=True, f_start=args.f_start, f_stop=args.f_stop, t='all')
        elif args.what_to_plot == "ank":
            plt.figure("Multiple diagnostic plots", figsize=(12, 9),facecolor='white')
            fil.plot_all(logged=True, f_start=args.f_start, f_stop=args.f_stop, t='all',kurtosis=False)

        if args.plt_filename != '':
            plt.savefig(args.plt_filename)

        if not args.save_only:
            if 'DISPLAY' in os.environ.keys():
                plt.show()
            else:
                print("No $DISPLAY available.")

if __name__ == "__main__":
    cmd_tool()


from blimpy import Waterfall
import numpy as np
from scipy.optimize import curve_fit
from .fluxcal import foldcal

def get_stokes(cross_dat, feedtype='l'):
    '''Output stokes parameters (I,Q,U,V) for a rawspec
    cross polarization filterbank file'''

    #Compute Stokes Parameters
    if feedtype=='l':
        #I = XX+YY
        I = cross_dat[:,0,:]+cross_dat[:,1,:]
        #Q = XX-YY
        Q = cross_dat[:,0,:]-cross_dat[:,1,:]
        #U = 2*Re(XY)
        U = 2*cross_dat[:,2,:]
        #V = -2*Im(XY)
        V = -2*cross_dat[:,3,:]

    elif feedtype=='c':
        #I = LL+RR
        I = cross_dat[:,0,:]+cross_dat[:,1,:]
        #Q = 2*Re(RL)
        Q = 2*cross_dat[:,2,:]
        #U = 2*Im(RL)
        U = -2*cross_dat[:,3,:]
        #V = RR-LL
        V = cross_dat[:,1,:]-cross_dat[:,0,:]
    else:
        raise ValueError('feedtype must be \'l\' (linear) or \'c\' (circular)')

    #Add middle dimension to match Filterbank format
    I = np.expand_dims(I,axis=1)
    Q = np.expand_dims(Q,axis=1)
    U = np.expand_dims(U,axis=1)
    V = np.expand_dims(V,axis=1)

    #Compute linear polarization
    #L=np.sqrt(np.square(Q)+np.square(U))

    return I,Q,U,V

def convert_to_coarse(data,chan_per_coarse):
    '''
    Converts a data array with length n_chans to an array of length n_coarse_chans
    by averaging over the coarse channels
    '''
    #find number of coarse channels and reshape array
    num_coarse = data.size/chan_per_coarse
    data_shaped = np.array(np.reshape(data,(num_coarse,chan_per_coarse)))

    #Return the average over each coarse channel
    return np.mean(data_shaped[:,2:-1],axis=1)

def phase_offsets(Idat,Qdat,Udat,Vdat,tsamp,chan_per_coarse,feedtype='l',**kwargs):
    '''
    Calculates phase difference between X and Y feeds given U and V (U and Q for circular basis)
    data from a noise diode measurement on the target
    '''
    #Fold noise diode data and calculate ON OFF diferences for U and V
    if feedtype=='l':
        U_OFF,U_ON = foldcal(Udat,tsamp,**kwargs)
        V_OFF,V_ON = foldcal(Vdat,tsamp,**kwargs)
        Udiff = U_ON-U_OFF
        Vdiff = V_ON-V_OFF
        poffset = np.arctan2(-1*Vdiff,Udiff)

    if feedtype=='c':
        U_OFF,U_ON = foldcal(Udat,tsamp,**kwargs)
        Q_OFF,Q_ON = foldcal(Qdat,tsamp,**kwargs)
        Udiff = U_ON-U_OFF
        Qdiff = Q_ON-Q_OFF
        poffset = np.arctan2(Udiff,Qdiff)

    coarse_p =  convert_to_coarse(poffset,chan_per_coarse)

    #Correct for problems created by discontinuity in arctan
    #Find whether phase offsets have increasing or decreasing slope
    y = coarse_p[:6]
    x = np.arange(y.size)
    m = np.polyfit(x,y,1)[0]

    for i in range(coarse_p.size-3):
        if (m>0 and coarse_p[i+1]<coarse_p[i]) or (m<0 and coarse_p[i+1]>coarse_p[i]):
            coarse_p[i+1] = 2*coarse_p[i+2]-coarse_p[i+3]    #Move problem point near the next

    return coarse_p

def gain_offsets(Idat,Qdat,Udat,Vdat,tsamp,chan_per_coarse,feedtype='l',**kwargs):
    '''
    Determines relative gain error in the X and Y feeds for an
    observation given I and Q (I and V for circular basis) noise diode data.
    '''
    if feedtype=='l':
        #Fold noise diode data and calculate ON OFF differences for I and Q
        I_OFF,I_ON = foldcal(Idat,tsamp,**kwargs)
        Q_OFF,Q_ON = foldcal(Qdat,tsamp,**kwargs)

        #Calculate power in each feed for noise diode ON and OFF
        XX_ON = (I_ON+Q_ON)/2
        XX_OFF = (I_OFF+Q_OFF)/2
        YY_ON = (I_ON-Q_ON)/2
        YY_OFF = (I_OFF-Q_OFF)/2

        #Calculate gain offset (divided by 2) as defined in Heiles (2001)
        G = (XX_OFF-YY_OFF)/(XX_OFF+YY_OFF)

    if feedtype=='c':
        #Fold noise diode data and calculate ON OFF differences for I and Q
        I_OFF,I_ON = foldcal(Idat,tsamp,**kwargs)
        V_OFF,V_ON = foldcal(Vdat,tsamp,**kwargs)

        #Calculate power in each feed for noise diode ON and OFF
        RR_ON = (I_ON+V_ON)/2
        RR_OFF = (I_OFF+V_OFF)/2
        LL_ON = (I_ON-V_ON)/2
        LL_OFF = (I_OFF-V_OFF)/2

        #Calculate gain offset (divided by 2) as defined in Heiles (2001)
        G = (RR_OFF-LL_OFF)/(RR_OFF+LL_OFF)

    return convert_to_coarse(G,chan_per_coarse)
def apply_Mueller(I,Q,U,V, gain_offsets, phase_offsets, chan_per_coarse, feedtype='l'):
    '''
    Returns calibrated Stokes parameters for an observation given an array
    of differential gains and phase differences.
    '''

    #Find shape of data arrays and calculate number of coarse channels
    shape = I.shape
    ax0 = I.shape[0]
    ax1 = I.shape[1]
    nchans = I.shape[2]
    ncoarse = nchans/chan_per_coarse

    #Reshape data arrays to separate coarse channels
    I = np.reshape(I,(ax0,ax1,ncoarse,chan_per_coarse))
    Q = np.reshape(Q,(ax0,ax1,ncoarse,chan_per_coarse))
    U = np.reshape(U,(ax0,ax1,ncoarse,chan_per_coarse))
    V = np.reshape(V,(ax0,ax1,ncoarse,chan_per_coarse))

    #Swap axes 2 and 3 to in order for broadcasting to work correctly
    I = np.swapaxes(I,2,3)
    Q = np.swapaxes(Q,2,3)
    U = np.swapaxes(U,2,3)
    V = np.swapaxes(V,2,3)

    #Apply top left corner of electronics chain inverse Mueller matrix
    a = 1/(1-gain_offsets**2)
    if feedtype=='l':
        Icorr = a*(I-gain_offsets*Q)
        Qcorr = a*(-1*gain_offsets*I+Q)
        I = None
        Q = None
    if feedtype=='c':
        Icorr = a*(I-gain_offsets*V)
        Vcorr = a*(-1*gain_offsets*I+V)
        I = None
        V = None

    #Apply bottom right corner of electronics chain inverse Mueller matrix
    if feedtype=='l':
        Ucorr = U*np.cos(phase_offsets)-V*np.sin(phase_offsets)
        Vcorr = U*np.sin(phase_offsets)+V*np.cos(phase_offsets)
        U = None
        V = None
    if feedtype=='c':
        Qcorr = Q*np.cos(phase_offsets)+U*np.sin(phase_offsets)
        Ucorr = -1*Q*np.sin(phase_offsets)+U*np.cos(phase_offsets)
        Q = None
        U = None

    #Reshape arrays to original shape
    Icorr = np.reshape(np.swapaxes(Icorr,2,3),shape)
    Qcorr = np.reshape(np.swapaxes(Qcorr,2,3),shape)
    Ucorr = np.reshape(np.swapaxes(Ucorr,2,3),shape)
    Vcorr = np.reshape(np.swapaxes(Vcorr,2,3),shape)

    #Return corrected data arrays
    return Icorr,Qcorr,Ucorr,Vcorr

def calibrate_pols(cross_pols,diode_cross,obsI=None,onefile=True,feedtype='l',**kwargs):
    '''
    Write Stokes-calibrated filterbank file for a given observation
    with a calibrator noise diode measurement on the source

    Parameters
    ----------
    cross_pols : string
        Path to cross polarization filterbank file (rawspec output) for observation to be calibrated
    diode_cross : string
        Path to cross polarization filterbank file of noise diode measurement ON the target
    obsI : string
        Path to Stokes I filterbank file of main observation (only needed if onefile=False)
    onefile : boolean
        True writes all calibrated Stokes parameters to a single filterbank file,
        False writes four separate files
    feedtype : 'l' or 'c'
        Basis of antenna dipoles. 'c' for circular, 'l' for linear
    '''
    #Obtain time sample length, frequencies, and noise diode data
    obs = Waterfall(diode_cross,max_load=150)
    cross_dat = obs.data
    tsamp = obs.header['tsamp']

    #Calculate number of coarse channels in the noise diode measurement (usually 8)
    dio_ncoarse = obs.calc_n_coarse_chan()
    dio_nchans = obs.header['nchans']
    dio_chan_per_coarse = dio_nchans/dio_ncoarse
    obs = None
    Idat,Qdat,Udat,Vdat = get_stokes(cross_dat,feedtype)
    cross_dat = None
    #Calculate differential gain and phase from noise diode measurements
    print('Calculating Mueller Matrix variables')
    gams = gain_offsets(Idat,Qdat,Udat,Vdat,tsamp,dio_chan_per_coarse,feedtype,**kwargs)
    psis = phase_offsets(Idat,Qdat,Udat,Vdat,tsamp,dio_chan_per_coarse,feedtype,**kwargs)

    #Clear data arrays to save memory
    Idat = None
    Qdat = None
    Udat = None
    Vdat = None

    #Get corrected Stokes parameters
    print('Opening '+cross_pols)
    cross_obs = Waterfall(cross_pols,max_load=150)
    obs_ncoarse = cross_obs.calc_n_coarse_chan()
    obs_nchans = cross_obs.header['nchans']
    obs_chan_per_coarse = obs_nchans/obs_ncoarse

    print('Grabbing Stokes parameters')
    I,Q,U,V = get_stokes(cross_obs.data,feedtype)

    print('Applying Mueller Matrix')
    I,Q,U,V = apply_Mueller(I,Q,U,V,gams,psis,obs_chan_per_coarse,feedtype)

    #Use onefile (default) to produce one filterbank file containing all Stokes information
    if onefile==True:
        cross_obs.data[:,0,:] = np.squeeze(I)
        cross_obs.data[:,1,:] = np.squeeze(Q)
        cross_obs.data[:,2,:] = np.squeeze(U)
        cross_obs.data[:,3,:] = np.squeeze(V)
        cross_obs.write_to_fil(cross_pols[:-15]+'.SIQUV.polcal.fil')
        print('Calibrated Stokes parameters written to '+cross_pols[:-15]+'.SIQUV.polcal.fil')
        return

    #Write corrected Stokes parameters to four filterbank files if onefile==False
    obs = Waterfall(obs_I,max_load=150)
    obs.data = I
    obs.write_to_fil(cross_pols[:-15]+'.SI.polcal.fil')   #assuming file is named *.cross_pols.fil
    print('Calibrated Stokes I written to '+cross_pols[:-15]+'.SI.polcal.fil')

    obs.data = Q
    obs.write_to_fil(cross_pols[:-15]+'.Q.polcal.fil')   #assuming file is named *.cross_pols.fil
    print('Calibrated Stokes Q written to '+cross_pols[:-15]+'.Q.polcal.fil')

    obs.data = U
    obs.write_to_fil(cross_pols[:-15]+'.U.polcal.fil')   #assuming file is named *.cross_pols.fil
    print('Calibrated Stokes U written to '+cross_pols[:-15]+'.U.polcal.fil')

    obs.data = V
    obs.write_to_fil(cross_pols[:-15]+'.V.polcal.fil')   #assuming file is named *.cross_pols.fil
    print('Calibrated Stokes V written to '+cross_pols[:-15]+'.V.polcal.fil')

def fracpols(str, **kwargs):
    '''Output fractional linear and circular polarizations for a
    rawspec cross polarization .fil file. NOT STANDARD USE'''

    I,Q,U,V,L=get_stokes(str, **kwargs)
    return L/I,V/I

def write_stokefils(str, str_I, Ifil=False, Qfil=False, Ufil=False, Vfil=False, Lfil=False, **kwargs):
    '''Writes up to 5 new filterbank files corresponding to each Stokes
    parameter (and total linear polarization L) for a given cross polarization .fil file'''

    I,Q,U,V,L=get_stokes(str, **kwargs)
    obs = Waterfall(str_I, max_load=150) #Load filterbank file to write stokes data to
    if Ifil:
        obs.data = I
        obs.write_to_fil(str[:-15]+'.I.fil')   #assuming file is named *.cross_pols.fil

    if Qfil:
        obs.data = Q
        obs.write_to_fil(str[:-15]+'.Q.fil')   #assuming file is named *.cross_pols.fil

    if Ufil:
        obs.data = U
        obs.write_to_fil(str[:-15]+'.U.fil')   #assuming file is named *.cross_pols.fil

    if Vfil:
        obs.data = V
        obs.write_to_fil(str[:-15]+'.V.fil')   #assuming file is named *.cross_pols.fil

    if Lfil:
        obs.data = L
        obs.write_to_fil(str[:-15]+'.L.fil')   #assuming file is named *.cross_pols.fil

def write_polfils(str, str_I, **kwargs):
    '''Writes two new filterbank files containing fractional linear and
    circular polarization data'''

    lin,circ=fracpols(str, **kwargs)
    obs = Waterfall(str_I, max_load=150)

    obs.data = lin
    obs.write_to_fil(str[:-15]+'.linpol.fil')   #assuming file is named *.cross_pols.fil

    obs.data = circ
    obs.write_to_fil(str[:-15]+'.circpol.fil')   #assuming file is named *.cross_pols.fil

#end module


"""
# utils.py
useful helper functions for common data manipulation tasks
"""
import numpy as np

def db(x):
    """ Convert linear to dB """
    return 10 * np.log10(x)

def lin(x):
    """ Convert dB to linear """
    return 10.0 ** (x / 10.0)

def closest(xarr, val):
    """ Return the index of the closest in xarr to value val """
    idx_closest = np.argmin(np.abs(np.array(xarr) - val))
    return idx_closest

def rebin(d, n_x, n_y=None):
    """ Rebin data by averaging bins together

    Args:
    d (np.array): data
    n_x (int): number of bins in x dir to rebin into one
    n_y (int): number of bins in y dir to rebin into one

    Returns:
    d: rebinned data with shape (n_x, n_y)
    """

    if d.ndim == 2:
        if n_y is None:
            n_y = 1
        if n_x is None:
            n_x = 1
        d = d[:int(d.shape[0] // n_x) * n_x, :int(d.shape[1] // n_y) * n_y]
        d = d.reshape((d.shape[0] // n_x, n_x, d.shape[1] // n_y, n_y))
        d = d.mean(axis=3)
        d = d.mean(axis=1)
    elif d.ndim == 1:
        d = d[:int(d.shape[0] // n_x) * n_x]
        d = d.reshape((d.shape[0] // n_x, n_x))
        d = d.mean(axis=1)
    else:
        raise RuntimeError("Only NDIM <= 2 supported")
    return d

def unpack(data, nbit):
    """upgrade data from nbits to 8bits

    Notes: Pretty sure this function is a little broken!
    """
    if nbit > 8:
        raise ValueError("unpack: nbit must be <= 8")
    if 8 % nbit != 0:
        raise ValueError("unpack: nbit must divide into 8")
    if data.dtype not in (np.uint8, np.int8):
        raise TypeError("unpack: dtype must be 8-bit")
    if nbit == 8:
        return data
    elif nbit == 4:
        data = unpack_4to8(data)
        return data
    elif nbit == 2:
        data = unpack_2to8(data)
        return data
    elif nbit == 1:
        data = unpack_1to8(data)
        return data

def unpack_1to8(data):
    """ Promote 1-bit unisgned data into 8-bit unsigned data.

    Args:
        data: Numpy array with dtype == uint8
    """
    return np.unpackbits(data)

def unpack_2to8(data):
    """ Promote 2-bit unisgned data into 8-bit unsigned data.

    Args:
        data: Numpy array with dtype == uint8

    Notes:
        DATA MUST BE LOADED as np.array() with dtype='uint8'.

        This works with some clever shifting and AND / OR operations.
        Data is LOADED as 8-bit, then promoted to 32-bits:
        /ABCD EFGH/ (8 bits of data)
        /0000 0000/0000 0000/0000 0000/ABCD EFGH/ (8 bits of data as a 32-bit word)

        Once promoted, we can do some shifting, AND and OR operations:
        /0000 0000/0000 ABCD/EFGH 0000/0000 0000/ (shifted << 12)
        /0000 0000/0000 ABCD/EFGH 0000/ABCD EFGH/ (bitwise OR of previous two lines)
        /0000 0000/0000 ABCD/0000 0000/0000 EFGH/ (bitwise AND with mask 0xF000F)
        /0000 00AB/CD00 0000/0000 00EF/GH00 0000/ (prev. line shifted << 6)
        /0000 00AB/CD00 ABCD/0000 00EF/GH00 EFGH/ (bitwise OR of previous two lines)
        /0000 00AB/0000 00CD/0000 00EF/0000 00GH/ (bitwise AND with 0x3030303)

        Then we change the view of the data to interpret it as 4x8 bit:
        [000000AB, 000000CD, 000000EF, 000000GH]  (change view from 32-bit to 4x8-bit)

        The converted bits are then mapped to values in the range [-40, 40] according to a lookup chart.
        The mapping is based on specifications in the breakthough docs:
        https://github.com/UCBerkeleySETI/breakthrough/blob/master/doc/RAW-File-Format.md

    """
    two_eight_lookup = {0: 40,
                        1: 12,
                        2: -12,
                        3: -40}

    tmp = data.astype(np.uint32)
    tmp = (tmp | (tmp << 12)) & 0xF000F
    tmp = (tmp | (tmp << 6)) & 0x3030303
    tmp = tmp.byteswap()
    tmp = tmp.view('uint8')
    mapped = np.array(tmp, dtype=np.int8)
    for k, v in two_eight_lookup.items():
        mapped[tmp == k] = v
    return mapped

def unpack_4to8(data):
    """ Promote 2-bit unisgned data into 8-bit unsigned data.

    Args:
        data: Numpy array with dtype == uint8

    Notes:
        # The process is this:
        # ABCDEFGH [Bits of one 4+4-bit value]
        # 00000000ABCDEFGH [astype(uint16)]
        # 0000ABCDEFGH0000 [<< 4]
        # 0000ABCDXXXXEFGH [bitwise 'or' of previous two lines]
        # 0000111100001111 [0x0F0F]
        # 0000ABCD0000EFGH [bitwise 'and' of previous two lines]
        # ABCD0000EFGH0000 [<< 4]
        # which effectively pads the two 4-bit values with zeros on the right
        # Note: This technique assumes LSB-first ordering
    """

    tmpdata = data.astype(np.int16)  # np.empty(upshape, dtype=np.int16)
    tmpdata = (tmpdata | (tmpdata << 4)) & 0x0F0F
    # tmpdata = tmpdata << 4 # Shift into high bits to avoid needing to sign extend
    updata = tmpdata.byteswap()
    return updata.view(data.dtype)


import pylab as plt
from .fluxcal import *
from .stokescal import *
from blimpy import Waterfall
import numpy as np

def get_diff(dio_cross,feedtype,**kwargs):
    '''
    Returns ON-OFF for all Stokes parameters given a cross_pols noise diode measurement
    '''
    #Get Stokes parameters, frequencies, and time sample length
    obs = Waterfall(dio_cross,max_load=150)
    freqs = obs.populate_freqs()
    tsamp = obs.header['tsamp']
    data = obs.data
    obs = None

    I,Q,U,V = get_stokes(data,feedtype)

    #Fold noise diode data
    I_OFF,I_ON = foldcal(I,tsamp,**kwargs)
    Q_OFF,Q_ON = foldcal(Q,tsamp,**kwargs)
    U_OFF,U_ON = foldcal(U,tsamp,**kwargs)
    V_OFF,V_ON = foldcal(V,tsamp,**kwargs)

    #Do ON-OFF subtraction
    Idiff = I_ON-I_OFF
    Qdiff = Q_ON-Q_OFF
    Udiff = U_ON-U_OFF
    Vdiff = V_ON-V_OFF

    return Idiff,Qdiff,Udiff,Vdiff,freqs

def plot_Stokes_diode(dio_cross,diff=True,feedtype='l',**kwargs):
    '''
    Plots the uncalibrated full stokes spectrum of the noise diode.
    Use diff=False to plot both ON and OFF, or diff=True for ON-OFF
    '''

    #If diff=True, get ON-OFF. If not get ON and OFF separately
    if diff==True:
        Idiff,Qdiff,Udiff,Vdiff,freqs = get_diff(dio_cross,feedtype,**kwargs)
    else:
        obs = Waterfall(dio_cross,max_load=150)
        freqs = obs.populate_freqs()
        tsamp = obs.header['tsamp']
        data = obs.data
        I,Q,U,V = get_stokes(data,feedtype)

        I_OFF,I_ON = foldcal(I,tsamp,**kwargs)
        Q_OFF,Q_ON = foldcal(Q,tsamp,**kwargs)
        U_OFF,U_ON = foldcal(U,tsamp,**kwargs)
        V_OFF,V_ON = foldcal(V,tsamp,**kwargs)

    #Plot spectra
    if diff==True:
        plt.plot(freqs,Idiff,'k-',label='I')
        plt.plot(freqs,Qdiff,'r-',label='Q')
        plt.plot(freqs,Udiff,'g-',label='U')
        plt.plot(freqs,Vdiff,'m-',label='V')
    else:
        plt.plot(freqs,I_ON,'k-',label='I ON')
        plt.plot(freqs,I_OFF,'k--',label='I OFF')
        plt.plot(freqs,Q_ON,'r-',label='Q ON')
        plt.plot(freqs,Q_OFF,'r--',label='Q OFF')
        plt.plot(freqs,U_ON,'g-',label='U ON')
        plt.plot(freqs,U_OFF,'g--',label='U OFF')
        plt.plot(freqs,V_ON,'m-',label='V ON')
        plt.plot(freqs,V_OFF,'m--',label='V OFF')

    plt.legend()
    plt.xlabel('Frequency (MHz)')
    plt.title('Uncalibrated Full Stokes Noise Diode Spectrum')
    plt.ylabel('Power (Counts)')

def plot_calibrated_diode(dio_cross,chan_per_coarse=8,feedtype='l',**kwargs):
    '''
    Plots the corrected noise diode spectrum for a given noise diode measurement
    after application of the inverse Mueller matrix for the electronics chain.
    '''
    #Get full stokes data for the ND observation
    obs = Waterfall(dio_cross,max_load=150)
    freqs = obs.populate_freqs()
    tsamp = obs.header['tsamp']
    data = obs.data
    obs = None
    I,Q,U,V = get_stokes(data,feedtype)
    data = None

    #Calculate Mueller Matrix variables for each coarse channel
    psis = phase_offsets(I,Q,U,V,tsamp,chan_per_coarse,feedtype,**kwargs)
    G = gain_offsets(I,Q,U,V,tsamp,chan_per_coarse,feedtype,**kwargs)

    #Apply the Mueller matrix to original noise diode data and refold
    I,Q,U,V = apply_Mueller(I,Q,U,V,G,psis,chan_per_coarse,feedtype)
    I_OFF,I_ON = foldcal(I,tsamp,**kwargs)
    Q_OFF,Q_ON = foldcal(Q,tsamp,**kwargs)
    U_OFF,U_ON = foldcal(U,tsamp,**kwargs)
    V_OFF,V_ON = foldcal(V,tsamp,**kwargs)

    #Delete data arrays for space
    I = None
    Q = None
    U = None
    V = None

    #Plot new ON-OFF spectra
    plt.plot(freqs,I_ON-I_OFF,'k-',label='I')
    plt.plot(freqs,Q_ON-Q_OFF,'r-',label='Q')
    plt.plot(freqs,U_ON-U_OFF,'g-',label='U')
    plt.plot(freqs,V_ON-V_OFF,'m-',label='V')

    plt.legend()
    plt.xlabel('Frequency (MHz)')
    plt.title('Calibrated Full Stokes Noise Diode Spectrum')
    plt.ylabel('Power (Counts)')

def plot_phase_offsets(dio_cross,chan_per_coarse=8,feedtype='l',ax1=None,ax2=None,legend=True,**kwargs):
    '''
    Plots the calculated phase offsets of each coarse channel along with
    the UV (or QU) noise diode spectrum for comparison
    '''
    #Get ON-OFF ND spectra
    Idiff,Qdiff,Udiff,Vdiff,freqs = get_diff(dio_cross,feedtype,**kwargs)
    obs = Waterfall(dio_cross,max_load=150)
    tsamp = obs.header['tsamp']
    data = obs.data
    obs = None
    I,Q,U,V = get_stokes(data,feedtype)

    #Get phase offsets and convert to degrees
    coarse_psis = phase_offsets(I,Q,U,V,tsamp,chan_per_coarse,feedtype,**kwargs)
    coarse_freqs = convert_to_coarse(freqs,chan_per_coarse)
    coarse_degs = np.degrees(coarse_psis)

    #Plot phase offsets
    if ax2==None:
        plt.subplot(211)
    else:
        axPsi = plt.axes(ax2)
        plt.setp(axPsi.get_xticklabels(),visible=False)
    plt.plot(coarse_freqs,coarse_degs,'ko',markersize=2,label='Coarse Channel $\psi$')
    plt.ylabel('Degrees')
    plt.grid(True)
    plt.title('Phase Offsets')
    if legend==True:
        plt.legend()

    #Plot U and V spectra
    if ax1==None:
        plt.subplot(212)
    else:
        axUV = plt.axes(ax1)
    plt.plot(freqs,Udiff,'g-',label='U')
    if feedtype=='l':
        plt.plot(freqs,Vdiff,'m-',label='V')
    if feedtype=='c':
        plt.plot(freqs,Qdiff,'r-',label='Q')
    plt.xlabel('Frequency (MHz)')
    plt.ylabel('Power (Counts)')
    plt.grid(True)
    if legend==True:
        plt.legend()

def plot_gain_offsets(dio_cross,dio_chan_per_coarse=8,feedtype='l',ax1=None,ax2=None,legend=True,**kwargs):
    '''
    Plots the calculated gain offsets of each coarse channel along with
    the time averaged power spectra of the X and Y feeds
    '''
    #Get ON-OFF ND spectra
    Idiff,Qdiff,Udiff,Vdiff,freqs = get_diff(dio_cross,feedtype,**kwargs)
    obs = Waterfall(dio_cross,max_load=150)
    tsamp = obs.header['tsamp']
    data = obs.data
    obs = None
    I,Q,U,V = get_stokes(data,feedtype)

    #Get phase offsets and convert to degrees
    coarse_G = gain_offsets(I,Q,U,V,tsamp,dio_chan_per_coarse,feedtype,**kwargs)
    coarse_freqs = convert_to_coarse(freqs,dio_chan_per_coarse)

    #Get X and Y spectra for the noise diode ON and OFF
    #If using circular feeds these correspond to LL and RR
    XX_OFF,XX_ON = foldcal(np.expand_dims(data[:,0,:],axis=1),tsamp,**kwargs)
    YY_OFF,YY_ON = foldcal(np.expand_dims(data[:,1,:],axis=1),tsamp,**kwargs)

    if ax1==None:
        plt.subplot(211)
    else:
        axG = plt.axes(ax1)
        plt.setp(axG.get_xticklabels(),visible=False)
    plt.plot(coarse_freqs,coarse_G,'ko',markersize=2)
    plt.ylabel(r'$\frac{\Delta G}{2}$',rotation=90)
    if feedtype=='l':
        plt.title('XY Gain Difference')
    if feedtype=='c':
        plt.title('LR Gain Difference')
    plt.grid(True)

    if ax2==None:
        plt.subplot(212)
    else:
        axXY = plt.axes(ax2,sharex=axG)
    if feedtype=='l':
        plt.plot(freqs,XX_OFF,'b-',label='XX')
        plt.plot(freqs,YY_OFF,'r-',label='YY')
    if feedtype=='c':
        plt.plot(freqs,XX_OFF,'b-',label='LL')
        plt.plot(freqs,YY_OFF,'r-',label='RR')
    plt.xlabel('Frequency (MHz)')
    plt.ylabel('Power (Counts)')
    if legend==True:
        plt.legend()

def plot_diode_fold(dio_cross,bothfeeds=True,feedtype='l',min_samp=-500,max_samp=7000,legend=True,**kwargs):
    '''
    Plots the calculated average power and time sampling of ON (red) and
    OFF (blue) for a noise diode measurement over the observation time series
    '''
    #Get full stokes data of ND measurement
    obs = Waterfall(dio_cross,max_load=150)
    tsamp = obs.header['tsamp']
    data = obs.data
    obs = None
    I,Q,U,V = get_stokes(data,feedtype)

    #Calculate time series, OFF and ON averages, and time samples for each
    tseriesI = np.squeeze(np.mean(I,axis=2))
    I_OFF,I_ON,OFFints,ONints = foldcal(I,tsamp,inds=True,**kwargs)
    if bothfeeds==True:
        if feedtype=='l':
            tseriesQ = np.squeeze(np.mean(Q,axis=2))
            tseriesX = (tseriesI+tseriesQ)/2
            tseriesY = (tseriesI-tseriesQ)/2
        if feedtype=='c':
            tseriesV = np.squeeze(np.mean(V,axis=2))
            tseriesR = (tseriesI+tseriesV)/2
            tseriesL = (tseriesI-tseriesV)/2

    stop = ONints[-1,1]

    #Plot time series and calculated average for ON and OFF
    if bothfeeds==False:
        plt.plot(tseriesI[0:stop],'k-',label='Total Power')
        for i in ONints:
            plt.plot(np.arange(i[0],i[1]),np.full((i[1]-i[0]),np.mean(I_ON)),'r-')
        for i in OFFints:
            plt.plot(np.arange(i[0],i[1]),np.full((i[1]-i[0]),np.mean(I_OFF)),'b-')
    else:
        if feedtype=='l':
            diff = np.mean(tseriesX)-np.mean(tseriesY)
            plt.plot(tseriesX[0:stop],'b-',label='XX')
            plt.plot(tseriesY[0:stop]+diff,'r-',label='YY (shifted)')
        if feedtype=='c':
            diff = np.mean(tseriesL)-np.mean(tseriesR)
            plt.plot(tseriesL[0:stop],'b-',label='LL')
            plt.plot(tseriesR[0:stop]+diff,'r-',label='RR (shifted)')

    #Calculate plotting limits
    if bothfeeds==False:
        lowlim = np.mean(I_OFF)-(np.mean(I_ON)-np.mean(I_OFF))/2
        hilim = np.mean(I_ON)+(np.mean(I_ON)-np.mean(I_OFF))/2
        plt.ylim((lowlim,hilim))

    plt.xlim((min_samp,max_samp))
    plt.xlabel('Time Sample Number')
    plt.ylabel('Power (Counts)')
    plt.title('Noise Diode Fold')
    if legend==True:
        plt.legend()

def plot_fullcalib(dio_cross,feedtype='l',**kwargs):
    '''
    Generates and shows five plots: Uncalibrated diode, calibrated diode, fold information,
    phase offsets, and gain offsets for a noise diode measurement. Most useful diagnostic plot to
    make sure calibration proceeds correctly.
    '''

    plt.figure("Multiple Calibration Plots", figsize=(12,9))
    left, width = 0.075,0.435
    bottom, height = 0.45,0.5
    width2 = 0.232
    bottom2, height2 = 0.115,0.0975

    rect_uncal = [left,bottom,width,height]
    rect_cal = [left+width+0.025,bottom,width,height]
    rect_fold = [left,bottom2,width2,0.22]
    rect_gain1 = [left+width2+0.1,bottom2,width2,height2]
    rect_phase1 = [left+width2*2+0.1*2,bottom2,width2,height2]
    rect_gain2 = [left+width2+0.1,bottom2+height2+0.025,width2,height2]
    rect_phase2 = [left+width2*2+0.1*2,bottom2+height2+0.025,width2,height2]

    #--------
    axFold = plt.axes(rect_fold)
    print('Plotting Diode Fold')
    plot_diode_fold(dio_cross,bothfeeds=False,feedtype=feedtype,min_samp=2000,max_samp=5500,legend=False,**kwargs)

    #--------
    print('Plotting Gain Offsets')
    plot_gain_offsets(dio_cross,feedtype=feedtype,ax1=rect_gain2,ax2=rect_gain1,legend=False,**kwargs)

    #--------
    print('Plotting Phase Offsets')
    plot_phase_offsets(dio_cross,feedtype=feedtype,ax1=rect_phase1,ax2=rect_phase2,legend=False,**kwargs)
    plt.ylabel('')

    #--------
    ax_uncal = plt.axes(rect_uncal)
    print('Plotting Uncalibrated Diode')
    plot_Stokes_diode(dio_cross,feedtype=feedtype,**kwargs)

    #--------
    ax_cal = plt.axes(rect_cal,sharey=ax_uncal)
    print('Plotting Calibrated Diode')
    plot_calibrated_diode(dio_cross,feedtype=feedtype,**kwargs)
    plt.ylabel('')
    plt.setp(ax_cal.get_yticklabels(),visible=False)

    plt.savefig(dio_cross[:-4]+'.stokescalib.png',dpi=2000)
    plt.show()

def plot_diodespec(ON_obs,OFF_obs,calflux,calfreq,spec_in,units='mJy',**kwargs):
    '''
    Plots the full-band Stokes I spectrum of the noise diode (ON-OFF)
    '''

    dspec = diode_spec(ON_obs,OFF_obs,calflux,calfreq,spec_in,**kwargs)
    obs = Waterfall(ON_obs,max_load=150)
    freqs = obs.populate_freqs()
    chan_per_coarse = obs.header['nchans']/obs.calc_n_coarse_chan()
    coarse_freqs = convert_to_coarse(freqs,chan_per_coarse)
    plt.ion()
    plt.figure()
    plt.plot(coarse_freqs,dspec)
    plt.xlabel('Frequency (MHz)')
    plt.ylabel('Flux Density ('+units+')')
    plt.title('Noise Diode Spectrum')


#!/usr/bin/env python
'''
    Script to dice data to course channel level. From BL FIL of HDF5 files, and outputs HDF5 with '_diced' appended to the file name.

    ..author: Greg Hellbourg (gregory.hellbourg@berkeley.edu)

    March 2018
'''

from .waterfall import Waterfall
import argparse
import math
import sys

#------
# Logging set up
import logging
logger = logging.getLogger(__name__)

level_log = logging.INFO

if level_log == logging.INFO:
    stream = sys.stdout
    format = '%(name)-15s %(levelname)-8s %(message)s'
else:
    stream =  sys.stderr
    format = '%%(relativeCreated)5d (name)-15s %(levelname)-8s %(message)s'

logging.basicConfig(format=format,stream=stream,level = level_log)
#------

def cmd_tool():
    '''Read input and output frequency, and output file name
    '''

    parser = argparse.ArgumentParser(description='Dices hdf5 or fil files and writes to hdf5 or fil.')
    parser.add_argument('-f', '--input_filename', action='store', default=None, dest='in_fname', type=str, help='Name of file to write from (HDF5 or FIL)')
    parser.add_argument('-b', action='store', default=None, dest='f_start', type=float, help='Start frequency in MHz')
    parser.add_argument('-e', action='store', default=None, dest='f_stop', type=float, help='Stop frequency in MHz')
    parser.add_argument('-x', '--output_file', action='store', default=None, dest='out_format', type=str, help='Output file format [.h5 or .fil].')
    parser.add_argument('-o', '--output_filename', action='store', default=None, dest='out_fname', type=str, help='Ouput file name to write (to HDF5 or FIL).')
    parser.add_argument('-l', action='store', default=None, dest='max_load', type=float,help='Maximum data limit to load. Default:1GB')

    args = parser.parse_args()

    if len(sys.argv) == 1:
            logger.error('Indicate file name and start and stop frequencies')
            sys.exit()

    if args.in_fname == None:
            logger.error('Need to indicate input file name')
            sys.exit()

    if args.out_fname == None:
        if (args.out_format == None) or (args.out_format == 'h5'):
            if args.in_fname[len(args.in_fname)-4:] == '.fil':
                args.out_fname = args.in_fname
                args.out_fname = args.out_fname.replace('.fil','_diced.h5')
            elif args.in_fname[len(args.in_fname)-3:] == '.h5':
                args.out_fname = args.in_fname
                args.out_fname = args.out_fname.replace('.h5','_diced.h5')
            else:
                logger.error('Input file not recognized')
                sys.exit()
        elif args.out_format == 'fil':
            if args.in_fname[len(args.in_fname)-4:] == '.fil':
                args.out_fname = args.in_fname
                args.out_fname = args.out_fname.replace('.fil','_diced.fil')
            elif args.in_fname[len(args.in_fname)-3:] == '.h5':
                args.out_fname = args.in_fname
                args.out_fname = args.out_fname.replace('.h5','_diced.fil')
            else:
                logger.error('input file not recognized.')
                sys.exit()

        else:
            logger.error('Must indicate either output file name or valid output file extension.')
            sys.exit()

    elif (args.out_fname[len(args.out_fname)-4:] == '.fil') and (args.out_format == 'h5'):
        logger.error('Output file extension does not match output file name')
        sys.exit()

    elif (args.out_fname[len(args.out_fname)-3:] == '.h5') and (args.out_format == 'fil'):
        logger.error('Output file extension does not match output file name.')
        sys.exit()

    if (args.out_fname[len(args.out_fname)-3:] != '.h5') and (args.out_fname[len(args.out_fname)-4:] != '.fil'):
        logger.error('Indicate output file name with extension, or simply output file extension.')
        sys.exit()

    if args.f_start == None and args.f_stop == None:
        logger.error('Please give either start and/or end frequencies. Otherwise use fil2h5 or h52fil functions.')
        sys.exit()

    if args.f_start == None:
        logger.warning('Lower frequency not given, setting to ' + str(f_min_file) + ' MHz to match file.')

    if args.f_stop == None:
        logger.warning('Higher frequency not given, setting to ' + str(f_max_file) + ' MHz to match file.')

    #Read start frequency and bandwidth from data set
    file_big = Waterfall(args.in_fname, max_load = args.max_load)
    f_min_file = file_big.header['fch1']
    f_max_file = file_big.header['fch1'] + file_big.header['nchans']*file_big.header['foff']

    if f_max_file < f_min_file:
            f_max_file,f_min_file = f_min_file,f_max_file

    FreqBWFile = f_max_file-f_min_file
    stdDF = FreqBWFile / float(file_big.calc_n_coarse_chan())   #stdDF = 2.9296875

    if args.f_stop < args.f_start:
            args.f_stop,args.f_start = args.f_start,args.f_stop

    if args.f_start < f_max_file and args.f_start > f_min_file and args.f_stop > f_max_file:
            args.f_stop = f_max_file
            logger.warning('Higher frequency set to ' + str(f_max_file) + ' MHz to match file.')

    if args.f_stop < f_max_file and args.f_stop > f_min_file and args.f_start < f_min_file:
            args.f_start = f_min_file
            logger.warning('Lower frequency set to ' + str(f_min_file) + ' MHz to match file.')

    if args.f_start < f_min_file and args.f_stop > f_max_file:
            args.f_start = f_min_file
            args.f_stop = f_max_file
            logger.warning('Lower frequency set to ' + str(f_min_file) + ' MHz and higher frequency set to ' + str(f_max_file) + ' MHz to match file.')
            # print '\nindicated frequencies include file frequency span - no need to dice\n'
            # sys.exit()

    if min(args.f_start,args.f_stop) < f_min_file or max(args.f_start,args.f_stop) > f_max_file:
            logger.error('Bandwidth to extract must be within ' + str(f_min_file) + ' MHz and ' + str(f_max_file) + ' MHz.')
            sys.exit()

    # calculate real coarse channel begin and end freqs
    f_start_real = math.floor((min(args.f_start,args.f_stop) - f_min_file)/stdDF)*stdDF + f_min_file
    f_stop_real = f_max_file - math.floor((f_max_file - max(args.f_start,args.f_stop))/stdDF)*stdDF

    # print
    # print "true start frequency is " + str(f_start_real)
    # print "true stop frequency is " + str(f_stop_real)

    logger.info('Writing to ' + args.out_fname)
    logger.info('Extacting from ' + str(f_start_real) + ' MHz to ' + str(f_stop_real) + ' MHz.')

    # create waterfall object
    file_small = Waterfall(args.in_fname, f_start = f_start_real, f_stop = f_stop_real, max_load = args.max_load)

    # write waterfall object
    if args.out_fname[len(args.out_fname)-4:] == '.fil':
        file_small.write_to_fil(args.out_fname)
    elif args.out_fname[len(args.out_fname)-3:] == '.h5':
        file_small.write_to_hdf5(args.out_fname)
    else:
        logger.error('Error in output file creation : verify output file name and extension.')
        sys.exit()

if __name__ == "__main__":

    cmd_tool()


#!/usr/bin/env python

from .filterbank import Filterbank
import h5py
try:
    HAS_BITSHUFFLE = True
    import bitshuffle.h5
except ImportError:
    HAS_BITSHUFFLE = False
import time
import os
import glob

import numpy as np

MAX_SIZE = 4e9

def cmd_tool(args=None):
    """ Command line utility for creating HDF5 blimpy files. """
    from argparse import ArgumentParser
    parser = ArgumentParser(description="Command line utility for creating HDF5 Filterbank files.")
    parser.add_argument('dirname', type=str, help='Name of directory to read')
    args = parser.parse_args()

    if not HAS_BITSHUFFLE:
        print("Error: the bitshuffle library is required to run this script.")
        exit()

    filelist = glob.glob(os.path.join(args.dirname, '*.fil'))

    for filename in filelist:
        if not os.path.exists(filename + '.h5'):
            t0 = time.time()
            print("\nReading %s header..." % filename)
            fb = Filterbank(filename, load_data=False)

            data_shape = (fb.n_ints_in_file, fb.header['nifs'], fb.header['nchans'])
            data_dtype = fb.data.dtype
            print(data_dtype)

            print("Creating new dataset, %s" % str(data_shape))
            block_size = 0
            h5 = h5py.File(filename + '.h5', 'w')

            h5.attrs['CLASS'] = 'FILTERBANK'

            dset = h5.create_dataset('data',
                              shape=data_shape,
                              compression=bitshuffle.h5.H5FILTER,
                              compression_opts=(block_size, bitshuffle.h5.H5_COMPRESS_LZ4),
                              dtype=data_dtype)

            dset_mask = h5.create_dataset('mask',
                                     shape=data_shape,
                                     compression=bitshuffle.h5.H5FILTER,
                                     compression_opts=(block_size, bitshuffle.h5.H5_COMPRESS_LZ4),
                                     dtype='uint8')

            dset.dims[0].label = "frequency"
            dset.dims[1].label = "feed_id"
            dset.dims[2].label = "time"

            dset_mask.dims[0].label = "frequency"
            dset_mask.dims[1].label = "feed_id"
            dset_mask.dims[2].label = "time"

            # Copy over header information as attributes
            for key, value in fb.header.items():
                dset.attrs[key] = value

            filesize = os.path.getsize(filename)

            if filesize >= MAX_SIZE:
                n_int_per_read = int(filesize / MAX_SIZE / 2)
                print("Filling in with data over %i reads..." % n_int_per_read)
                for ii in range(0, n_int_per_read):
                    print("Reading %i of %i" % (ii + 1, n_int_per_read))
                    #print  ii*n_int_per_read, (ii+1)*n_int_per_read
                    fb = Filterbank(filename, t_start=ii*n_int_per_read, t_stop=(ii+1) * n_int_per_read)
                    dset[ii*n_int_per_read:(ii+1)*n_int_per_read] = fb.data[:]
            else:
                fb = Filterbank(filename)
                print(dset.shape, " -> ", fb.data.shape)
                dset[:] = fb.data[:]

            h5.close()

            t1 = time.time()
            print("Conversion time: %2.2fs" % (t1- t0))

if __name__ == "__main__":
    cmd_tool()