# -*- coding: utf-8 -*-
"""
Defines a 'wire' format for transmitting or saving image frame data, optionally with Huffman compression and/or
sqrt quantization. The combination of quantization and Huffman coding allows compression ratios of 6-10 fold on typical
microscopy data. Using compression or quantization requires the pyme-compress companion library to be installed which
has optimized c code for performing the compression and quantization. If pyme-compress is compiled and installed on an
AVX capable processor, a throughput in excess of 800MB/s can be achieved.

Most users will just want the :func:`dumps` and :func:`loads` functions
"""

import numpy as np
import six
import zstandard as zstd

if six.PY2:
    _ord = ord
else:
    _ord = lambda a: a

from multiprocessing.pool import ThreadPool

import threading
import logging

logging.basicConfig(level=logging.DEBUG)

try:
    from pymecompress import bcl
except ImportError:
    logging.warning('''Could not import pymecompress library - saving or loading compressed PZF will fail
    (the library is installable from david_baddely conda channel, but requires an AVX capable processor)''')

NUM_COMP_THREADS = 2  # cpu_count()

compPool = ThreadPool(NUM_COMP_THREADS)


def ChunkedHuffmanCompress(data, quantization=None):
    num_chunks = NUM_COMP_THREADS

    chunk_size = int(np.ceil(float(len(data)) / num_chunks))
    raw_chunks = [data[j * chunk_size:(j + 1) * chunk_size].data for j in range(num_chunks)]

    # compPool = ThreadPool(NUM_COMP_THREADS)

    comp_chunk_d = {}

    if quantization is None:
        def _compChunk(c, j):
            comp_chunk_d[j] = bcl.HuffmanCompress(c)

        threads = [threading.Thread(target=_compChunk, args=(rc, j)) for j, rc in enumerate(raw_chunks)]
    else:
        def _compChunk(c, j):
            comp_chunk_d[j] = bcl.HuffmanCompressQuant(c, *quantization)

        threads = [threading.Thread(target=_compChunk, args=(rc, j)) for j, rc in enumerate(raw_chunks)]

    for p in threads:
        # print p
        p.start()

    for p in threads:
        p.join()

    # comp_chunks = compPool.map(bcl.HuffmanCompress, raw_chunks)

    s = np.array([num_chunks], 'u2').tostring()

    for j, r in enumerate(raw_chunks):
        c = comp_chunk_d[j]
        s += np.array([len(c), len(r)], 'u4').tostring()
        s += c.tostring()

    return s


def ChunkedHuffmanCompress_o(data):
    num_chunks = NUM_COMP_THREADS

    chunk_size = int(np.ceil(float(len(data)) / num_chunks))
    raw_chunks = [data[j * chunk_size:(j + 1) * chunk_size].data for j in range(num_chunks)]

    # compPool = ThreadPool(NUM_COMP_THREADS)

    comp_chunks = compPool.map(bcl.HuffmanCompress, raw_chunks)

    s = np.array([num_chunks], 'u2').tostring()

    for c, r in zip(comp_chunks, raw_chunks):
        s += np.array([len(c), len(r)], 'u4').tostring()
        s += c.tostring()

    return s


def _chunkDecompress(args):
    chunk, length = args
    return bcl.HuffmanDecompress(np.fromstring(chunk, 'u1'), length)


def ChunkedHuffmanDecompress(datastring):
    num_chunks = np.fromstring(datastring[:2], 'u2')

    # compPool = ThreadPool(NUM_COMP_THREADS)

    sp = 2

    comp_chunks = []
    for i in range(num_chunks):
        chunk_len, raw_len = np.fromstring(datastring[sp:(sp + 8)], 'u4')
        sp += 8
        comp_chunks.append((datastring[sp:(sp + chunk_len)], raw_len))
        sp += chunk_len

    decomp_chunks = compPool.map(_chunkDecompress, comp_chunks)

    data = np.hstack(decomp_chunks)

    # print data.shape #, comp_chunks
    return data


def ZstdDecompress(data_s, data_format):
    dctx = zstd.ZstdDecompressor()
    data = dctx.decompress(data_s)
    return np.frombuffer(data, dtype=data_format)


