import base64
import zlib

import numpy as np

import six

if six.PY2:
    decode_base64 = base64.decodestring
else:
    decode_base64 = base64.decodebytes


COMPRESSION_NONE = 'none'
COMPRESSION_ZLIB = 'zlib'
COMPRESSION_LINEAR = 'linear'
COMPRESSION_DELTA = 'delta'


encoding_map = {
    32: np.float32,
    64: np.float64,
    '64-bit integer': np.int64,
    'MS:1000522': np.int64,
    'MS:1000519': np.int32,
    '32-bit integer': np.int32,
    'MS:1000520': np.float16,
    '16-bit float': np.float16,
    'MS:1000521': np.float32,
    '32-bit float': np.float32,
    'MS:1000523': np.float64,
    '64-bit float': np.float64,
    'MS:1001479': np.bytes_,
    'null-terminated ASCII string': np.bytes_,
    float: np.float64,
    int: np.int32,
    None: np.float32
}


dtype_to_encoding = {
    np.int64: '64-bit integer',
    np.int32: '32-bit integer',
    np.float16: '16-bit float',
    np.float32: '32-bit float',
    np.float64: '64-bit float',
    np.bytes_: 'null-terminated ASCII string',
}


compression_map = {
    COMPRESSION_ZLIB: "zlib compression",
    COMPRESSION_NONE: 'no compression',
    None: 'no compression',
    False: 'no compression',
    True: "zlib compression",
    COMPRESSION_LINEAR: "linear prediction",
    COMPRESSION_DELTA: "delta prediction"
}


for dtype in list(encoding_map.values()):
    encoding_map[dtype] = dtype


def encode_array(array, compression=COMPRESSION_NONE, dtype=np.float32):
    if compression == COMPRESSION_LINEAR or compression == COMPRESSION_DELTA:
        array = np.asanyarray(array).astype(dtype, copy=False)
        if compression == COMPRESSION_LINEAR:
            array = linear_encode(array, copy=True)
        else:
            array = delta_encode(array, copy=True)
        bytestring = array.tobytes()
        # Now don't apply any additional encoding to the generated bytestring
        compression = COMPRESSION_NONE
    else:
        bytestring = np.asanyarray(array).astype(dtype, copy=False).tobytes()
    if compression == COMPRESSION_NONE:
        bytestring = bytestring
    elif compression == COMPRESSION_ZLIB:
        bytestring = zlib.compress(bytestring)
    else:
        raise ValueError("Unknown compression: %s" % compression)
    encoded_string = base64.standard_b64encode(bytestring)
    return encoded_string


def encode_array_direct(array, compression=COMPRESSION_NONE, dtype=np.float32):
    array = np.asanyarray(array).astype(dtype, copy=False)
    if compression == COMPRESSION_LINEAR:
        array = linear_encode(array, copy=True)
    elif compression == COMPRESSION_DELTA:
        array = delta_encode(array, copy=True)
    return array


def decode_array(bytestring, compression=COMPRESSION_NONE, dtype=np.float32):
    try:
        decoded_string = bytestring.encode("ascii")
    except AttributeError:
        decoded_string = bytestring
    decoded_string = decode_base64(decoded_string)
    if compression == COMPRESSION_NONE:
        decoded_string = decoded_string
    elif compression == COMPRESSION_ZLIB:
        decoded_string = zlib.decompress(decoded_string)
    else:
        raise ValueError("Unknown compression: %s" % compression)
    array = np.frombuffer(decoded_string, dtype=dtype)
    return array


def delta_predict(data, copy=True):
    '''Reverse the lossy transformation of the delta compression
    helper.

    Parameters
    ----------
    data : :class:`numpy.ndarray`
        The data to transform
    copy : bool
        Whether to make a copy of the data array or transform it in-place.

    Returns
    -------
    :class:`numpy.ndarray`
        The transformed data array
    '''
    if copy:
        out = data.copy()
    else:
        out = data
    for i in range(2, len(data)):
        out[i] = out[i] + out[i - 1] - out[0]
    return out


def linear_predict(data, copy=True):
    '''Reverse the lossy transformation of the linear interpolation compression
    helper.

    Parameters
    ----------
    data : :class:`numpy.ndarray`
        The data to transform
    copy : bool
        Whether to make a copy of the data array or transform it in-place.

    Returns
    -------
    :class:`numpy.ndarray`
        The transformed data array
    '''
    if copy:
        out = data.copy()
    else:
        out = data
    for i in range(2, len(data)):
        out[i] = out[i] + 2 * out[i - 1] - out[i - 2] - out[1]
    return out


def linear_encode(data, copy=True):
    if copy:
        out = data.copy()
    else:
        out = data
    n = len(data)
    if n < 3:
        return data
    prev2 = out[0]
    prev1 = out[1]
    offset = out[1]
    for i in range(2, n):
        out[i] = offset + out[i] - 2 * prev1 + prev2
        tmp = prev1
        prev1 = out[i] + 2 * prev1 - prev2 - offset
        prev2 = tmp
    return out


def delta_encode(data, copy=True):
    if copy:
        out = data.copy()
    else:
        out = data
    n = len(data)
    if n < 2:
        return data
    prev = out[0]
    offset = out[0]
    for i in range(1, n):
        tmp = out[i]
        out[i] = offset + out[i] - prev
        prev = tmp
    return out
