"""Input-output for PyConTurb files
"""
import os
from struct import pack, unpack
import time

import numpy as np
import pandas as pd

from pyconturb._version import __version__


_HAWC2_BIN_FMT = '<f'  # HAWC2 binary turbulence datatype
_TS_BIN_FMT = '<h4l12fl'  # TurbSim binary format


def bts_to_df(path):
    """Load TurbSim-style .bts file to pyconturb dataframe"""
    if path.endswith('.bts'):  # remove file extension, will be added back later
        path = path[:-4]
    u_scl = np.zeros(3, np.float32)
    u_off = np.zeros(3, np.float32)
    with open(path + '.bts', 'rb') as fl:
        (_,  # periodic (7) or non periodic (8), unused
         nz,
         ny,
         _,  # no. tower ptsbelow box, unused
         nt,
         dz,
         dy,
         dt,
         uhub,
         zhub,
         z0,
         u_scl[0],
         u_off[0],
         u_scl[1],
         u_off[1],
         u_scl[2],
         u_off[2],
         strlen) = unpack(_TS_BIN_FMT, fl.read(70))
        desc_str = fl.read(strlen)  # skip these bytes
        nbt = 3 * ny * nz * nt
        dat = np.rollaxis(np.frombuffer(fl.read(2 * nbt), dtype=np.int16).astype(
            np.float32).reshape([3, ny, nz, nt], order='F'), 2, 1)  # 3 x ny x nz x nt
    dat -= u_off[:, None, None, None]
    dat /= u_scl[:, None, None, None]
    # create pyconturb-flavor dataframe
    t = np.arange(nt) * dt
    turb_df = pd.DataFrame(index=t)
    for i, c in enumerate('uvw'):
        arr = np.transpose(dat[i], (2, 0, 1))  # nt x ny x nz
        comp_df = pd.DataFrame(arr.reshape(nt, ny*nz),
                               index=t).add_prefix(f'{c}_p')
        turb_df = turb_df.join(comp_df, how='outer')
    return turb_df


