# -*- coding: utf-8 -*-

# Python3 compatiability
from __future__ import division
import sys
if sys.version_info > (3,):
    long = int
    
"""
Module to write VDIF frames.  The implementation of this module is similar
to that of lsl.sim.tbw in that the primary element defined in this module is
a Frame object which as attribute functions that can create a numpy 
representation of the raw frame and write that raw frame to an open file-
handle.

.. seealso::
    :mod:`lsl.sim.tbw`
"""

import numpy
import ephem

from lsl.common import dp as dp_common
import lsl.astro as astro

__version__ = '0.1'
__revision__ = '$Rev: 2577 $'
__all__ = ['Frame', '__version__', '__revision__', '__all__']

vdifEpoch = ephem.Date('2000/01/01 00:00:00.00')
unixEpoch = ephem.Date('1970/01/01 00:00:00.00')


class Frame(object):
    """
    Object to create and write a VDIF (VLBI Data Interchange Format) 
    frame (version 1, June 26, 2009).
    """
    
    def __init__(self, stand=0, time=0, bits=8, data=None, sampleRate=dp_common.fS):
        self.stand = stand
        self.time = time
        self.bits = bits
        self.data = numpy.squeeze(data)
        self.sampleRate = sampleRate
        if data.dtype.kind == 'c':
            self.dataReal = False
        else:
            self.dataReal = True
            
        self.seconds = 0
        self.frame = 0
        self.epoch = 0
        
        # Validate
        if self.bits not in (1, 2, 4, 8):
            raise ValueError("Unsupported output bitdepth: %i" % self.bits)
                    
        # Convert the time from UNIX to epoch and make the data ready to 
        # be written to the disk.
        self.__dataAdjust()
        self.__setEpoch()
        
    def __dataAdjust(self):
        """
        Adjust the data to the number of bits required for the frame 
        and convert to unsigned 32-bit integers.  Also, interlace complex
        arrays.
        """
        
        if self.dataReal:
            interlaced = self.data
        else:
            interlaced = numpy.zeros(2*len(self.data))
            interlaced[0::2] = self.data.real
            interlaced[1::2] = self.data.imag
        
        biased = interlaced + 2**(self.bits-1)
        biased = biased.astype(numpy.uint32) & (2**self.bits-1)
        
        self.data = biased
        
    def __setEpoch(self):
        """
        Convert a UNIX timestap in seconds since 1970 to seconds since
        January 1st, 2000 using ephem.Data objects.  Also, calculate the 
        frame number since the beginning of the second.
        """
        
        # UNIX time to seconds since DJD  = 0
        curEpoch = float(unixEpoch)*astro.SECS_IN_DAY + self.time
        self.epoch = 0
        # Seconds since the VDIF epoch
        epochSeconds = curEpoch - float(vdifEpoch)*astro.SECS_IN_DAY
        # Integer seconds
        self.seconds = long(epochSeconds)
        
        # Compute the frames since the beginning of the second
        frame = (epochSeconds - self.seconds) * (self.sampleRate/len(self.data))
        self.frame = long(round(frame))

    def createRawFrame(self):
        """
        Using the data and information stored in the object, create a
        numpy array on unsigned 1-byte integers that represents the frame.
        """
        
        # Find out how many samples can be packed into a work and build
        # the corresponding numpy array.
        samplesPerWord = int(32 // self.bits)
        raw = numpy.zeros(32 + 4*len(self.data)//samplesPerWord, dtype=numpy.uint8)

        # Valid data, standard (not legacy) 32-bit header, and seconds since 
        # the 01/01/2000 epoch.
        raw[3] = (0 << 7) | (0 << 6) | ((self.seconds >> 24) & 63)
        raw[2] = (self.seconds >> 16) & 255
        raw[1] = (self.seconds >> 8) & 255
        raw[0] = self.seconds & 255

        # Refernce epoch (0 == 01/01/2000) and frame count
        raw[7] = self.epoch & 63
        raw[6] = (self.frame >> 16) & 255
        raw[5] = (self.frame >> 8) & 255
        raw[4] = self.frame & 255

        # VDIF version number, number of channels (just 1), and data frame 
        # length in units to 8-bytes (8 raw array elements)
        raw[11] = (1 << 6) | (0 & 31)
        raw[10] = ((len(raw) // 8) >> 16) & 255
        raw[9] = ((len(raw) // 8) >> 8) & 255
        raw[8] = (len(raw) // 8) & 255

        # Data type, bits per sample, thread ID, and station ID
        # NB:  The thread ID is fixed at `0'
        if self.dataReal:
            raw[15] = (0 << 7) | ((self.bits-1 & 31) << 2) | ((0 >> 8) & 3)
        else:
            raw[15] = (1 << 7) | ((self.bits-1 & 31) << 2) | ((0 >> 8) & 3)
        raw[14] = 0 & 255
        raw[13] = (self.stand >> 8) & 255
        raw[12] = self.stand & 255
        
        # User-defined words 4 through 7 (array entries 16 to 31)
        ## NICT format so that we can put in the sample rate
        sampleRatekHz = int(self.sampleRate / 1000.0)
        raw[19] = 0x01
        raw[18] = (0 << 7) | ((sampleRatekHz >> 16) & 0xFF)
        raw[17] = (sampleRatekHz >> 8) & 0xFF
        raw[16] = sampleRatekHz & 0xFF
        raw[20:24] = 0
        raw[24:28] = 0
        raw[28:32] = 0
        
        # Data values
        for f in range(32, len(raw), 4):
            i = (f - 32) // 4
            word = 0
            for p in range(samplesPerWord):
                word |= (self.data[i*samplesPerWord + p] << (self.bits*p))
            raw[f+3] = (word >> 24) & 255
            raw[f+2] = (word >> 16) & 255
            raw[f+1] = (word >> 8) & 255
            raw[f+0] = word & 255
            
        return raw

    def writeRawFrame(self, fh):
        """
        Create a numpy representation of the VDIF frame and then write 
        it to the specified file handle.
        """
        
        rawFrame = self.createRawFrame()
        rawFrame.tofile(fh)