FILE_FORMAT_ID = b'BD'
FORMAT_VERSION = 3

DATA_FMT_UINT8 = 0
DATA_FMT_UINT16 = 1
DATA_FMT_FLOAT32 = 2

DATA_FMTS = ['u1', 'u2', 'f4']
DATA_FMTS_SIZES = [1, 2, 4]

DATA_COMP_RAW = 0
DATA_COMP_HUFFCODE = 1
DATA_COMP_HUFFCODE_CHUNKS = 2
DATA_COMP_ZSTD = 3
DATA_COMP_ZSTD_CHUNKS = 4

DATA_QUANT_NONE = 0
DATA_QUANT_SQRT = 1

# header_dtype = [('ID', 'S2'), ('Version', 'u1') , ('DataFormat', 'u1'), ('DataCompression', 'u1'), ('RESERVED0', 'S3'), ('SequenceID', 'i8'),
#                ('FrameNum', 'u4'), ('Width', 'u4'), ('Height', 'u4'), ('Depth', 'u4'),
#                ('FrameTimestamp', 'u8'), ('RESERVED1', 'u8')]

header_dtype = [('ID', 'S2'), ('Version', 'u1'), ('DataFormat', 'u1'), ('DataCompression', 'u1'),
                ('DataQuantization', 'u1'), ('DimOrder', 'S1'), ('RESERVED0', 'S1'), ('SequenceID', 'i8'),
                ('FrameNum', 'u4'), ('Width', 'u4'), ('Height', 'u4'), ('Depth', 'u4'),
                ('FrameTimestamp', 'u8'), ('QuantOffset', 'f4'), ('QuantScale', 'f4')]

# v3 increases header size to support arbitrary offsets to data so that data can be aligned.
header_dtype_v3 = [('ID', 'S2'), ('Version', 'u1'), ('DataFormat', 'u1'), ('DataCompression', 'u1'),
                   ('DataQuantization', 'u1'), ('DimOrder', 'S1'), ('RESERVED0', 'S1'), ('SequenceID', 'i8'),
                   ('FrameNum', 'u4'), ('Width', 'u4'), ('Height', 'u4'), ('Depth', 'u4'),
                   ('FrameTimestamp', 'u8'), ('QuantOffset', 'f4'), ('QuantScale', 'f4'), ('DataOffset', 'u4'),
                   ('RESERVED1', 'S12')]

"""
numpy dtype used to define the file header struct.

Most of the entries should be fairly self explanatory, with the following
deserving a bit more explanation:

:ID: a 2-character string that we can test to see if the file type is consistent
:Version: the version of this format the file uses
:DataFormat: what the data type of individual pixels is
:DataCompression: whether the data is compressed, and which algorithm is used
:SequenceID: A unique identifier for the sequence to which this frame belongs.
    The most important property of this number is that it is unique to
    each sequence. A reasonable method of generation would be to use
    a unix-format integer timestamp for the first dword, and a random
    integer for the second. A hash of the first n image pixels could
    also be used.
:FrameNum: The position of this frame within the sequence
:FrameTimestamp: Space to save camera derived frame timestamps, if available
:Depth: As envisaged, the format is expected to contain individual 2D frames, with
    multiple frames being pulled together in a higher level container to
    construct a sequence or stack. Depth is included just because it doesn't
    take a significant ammount of extra space, but gives us flexibility for
    the future.
"""

HEADER_LENGTH = np.zeros(1, header_dtype).nbytes

HEADER_LENGTH_V3 = np.zeros(1, header_dtype_v3).nbytes


