__all__ = ['write_point']
__doc__ = """
.. _Write
:mod:`Write` -- CAMx point_source writer
============================================

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

# Distribution packages
import operator
from functools import reduce

# Site-Packages
import numpy as np

# This Package modules
from PseudoNetCDF.camxfiles.timetuple import timeadd, timerange
from PseudoNetCDF.camxfiles.FortranFileUtil import writeline, Asc2Int
from PseudoNetCDF._getwriter import registerwriter

_emiss_hdr_fmt = np.dtype(dict(
    names=['SPAD', 'name', 'note', 'itzon', 'nspec', 'ibdate', 'btime',
           'iedate', 'etime', 'EPAD'],
    formats=['>i', '(10,4)>S1', '(60,4)>S1', '>i', '>i', '>i', '>f', '>i',
             '>f', '>i']))

_grid_hdr_fmt = np.dtype(dict(
    names=['SPAD', 'plon', 'plat', 'iutm', 'xorg', 'yorg', 'delx', 'dely',
           'nx', 'ny', 'nz', 'iproj', 'istag', 'tlat1', 'tlat2', 'rdum5',
           'EPAD'],
    formats=['>i', '>f', '>f', '>i', '>f', '>f', '>f', '>f', '>i', '>i',
             '>i', '>i', '>i', '>f', '>f', '>f', '>i']))

_cell_hdr_fmt = np.dtype(dict(
    names=['SPAD', 'ione1', 'ione2', 'nx', 'ny', 'EPAD'],
    formats=['>i', '>i', '>i', '>i', '>i', '>i']))

_time_hdr_fmt = np.dtype(dict(
    names=['SPAD', 'ibdate', 'btime', 'iedate', 'etime', 'EPAD'],
    formats=['>i', '>i', '>f', '>i', '>f', '>i']))


_nstk_hdr_fmt = np.dtype(dict(
    names=['SPAD', 'ione', 'nstk', 'EPAD'],
    formats=['>i', '>i', '>i', '>i']))


def ncf2point_source(ncffile, outpath):
    outfile = open(outpath, 'wb')
    emiss_hdr = np.zeros((1,), dtype=_emiss_hdr_fmt)
    emiss_hdr['SPAD'] = emiss_hdr['EPAD'] = emiss_hdr.dtype.itemsize - 8
    emiss_hdr[0]['name'][:, :] = ' '
    emiss_hdr[0]['name'][:, 0] = np.array(
        [c for c in ncffile.NAME], dtype='>S1')
    emiss_hdr[0]['note'][:, :] = ' '
    emiss_hdr[0]['note'][:, 0] = np.array(
        [c for c in ncffile.NOTE], dtype='>S1')
    # gdtype = getattr(ncffile, 'GDTYPE', -999)
    emiss_hdr['itzon'][0] = ncffile.ITZON
    tmps1 = np.array([c for c in ncffile.SPC_NAMES], dtype='S1')
    tmps1 = tmps1.reshape(-1, 16)[:, :10].copy().view('>S10')[:, 0]
    SPC_NAMES = [s.decode() for s in tmps1]
    nspec = len(SPC_NAMES)
    emiss_hdr['nspec'] = nspec

    grid_hdr = np.zeros((1,), dtype=_grid_hdr_fmt)
    grid_hdr['SPAD'] = grid_hdr['EPAD'] = grid_hdr.dtype.itemsize - 8
    grid_hdr['plon'] = ncffile.PLON
    grid_hdr['plat'] = ncffile.PLAT
    grid_hdr['iutm'][0] = ncffile.IUTM
    grid_hdr['xorg'] = ncffile.XORIG
    grid_hdr['yorg'] = ncffile.YORIG
    grid_hdr['delx'] = ncffile.XCELL
    grid_hdr['dely'] = ncffile.YCELL
    grid_hdr['nx'] = ncffile.NCOLS
    grid_hdr['ny'] = ncffile.NROWS
    grid_hdr['nz'] = ncffile.NLAYS
    grid_hdr['iproj'] = ncffile.CPROJ
    grid_hdr['tlat1'] = ncffile.TLAT1
    grid_hdr['tlat2'] = ncffile.TLAT2
    grid_hdr['istag'] = ncffile.ISTAG
    grid_hdr['rdum5'] = 100.

    cell_hdr = np.zeros((1,), dtype=_cell_hdr_fmt)
    cell_hdr['SPAD'] = cell_hdr['EPAD'] = cell_hdr.dtype.itemsize - 8
    cell_hdr['ione1'] = 1
    cell_hdr['ione2'] = 1
    cell_hdr['nx'] = ncffile.NCOLS
    cell_hdr['ny'] = ncffile.NROWS
    time_hdr = np.zeros(
        shape=(len(ncffile.dimensions['TSTEP']),), dtype=_time_hdr_fmt)
    time_hdr['SPAD'] = time_hdr['EPAD'] = time_hdr.dtype.itemsize - 8
    date, time = ncffile.variables['TFLAG'][:, 0].T
    edate, etime = ncffile.variables['ETFLAG'][:, 0].T
    time = time.astype('>f') / 10000.
    date = date % (date // 100000 * 100000)
    etime = etime.astype('>f') / 10000.
    edate = edate % (edate // 100000 * 100000)
    time_hdr['ibdate'] = date
    time_hdr['btime'] = time
    time_hdr['iedate'] = edate
    time_hdr['etime'] = etime
    # time_hdr['iedate'] += time_hdr['etime'] // 24
    # time_hdr['etime'] -= (time_hdr['etime'] // 24) * 24

    emiss_hdr['ibdate'] = time_hdr[0]['ibdate']
    emiss_hdr['btime'] = time_hdr[0]['btime']
    emiss_hdr['iedate'] = ncffile.EDATE
    emiss_hdr['etime'] = ncffile.ETIME

    nspec = emiss_hdr['nspec'][0]
    _spc_hdr_fmt = np.dtype(dict(names=['SPAD', 'NAME', 'EPAD'], formats=[
                            '>i', '(%d,10,4)>S1' % nspec, '>i']))
    spc_hdr = np.zeros((1,), dtype=_spc_hdr_fmt)
    spc_hdr['NAME'][:] = ' '
    spc_hdr['SPAD'] = spc_hdr['EPAD'] = spc_hdr.dtype.itemsize - 8
    for si, spc in enumerate(SPC_NAMES):
        spc_hdr[0]['NAME'][si, :, 0] = np.array(
            [c for c in spc.ljust(10)], dtype='>S1')

    nstk_hdr = np.ones((1,), dtype=_nstk_hdr_fmt)
    nstk_hdr['SPAD'] = nstk_hdr['EPAD'] = nstk_hdr.dtype.itemsize - 8
    nstk_hdr['nstk'] = len(ncffile.dimensions['NSTK'])

    nstk = nstk_hdr['nstk'][0]
    # (xstk(n),ystk(n),hstk(n),dstk(n),tstk(n),vstk(n),n=1,nstk)
    stk_prop_fmt = np.dtype(dict(
        names=['xstk', 'ystk', 'hstk', 'dstk', 'tstk', 'vstk'],
        formats=['>f', '>f', '>f', '>f', '>f', '>f']))
    stk_prop = np.ones((1,), dtype=np.dtype([
        ('SPAD', '>i', (1,)), ('DATA', stk_prop_fmt, nstk),
        ('EPAD', '>i', (1,))
    ]))
    stk_prop['SPAD'] = stk_prop['EPAD'] = stk_prop.dtype.itemsize - 8
    stk_prop['DATA']['xstk'] = ncffile.variables['XSTK'][:]
    stk_prop['DATA']['ystk'] = ncffile.variables['YSTK'][:]
    stk_prop['DATA']['hstk'] = ncffile.variables['HSTK'][:]
    stk_prop['DATA']['dstk'] = ncffile.variables['DSTK'][:]
    stk_prop['DATA']['tstk'] = ncffile.variables['TSTK'][:]
    stk_prop['DATA']['vstk'] = ncffile.variables['VSTK'][:]

    # Start writing
    emiss_hdr.tofile(outfile)
    grid_hdr.tofile(outfile)
    cell_hdr.tofile(outfile)
    spc_hdr.tofile(outfile)
    nstk_hdr.tofile(outfile)
    spc_names = [spc[:, 0].copy().view('>S10')[0].decode().strip()
                 for spc in spc_hdr[0]['NAME']]
    stk_prop.tofile(outfile)
    ione = np.array([1]).astype('>i')
    for di, (d, t) in enumerate(ncffile.variables['TFLAG'][:, 0]):
        tempout = b''
        tempout += time_hdr[di].tobytes()
        tempout += np.array([8, 1, nstk, 8], dtype='>i').tobytes()
        # (idum,idum,kcell(n),flow(n),plmht(n),n=1,nstk)
        time_prop_fmt = np.dtype(dict(
            names=['ione1', 'ione2', 'kcell', 'flow', 'plmht'],
            formats=['>i', '>i', '>i', '>f', '>f']))
        props = np.ones((1,), dtype=np.dtype(
            [('SPAD', '>i', (1,)),
             ('DATA', time_prop_fmt, (nstk,)),
             ('EPAD', '>i', (1,))]))
        props['SPAD'] = props['EPAD'] = props.dtype.itemsize - 8
        props['DATA']['ione1'] = ncffile.variables['IONE'][di]
        props['DATA']['ione2'] = ncffile.variables['ITWO'][di]
        props['DATA']['kcell'] = ncffile.variables['KCELL'][di]
        props['DATA']['flow'] = ncffile.variables['FLOW'][di]
        props['DATA']['plmht'] = ncffile.variables['PLMHT'][di]
        tempout += props.tobytes()
        for spc_key, spc_name in zip(spc_names, spc_hdr[0]['NAME']):
            var = ncffile.variables[spc_key]
            data = var[di].astype('>f')
            buf = np.array([4 + 40 + data.size * 4]).astype('>i')
            tempout += buf.tobytes()
            tempout += ione.tobytes()
            tempout += spc_name.tobytes()
            tempout += data.tobytes()
            tempout += buf.tobytes()
        outfile.write(tempout)
    outfile.flush()
    return outfile


registerwriter('camxfiles.point_source', ncf2point_source)
registerwriter('point_source', ncf2point_source)


def write_point(start_date, start_time, time_step, hdr, vals):
    # Internalize header
    hdr = [h for h in hdr]
    species = hdr[-4]
    nstk = hdr[-3][1]
    timeprops = hdr.pop()

    # initialize hdr_fmts with species count
    hdr_fmts = ["10i60i3ifif", "ffiffffiiiiifff",
                "iiii", "10i" * len(species), "ii", "ffffff" * nstk]

    # initialize output variable
    pt_string = ''

    # Change name and note
    hdr[0] = list(hdr[0][0]) + list(hdr[0][1]) + list(hdr[0][2:])

    # Reducing stk props
    stkprops = hdr[-1]
    hdr[-1] = []
    for stk in stkprops:
        hdr[-1].extend(stk)
    stk_time_prop_fmt = "iiiff" * nstk

    stk_time_props = []
    for time in timeprops:
        stk_time_props.append([])
        for stk in time:
            stk_time_props[-1].extend(stk)

    # Change species names to array of characters
    hdr[-3] = reduce(operator.concat, [Asc2Int(s) for s in hdr[-3]])

    # for each item in the header, write it to output
    for i, (h, f) in enumerate(zip(hdr, hdr_fmts)):
        pt_string += writeline(h, f, False)

    # create value format
    valfmt = 'i10i' + ('f' * nstk)

    # Get end date
    (end_date, end_time) = timeadd(
        (start_date, start_time), (0, time_step * vals.shape[0]))

    # Write out values
    for ti, (d, t) in enumerate(timerange((start_date, start_time),
                                          (end_date, end_time),
                                          time_step)):
        ed, et = timeadd((d, t), (0, time_step))
        pt_string += writeline((d, t, ed, et), 'ifif', False)
        pt_string += writeline((1, nstk), 'ii', False)
        pt_string += writeline(stk_time_props[ti], stk_time_prop_fmt, False)
        for si, spc in enumerate(species):
            # Dummy variable,spc characters and values flattened
            temp = [1]
            temp.extend(Asc2Int(spc))
            temp.extend(vals[ti, si, ...])
            pt_string += writeline(temp, valfmt, False)

    return pt_string
