#!/usr/bin/env python
""" This modele handles file types.
"""

import os
import sys
import numpy as np
import h5py
import six

from astropy.coordinates import Angle

try:
    from . import sigproc
except:
    import sigproc

# import pdb;# pdb.set_trace()

import logging
logger = logging.getLogger(__name__)

level_log = logging.INFO

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

logging.basicConfig(format=lformat, stream=stream, level=level_log)

# Max size of data array to load into memory (1GB in bytes)
MAX_DATA_ARRAY_SIZE_UNIT = 1024 * 1024 * 1024.0

class Reader(object):
    """ Basic reader object """

    def __init__(self):

        self.t_begin = 0
        self.t_end   = 0

    def _setup_selection_range(self, f_start=None, f_stop=None, t_start=None, t_stop=None, init=False):
        """Making sure the selection if time and frequency are within the file limits.

        Args:
            init (bool): If call during __init__
        """

        # This avoids resetting values
        if init is True:
            if t_start is None:
                t_start = self.t_begin
            if t_stop is None:
                t_stop = self.t_end
            if f_start is None:
                f_start = self.f_begin
            if f_stop is None:
                f_stop = self.f_end
        else:
            if f_start is None:
                f_start = self.f_start
            if f_stop is None:
                f_stop = self.f_stop
            if t_start is None:
                t_start = self.t_start
            if t_stop is None:
                t_stop = self.t_stop

        # By now, all values start/stop are populated.

        if t_stop >= 0 and t_start >= 0 and t_stop < t_start:
            t_stop, t_start = t_start,t_stop
            logger.warning('Given t_stop < t_start, assuming reversed values.')
        if f_stop and f_start and f_stop < f_start:
            f_stop, f_start = f_start,f_stop
            logger.warning('Given f_stop < f_start, assuming reversed values.')

        if t_start >= self.t_begin and t_start < self.t_end:
            self.t_start = int(t_start)
        else:
            if init is False or t_start != None:
                logger.warning('Setting t_start = %f, since t_start not given or not valid.'%self.t_begin)
            self.t_start = self.t_begin

        if t_stop <= self.t_end  and t_stop > self.t_begin:
            self.t_stop = int(t_stop)
        else:
            if init is False or t_stop:
                logger.warning('Setting t_stop = %f, since t_stop not given or not valid.'%self.t_end)
            self.t_stop = self.t_end

        if f_start >= self.f_begin and f_start < self.f_end:
            self.f_start = f_start
        else:
            if init is False or f_start:
                logger.warning('Setting f_start = %f, since f_start not given or not valid.'%self.f_begin)
            self.f_start = self.f_begin

        if f_stop <= self.f_end and f_stop > self.f_begin:
            self.f_stop = f_stop
        else:
            if init is False or f_stop:
                logger.warning('Setting f_stop = %f, since f_stop not given or not valid.'%self.f_end)
            self.f_stop = self.f_end

        # Now we have setup bounds, we can calculate shape of selection
        self.selection_shape = self._calc_selection_shape()

    def _init_empty_selection(self):
        """
        """

        self.data = np.array([0],dtype=self._d_type)

    def _setup_dtype(self):
        """Calculating dtype
        """

        #Set up the data type
        if self._n_bytes  == 4:
            return b'float32'
        elif self._n_bytes  == 2:
            return b'uint16'
        elif self._n_bytes  == 1:
            return b'uint8'
        else:
            logger.warning('Having trouble setting dtype, assuming float32.')
            return b'float32'

    def _calc_selection_size(self):
        """Calculate size of data of interest.
        """

        #Check to see how many integrations requested
        n_ints = self.t_stop - self.t_start
        #Check to see how many frequency channels requested
        n_chan = (self.f_stop - self.f_start) / abs(self.header[b'foff'])

        n_bytes  = self._n_bytes
        selection_size = int(n_ints*n_chan*n_bytes)

        return selection_size

    def _calc_selection_shape(self):
        """Calculate shape of data of interest.
        """

        #Check how many integrations requested
        n_ints = int(self.t_stop - self.t_start)
        #Check how many frequency channels requested
        n_chan = int(np.round((self.f_stop - self.f_start) / abs(self.header[b'foff'])))

        selection_shape = (n_ints,int(self.header[b'nifs']),n_chan)

        return selection_shape

    def _setup_chans(self):
        """Setup channel borders
        """

        if self.header[b'foff'] < 0:
            f0 = self.f_end
        else:
            f0 = self.f_begin

        i_start, i_stop = 0, self.n_channels_in_file
        if self.f_start:
            i_start = np.round((self.f_start - f0) / self.header[b'foff'])
        if self.f_stop:
            i_stop  = np.round((self.f_stop - f0)  / self.header[b'foff'])

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

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

        self.chan_start_idx =  chan_start_idx
        self.chan_stop_idx = chan_stop_idx

    def _setup_freqs(self):
        """Updating frequency borders from channel values
        """

        if self.header[b'foff'] > 0:
            self.f_start = self.f_begin + self.chan_start_idx*abs(self.header[b'foff'])
            self.f_stop = self.f_begin + self.chan_stop_idx*abs(self.header[b'foff'])
        else:
            self.f_start = self.f_end - self.chan_stop_idx*abs(self.header[b'foff'])
            self.f_stop = self.f_end - self.chan_start_idx*abs(self.header[b'foff'])

    def populate_timestamps(self,update_header=False):
        """  Populate time axis.
            IF update_header then only return tstart
        """

        #Check to see how many integrations requested
        ii_start, ii_stop = 0, self.n_ints_in_file
        if self.t_start:
            ii_start = self.t_start
        if self.t_stop:
            ii_stop = self.t_stop

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

        if update_header:
            timestamps = ii_start * t_delt / 24./60./60. + t0
        else:
            timestamps = np.arange(ii_start, ii_stop) * t_delt / 24./60./60. + t0

        return timestamps

    def populate_freqs(self):
        """
         Populate frequency axis
        """

        if self.header[b'foff'] < 0:
            f0 = self.f_end
        else:
            f0 = self.f_begin

        self._setup_chans()

        #create freq array
        i_vals = np.arange(self.chan_start_idx, self.chan_stop_idx)
        freqs = self.header[b'foff'] * i_vals + f0

        return freqs

    def calc_n_coarse_chan(self, chan_bw=None):
        """ This makes an attempt to calculate the number of coarse channels in a given file.

            Note:
                This is unlikely to work on non-Breakthrough Listen data, as a-priori knowledge of
                the digitizer system is required.

        """
        nchans = int(self.header[b'nchans'])
        # Do we have a file with enough channels that it has coarse channelization?
        if chan_bw is not None:
            bandwidth = abs(self.f_stop - self.f_start)
            n_coarse_chan = int(bandwidth / chan_bw)
            return n_coarse_chan
        elif nchans >= 2**20:
            # Does the common FFT length of 2^20 divide through without a remainder?
            # This should work for most GBT and all Parkes hires data
            if nchans % 2**20 == 0:
                n_coarse_chan = nchans // 2**20
                return n_coarse_chan
            # Early GBT data has non-2^N FFT length, check if it is GBT data
            elif self.header[b'telescope_id'] == 6:
                coarse_chan_bw = 2.9296875
                bandwidth = abs(self.f_stop - self.f_start)
                n_coarse_chan = int(bandwidth / coarse_chan_bw)
                return n_coarse_chan
            else:
                logger.warning("Couldn't figure out n_coarse_chan")
        elif self.header[b'telescope_id'] == 6 and nchans < 2**20:
            #For GBT non-hires data
            coarse_chan_bw = 2.9296875
            bandwidth = abs(self.f_stop - self.f_start)
            n_coarse_chan = int(bandwidth / coarse_chan_bw)
            return n_coarse_chan
        else:
            logger.warning("This function currently only works for hires BL Parkes or GBT data.")

    def calc_n_blobs(self, blob_dim):
        """ Given the blob dimensions, calculate how many fit in the data selection.
        """

        n_blobs = int(np.ceil(1.0 * np.prod(self.selection_shape) / np.prod(blob_dim)))

        return n_blobs

    def isheavy(self):
        """ Check if the current selection is too large.
        """

        selection_size_bytes = self._calc_selection_size()

        if selection_size_bytes > self.MAX_DATA_ARRAY_SIZE:
            return True
        else:
            return False