def dumps(data, sequenceID=0, frameNum=0, frameTimestamp=0, compression=DATA_COMP_RAW, quantization=DATA_QUANT_NONE,
          quantizationOffset=0, quantizationScale=1):
    """Dump an image frame (supplied as a numpy array) into a string in PZF format.

    Parameters
    ==========

    data:  ndarray
            The frame as a 2D (or optionally 3D) numpy array

    sequenceID:  int
            A unique identifier for the sequence to which this frame belongs.
            This will let us connect the frame with it's metadata even if
            they end up in different directories etc ...

    frameNum:   int
            The position of this frame within the sequence

    frameTimestamp:  float
            A timestamp for the frame (if provided by the camera)

    compression:  int (enum)
            compression method to use - one of: `PZFFormat.DATA_COMP_RAW`,
            `PZFFormat.DATA_COMP_HUFFCODE`, or `PZFFormat.DATA_COMP_HUFFCODE_CHUNKS`
            Where raw stores the data with no compression, huffcode uses
            Huffman coding, and huffcode chunks breaks the data into chunks
            first, with each chunk meing encodes by a separate thread.

    quantization: int (enum)
            Whether or not the data is quantized before saving.
            One of `DATA_QUANT_NONE` or `DATA_QUANT_SQRT`. If `DATA_QUANT_SQRT`
            is selected, then the data is quantized as follows prior to
            compression:

            .. math:: data_{quant} =  \\frac{\\sqrt{data - quantizationOffset}}{quantizationScale}
    """

    header = np.zeros(1, header_dtype_v3)

    header['ID'] = FILE_FORMAT_ID
    header['Version'] = FORMAT_VERSION

    header['FrameNum'] = frameNum
    header['SequenceID'] = sequenceID
    header['FrameTimestamp'] = frameTimestamp

    header['DataOffset'] = HEADER_LENGTH_V3  # don't support padding on save yet

    if data.dtype == 'uint8':
        header['DataFormat'] = DATA_FMT_UINT8
    elif data.dtype == 'uint16':
        header['DataFormat'] = DATA_FMT_UINT16
    elif data.dtype == 'float32':
        header['DataFormat'] = DATA_FMT_FLOAT32
    else:
        raise RuntimeError('Unsupported data type')

    header['DimOrder'] = 'C'
    header['Width'] = data.shape[0]
    header['Height'] = data.shape[1]
    if data.ndim > 2:
        header['Depth'] = data.shape[2]
    else:
        header['Depth'] = 1

    if quantization == DATA_QUANT_SQRT:
        # sqrt- quantize the data
        header['DataQuantization'] = DATA_QUANT_SQRT
        header['QuantOffset'] = quantizationOffset
        header['QuantScale'] = quantizationScale

    #   qs = 1.0/quantizationScale
    #   data = (np.sqrt(np.maximum(data-quantizationOffset,0))*qs).astype('uint8')

    if data.flags['F_CONTIGUOUS']:
        header['DimOrder'] = 'F'
        d1 = data
        # d1 = np.frombuffer(data, data.dtype) #this will flatten without respecting order - we re-order later
    else:
        # if the array is c-contiguous this will just pass through. If it is non contiguous (i.e. a wierd slice)
        # it will force it to be c-contiguous
        d1 = np.ascontiguousarray(data)

    if compression == DATA_COMP_HUFFCODE:
        header['DataCompression'] = DATA_COMP_HUFFCODE

        if quantization:
            dataString = bcl.HuffmanCompressQuant(d1, quantizationOffset, quantizationScale).tostring()
        else:
            d2 = bcl.HuffmanCompress(d1)
            dataString = d2.tostring()
    elif compression == DATA_COMP_HUFFCODE_CHUNKS:
        header['DataCompression'] = DATA_COMP_HUFFCODE_CHUNKS

        dataString = ChunkedHuffmanCompress(d1)
    else:
        # print('saving raw')
        # print(header['DimOrder'][0])
        dataString = d1.tostring(order=header['DimOrder'][0])

    return header.tostring() + dataString


def load_header(datastring):
    if (_ord(datastring[2]) >= 3):
        return np.fromstring(datastring[:HEADER_LENGTH_V3], header_dtype_v3)
    else:
        return np.fromstring(datastring[:HEADER_LENGTH], header_dtype)