def df_to_bts(turb_df, spat_df, path, uzhub=None, periodic=True):
    """pyconturb-style turbulence dataframe to TurbSim-style binary file
    Code modified based on `turbsim` in PyTurbSim:
        https://github.com/lkilcher/pyTurbSim/blob/master/pyts/io/write.py

    Notes
    -----
    * The turbulence must have been generated on a y-z grid.
    * The naming convention must be 'u_p0', 'v_p0, 'w_p0', 'u_p1', etc.,
        where the point indices proceed vertically along z before horizontally
        along y.
    """
    if path.endswith('.bts'):  # remove file extension, will be added back later
        path = path[:-4]
    # format-specific constants
    intmin = -32768  # minimum integer
    intrng = 65535  # range of integers
    # calculate intermediate parameters
    y = np.sort(np.unique(spat_df.loc['y'].values))
    z = np.sort(np.unique(spat_df.loc['z'].values))
    ny = y.size  # no. of y points in grid
    nz = z.size  # no. of z points in grif
    nt = turb_df.shape[0]  # no. of time steps
    if y.size == 1:
        dy = 0
    else:
        dy = np.mean(y[1:] - y[:-1])  # hopefully will reduce possible errors
    if z.size == 1:
        dz = 0
    else:
        dz = np.mean(z[1:] - z[:-1])  # hopefully will reduce possible errors
    dt = turb_df.index[-1] / (turb_df.shape[0] - 1)  # time step
    if uzhub is None:  # default is center of grid
        zhub = z[z.size // 2]  # halfway up
        u_df = turb_df.filter(regex=f'u_', axis=1).values
        uhub = u_df[:, u_df.shape[1] // 2].mean()  # mean of center of grid
    else:
        uhub, zhub = uzhub
    # convert pyconturb dataframe to pyturbsim format (3 x ny x nz x nt)
    ts = np.empty((3, ny, nz, nt))
    for i, c in enumerate('uvw'):
        arr = turb_df.filter(regex=f'{c}_', axis=1).values.reshape((nt, ny, nz))
        ts[i] = np.transpose(arr, (1, 2, 0))  # reorder to ny, nz, nt
    # initialize output arrays
    u_minmax = np.empty((3, 2), dtype=np.float32)  # mins and maxs of time series
    u_off = np.empty((3), dtype=np.float32)  # offsets of each time series
    u_scl = np.empty((3), dtype=np.float32)  # scales of each time series
    desc_str = 'generated by PyConTurb v%s, %s.' % (
        __version__,
        time.strftime('%b %d, %Y, %H:%M (%Z)', time.localtime()))  # description
    # calculate the scales and offsets of each time series
    out = np.empty((3, ny, nz, nt), dtype=np.int16)
    for ind in range(3):
        u_minmax[ind] = ts[ind].min(), ts[ind].max()
        if np.isclose(u_minmax[ind][0], u_minmax[ind][1]):
            u_scl[ind] = 1
        else:
            u_scl[ind] = intrng / np.diff(u_minmax[ind])
        u_off[ind] = intmin - u_scl[ind] * u_minmax[ind, 0]
        out[ind] = (ts[ind] * u_scl[ind] + u_off[ind]).astype(np.int16)
    with open(path + '.bts', 'wb') as fl:
        # write the header
        fl.write(pack(_TS_BIN_FMT,
                      [7, 8][periodic],  # 7 is not periodic, 8 is periodic
                      nz,
                      ny,
                      0,  # assuming 0 tower points below grid
                      nt,
                      dz,
                      dy,
                      dt,
                      uhub,
                      zhub,
                      z[0],
                      u_scl[0],
                      u_off[0],
                      u_scl[1],
                      u_off[1],
                      u_scl[2],
                      u_off[2],
                      len(desc_str)))
        fl.write(desc_str.encode(encoding='UTF-8'))
        # Swap the y and z indices so that fortran-order writing agrees with the file format.
        # Also, we swap the order of z-axis to agree with the file format.
        # Write the data so that the first index varies fastest (F order).
        # The indexes vary in the following order:
        # component (fastest), y-index, z-index, time (slowest).
        fl.write(np.rollaxis(out, 2, 1).tobytes(order='F'))


def df_to_h2turb(turb_df, spat_df, path, prefix=''):
    """pyconturb-style turbulence dataframe to binary files for hawc2

    Notes
    -----
    * The turbulence must have been generated on a y-z grid.
    * The naming convention must be 'u_p0', 'v_p0, 'w_p0', 'u_p1', etc.,
       where the point indices proceed vertically along z before horizontally
       along y.
    """
    nx = turb_df.shape[0]  # turbulence dimensions for reshaping
    ny = len(set(spat_df.loc['y'].values))
    nz = len(set(spat_df.loc['z'].values))
    # make and save binary files for all three components
    for c in 'uvw':
        arr = turb_df.filter(regex=f'{c}_', axis=1).values.reshape((nx, ny, nz))
        bin_path = os.path.join(path, f'{prefix}{c}.bin')
        with open(bin_path, 'wb') as bin_fid:
            arr.astype(np.dtype(_HAWC2_BIN_FMT)).tofile(bin_fid)
    return


def h2turb_to_arr(spat_df, path):
    """Raw-load a hawc2 turbulent binary file to float-64 numeric array"""
    ny, nz = pd.unique(spat_df.loc['y']).size, pd.unique(spat_df.loc['z']).size
    bin_arr = np.fromfile(path, dtype=np.dtype(_HAWC2_BIN_FMT))
    nx = bin_arr.size // (ny * nz)
    if (nx * ny * nz) != bin_arr.size:
        raise ValueError('Binary file size does not match spat_df!')
    bin_arr.shape = (nx, ny, nz)
    return bin_arr


def h2turb_to_df(spat_df, path, nt=600, dt=1, prefix=''):
    """load a hawc2 binary file into a float-64 pandas datafram with transform to uvw"""
    t = np.arange(nt) * dt
    turb_df = pd.DataFrame(index=t)
    for c in 'uvw':
        comp_path = os.path.join(path, f'{prefix}{c}.bin')
        arr = h2turb_to_arr(spat_df, comp_path)
        nx, ny, nz = arr.shape
        comp_df = pd.DataFrame(arr.reshape(nx, ny*nz),
                               index=t).add_prefix(f'{c}_p')
        turb_df = turb_df.join(comp_df, how='outer')
    turb_df = turb_df[[f'{c}_p{i}' for i in range(ny*nz) for c in 'uvw']]
    return turb_df
