import os
import base64
import zlib

try:
    from isal import isal_zlib as zlib
except ImportError:
    pass

try:
    import pyzstd
    zstd_compress = pyzstd.compress
    zstd_decompress = pyzstd.decompress
except ImportError:
    zstd_compress = zstd_decompress = None


import numpy as np

import six

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


try:
    if os.environ.get("PSIMS_NO_PYNUMPRESS"):
        pynumpress = False
    else:
        import pynumpress
except ImportError:
    pynumpress = None


COMPRESSION_NONE = 'none'
COMPRESSION_ZLIB = 'zlib'
COMPRESSION_ZSTD = 'zstd'
COMPRESSION_LINEAR = 'linear'
COMPRESSION_DELTA = 'delta'
COMPRESSION_NUMPRESS_LINEAR_PREDICTION = "MS-Numpress linear prediction compression"
COMPRESSION_NUMPRESS_POSITIVE_INTEGER = "MS-Numpress positive integer compression"
COMPRESSION_NUMPRESS_SHORT_LOGGED_FLOAT = "MS-Numpress short logged float compression"


class Namespace(object):
    def __init__(self, d):
        self.__dict__.update(d)

    def __getitem__(self, key):
        return self.__dict__[key]

    def __contains__(self, key):
        return key in self.__dict__

    def items(self):
        return self.__dict__.items()

    def keys(self):
        return self.__dict__.keys()

    def values(self):
        return self.__dict__.values()

    def __repr__(self):
        return "{self.__class__.__name__}({self.__dict__!r})".format(self=self)


compressors = Namespace({k.replace("COMPRESSION_", '').lower(): v for k,
               v in locals().items() if k.startswith("COMPRESSION_")})


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',
    COMPRESSION_ZSTD: 'zstd compression',
    'no compression': 'no compression',
    None: 'no compression',
    False: 'no compression',
    True: "zlib compression",
    COMPRESSION_LINEAR: "MS:1003090",
    COMPRESSION_DELTA: "MS:1003089"
}

if pynumpress is not None:
    compression_map[COMPRESSION_NUMPRESS_LINEAR_PREDICTION] = COMPRESSION_NUMPRESS_LINEAR_PREDICTION
    compression_map[COMPRESSION_NUMPRESS_POSITIVE_INTEGER] = COMPRESSION_NUMPRESS_POSITIVE_INTEGER
    compression_map[COMPRESSION_NUMPRESS_SHORT_LOGGED_FLOAT] = COMPRESSION_NUMPRESS_SHORT_LOGGED_FLOAT


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


def coerce_array(array, dtype):
    array = np.asanyarray(array).astype(dtype, copy=False)
    if dtype == np.bytes_:
        array = build_null_terminated_string_byte_array(array)
    return array


def build_null_terminated_string_byte_array(array):
    buff = bytearray()
    for t in array:
        buff.extend(t)
        buff.extend(b"\0")
    return np.asanyarray(buff).astype(dtype=np.uint8, copy=False)


def encode_array(array, compression=COMPRESSION_NONE, dtype=np.float32):
    if compression is None:
        compression = COMPRESSION_NONE
    if compression == COMPRESSION_LINEAR or compression == COMPRESSION_DELTA:
        array = coerce_array(array, dtype)
        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
    elif compression.startswith("MS-Numpress"):
        if pynumpress is None:
            raise ImportError("pynumpress is not installed, required for MS-Numpress compression")
        array = coerce_array(array, dtype)
        if compression == COMPRESSION_NUMPRESS_SHORT_LOGGED_FLOAT:
            fixed_point = pynumpress.optimal_slof_fixed_point(array)
            array = pynumpress.encode_slof(array, fixed_point)
        elif compression == COMPRESSION_NUMPRESS_POSITIVE_INTEGER:
            array = pynumpress.encode_pic(array)
        elif compression == COMPRESSION_NUMPRESS_LINEAR_PREDICTION:
            fixed_point = pynumpress.optimal_linear_fixed_point(array)
            array = pynumpress.encode_linear(array, fixed_point)
        bytestring = array.tobytes()
        # Now don't apply any additional encoding to the generated bytestring
        compression = COMPRESSION_NONE
    else:
        bytestring = coerce_array(array, dtype).tobytes()
    if compression == COMPRESSION_NONE:
        bytestring = bytestring
    elif compression == COMPRESSION_ZLIB:
        bytestring = zlib.compress(bytestring)
    elif compression == COMPRESSION_ZSTD:
        bytestring = zstd_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 = coerce_array(array, dtype)
    if compression == COMPRESSION_LINEAR:
        array = linear_encode(array, copy=True)
    elif compression == COMPRESSION_DELTA:
        array = delta_encode(array, copy=True)
    elif compression == COMPRESSION_NUMPRESS_LINEAR_PREDICTION:
        array = pynumpress.encode_linear(array, pynumpress.optimal_linear_fixed_point(array))
    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
