__all__ = ['point_source']
__doc__ = """
.. _Read
:mod:`Read` -- point_source Read interface
============================================

.. module:: Read
   :platform: Unix, Windows
   :synopsis: Provides :ref:`PseudoNetCDF` random access read for CAMx
              point_source files.  See PseudoNetCDF.sci_var.PseudoNetCDFFile
              for interface details
.. moduleauthor:: Barron Henderson <barronh@unc.edu>
"""

# Distribution packages
import unittest
import struct

# Site-Packages
from numpy import zeros, array

# This Package modules
from PseudoNetCDF.camxfiles.timetuple import timediff, timerange
from PseudoNetCDF.camxfiles.FortranFileUtil import OpenRecordFile, read_into
from PseudoNetCDF.camxfiles.FortranFileUtil import Int2Asc
from PseudoNetCDF.sci_var import PseudoNetCDFFile, PseudoNetCDFVariables


# for use in identifying uncaught nan
listnan = struct.unpack('>f', b'\xff\xc0\x00\x00')[0]
checkarray = zeros((1,), 'f')
checkarray[0] = listnan
array_nan = checkarray[0]


class point_source(PseudoNetCDFFile):
    """
    point_source provides a PseudoNetCDF interface for CAMx
    point_source files.  Where possible, the inteface follows
    IOAPI conventions (see www.baronams.com).

    ex:
        >>> point_source_path = 'camx_point_source.bin'
        >>> rows,cols = 65,83
        >>> point_sourcefile = point_source(point_source_path,rows,cols)
        >>> point_sourcefile.variables.keys()
        ['TFLAG', 'ETFLAG', 'TFLAG', 'XSTK', 'YSTK', 'HSTK', 'DSTK', 'TSTK',
         'VSTK', 'KCELL', 'FLOW', 'PLMHT', 'NSTKS', 'NO', 'NO2', ...]
        >>> tflag = point_sourcefile.variables['TFLAG']
        >>> tflag.dimensions
        ('TSTEP', 'VAR', 'DATE-TIME')
        >>> tflag[0,0,:]
        array([2005185,       0])
        >>> tflag[-1,0,:]
        array([2005185,  240000])
        >>> v = point_sourcefile.variables['XSTK']
        >>> v.dimensions
        ('NSTK',)
        >>> v.shape
        (38452,)
        >>> v = point_sourcefile.variables['NO2']
        >>> v.dimensions
        ('TSTEP', 'NSTK')
        >>> v.shape
        (25, 38452)
        >>> point_sourcefile.dimensions
        {'TSTEP': 25, 'NSTK': 38452}
    """

    emiss_hdr_fmt = "10i60i3ifif"
    grid_hdr_fmt = "ffiffffiiiiifff"
    cell_hdr_fmt = "iiii"
    time_hdr_fmt = "ifif"
    spc_fmt = "10i"
    nstk_hdr_fmt = "ii"
    padded_nstk_hdr_size = struct.calcsize("ii" + nstk_hdr_fmt)
    padded_time_hdr_size = struct.calcsize("ii" + time_hdr_fmt)
    stk_hdr_fmt = "ffffff"
    id_fmt = "i" + spc_fmt
    id_size = struct.calcsize(id_fmt)
    data_fmt = "f"
    stkprops = ['XSTK', 'YSTK', 'HSTK', 'DSTK', 'TSTK', 'VSTK']
    stktimeprops = ['KCELL', 'FLOW', 'PLMHT']

    def __init__(self, rf):
        """
        Initialization included reading the header and learning
        about the format.

        see __readheader and __gettimestep() for more info
        """
        self.rffile = OpenRecordFile(rf)
        self.padded_time_hdr_size = struct.calcsize(self.time_hdr_fmt + "ii")
        self.__readheader()
        self.__gettimestep()
        self.__gettimeprops()
        self.createDimension('TSTEP', self.time_step_count)
        self.createDimension('STK', self.nstk)
        varkeys = (['XSTK', 'YSTK', 'HSTK', 'DSTK', 'TSTK', 'VSTK',
                    'KCELL', 'FLOW', 'PLMHT'] +
                   [i.strip() for i in self.spcnames])
        self.variables = PseudoNetCDFVariables(self.__var_get, varkeys)

    def __var_get(self, key):
        values = self.__variables(key)
        if key in self.stkprops:
            var = self.createVariable(key, 'f', ('STK',))
        else:
            var = self.createVariable(key, 'f', ('TSTEP', 'STK'))
        var[:] = values
        setattr(var, 'notread', 1)
        return var

    def __variables(self, k):
        if k in self.stkprops:
            return array(self.stk_props)[:, self.stkprops.index(k)]
        elif k in self.stktimeprops:
            stkps = array(self.stk_time_props)[:, :, 2:]
            return stkps[:, :, self.stktimeprops.index(k)]
        else:
            return self.getArray()[:, self.spcnames.index(k.ljust(10)), :]

    def header(self):
        rdum = 0.
        idum = 0
        ione = 1
        return [
            [self.name, self.note, ione, self.nspec, self.start_date,
                self.start_time, self.end_date, self.end_time],
            [rdum, rdum, self.iutm, self.xorg, self.yorg, self.delx, self.dely,
                self.nx, self.ny, self.nz, idum, idum, rdum, rdum, rdum],
            [ione, ione, self.nx, self.ny],
            self.spcnames,
            [ione, self.nstk],
            self.stk_props,
            self.stk_time_props
        ]

    def __readheader(self):
        """
        __readheader reads the header section of the ipr file
        it initializes each header field (see CAMx Users Manual for a list)
        as properties of the ipr class
        """
        vals = self.rffile.read(self.emiss_hdr_fmt)
        self.name = vals[0:10]
        self.note = vals[10:70]
        ione = vals[70]
        self.nspec = vals[71]
        self.start_date = vals[72]
        self.start_time = vals[73]
        self.end_date = vals[74]
        self.end_time = vals[75]

        vals = self.rffile.read(self.grid_hdr_fmt)
        rdum, rdum, self.iutm = vals[0:3]
        self.xorg, self.yorg, self.delx, self.dely = vals[3:7]
        self.nx, self.ny, self.nz = vals[7:10]
        idum, idum, rdum, rdum, rdum = vals[10:]
        if self.nz == 0:
            # Special case of gridded emissions
            # Seems to be same as avrg
            self.nlayers = 1
        else:
            self.nlayers = self.nz
        ione, ione, nx, ny = self.rffile.read(self.cell_hdr_fmt)
        if not (self.nx, self.ny) == (nx, ny):
            raise ValueError(("nx, ny defined first as %i, %i and then as " +
                              "%i, %i") % (self.nx, self.ny, nx, ny))
        species_temp = self.rffile.read(self.nspec * self.spc_fmt)
        self.spcnames = []
        for i in range(0, self.nspec * 10, 10):
            self.spcnames.append(Int2Asc(species_temp[i:i + 10]))

        ione, self.nstk = self.rffile.read(self.nstk_hdr_fmt)

        stkprms = zeros((self.nstk * len(self.stk_hdr_fmt),), 'f')
        read_into(self.rffile, stkprms, '')
        self.rffile.next()
        # self.rffile.previous()
        # self.tmplist=self.rffile.read('ffffff' * self.nstk)

        stkprms = stkprms.reshape((self.nstk, len(self.stk_hdr_fmt)))
        for i in range(stkprms.shape[0]):
            if stkprms[i, -1] == array_nan:
                stkprms[i, -1] = float('-nan')
        self.stk_props = stkprms.tolist()
        self.data_start_byte = self.rffile.record_start
        self.start_date, self.start_time, end_date, end_time = \
            self.rffile.read(self.time_hdr_fmt)

        self.time_step = timediff(
            (self.start_date, self.start_time), (end_date, end_time))
        # self.end_time += self.time_step
        mydayhrs = (2400, 24)[int(self.time_step % 2)]
        self.time_step_count = int(timediff((self.start_date, self.start_time),
                                            (self.end_date, self.end_time),
                                            mydayhrs) / self.time_step)

        self.stk_time_prop_fmt = "" + ("iiiff" * self.nstk)
        self.padded_stk_time_prop_size = struct.calcsize(
            "ii" + self.stk_time_prop_fmt)

        self.record_fmt = ("i10i") + self.data_fmt * (self.nstk)
        self.record_size = struct.calcsize(self.record_fmt)
        self.padded_size = self.record_size + 8

    def __gettimestep(self):
        """
        this is taken care of in the readheader routine
        record format provides start and end for each hour,
        which translates to t1 and t2
        """
        pass

    def __gettimeprops(self):
        self.stk_time_props = []
        dates = timerange((self.start_date, self.start_time),
                          (self.end_date, self.end_time),
                          self.time_step,
                          (2400, 24)[int(self.time_step % 2)])
        for ti, (d, t) in enumerate(dates):
            tmpprop = zeros((len(self.stk_time_prop_fmt)), 'f')
            tmpprop[...] = self.seekandread(
                d, t, 1, True, self.stk_time_prop_fmt)
            tmpprop = tmpprop.reshape(self.nstk, 5)
            for i in range(tmpprop.shape[0]):
                if tmpprop[i, -2] == array_nan:
                    tmpprop[i, -2] = float('-nan')

            self.stk_time_props.append(tmpprop.tolist())

    def __timerecords(self, dt):
        """
        Calculate the number of records to increment to reach time (d,t)
        """
        d, t = dt
        nsteps = int(timediff((self.start_date, self.start_time),
                              (d, t), (2400, 24)[int(self.time_step % 2)]))
        nspec = self.__spcrecords(self.nspec + 1)
        return nsteps * (nspec)

    def __spcrecords(self, spc):
        """
        Calculated number of records before spc
        """

        return spc - 1

    def __recordposition(self, date, time, spc, offset=False):
        """
        Use time (d,t), spc, and k to calculate number of records before
        desired record

        date - integer julian
        time - float
        spc - integer
        """
        ntime = self.__timerecords((date, time))
        nhdr = ((ntime // self.__spcrecords(self.nspec + 1)) + 1)
        nspc = self.__spcrecords(spc)
        noffset = -abs(int(offset))
        byte = self.data_start_byte
        byte += nhdr * (self.padded_time_hdr_size +
                        self.padded_nstk_hdr_size +
                        self.padded_stk_time_prop_size)
        byte += (ntime + nspc) * self.padded_size
        byte += noffset * self.padded_stk_time_prop_size
        return byte

    def seek(self, date=None, time=None, spc=-1, offset=False):
        """
        Move file cursor to the beginning of the specified record
        see __recordposition for parameter definitions
        """
        seekto = self.__recordposition(date, time, spc, offset)
        self.rffile._newrecord(seekto)

    def read(self, fmt=None):
        """
        Provide direct access to record file read
        """
        if fmt is None:
            fmt = self.record_fmt
        return self.rffile.read(fmt)

    def read_into(self, dest):
        """
        Transfer values from current record to dest
        dest - numeric or numpy array
        """

        return read_into(self.rffile, dest, self.id_fmt, self.data_fmt)

    def seekandreadinto(self, dest, date=None, time=None, spc=1):
        """
        see seek and read_into
        """

        self.seek(date, time, spc)
        self.read_into(dest)

    def seekandread(self, date=None, time=None, spc=1, offset=False, fmt=None):
        """
        see seek and read
        """
        self.seek(date, time, spc, offset)
        return self.read(fmt)

    def values(self):
        for d, t, spc in self.__iter__():
            yield self.seekandread(d, t, spc)

    def items(self):
        for d, t, spc in self.__iter__():
            yield d, t, spc, self.seekandread(d, t, spc)

    def keys(self):
        for ti, (d, t) in enumerate(self.timerange()):
            for spc in range(1, len(self.spcnames) + 1):
                yield d, t, spc

    __iter__ = keys

    def getArray(self):
        a = zeros((self.time_step_count, self.nspec, self.nstk), 'f')
        for ti, (d, t) in enumerate(self.timerange()):
            for spc in range(1, len(self.spcnames) + 1):
                self.seekandreadinto(a[ti, spc - 1, ...], d, t, spc)
        return a.copy()

    def timerange(self):
        return timerange((self.start_date, self.start_time),
                         (self.end_date, self.end_time), self.time_step,
                         eod=24)


class TestRead(unittest.TestCase):
    def runTest(self):
        pass

    def setUp(self):
        pass

    def testPT(self):
        import PseudoNetCDF.testcase
        emissfile = point_source(
            PseudoNetCDF.testcase.camxfiles_paths['point_source'])
        v = emissfile.variables['NO2']
        self.assertTrue((v[:] == array(
            [0.00000000e+00, 3.12931000e+02, 1.23599997e+01, 0.00000000e+00,
             5.27999992e+01, 0.00000000e+00, 3.12931000e+02, 1.23599997e+01,
             0.00000000e+00, 5.27999992e+01], dtype='f').reshape(2, 5)).all())


if __name__ == '__main__':
    unittest.main()