def display_header(datastring):
    header = load_header(datastring)

    if header['Version'] >= 2:
        dim_order = header['DimOrder'][0]
    else:
        dim_order = 'C'

    w, h, d = header['Width'][0], header['Height'][0], header['Depth'][0]

    if header['Version'] < 3:
        data_offset = HEADER_LENGTH
    else:
        data_offset = int(header['DataOffset'])

    if header['DataCompression'] == DATA_COMP_RAW:
        compression = 'None'
    elif header['DataCompression'] == DATA_COMP_HUFFCODE:
        compression = 'Huffmancode'
    elif header['DataCompression'] == DATA_COMP_HUFFCODE_CHUNKS:
        compression = 'Chunked Huffmancode'
    elif header['DataCompression'] == DATA_COMP_ZSTD:
        compression = 'Zstd'
    elif header['DataCompression'] == DATA_COMP_ZSTD_CHUNKS:
        compression = 'Chunked Zstd'
    else:
        raise RuntimeError('Compression type not understood')

    header_json = {
        'ID': str(header['ID'][0]),
        'Version': header['Version'][0],
        'Format': DATA_FMTS[header['DataFormat'][0]],
        'Compression': compression,
        'Quantisation': 'SQRT' if header['DataQuantization'] else 'None',
        'Order': str(dim_order),
        'SequenceId': header['SequenceID'][0],
        'Frame no': header['FrameNum'][0],
        'Width': w,
        'Height': h,
        'Depth': d,
        'Data offset': data_offset,
    }

    return header_json


def loads(datastring):
    """
    Loads image data from a string in PZF format.

    Parameters
    ----------
    datastring : string / bytes
        The encoded data

    Returns
    -------

    data : ndarray
        The image data as a numpy array

    header : recarray
        The image header, as a numpy record array with the :const:`header_dtype` dtype.

    """
    header = load_header(datastring)

    if not header['ID'] == FILE_FORMAT_ID:
        raise RuntimeError("Invalid format: This doesn't appear to be a PZF file")

    if header['Version'] >= 2:
        dimOrder = header['DimOrder'][0]
    else:
        dimOrder = 'C'

    # print(dimOrder)

    w, h, d = header['Width'][0], header['Height'][0], header['Depth'][0]

    if header['DataQuantization'] == DATA_QUANT_SQRT:
        # quantized data is always 8 bit
        outsize = w * h * d
    else:
        outsize = w * h * d * DATA_FMTS_SIZES[int(header['DataFormat'])]

    if header['Version'] < 3:
        data_offset = HEADER_LENGTH
    else:
        data_offset = int(header['DataOffset'])

    data_s = datastring[data_offset:]
    data_format = DATA_FMTS[header['DataFormat'][0]]

    # logging.debug('About to decompress')
    # logging.debug({k:header[0][k] for k in header.dtype.names})

    # logging.debug('Compressed size: %s' % len(data_s))

    if header['DataCompression'] == DATA_COMP_RAW:
        # no need to decompress
        data = np.fromstring(data_s, 'u1')
    elif header['DataCompression'] == DATA_COMP_HUFFCODE:
        # logging.debug('Decompressing ...')
        data = bcl.HuffmanDecompress(np.fromstring(data_s, 'u1'), outsize)
    elif header['DataCompression'] == DATA_COMP_HUFFCODE_CHUNKS:
        data = ChunkedHuffmanDecompress(data_s)
    elif header['DataCompression'] == DATA_COMP_ZSTD:
        data = ZstdDecompress(data_s, data_format)
    else:
        raise RuntimeError('Compression type not understood')

    # logging.debug('Uncompressed shape: %s, %s, (%d, %d, %d)' % (data.shape, w * h * d, w, h, d))

    if header['DataQuantization'] == DATA_QUANT_SQRT:
        # un-quantize data
        # logging.debug('Dequantizing')
        #         print(data.max())
        #         data = data.astype('f') * header['QuantScale']
        data = data.astype('f') * 1.0
        #         print('data dtype: %s' % data.dtype)
        #         print('data shape: %s' % data.shape)
        data = data * data
    #         data = (data * data + header['QuantOffset'])

    #     print(dimOrder, [w, h, d])
    data = data.reshape([d, h, w], order=dimOrder).T

    return data, header