class H5Reader(Reader):
    """ This class handles .h5 files.
    """

    def __init__(self, filename, f_start=None, f_stop=None, t_start=None, t_stop=None, load_data=True, max_load=1.):
        """ Constructor.

        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 time bin
            t_stop (int): stop time bin
        """
        super(H5Reader, self).__init__()

        if filename and os.path.isfile(filename) and h5py.is_hdf5(filename):

            #These values may be modified once code for multi_beam and multi_stokes observations are possible.
            self.freq_axis = 2
            self.time_axis = 0
            self.beam_axis = 1  # Place holder
            self.stokes_axis = 4  # Place holder

            self.filename = filename
            self.filestat = os.stat(filename)
            self.filesize = self.filestat.st_size/(1024.0**2)
            self.load_data = load_data
            self.h5 = h5py.File(self.filename)
            self.read_header()
            self.file_size_bytes = os.path.getsize(self.filename)  # In bytes
            self.n_ints_in_file = self.h5["data"].shape[self.time_axis] #
            self.n_channels_in_file = self.h5["data"].shape[self.freq_axis] #
            self.n_beams_in_file = self.header[b'nifs'] #Placeholder for future development.
            self.n_pols_in_file = 1 #Placeholder for future development.
            self._n_bytes = int(self.header[b'nbits'] / 8)  #number of bytes per digit.
            self._d_type = self._setup_dtype()
            self.file_shape = (self.n_ints_in_file,self.n_beams_in_file,self.n_channels_in_file)

            if self.header[b'foff'] < 0:
                self.f_end = self.header[b'fch1']
                self.f_begin = self.f_end + self.n_channels_in_file*self.header[b'foff']
            else:
                self.f_begin = self.header[b'fch1']
                self.f_end = self.f_begin + self.n_channels_in_file*self.header[b'foff']

            self.t_begin = 0
            self.t_end = self.n_ints_in_file

            #Taking care all the frequencies are assigned correctly.
            self._setup_selection_range(f_start=f_start, f_stop=f_stop, t_start=t_start, t_stop=t_stop, init=True)
            #Convert input frequencies into what their corresponding channel number would be.
            self._setup_chans()
            #Update frequencies ranges from channel number.
            self._setup_freqs()

            #Applying data size limit to load.
            if max_load is not None:
                if max_load > 1.0:
                    logger.warning('Setting data limit > 1GB, please handle with care!')
                self.MAX_DATA_ARRAY_SIZE = max_load * MAX_DATA_ARRAY_SIZE_UNIT
            else:
                self.MAX_DATA_ARRAY_SIZE = MAX_DATA_ARRAY_SIZE_UNIT

            if self.file_size_bytes > self.MAX_DATA_ARRAY_SIZE:
                self.large_file = True
            else:
                self.large_file = False

            if self.load_data:
                if self.large_file:
                    #Only checking the selection, if the file is too large.
                    if self.f_start or self.f_stop or self.t_start or self.t_stop:
                        if self.isheavy():
                            logger.warning("Selection size of %.2f GB, exceeding our size limit %.2f GB. Instance created, header loaded, but data not loaded, please try another (t,v) selection." % (self._calc_selection_size() / (1024. ** 3), self.MAX_DATA_ARRAY_SIZE / (1024. ** 3)))
                            self._init_empty_selection()
                        else:
                            self.read_data()
                    else:
                        logger.warning("The file is of size %.2f GB, exceeding our size limit %.2f GB. Instance created, header loaded, but data not loaded. You could try another (t,v) selection."%(self.file_size_bytes/(1024.**3), self.MAX_DATA_ARRAY_SIZE/(1024.**3)))
                        self._init_empty_selection()
                else:
                    self.read_data()
            else:
                logger.info("Skipping loading data ...")
                self._init_empty_selection()
        else:
            raise IOError("Need a file to open, please give me one!")

    def read_header(self):
        """ Read header and return a Python dictionary of key:value pairs
        """

        self.header = {}

        for key, val in self.h5['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

        return self.header

    def _find_blob_start(self, blob_dim, n_blob):
        """Find first blob from selection.
        """

        #Convert input frequencies into what their corresponding channel number would be.
        self._setup_chans()

        #Check which is the blob time offset
        blob_time_start = self.t_start + blob_dim[self.time_axis]*n_blob

        #Check which is the blob frequency offset (in channels)
        blob_freq_start = self.chan_start_idx + (blob_dim[self.freq_axis]*n_blob)%self.selection_shape[self.freq_axis]

        blob_start = np.array([blob_time_start, 0, blob_freq_start])

        return blob_start

    def read_data(self, f_start=None, f_stop=None,t_start=None, t_stop=None):
        """ Read data
        """

        self._setup_selection_range(f_start=f_start, f_stop=f_stop, t_start=t_start, t_stop=t_stop)

        #check if selection is small enough.
        if self.isheavy():
            logger.warning("Selection size of %.2f GB, exceeding our size limit %.2f GB. Instance created, header loaded, but data not loaded, please try another (t,v) selection." % (self._calc_selection_size() / (1024. ** 3), self.MAX_DATA_ARRAY_SIZE / (1024. ** 3)))
            self.data = np.array([0],dtype=self._d_type)
            return None

        #Convert input frequencies into what their corresponding channel number would be.
        self._setup_chans()
        #Update frequencies ranges from channel number.
        self._setup_freqs()

        self.data = self.h5["data"][self.t_start:self.t_stop,:,self.chan_start_idx:self.chan_stop_idx]

    def read_blob(self,blob_dim,n_blob=0):
        """Read blob from a selection.
        """

        n_blobs = self.calc_n_blobs(blob_dim)
        if n_blob > n_blobs or n_blob < 0:
            raise ValueError('Please provide correct n_blob value. Given %i, but max values is %i'%(n_blob,n_blobs))

        #This prevents issues when the last blob is smaller than the others in time
        if blob_dim[self.time_axis]*(n_blob+1) > self.selection_shape[self.time_axis]:
            updated_blob_dim = (self.selection_shape[self.time_axis] - blob_dim[self.time_axis]*n_blob, 1, blob_dim[self.freq_axis])
        else:
            updated_blob_dim = [int(i) for i in blob_dim]

        blob_start = self._find_blob_start(blob_dim, n_blob)
        blob_end = blob_start + np.array(updated_blob_dim)

        blob = self.h5["data"][int(blob_start[self.time_axis]):int(blob_end[self.time_axis]),
                               :,
                               int(blob_start[self.freq_axis]):int(blob_end[self.freq_axis])
                               ]

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

        return blob

class FilReader(Reader):
    """ This class handles .fil files.
    """

    def __init__(self, filename,f_start=None, f_stop=None,t_start=None, t_stop=None, load_data=True, max_load=1.):
        """ Constructor.

        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 time bin
            t_stop (int): stop time bin
        """
        super(FilReader, self).__init__()

        self.header_keywords_types = sigproc.header_keyword_types

        if filename and os.path.isfile(filename):
            self.filename = filename
            self.load_data = load_data
            self.header = self.read_header()
            self.file_size_bytes = os.path.getsize(self.filename)
            self.idx_data = sigproc.len_header(self.filename)
            self.n_channels_in_file  = self.header[b'nchans']
            self.n_beams_in_file = self.header[b'nifs'] #Placeholder for future development.
            self.n_pols_in_file = 1 #Placeholder for future development.
            self._n_bytes = int(self.header[b'nbits'] / 8)  #number of bytes per digit.
            self._d_type = self._setup_dtype()
            self._setup_n_ints_in_file()
            self.file_shape = (self.n_ints_in_file,self.n_beams_in_file,self.n_channels_in_file)

            if self.header[b'foff'] < 0:
                self.f_end  = self.header[b'fch1']
                self.f_begin  = self.f_end + self.n_channels_in_file*self.header[b'foff']
            else:
                self.f_begin  = self.header[b'fch1']
                self.f_end  = self.f_begin + self.n_channels_in_file*self.header[b'foff']

            self.t_begin = 0
            self.t_end = self.n_ints_in_file

            #Taking care all the frequencies are assigned correctly.
            self._setup_selection_range(f_start=f_start, f_stop=f_stop, t_start=t_start, t_stop=t_stop, init=True)
            #Convert input frequencies into what their corresponding channel number would be.
            self._setup_chans()
            #Update frequencies ranges from channel number.
            self._setup_freqs()

            self.freq_axis = 2
            self.time_axis = 0
            self.beam_axis = 1  # Place holder

#EE ie.
#           spec = np.squeeze(fil_file.data)
            # set start of data, at real length of header  (future development.)
#            self.datastart=self.hdrraw.find('HEADER_END')+len('HEADER_END')+self.startsample*self.channels

            #Applying data size limit to load.
            if max_load is not None:
                if max_load > 1.0:
                    logger.warning('Setting data limit != 1GB, please handle with care!')
                self.MAX_DATA_ARRAY_SIZE = max_load * MAX_DATA_ARRAY_SIZE_UNIT
            else:
                self.MAX_DATA_ARRAY_SIZE = MAX_DATA_ARRAY_SIZE_UNIT

            if self.file_size_bytes > self.MAX_DATA_ARRAY_SIZE:
                self.large_file = True
            else:
                self.large_file = False

            if self.load_data:
                if self.large_file:
                    if self.f_start or self.f_stop or self.t_start or self.t_stop:
                        if self.isheavy():
                            logger.warning("Selection size of %.2f GB, exceeding our size limit %.2f GB. Instance created, header loaded, but data not loaded, please try another (t,v) selection." % (self._calc_selection_size() / (1024. ** 3), self.MAX_DATA_ARRAY_SIZE / (1024. ** 3)))
                            self._init_empty_selection()
                        else:
                            self.read_data()
                    else:
                        logger.warning("The file is of size %.2f GB, exceeding our size limit %.2f GB. Instance created, header loaded, but data not loaded. You could try another (t,v) selection."%(self.file_size_bytes/(1024.**3), self.MAX_DATA_ARRAY_SIZE/(1024.**3)))
                        self._init_empty_selection()
                else:
                    self.read_data()
            else:
                logger.info("Skipping loading data ...")
                self._init_empty_selection()
        else:
            raise IOError("Need a file to open, please give me one!")

    def _setup_n_ints_in_file(self):
        """ Calculate the number of integrations in the file. """
        self.n_ints_in_file = sigproc.calc_n_ints_in_file(self.filename)

    def read_header(self, return_idxs=False):
        """ Read blimpy header and return a Python dictionary of key:value pairs

        Args:
            filename (str): name of file to open

        Optional args:
            return_idxs (bool): Default False. If true, returns the file offset indexes
                                for values

        Returns:
            Python dict of key:value pairs, OR returns file offset indexes for values.

        """
        self.header = sigproc.read_header(self.filename, return_idxs=return_idxs)
        return self.header

    def read_data(self, f_start=None, f_stop=None,t_start=None, t_stop=None):
        """ Read data.
        """

        self._setup_selection_range(f_start=f_start, f_stop=f_stop, t_start=t_start, t_stop=t_stop)

        #check if selection is small enough.
        if self.isheavy():
            logger.warning("Selection size of %.2f GB, exceeding our size limit %.2f GB. Instance created, "
                           "header loaded, but data not loaded, please try another (t,v) selection." % (self._calc_selection_size() / (1024. ** 3), self.MAX_DATA_ARRAY_SIZE / (1024. ** 3)))
            self.data = np.array([0],dtype=self._d_type)
            return None

        #Convert input frequencies into what their corresponding channel number would be.
        self._setup_chans()
        #Update frequencies ranges from channel number.
        self._setup_freqs()

        n_chans = self.header[b'nchans']
        n_chans_selected = self.selection_shape[self.freq_axis]
        n_ifs   = self.header[b'nifs']

        # Load binary data
        f = open(self.filename, 'rb')
        f.seek(int(self.idx_data))

        # now check to see how many integrations requested
        n_ints = self.t_stop - self.t_start

        # Seek to first integration
        f.seek(int(self.t_start * self._n_bytes  * n_ifs * n_chans), 1)

        #Loading  data
        self.data = np.zeros((n_ints, n_ifs, n_chans_selected), dtype=self._d_type)

        for ii in range(n_ints):
            for jj in range(n_ifs):
                f.seek(int(self._n_bytes  * self.chan_start_idx), 1) # 1 = from current location
                dd = np.fromfile(f, count=n_chans_selected, dtype=self._d_type)

                # Reverse array if frequency axis is flipped
#                     if self.header[b'foff'] < 0:
#                         dd = dd[::-1]

                self.data[ii, jj] = dd

                f.seek(int(self._n_bytes  * (n_chans - self.chan_stop_idx)), 1)  # Seek to start of next block

    def _find_blob_start(self):
        """Find first blob from selection.
        """

        # Convert input frequencies into what their corresponding channel number would be.
        self._setup_chans()

        # Check which is the blob time offset
        blob_time_start = self.t_start

        # Check which is the blob frequency offset (in channels)
        blob_freq_start = self.chan_start_idx

        blob_start = blob_time_start * self.n_channels_in_file + blob_freq_start

        return blob_start

    def read_blob(self,blob_dim,n_blob=0):
        """Read blob from a selection.
        """

        n_blobs = self.calc_n_blobs(blob_dim)
        if n_blob > n_blobs or n_blob < 0:
            raise ValueError('Please provide correct n_blob value. Given %i, but max values is %i'%(n_blob,n_blobs))

        # This prevents issues when the last blob is smaller than the others in time.
        if blob_dim[self.time_axis]*(n_blob+1) > self.selection_shape[self.time_axis]:
            updated_blob_dim = (int(self.selection_shape[self.time_axis] - blob_dim[self.time_axis]*n_blob), 1, int(blob_dim[self.freq_axis]))
        else:
            updated_blob_dim = [int(i) for i in blob_dim]

        blob_start = self._find_blob_start()
        blob = np.zeros(updated_blob_dim, dtype=self._d_type)

        # EE: For now; also assuming one polarization and one beam.

        # Assuming the blob will loop over the whole frequency range.
        if self.f_start == self.f_begin and self.f_stop == self.f_end:

            blob_flat_size = np.prod(blob_dim)
            updated_blob_flat_size = np.prod(updated_blob_dim)

            # Load binary data
            with open(self.filename, 'rb') as f:
                f.seek(int(self.idx_data + self._n_bytes  * (blob_start + n_blob*blob_flat_size)))
                dd = np.fromfile(f, count=updated_blob_flat_size, dtype=self._d_type)

            if dd.shape[0] == updated_blob_flat_size:
                blob = dd.reshape(updated_blob_dim)
            else:
                logger.info('DD shape != blob shape.')
                blob = dd.reshape((int(dd.shape[0]/blob_dim[self.freq_axis]),blob_dim[self.beam_axis],blob_dim[self.freq_axis]))
        else:

            for blobt in range(updated_blob_dim[self.time_axis]):

                #Load binary data
                with open(self.filename, 'rb') as f:
                    f.seek(int(self.idx_data + self._n_bytes * (blob_start + n_blob*blob_dim[self.time_axis]*self.n_channels_in_file + blobt*self.n_channels_in_file)))
                    dd = np.fromfile(f, count=blob_dim[self.freq_axis], dtype=self._d_type)

                blob[blobt] = dd

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

        return blob

    def read_all(self,reverse=True):
        """ read all the data.
            If reverse=True the x axis is flipped.
        """
        raise NotImplementedError('To be implemented')

        # go to start of the data
        self.filfile.seek(int(self.datastart))
        # read data into 2-D numpy array
#        data=np.fromfile(self.filfile,dtype=self.dtype).reshape(self.channels,self.blocksize,order='F')
        data=np.fromfile(self.filfile,dtype=self.dtype).reshape(self.blocksize, self.channels)
        if reverse:
            data = data[:,::-1]
        return data

    def read_row(self,rownumber,reverse=True):
        """ Read a block of data. The number of samples per row is set in self.channels
            If reverse=True the x axis is flipped.
        """
        raise NotImplementedError('To be implemented')

        # go to start of the row
        self.filfile.seek(int(self.datastart+self.channels*rownumber*(int(self.nbits/8))))
        # read data into 2-D numpy array
        data=np.fromfile(self.filfile,count=self.channels,dtype=self.dtype).reshape(1, self.channels)
        if reverse:
            data = data[:,::-1]
        return data

    def read_rows(self,rownumber,n_rows,reverse=True):
        """ Read a block of data. The number of samples per row is set in self.channels
            If reverse=True the x axis is flipped.
        """
        raise NotImplementedError('To be implemented')

        # go to start of the row
        self.filfile.seek(int(self.datastart+self.channels*rownumber*(int(self.nbits/8))))
        # read data into 2-D numpy array
        data=np.fromfile(self.filfile,count=self.channels*n_rows,dtype=self.dtype).reshape(n_rows, self.channels)
        if reverse:
            data = data[:,::-1]
        return data

def open_file(filename, f_start=None, f_stop=None,t_start=None, t_stop=None,load_data=True,max_load=1.):
    """Open a HDF5 or filterbank file

    Returns instance of a Reader to read data from file.

    ================== ==================================================
    Filename extension File type
    ================== ==================================================
    h5, hdf5           HDF5 format
    fil                fil format
    *other*            Will raise NotImplementedError
    ================== ==================================================

    """
    if not os.path.isfile(filename):
        type(filename)
        print(filename)
        raise IOError("No such file or directory: " + filename)

    filename = os.path.expandvars(os.path.expanduser(filename))
    # Get file extension to determine type
    ext = filename.split(".")[-1].strip().lower()

    if six.PY3:
        ext = bytes(ext, 'ascii')

    if h5py.is_hdf5(filename):
        # Open HDF5 file
        return H5Reader(filename, f_start=f_start, f_stop=f_stop, t_start=t_start, t_stop=t_stop,
                        load_data=load_data, max_load=max_load)
    elif sigproc.is_filterbank(filename):
        # Open FIL file
        return FilReader(filename, f_start=f_start, f_stop=f_stop, t_start=t_start, t_stop=t_stop, load_data=load_data, max_load=max_load)
    else:
        raise NotImplementedError('Cannot open this type of file with Waterfall')


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

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

This provides a class, Waterfall(), which can be used to read a blimpy file (.fil or .h5):

    fil = Waterfall("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 os
import sys
import time

from blimpy.filterbank import Filterbank
from blimpy import file_wrapper as fw
from blimpy.sigproc import *

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

try:
    HAS_BITSHUFFLE = True
    import bitshuffle.h5
except ImportError:
    try:
        import hdf5plugin
        logger.warning('Could not import bitshuffle, using hdf5plugin failover. Write access disabled.')
    except ImportError:
        pass
    HAS_BITSHUFFLE = False
    pass

try:
    import h5py
    HAS_HDF5 = True
except ImportError:
    HAS_HDF5 = 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

###
# 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_HEADER_BLOCKS   = 100                    # Max size of header (in 512-byte blocks)
MAX_BLOB_MB         = 1024                   # Max size of blob in MB

###
# Main blimpy class
###

class Waterfall(Filterbank):
    """ Class for loading and writing blimpy data (.fil, .h5) """

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

        This class parses the blimpy file and stores the header and data
        as objects:
            fb = Waterfall('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.
            max_load (float): maximum data to load in GB. Default: 1GB.
                              e.g. 0.1 is 100 MB
            header_dict (dict): Create blimpy from header dictionary + data array
            data_array (np.array): Create blimpy from header dict + data array
        """

##EE        super(Waterfall, self).__init__()

        if filename:
            self.filename = filename
            self.ext = os.path.splitext(filename)[-1].lower()
            self.container = fw.open_file(filename, f_start=f_start, f_stop=f_stop, t_start=t_start, t_stop=t_stop,
                                          load_data=load_data, max_load=max_load)
            self.file_header = self.container.header
            self.header = self.file_header
            self.n_ints_in_file = self.container.n_ints_in_file
            self.file_shape = self.container.file_shape
            self.file_size_bytes = self.container.file_size_bytes
            self.selection_shape = self.container.selection_shape
            self.n_channels_in_file = self.container.n_channels_in_file

            # These values will be modified once code for multi_beam and multi_stokes observations are possible.
            self.freq_axis = 2
            self.time_axis = 0
            self.beam_axis = 1  # Place holder
            self.stokes_axis = 4  # Place holder

            self.__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:
            self.filename = b''

    def __load_data(self):
        """ Helper for loading data form a container. Should not be called manually. """

        self.data = self.container.data

    def read_data(self, f_start=None, f_stop=None,t_start=None, t_stop=None):
        """ Reads data selection if small enough.
        """

        self.container.read_data(f_start=f_start, f_stop=f_stop,t_start=t_start, t_stop=t_stop)

        self.__load_data()

    def __update_header(self):
        """ Updates the header information from the original file to the selection.
        """

        #Updating frequency of first channel from selection
        if self.header[b'foff'] < 0:
            self.header[b'fch1'] = self.container.f_stop
        else:
            self.header[b'fch1'] = self.container.f_start

        #Updating number of coarse channels.
        self.header[b'nchans'] = self.container.selection_shape[self.freq_axis]

        #Updating time stamp for first time bin from selection
        self.header[b'tstart'] = self.container.populate_timestamps(update_header=True)

    def populate_freqs(self):
        """
        """

        return self.container.populate_freqs()

    def populate_timestamps(self):
        """
        """

        return self.container.populate_timestamps()

    def info(self):
        """ Print header information and other derived information. """

        print("\n--- File Info ---")

        for key, val in self.file_header.items():
            if key == 'src_raj':
                val = val.to_string(unit=u.hour, sep=':')
            if key == 'src_dej':
                val = val.to_string(unit=u.deg, sep=':')
            print("%16s : %32s" % (key, val))

        print("\n%16s : %32s" % ("Num ints in file", self.n_ints_in_file))
        print("%16s : %32s" % ("File shape", self.file_shape))
        print("--- Selection Info ---")
        print("%16s : %32s" % ("Data selection shape", self.selection_shape))
        print("%16s : %32s" % ("Minimum freq (MHz)", self.container.f_start))
        print("%16s : %32s" % ("Maximum freq (MHz)", self.container.f_stop))

    def write_to_fil(self, filename_out, *args, **kwargs):
        """ Write data to .fil file.
            It check the file size then decides how to write the file.

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

        #For timing how long it takes to write a file.
        t0 = time.time()

        #Update header
        self.__update_header()

        if self.container.isheavy():
            self.__write_to_fil_heavy(filename_out)
        else:
            self.__write_to_fil_light(filename_out)

        t1 = time.time()
        logger.info('Conversion time: %2.2fsec' % (t1- t0))

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

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

        #Note that a chunk is not a blob!!
        chunk_dim = self.__get_chunk_dimensions()
        blob_dim = self.__get_blob_dimensions(chunk_dim)
        n_blobs = self.container.calc_n_blobs(blob_dim)

        #Write header of .fil file
        n_bytes  = self.header[b'nbits'] / 8
        with open(filename_out, "wb") as fileh:
            fileh.write(generate_sigproc_header(self)) #generate_sigproc_header comes from sigproc.py

        logger.info('Using %i n_blobs to write the data.'% n_blobs)
        for ii in range(0, n_blobs):
            logger.info('Reading %i of %i' % (ii + 1, n_blobs))

            bob = self.container.read_blob(blob_dim,n_blob=ii)

            #Write data of .fil file.
            with open(filename_out, "a") as fileh:
                j = bob
                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_fil_light(self, filename_out, *args, **kwargs):
        """ Write data to .fil file.

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

        n_bytes  = self.header[b'nbits'] / 8
        with open(filename_out, "wb") as fileh:
            fileh.write(generate_sigproc_header(self)) #generate_sigproc_header comes from sigproc.py
            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.
            It check the file size then decides how to write the file.

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

        #For timing how long it takes to write a file.
        t0 = time.time()

        #Update header
        self.__update_header()

        if self.container.isheavy():
            self.__write_to_hdf5_heavy(filename_out)
        else:
            self.__write_to_hdf5_light(filename_out)

        t1 = time.time()
        logger.info('Conversion time: %2.2fsec' % (t1- t0))

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

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

        block_size = 0

        #Note that a chunk is not a blob!!
        chunk_dim = self.__get_chunk_dimensions()
        blob_dim = self.__get_blob_dimensions(chunk_dim)
        n_blobs = self.container.calc_n_blobs(blob_dim)

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

            h5.attrs[b'CLASS'] = b'FILTERBANK'
            h5.attrs[b'VERSION'] = b'1.0'

            if HAS_BITSHUFFLE:
                bs_compression = bitshuffle.h5.H5FILTER
                bs_compression_opts = (block_size, bitshuffle.h5.H5_COMPRESS_LZ4)
            else:
                bs_compression = None
                bs_compression_opts = None
                logger.warning("Warning: bitshuffle not found. No compression applied.")

            dset = h5.create_dataset('data',
                            shape=self.selection_shape,
                            chunks=chunk_dim,
                            compression=bs_compression,
                            compression_opts=bs_compression_opts,
                            dtype=self.data.dtype)

            dset_mask = h5.create_dataset('mask',
                            shape=self.selection_shape,
                            chunks=chunk_dim,
                            compression=bs_compression,
                            compression_opts=bs_compression_opts,
                            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

            if blob_dim[self.freq_axis] < self.selection_shape[self.freq_axis]:

                logger.info('Using %i n_blobs to write the data.'% n_blobs)
                for ii in range(0, n_blobs):
                    logger.info('Reading %i of %i' % (ii + 1, n_blobs))

                    bob = self.container.read_blob(blob_dim,n_blob=ii)

                    #-----
                    #Using channels instead of frequency.
                    c_start = self.container.chan_start_idx + ii*blob_dim[self.freq_axis]
                    t_start = self.container.t_start + (c_start/self.selection_shape[self.freq_axis])*blob_dim[self.time_axis]
                    t_stop = t_start + blob_dim[self.time_axis]

                    # Reverse array if frequency axis is flipped
#                     if self.header['foff'] < 0:
#                         c_stop = self.selection_shape[self.freq_axis] - (c_start)%self.selection_shape[self.freq_axis]
#                         c_start = c_stop - blob_dim[self.freq_axis]
#                     else:
                    c_start = (c_start)%self.selection_shape[self.freq_axis]
                    c_stop = c_start + blob_dim[self.freq_axis]
                    #-----

                    logger.debug(t_start,t_stop,c_start,c_stop)

                    dset[t_start:t_stop,0,c_start:c_stop] = bob[:]

            else:

                logger.info('Using %i n_blobs to write the data.'% n_blobs)
                for ii in range(0, n_blobs):
                    logger.info('Reading %i of %i' % (ii + 1, n_blobs))

                    bob = self.container.read_blob(blob_dim,n_blob=ii)
                    t_start = self.container.t_start + ii*blob_dim[self.time_axis]

                    #This prevents issues when the last blob is smaller than the others in time
                    if (ii+1)*blob_dim[self.time_axis] > self.n_ints_in_file:
                        t_stop = self.n_ints_in_file
                    else:
                        t_stop = (ii+1)*blob_dim[self.time_axis]

                    dset[t_start:t_stop] = bob[:]

    def __write_to_hdf5_light(self, filename_out, *args, **kwargs):
        """ Write data to HDF5 file in one go.

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

        block_size = 0

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

            h5.attrs[b'CLASS']   = b'FILTERBANK'
            h5.attrs[b'VERSION'] = b'1.0'

            if HAS_BITSHUFFLE:
                bs_compression = bitshuffle.h5.H5FILTER
                bs_compression_opts = (block_size, bitshuffle.h5.H5_COMPRESS_LZ4)
            else:
                bs_compression = None
                bs_compression_opts = None
                logger.warning("Warning: bitshuffle not found. No compression applied.")

            dset = h5.create_dataset('data',
                        data=self.data,
#                          compression='lzf')
                        compression=bs_compression,
                        compression_opts=bs_compression_opts)

            dset_mask = h5.create_dataset('mask',
                        shape=self.file_shape,
#                                 compression='lzf',
                        compression=bs_compression,
                        compression_opts=bs_compression_opts,
                        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 __get_blob_dimensions(self, chunk_dim):
        """ Sets the blob dimmentions, trying to read around 1024 MiB at a time.
            This is assuming a chunk is about 1 MiB.
        """

        #Taking the size into consideration, but avoiding having multiple blobs within a single time bin.
        if self.selection_shape[self.freq_axis] > chunk_dim[self.freq_axis]*MAX_BLOB_MB:
            freq_axis_size = self.selection_shape[self.freq_axis]
#             while freq_axis_size > chunk_dim[self.freq_axis]*MAX_BLOB_MB:
#                 freq_axis_size /= 2
            time_axis_size = 1
        else:
            freq_axis_size = self.selection_shape[self.freq_axis]
            time_axis_size = np.min([chunk_dim[self.time_axis] * MAX_BLOB_MB * chunk_dim[self.freq_axis] / freq_axis_size, self.selection_shape[self.time_axis]])

        blob_dim = (int(time_axis_size), 1, freq_axis_size)

        return blob_dim

    def __get_chunk_dimensions(self):
        """ Sets the chunking dimmentions depending on the file type.
        """

        #Usually '.0000.' is in self.filename
        if np.abs(self.header[b'foff']) < 1e-5:
            logger.info('Detecting high frequency resolution data.')
            chunk_dim = (1,1,1048576) #1048576 is the number of channels in a coarse channel.
            return chunk_dim
        #Usually '.0001.' is in self.filename
        elif np.abs(self.header[b'tsamp']) < 1e-3:
            logger.info('Detecting high time resolution data.')
            chunk_dim = (2048,1,512) #512 is the total number of channels per single band (ie. blc00)
            return chunk_dim
        #Usually '.0002.' is in self.filename
        elif np.abs(self.header[b'foff']) < 1e-2 and np.abs(self.header[b'foff'])  >= 1e-5:
            logger.info('Detecting intermediate frequency and time resolution data.')
            chunk_dim = (10,1,65536)  #65536 is the total number of channels per single band (ie. blc00)
#            chunk_dim = (1,1,65536/4)
            return chunk_dim
        else:
            logger.warning('File format not known. Will use minimum chunking. NOT OPTIMAL.')
            chunk_dim = (1,1,512)
            return chunk_dim

    def calc_n_coarse_chan(self, chan_bw=None):
        """ This makes an attempt to calculate the number of coarse channels in a given freq selection.
            It assumes for now that a single coarse channel is 2.9296875 MHz
        """

        n_coarse_chan = self.container.calc_n_coarse_chan(chan_bw)

        return n_coarse_chan

    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
        """

        self.freqs = self.populate_freqs()
        self.timestamps = self.populate_timestamps()

        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 cmd_tool(args=None):
    """ Command line tool for plotting and viewing info on blimpy files """

    from argparse import ArgumentParser

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

    parser.add_argument('filename', type=str,
                        help='Name of file to read')
    parser.add_argument('-p', action='store',  default='a', 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('-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, inclusive) ID ')
    parser.add_argument('-E', action='store', default=None, dest='t_stop', type=int,
                        help='Stop integration (end, exclusive) 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('-H', action='store_true', default=False, dest='to_hdf5',
                       help='Write file to hdf5 format.')
    parser.add_argument('-F', action='store_true', default=False, dest='to_fil',
                       help='Write file to .fil format.')
    parser.add_argument('-o', action='store', default=None, dest='filename_out', type=str,
                        help='Filename output (if not probided, the name will be the same but with apropiate extension).')
    parser.add_argument('-l', action='store', default=None, dest='max_load', type=float,
                        help='Maximum data limit to load. Default:1GB')

    if args is None:
        args = sys.argv[1:]

    parse_args = parser.parse_args(args)

    # Open blimpy data
    filename = parse_args.filename
    load_data = not parse_args.info_only
    info_only = parse_args.info_only
    filename_out = parse_args.filename_out

    fil = Waterfall(filename, f_start=parse_args.f_start, f_stop=parse_args.f_stop, t_start=parse_args.t_start, t_stop=parse_args.t_stop, load_data=load_data, max_load=parse_args.max_load)
    fil.info()

    #Check the size of selection.
    if fil.container.isheavy() or parse_args.to_hdf5 or parse_args.to_fil:
        info_only = True

    # And if we want to plot data, then plot data.

    if not info_only:
        print('')

        if parse_args.blank_dc:
            logger.info("Blanking DC bin")
            n_coarse_chan = fil.calc_n_coarse_chan()
            fil.blank_dc(n_coarse_chan)

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

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

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

    else:

        if parse_args.to_hdf5 and parse_args.to_fil:
            raise Warning('Either provide to_hdf5 or to_fil, but not both.')

        if parse_args.to_hdf5:
            if not filename_out:
                filename_out = filename.replace('.fil','.h5')
            elif '.h5' not in filename_out:
                filename_out = filename_out.replace('.fil','')+'.h5'

            logger.info('Writing file : %s'%(filename_out))
            fil.write_to_hdf5(filename_out)
            logger.info('File written.')

        if parse_args.to_fil:
            if not filename_out:
                filename_out = filename.replace('.h5','.fil')
            elif '.fil' not in filename_out:
                filename_out = filename_out.replace('.h5','')+'.fil'

            logger.info('Writing file : %s'%(filename_out))
            fil.write_to_fil(filename_out)
            logger.info('File written.')

if __name__ == "__main__":
    cmd_tool()


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

A python file handler for guppi RAW files from the GBT.

The guppi raw format consists of a FITS-like header, followed by a block of data,
and repeated over and over until the end of the file.
"""

import numpy as np
import os
import time
from pprint import pprint
from astropy.coordinates import Angle

from .utils import unpack, rebin
import sys

PYTHON3 = sys.version_info >= (3, 0)

# Check if $DISPLAY is set (for handling plotting on remote machines with no X-forwarding)
if 'DISPLAY' in os.environ.keys():

    try:
        import matplotlib
        # matplotlib.use('Qt5Agg')
    except ImportError:
        pass
    import pylab as plt
else:
    import matplotlib

    matplotlib.use('Agg')
    import pylab as plt

###
# Config values
###

MAX_PLT_POINTS = 65536 * 4  # 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  # Max size of data array to load into memory

class EndOfFileError(Exception):
    pass

class GuppiRaw(object):
    """ Python class for reading Guppi raw files

    Args:
        filename (str): name of the .raw file to open

    Optional args:
        n_blocks (int): if number of blocks to read is known, set it here.
                        This saves seeking through the file to check how many
                        integrations there are in the file.
    """

    def __init__(self, filename, n_blocks=None):
        self.filename = filename
        if PYTHON3:
            self.file_obj = open(filename, 'rb')
        else:
            self.file_obj = open(filename, 'r')
        self.filesize = os.path.getsize(filename)

        if not n_blocks:
            self.n_blocks = self.find_n_data_blocks()

        else:
            self.n_blocks = n_blocks

        self._d = np.zeros(1, dtype='complex64')
        self._d_x = np.zeros(1, dtype='int8')
        self._d_y = np.zeros(1, dtype='int8')

    def __enter__(self):
        """
        reopen the file each time a `with` block is entered
        :return:
        """
        if PYTHON3:
            self.file_obj = open(self.filename, 'rb')
        else:
            self.file_obj = open(self.filename, 'r')
        self.filesize = os.path.getsize(self.filename)
        return self

    def __exit__(self, exception_type, exception_value, traceback):
        """
        closes the file after `with` block has exited
        :param exception_type:
        :param exception_value:
        :param traceback:
        :return:
        """
        self.file_obj.close()

    def __repr__(self):
        return "<GuppiRaw file handler for %s>" % self.filename

    def read_header(self):
        """ Read next header (multiple headers in file)

        Returns:
            (header, data_idx) - a dictionary of keyword:value header data and
            also the byte index of where the corresponding data block resides.
        """
        start_idx = self.file_obj.tell()
        key, val = '', ''

        header_dict = {}
        keep_reading = True

        first_line = self.file_obj

        try:
            while keep_reading:
                if start_idx + 80 > self.filesize:
                    keep_reading = False
                    raise EndOfFileError("End Of Data File")
                line = self.file_obj.read(80)
                if PYTHON3:
                    line = line.decode("utf-8")
                # print line
                if line.startswith('END'):
                    keep_reading = False
                    break
                else:
                    key, val = line.split('=')
                    key, val = key.strip(), val.strip()

                    if "'" in val:
                        # Items in quotes are strings
                        val = str(val.strip("'").strip())
                    elif "." in val:
                        # Items with periods are floats (if not a string)
                        val = float(val)
                    else:
                        # Otherwise it's an integer
                        val = int(val)

                header_dict[key] = val
        except ValueError:
            print("CURRENT LINE: ", line)
            print("BLOCK START IDX: ", start_idx)
            print("FILE SIZE: ", self.filesize)
            print("NEXT 512 BYTES: \n")
            print(self.file_obj.read(512))
            raise

        data_idx = self.file_obj.tell()

        # Seek past padding if DIRECTIO is being used
        if "DIRECTIO" in header_dict.keys():
            if int(header_dict["DIRECTIO"]) == 1:
                if data_idx % 512:
                    data_idx += (512 - data_idx % 512)

        self.file_obj.seek(start_idx)
        return header_dict, data_idx

    def read_first_header(self):
        """ Read first header in file

        Returns:
            header (dict): keyword:value pairs of header metadata
        """
        self.file_obj.seek(0)
        header_dict, pos = self.read_header()
        self.file_obj.seek(0)
        return header_dict

    def read_next_data_block_shape(self):
        header, data_idx = self.read_header()
        n_chan = int(header['OBSNCHAN'])
        n_pol = int(header['NPOL'])
        n_bit = int(header['NBITS'])
        n_samples = int(int(header['BLOCSIZE']) / (n_chan * n_pol * (n_bit / 8)))

        is_chanmaj = False
        if 'CHANMAJ' in header.keys():
            if int(header['CHANMAJ']) == 1:
                is_chanmaj = True
        if is_chanmaj:
            dshape = (int(n_samples), n_chan, n_pol)
        else:
            dshape = (n_chan, int(n_samples), n_pol)
        return dshape

    def get_data(self):
        """
        returns a generator object that reads data a block at a time;
        the generator prints "File depleted" and returns nothing when all data in the file has been read.
        :return:
        """
        with self as gr:
            while True:
                try:
                    yield gr.read_next_data_block_int8()
                except Exception as e:
                    print("File depleted")
                    yield None, None, None

    def read_next_data_block_int8(self):
        """ Read the next block of data and its header

        Returns: (header, data)
            header (dict): dictionary of header metadata
            data (np.array): Numpy array of data, converted into to complex64.
        """
        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)

        # Read data and reshape

        n_chan = int(header['OBSNCHAN'])
        n_pol = int(header['NPOL'])
        n_bit = int(header['NBITS'])
        n_samples = int(int(header['BLOCSIZE']) / (n_chan * n_pol * (n_bit / 8)))

        d = np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8')

        # Handle 2-bit and 4-bit data
        if n_bit != 8:
            d = unpack(d, n_bit)

        d = d.reshape((n_chan, n_samples, n_pol))  # Real, imag

        if self._d_x.shape != d[..., 0:2].shape:
            self._d_x = np.ascontiguousarray(np.zeros(d[..., 0:2].shape, dtype='int8'))
            self._d_y = np.ascontiguousarray(np.zeros(d[..., 2:4].shape, dtype='int8'))

        self._d_x[:] = d[..., 0:2]
        self._d_y[:] = d[..., 2:4]
        return header, self._d_x, self._d_y

    def read_next_data_block_int8_2x(self):
        """ Read the next block of data and its header

        Returns: (header, data)
            header (dict): dictionary of header metadata
            data (np.array): Numpy array of data, converted into to complex64.
        """
        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)

        # Read data and reshape

        n_chan = int(header['OBSNCHAN'])
        n_pol = int(header['NPOL'])
        n_bit = int(header['NBITS'])
        n_samples = int(int(header['BLOCSIZE']) / (n_chan * n_pol * (n_bit / 8)))

        d = np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8')

        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)
        d2 = np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8')

        # Handle 2-bit and 4-bit data
        if n_bit != 8:
            d = unpack(d, n_bit)

        d = d.reshape((n_chan, n_samples, n_pol))  # Real, imag
        d2 = d2.reshape((n_chan, n_samples, n_pol))
        d = np.concatenate((d, d2), axis=1)
        print(d.shape)

        if self._d_x.shape != (n_chan, n_samples * 2, n_pol):
            self._d_x = np.ascontiguousarray(np.zeros(d[..., 0:2].shape, dtype='int8'))
            self._d_y = np.ascontiguousarray(np.zeros(d[..., 2:4].shape, dtype='int8'))

        self._d_x[:] = d[..., 0:2]
        self._d_y[:] = d[..., 2:4]
        return header, self._d_x, self._d_y

    def read_next_data_block(self):
        """ Read the next block of data and its header

        Returns: (header, data)
            header (dict): dictionary of header metadata
            data (np.array): Numpy array of data, converted into to complex64.
        """
        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)

        # Read data and reshape

        n_chan = int(header['OBSNCHAN'])
        n_pol = int(header['NPOL'])
        n_bit = int(header['NBITS'])
        n_samples = int(int(header['BLOCSIZE']) / (n_chan * n_pol * (n_bit / 8)))

        d = np.ascontiguousarray(np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8'))

        # Handle 2-bit and 4-bit data
        if n_bit != 8:
            d = unpack(d, n_bit)

        dshape = self.read_next_data_block_shape()

        d = d.reshape(dshape)  # Real, imag

        if self._d.shape != d.shape:
            self._d = np.zeros(d.shape, dtype='float32')

        self._d[:] = d

        return header, self._d[:].view('complex64')

    def find_n_data_blocks(self):
        """ Seek through the file to find how many data blocks there are in the file

        Returns:
            n_blocks (int): number of data blocks in the file
        """
        self.file_obj.seek(0)
        header0, data_idx0 = self.read_header()

        self.file_obj.seek(data_idx0)
        block_size = int(header0['BLOCSIZE'])
        n_bits = int(header0['NBITS'])
        self.file_obj.seek(int(header0['BLOCSIZE']), 1)
        n_blocks = 1
        end_found = False
        while not end_found:
            try:
                header, data_idx = self.read_header()
                self.file_obj.seek(data_idx)
                self.file_obj.seek(header['BLOCSIZE'], 1)
                n_blocks += 1
            except EndOfFileError:
                end_found = True
                break

        self.file_obj.seek(0)
        return n_blocks

    def reset_index(self):
        """ Return file_obj seek to start of file """
        self.file_obj.seek(0)

    def print_stats(self):
        """ Compute some basic stats on the next block of data """

        header, data = self.read_next_data_block()
        data = data.view('float32')

        print("AVG: %2.3f" % data.mean())
        print("STD: %2.3f" % data.std())
        print("MAX: %2.3f" % data.max())
        print("MIN: %2.3f" % data.min())

        import pylab as plt

    def plot_histogram(self, filename=None):
        """ Plot a histogram of data values """
        header, data = self.read_next_data_block()
        data = data.view('float32')

        plt.figure("Histogram")
        plt.hist(data.flatten(), 65, facecolor='#cc0000')
        if filename:
            plt.savefig(filename)
        plt.show()

    def plot_spectrum(self, filename=None, plot_db=True):
        """ Do a (slow) numpy FFT and take power of data """
        header, data = self.read_next_data_block()

        print("Computing FFT...")
        d_xx_fft = np.abs(np.fft.fft(data[..., 0]))
        d_xx_fft = d_xx_fft.flatten()

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

        d_xx_fft = rebin(d_xx_fft, dec_fac_x)

        print("Plotting...")
        if plot_db:
            plt.plot(10 * np.log10(d_xx_fft))
            plt.ylabel("Power [dB]")
        else:
            plt.plot(d_xx_fft)
            plt.ylabel("Power")
            plt.xlabel("Channel")
        plt.title(self.filename)
        if filename:
            plt.savefig(filename)
        plt.show()

    def generate_filterbank_header(self, nchans=1, ):
        """ Generate a blimpy header dictionary """
        gp_head = self.read_first_header()
        fb_head = {}

        telescope_str = gp_head.get("TELESCOP", "unknown")
        if telescope_str in ('GBT', 'GREENBANK'):
            fb_head["telescope_id"] = 6
        elif telescope_str in ('PKS', 'PARKES'):
            fb_head["telescop_id"] = 7
        else:
            fb_head["telescop_id"] = 0

        # Using .get() method allows us to fill in default values if not present
        fb_head["source_name"] = gp_head.get("SRC_NAME", "unknown")
        fb_head["az_start"] = gp_head.get("AZ", 0)
        fb_head["za_start"] = gp_head.get("ZA", 0)

        fb_head["src_raj"] = Angle(str(gp_head.get("RA", 0.0)) + "hr")
        fb_head["src_dej"] = Angle(str(gp_head.get("DEC", 0.0)) + "deg")
        fb_head["rawdatafile"] = self.filename

        # hardcoded
        fb_head["machine_id"] = 20
        fb_head["data_type"] = 1  # blio datatype
        fb_head["barycentric"] = 0
        fb_head["pulsarcentric"] = 0
        fb_head["nbits"] = 32

        # TODO - compute these values. Need to figure out the correct calcs
        fb_head["tstart"] = 0.0
        fb_head["tsamp"] = 1.0
        fb_head["fch1"] = 0.0
        fb_head["foff"] = 187.5 / nchans

        # Need to be updated based on output specs
        fb_head["nchans"] = nchans
        fb_head["nifs"] = 1
        fb_head["nbeams"] = 1

        return fb_head

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

    from argparse import ArgumentParser

    parser = ArgumentParser(description="Command line utility for creating spectra from GuppiRaw files.")

    parser.add_argument('filename', type=str, help='Name of file to read')
    parser.add_argument('-o', dest='outdir', type=str, default='./', help='output directory for PNG files')
    args = parser.parse_args()

    r = GuppiRaw(args.filename)

    r.print_stats()
    bname = os.path.splitext(os.path.basename(args.filename))[0]
    bname = os.path.join(args.outdir, bname)
    r.plot_histogram(filename="%s_hist.png" % bname)
    r.plot_spectrum(filename="%s_spec.png" % bname)

if __name__ == "__main__":
    cmd_tool()


#!/usr/bin/env python

from optparse import OptionParser
import socket
import subprocess
import sys
import os
try:
    from .waterfall import Waterfall
except:
    from blimpy import Waterfall

#import pdb #pdb.set_trace()

local_host = socket.gethostname()

def reset_outs():
    '''
    '''

    return None,None

def make_batch_script():
    """ This creates a batch script for the heavy lifting"""

    script_text  = '#! /bin/bash\n# This script calculates the md5sum for the last number of bits (nbits).# Usage :\n#     ./tail_sum file bnits\n#\n\nfile_name=$1\nbit_num=$2\n\ntail -c $bit_num $file_name | md5sum\n'

    with open('tail_sum.sh', 'w') as batch_script:
        batch_script.write(script_text)

    os.chmod('tail_sum.sh', 0o775)

def find_header_size(filename):
    ''' Script to find the header size of a filterbank file'''

    # open datafile
    filfile=open(filename,'rb')
    # go to the start of the file
    filfile.seek(0)
    #read some region larger than the header.
    round1 = filfile.read(1000)
    headersize = round1.find('HEADER_END')+len('HEADER_END')

    return headersize

def cmd_tool(args=None):
    """ Command line tool to make a md5sum comparison of two .fil files. """

    if 'bl' in local_host:
        header_loc = '/usr/local/sigproc/bin/header' #Current location of header command in GBT.
    else:
        raise IOError('Script only able to run in BL systems.')

    p = OptionParser()
    p.set_usage('matchfils <FIL_FILE1> <FIL_FILE2>')
    opts, args = p.parse_args(sys.argv[1:])

    file1 = args[0]
    file2 = args[1]

    #------------------------------------
    #Create batch script

    make_batch_script()

    #------------------------------------
    #First checksum

    headersize1 = find_header_size(file1)
    file_size1 = os.path.getsize(file1)

    #Strip header from file, and calculate the md5sum of the rest.
    #command=['tail','-c',str(file_size1-headersize1),file1,'|','md5sum']
    command=['./tail_sum.sh',file1,str(file_size1-headersize1)]
    print('[matchfils] '+' '.join(command))

    proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (out, err) = proc.communicate()

    check_sum1 = out.split()[0]
    print('[matchfils] Checksum is:', check_sum1)

    if err:
        raise Error('There is an error.')

    #---
    out,err = reset_outs()

    command=[header_loc,file1]
    print('[matchfils] Header information:')

    proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (out, err) = proc.communicate()

    header1 = out
    print(header1)

    #------------------------------------
    #Second checksum

    out,err = reset_outs()

    headersize2 = find_header_size(file2)
    file_size2 = os.path.getsize(file2)

    #Strip header from file, and calculate the md5sum of the rest.
    command=['./tail_sum.sh',file2,str(file_size2-headersize2)]
    print('[matchfils] '+' '.join(command))

    proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (out, err) = proc.communicate()

    check_sum2 = out.split()[0]
    print('[matchfils] Checksum is:', check_sum2)

    if err:
        raise Error('There is an error.')

    #---
    out,err = reset_outs()

    command=[header_loc,file2]
    print('[matchfils] Header information:')

    proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (out, err) = proc.communicate()

    header2 = out
    print(header2)

    #------------------------------------
    #check the checksums

    if check_sum1 != check_sum2:
        print('[matchfils] Booo! Checksum does not match between files.')
    else:
        print('[matchfils] Hooray! Checksum matches between files.')

    #------------------------------------
    #Remove batch script

    os.remove('tail_sum.sh')

if __name__ == "__main__":
    cmd_tool()


#!/usr/bin/env python
'''
    Simple script for quicly making an h5 file from a .fil.

    ..author: Emilio Enriquez (jeenriquez@gmail.com)

    July 28th 2017
'''

try:
    from .waterfall import Waterfall
except:
    from waterfall import Waterfall

from optparse import OptionParser
import sys
import os

#------
# 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 make_h5_file(filename,out_dir='./', new_filename = None, max_load = None):
    ''' Converts file to HDF5 (.h5) format. Default saves output in current dir.
    '''

    fil_file = Waterfall(filename, max_load = max_load)
    if not new_filename:
        new_filename = out_dir+filename.replace('.fil','.h5').split('/')[-1]

    if '.h5' not in new_filename:
        new_filename = new_filename+'.h5'

    fil_file.write_to_hdf5(new_filename)

def cmd_tool():
    '''
    '''

    p = OptionParser()
    p.set_usage('Command line utility for converting Sigproc filterbank (.fil) to  HDF5 (.h5) format  \n >>fil2h5 <FULL_PATH_TO_FIL_FILE> [options]')
    p.add_option('-o', '--out_dir', dest='out_dir', type='str', default='./', help='Location for output files. Default: local dir. ')
    p.add_option('-n', '--new_filename', dest='new_filename', type='str', default='', help='New name. Default: replaces extention to .h5')
    p.add_option('-d', '--delete_input', dest='delete_input', action='store_true', default=False, help='This option deletes the input file after conversion.')
    p.add_option('-l', action='store', default=None, dest='max_load', type=float,help='Maximum data limit to load. Default:1GB')

    opts, args = p.parse_args(sys.argv[1:])

    if len(args)!=1:
        logger.info('Please specify a file name \nExiting.')
        sys.exit()
    else:
        filename = args[0]

    make_h5_file(filename, out_dir = opts.out_dir, new_filename = opts.new_filename, max_load = opts.max_load)

    if opts.delete_input:
        logger.info("'Deleting input file: %s"%(filename))
        os.remove(filename)

if __name__ == "__main__":

    cmd_tool()


#!/usr/bin/env python

from .guppi import GuppiRaw
import h5py
try:
    import bitshuffle.h5
    HAS_BITSHUFFLE = True
except ImportError:
    HAS_BITSHUFFLE = False

import time
import os
import glob
import numpy as np

def cmd_tool(args=None):
    """ Command line tool for converting guppi raw into HDF5 versions of guppi raw """
    from argparse import ArgumentParser

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

    parser = ArgumentParser(description="Command line utility for creating HDF5 Raw files.")
    parser.add_argument('filename', type=str, help='Name of filename to read')
    args = parser.parse_args()

    fileroot = args.filename.split('.0000.raw')[0]

    filelist = glob.glob(fileroot + '*.raw')
    filelist = sorted(filelist)

    # Read first file
    r = GuppiRaw(filelist[0])
    header, data = r.read_next_data_block()
    dshape = data.shape #r.read_next_data_block_shape()
    print(dshape)

    n_blocks_total = 0
    for filename in filelist:
        print(filename)
        r = GuppiRaw(filename)
        n_blocks_total += r.n_blocks
    print(n_blocks_total)

    full_dshape = np.concatenate(((n_blocks_total,), dshape))

    # Create h5py file
    h5 = h5py.File(fileroot + '.h5', 'w')
    h5.attrs['CLASS'] = 'GUPPIRAW'
    block_size = 0      # This is chunk block size
    dset = h5.create_dataset('data',
                  shape=full_dshape,
                  #compression=bitshuffle.h5.H5FILTER,
                  #compression_opts=(block_size, bitshuffle.h5.H5_COMPRESS_LZ4),
                  dtype=data.dtype)

    h5_idx = 0
    for filename in filelist:
        print("\nReading %s header..." % filename)
        r = GuppiRaw(filename)
        h5 = h5py.File(filename + '.h5', 'w')

        header, data = r.read_next_data_block()

        for ii in range(0, r.n_blocks):
            t0 = time.time()
            print("Reading block %i of %i" % (h5_idx+1, full_dshape[0]))
            header, data = r.read_next_data_block()
            t1 = time.time()

            t2 = time.time()
            print("Writing block %i of %i" % (h5_idx+1, full_dshape[0]))
            dset[h5_idx, :] = data
            t3 = time.time()
            print("Read: %2.2fs, Write %2.2fs" % ((t1-t0), (t3-t2)))

            h5_idx += 1

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

        h5.close()

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

if __name__ == "__main__":
    cmd_tool()


from blimpy import Waterfall
import numpy as np
import pylab as plt
from scipy.integrate import trapz

def foldcal(data,tsamp, diode_p=0.04,numsamps=1000,switch=False,inds=False):
    '''
    Returns time-averaged spectra of the ON and OFF measurements in a
    calibrator measurement with flickering noise diode

    Parameters
    ----------
    data : 2D Array object (float)
        2D dynamic spectrum for data (any Stokes parameter) with flickering noise diode.
    tsamp : float
        Sampling time of data in seconds
    diode_p : float
        Period of the flickering noise diode in seconds
    numsamps : int
        Number of samples over which to average noise diode ON and OFF
    switch : boolean
        Use switch=True if the noise diode "skips" turning from OFF to ON once or vice versa
    inds : boolean
        Use inds=True to also return the indexes of the time series where the ND is ON and OFF
    '''

    halfper = diode_p/2.0

    foldt = halfper/tsamp   #number of time samples per diode switch

    onesec = 1/tsamp    #number of time samples in the first second

    #Find diode switches in units of time samples and round down to the nearest int
    ints = np.arange(0,numsamps)
    t_switch = (onesec+ints*foldt)
    t_switch = t_switch.astype('int')

    ONints = np.array(np.reshape(t_switch[:],(numsamps/2,2)))
    ONints[:,0] = ONints[:,0]+1   #Find index ranges of ON time samples

    OFFints = np.array(np.reshape(t_switch[1:-1],(numsamps/2-1,2)))
    OFFints[:,0] = OFFints[:,0]+1   #Find index ranges of OFF time samples

    av_ON = []
    av_OFF = []

    #Average ON and OFF spectra separately with respect to time
    for i in ONints:
        if i[1]!=i[0]:
            av_ON.append(np.sum(data[i[0]:i[1],:,:],axis=0)/(i[1]-i[0]))

    for i in OFFints:
        if i[1]!=i[0]:
            av_OFF.append(np.sum(data[i[0]:i[1],:,:],axis=0)/(i[1]-i[0]))

    #If switch=True, flip the return statement since ON is actually OFF
    if switch==False:
        if inds==False:
            return np.squeeze(np.mean(av_ON,axis=0)), np.squeeze(np.mean(av_OFF,axis=0))
        else:
            return np.squeeze(np.mean(av_ON,axis=0)), np.squeeze(np.mean(av_OFF,axis=0)),ONints,OFFints
    if switch==True:
        if inds==False:
            return np.squeeze(np.mean(av_OFF,axis=0)), np.squeeze(np.mean(av_ON,axis=0))
        else:
            return np.squeeze(np.mean(av_OFF,axis=0)), np.squeeze(np.mean(av_ON,axis=0)),OFFints,ONints

def integrate_chans(spec,freqs,chan_per_coarse):
    '''
    Integrates over each core channel of a given spectrum.
    Important for calibrating data with frequency/time resolution different from noise diode data

    Parameters
    ----------
    spec : 1D Array (float)
        Spectrum (any Stokes parameter) to be integrated
    freqs : 1D Array (float)
        Frequency values for each bin of the spectrum
    chan_per_coarse: int
        Number of frequency bins per coarse channel
    '''

    num_coarse = spec.size/chan_per_coarse     #Calculate total number of coarse channels

    #Rearrange spectrum by coarse channel
    spec_shaped = np.array(np.reshape(spec,(num_coarse,chan_per_coarse)))
    freqs_shaped = np.array(np.reshape(freqs,(num_coarse,chan_per_coarse)))

    #Average over coarse channels
    return np.mean(spec_shaped[:,1:-1],axis=1)

def integrate_calib(name,chan_per_coarse,fullstokes=False,**kwargs):
    '''
    Folds Stokes I noise diode data and integrates along coarse channels

    Parameters
    ----------
    name : str
        Path to noise diode filterbank file
    chan_per_coarse : int
        Number of frequency bins per coarse channel
    fullstokes : boolean
        Use fullstokes=True if data is in IQUV format or just Stokes I, use fullstokes=False if
        it is in cross_pols format
    '''
    #Load data
    obs = Waterfall(name,max_load=150)
    data = obs.data

    #If the data has cross_pols format calculate Stokes I
    if fullstokes==False and data.shape[1]>1:
        data = data[:,0,:]+data[:,1,:]
        data = np.expand_dims(data,axis=1)
    #If the data has IQUV format get Stokes I
    if fullstokes==True:
        data = data[:,0,:]
        data = np.expand_dims(data,axis=1)

    tsamp = obs.header['tsamp']

    #Calculate ON and OFF values
    OFF,ON = foldcal(data,tsamp,**kwargs)

    freqs = obs.populate_freqs()

    #Find ON and OFF spectra by coarse channel
    ON_int = integrate_chans(ON,freqs,chan_per_coarse)
    OFF_int = integrate_chans(OFF,freqs,chan_per_coarse)

    #If "ON" is actually "OFF" switch them
    if np.sum(ON_int)<np.sum(OFF_int):
        temp = ON_int
        ON_int = OFF_int
        OFF_int = temp

    #Return coarse channel spectrum of OFF and ON
    return OFF_int,ON_int

def get_calfluxes(calflux,calfreq,spec_in,centerfreqs,oneflux):
    '''
    Given properties of the calibrator source, calculate fluxes of the source
    in a particular frequency range

    Parameters
    ----------
    calflux : float
        Known flux of calibrator source at a particular frequency
    calfreq : float
        Frequency where calibrator source has flux calflux (see above)
    spec_in : float
        Known power-law spectral index of calibrator source. Use convention flux(frequency) = constant * frequency^(spec_in)
    centerfreqs : 1D Array (float)
        Central frequency values of each coarse channel
    oneflux : boolean
        Use oneflux to choose between calculating the flux for each core channel (False)
        or using one value for the entire frequency range (True)
    '''

    const = calflux/np.power(calfreq,spec_in)
    if oneflux==False:
        return const*np.power(centerfreqs,spec_in)
    else:
        return const*np.power(np.mean(centerfreqs),spec_in)

def get_centerfreqs(freqs,chan_per_coarse):
    '''
    Returns central frequency of each coarse channel

    Parameters
    ----------
    freqs : 1D Array (float)
        Frequency values for each bin of the spectrum
    chan_per_coarse: int
        Number of frequency bins per coarse channel
    '''

    num_coarse = freqs.size/chan_per_coarse
    freqs = np.reshape(freqs,(num_coarse,chan_per_coarse))
    return np.mean(freqs,axis=1)

def f_ratios(calON_obs,calOFF_obs,chan_per_coarse,**kwargs):
    '''
    Calculate f_ON, and f_OFF as defined in van Straten et al. 2012 equations 2 and 3

    Parameters
    ----------
    calON_obs : str
        Path to filterbank file (any format) for observation ON the calibrator source
    calOFF_obs : str
        Path to filterbank file (any format) for observation OFF the calibrator source
    '''
    #Calculate noise diode ON and noise diode OFF spectra (H and L) for both observations
    L_ON,H_ON = integrate_calib(calON_obs,chan_per_coarse,**kwargs)
    L_OFF,H_OFF = integrate_calib(calOFF_obs,chan_per_coarse,**kwargs)

    f_ON = H_ON/L_ON-1
    f_OFF = H_OFF/L_OFF-1

    return f_ON, f_OFF

def diode_spec(calON_obs,calOFF_obs,calflux,calfreq,spec_in,average=True,oneflux=False,**kwargs):
    '''
    Calculate the coarse channel spectrum and system temperature of the noise diode in Jy given two noise diode
    measurements ON and OFF the calibrator source with the same frequency and time resolution

    Parameters
    ----------
    calON_obs : str
        (see f_ratios() above)
    calOFF_obs : str
        (see f_ratios() above)
    calflux : float
        Known flux of calibrator source at a particular frequency
    calfreq : float
        Frequency where calibrator source has flux calflux (see above)
    spec_in : float
        Known power-law spectral index of calibrator source. Use convention flux(frequency) = constant * frequency^(spec_in)
    average : boolean
        Use average=True to return noise diode and Tsys spectra averaged over frequencies
    '''
    #Load frequencies and calculate number of channels per coarse channel
    obs = Waterfall(calON_obs,max_load=150)
    freqs = obs.populate_freqs()
    ncoarse = obs.calc_n_coarse_chan()
    nchans = obs.header['nchans']
    chan_per_coarse = nchans/ncoarse

    f_ON, f_OFF = f_ratios(calON_obs,calOFF_obs,chan_per_coarse,**kwargs)

    #Obtain spectrum of the calibrator source for the given frequency range
    centerfreqs = get_centerfreqs(freqs,chan_per_coarse)
    calfluxes = get_calfluxes(calflux,calfreq,spec_in,centerfreqs,oneflux)

    #C_o and Tsys as defined in van Straten et al. 2012
    C_o = calfluxes/(1/f_ON-1/f_OFF)
    Tsys = C_o/f_OFF

    #return coarse channel diode spectrum
    if average==True:
        return np.mean(C_o),np.mean(Tsys)
    else:
        return C_o,Tsys

def get_Tsys(calON_obs,calOFF_obs,calflux,calfreq,spec_in,oneflux=False,**kwargs):
    '''
    Returns frequency dependent system temperature given observations on and off a calibrator source

    Parameters
    ----------
    (See diode_spec())
    '''
    return diode_spec(calON_obs,calOFF_obs,calflux,calfreq,spec_in,average=False,oneflux=False,**kwargs)[1]

def calibrate_fluxes(main_obs_name,dio_name,dspec,Tsys,fullstokes=False,**kwargs):
    '''
    Produce calibrated Stokes I for an observation given a noise diode
    measurement on the source and a diode spectrum with the same number of
    coarse channels

    Parameters
    ----------
    main_obs_name : str
        Path to filterbank file containing final data to be calibrated
    dio_name : str
        Path to filterbank file for observation on the target source with flickering noise diode
    dspec : 1D Array (float) or float
        Coarse channel spectrum (or average) of the noise diode in Jy (obtained from diode_spec())
    Tsys : 1D Array (float) or float
        Coarse channel spectrum (or average) of the system temperature in Jy
    fullstokes: boolean
        Use fullstokes=True if data is in IQUV format or just Stokes I, use fullstokes=False if
        it is in cross_pols format
    '''

    #Find folded spectra of the target source with the noise diode ON and OFF
    main_obs = Waterfall(main_obs_name,max_load=150)
    ncoarse = main_obs.calc_n_coarse_chan()
    dio_obs = Waterfall(dio_name,max_load=150)
    dio_chan_per_coarse = dio_obs.header['nchans']/ncoarse
    dOFF,dON = integrate_calib(dio_name,dio_chan_per_coarse,fullstokes,**kwargs)

    #Find Jy/count for each coarse channel using the diode spectrum
    main_dat = main_obs.data
    scale_facs = dspec/(dON-dOFF)
    print(scale_facs)

    nchans = main_obs.header['nchans']
    obs_chan_per_coarse = nchans/ncoarse

    ax0_size = np.size(main_dat,0)
    ax1_size = np.size(main_dat,1)

    #Reshape data array of target observation and multiply coarse channels by the scale factors
    main_dat = np.reshape(main_dat,(ax0_size,ax1_size,ncoarse,obs_chan_per_coarse))
    main_dat = np.swapaxes(main_dat,2,3)

    main_dat = main_dat*scale_facs
    main_dat = main_dat-Tsys
    main_dat = np.swapaxes(main_dat,2,3)
    main_dat = np.reshape(main_dat,(ax0_size,ax1_size,nchans))

    #Write calibrated data to a new filterbank file with ".fluxcal" extension

    main_obs.data = main_dat
    main_obs.write_to_filterbank(main_obs_name[:-4]+'.fluxcal.fil')
    print('Finished: calibrated product written to ' + main_obs_name[:-4]+'.fluxcal.fil')

#end module