#ifndef __COLUMN_QUANTIZER_H__
#define __COLUMN_QUANTIZER_H__
#include "ValueQuantizer.h"
#include <math.h>

#pragma warning(disable : 4127) // conditional expression is constant

namespace Microsoft { namespace MSR { namespace CNTK {

#define ColMIDX(i, j, numRow) (((j) * (numRow)) + (i)) // 0 based indexing for column major

// ---------------------------------------------------------------------------
// Class to perform columnwise quantization/unquantization
//
// The quantization of a column is performed in 2 steps
// a) Compute the values used for unquantizing/reconstructing the quantized values. This is done by computing a pair of
//    values that specify the range of reconstructed unquantized values, such that the aggregate qunatization error is minimized.
// b) Perform the actual quantization by quantizing each value in the column to an integer of the size of
//    the specified number of bits and then packing these integer bits into the quantized matrix storage
// ---------------------------------------------------------------------------

template <class ElemType>
class ColumnQuantizer
{
    typedef typename ValueQuantizer<ElemType>::QWord QWord;
    typedef typename ValueQuantizer<ElemType>::QWordVal QWordVal;
    static const size_t QWordNumBits = ValueQuantizer<ElemType>::QWordNumBits;

public:
    cudacode ColumnQuantizer(size_t logNbits, ElemType lower, ElemType upper)
        : valQ(logNbits, lower, upper)
    {
    }

    // compute #QWords per column of a given height
    static size_t QWordsPerCol(size_t rows, size_t Nbits)
    {
        const size_t valsPerQWord = QWordNumBits / Nbits;
        return (rows + valsPerQWord - 1) / valsPerQWord;
    }

    size_t QWordsPerCol(size_t rows) const
    {
        return QWordsPerCol(rows, valQ.NBits());
    }

    // quantize a matrix column into qcoldata
    //  The current value of 'inResidual' is added to the matrix, and 'outResidual' gets updated with the new residual;
    //  inResidual = outResidual is allowed (intended)
    template <bool ZeroThresholdFor1Bit>
    cudacode void Quantize(const ElemType* inMat, const ElemType* inResidual, long M, size_t j, QWord* qColBits, ElemType* outResidual) const
    {
        // we loop over QWord values
        // E.g. there are 35 ints for a 1100-dim column (at 1-bit quantization).
        // For better CUDA memory collating, we interleave memory such that computing consecutive ints triggers consecutive memory accesses
        // (although for the CPU side, it breaks caching; we could do in-place op)
        // E.g., int  0 accesses elements 0, 35, 70, etc.
        // while int  1 accesses elements 1, 36, 71, etc
        // up to int 34 accesses elements 34, 69, 104, etc.
        const size_t numQWordsPerCol = QWordsPerCol(M);
        for (size_t iQWord = 0; iQWord < numQWordsPerCol; iQWord++)
        {
            qColBits[iQWord] = QuantizeOneQWord<ZeroThresholdFor1Bit>(inMat, inResidual, M, iQWord, M, numQWordsPerCol, j, outResidual);
        }
    }

    // unquantize a matrix column from qcoldata
    // If 'add' then add to the existing content of the matrix (this is a common thing to do; saves a buffer).
    cudacode void Unquantize(ElemType* outMat, long M, size_t j, const QWord* qColBits, bool add) const
    {
        // loop over QWord values
        const size_t numQWordsPerCol = QWordsPerCol(M);
        for (size_t iQWord = 0; iQWord < numQWordsPerCol; iQWord++)
        {
            UnquantizeOneQWord(outMat, M, iQWord, M, numQWordsPerCol, j, qColBits[iQWord], add);
        }
    }

    // workaround for not being able to declare a default argument for lambda parameters
    template <bool ZeroThresholdFor1Bit>
    static cudacode void ComputeRangeStatColj(const ElemType* inMat, const ElemType* inResidual, long M, size_t j, size_t bits, ElemType& lower, ElemType& upper)
    {
        /*dummy reducers do nothing in linear CPU version*/
        ComputeRangeStatColjSubset<ZeroThresholdFor1Bit>(inMat, inResidual, M, j, bits, lower, upper, 0, 1, [](ElemType&){}, [](unsigned int&){});
    }

public:
    // quantize the value in  inMat[rowStart,colIdx],  inMat[rowStart + rowStride,colIdx],inMat[rowStart + rowStride*,colIdx]  ... and pack them into a QWord
    // Question: note that it is somewhat un-intuitional, but this memory access pattern is efficient for GPU?
    template <bool ZeroThresholdFor1Bit>
    cudacode QWord QuantizeOneQWord(
        const ElemType* inMat, const ElemType* inResidual,
        long M,
        size_t rowStart, size_t rowEnd, size_t rowStride,
        size_t colIdx,
        ElemType* outResidual) const
    {
        QWord bitBuf = 0;

        if ((valQ.NBits() == 1) && (inResidual == outResidual) /*in-place*/)
        {
            ElemType val0 = valQ.Unquantize(0);
            ElemType val1 = valQ.Unquantize(1);
            size_t ij = ColMIDX(rowStart, colIdx, M);
            const ElemType* usibj = inMat + ij;
            const ElemType* usibjend = usibj + (rowEnd - rowStart);
            ElemType* resibj = outResidual + ij;
            // we know that the range covers at most the number of bits in a 'QWord'
            for (QWord bitmask = 1; usibj < usibjend; bitmask <<= 1, usibj += rowStride, resibj += rowStride)
            {
                // quantize   --we access element (i,j) through the three increasing pointers
                ElemType val = *usibj + *resibj;

                // Explicit use of 'template' keyword is needed to compile with GCC
                bool qval = valQ.template Quantize1<ZeroThresholdFor1Bit>(val);
                if (qval)
                {
                    bitBuf |= bitmask;
                }

                // compute residual
                ElemType uval = valQ.Unquantize1(qval, val0, val1);
                *resibj = val - uval;
            }
        }
        else
        {
            // number of bits in a QWord
            size_t i = rowStart;
            for (size_t k = 0; (k < QWordNumBits) && (i < rowEnd); k += valQ.NBits(), i += rowStride)
            {
                // quantize
                size_t ij = ColMIDX(i, colIdx, M);
                ElemType val = inMat[ij] + inResidual[ij];

                // 'template' keyword to compile with GCC
                QWordVal qval = valQ.template Quantize<ZeroThresholdFor1Bit>(val);

                // compute residual
                ElemType uval = valQ.Unquantize(qval);
                ElemType r = val - uval;
                outResidual[ij] = r;
                bitBuf = bitBuf | (qval << k);
            }
        }
        return bitBuf;
    }

    // unquantize one QWord of a quantized matrix column
    cudacode void UnquantizeOneQWord(
        ElemType* us, long M,
        size_t rowStart, size_t rowEnd, size_t rowStride,
        size_t colIdx, QWord bitBuf, bool add) const
    {
        // special case for 1 bit
        if (valQ.NBits() == 1)
        {
            ElemType val0 = valQ.Unquantize(0);
            ElemType val1 = valQ.Unquantize(1);
            size_t ij = ColMIDX(rowStart, colIdx, M);
            ElemType* usibj = us + ij;
            const ElemType* usibjend = usibj + (rowEnd - rowStart);
            for (; usibj < usibjend; usibj += rowStride)
            {
                // get value
                // bitbuf is shifted in-place
                bool qval = (bitBuf & 1) != 0;

                // and get bitbuf into next position
                bitBuf >>= 1;

                // unquantize
                ElemType val = ValueQuantizer<ElemType>::Unquantize1(qval, val0, val1);
                if (add)
                {
                    val += *usibj;
                }

                *usibj = val;
            }
        }
        else
        {
            // (rangeend MUST be a power of two; ensured by constructing off ldNbits)
            const QWordVal bitmask = valQ.QuanRangeEnd() - 1;
            size_t i = rowStart;
            for (size_t k = 0; (k < QWordNumBits) && (i < rowEnd); k += valQ.NBits(), i += rowStride)
            {
                // get value
                const QWordVal qval = (bitBuf >> k) & bitmask; // % 2^Nbits

                // unquantize
                ElemType val = valQ.Unquantize(qval);
                size_t ij = ColMIDX(i, colIdx, M);
                if (add)
                {
                    val += us[ij];
                }

                us[ij] = val;
            }
        }
    }

    // determine quantization range of one column
    // This code is written so that it can run in parallel threads on CUDA for collated memory access;
    // set 'subsets' to >1 and pass cross-thread reducer functions for 'float' and 'size_t' (which would reduce through using CUDA __shared__ memory).
    // TODO: further opportunity for speed-up: use 'mean' from last round for 1-bit and stddev calc
    template <bool ZeroThresholdFor1Bit, class F1, class F2>
    static cudacode void ComputeRangeStatColjSubset(
        const ElemType* inMat,
        const ElemType* inResidual, long M,
        size_t j,
        size_t bits,
        ElemType& lower, ElemType& upper,
        size_t subset, size_t subsets,
        F1 allReduceElem, F2 allReduceUint)
    {
        // quantization range, cut off after how many standard deviations (make this a parameter if we care)
        size_t rows = M;

        // compute mean
        // computing the mean is expensive; we assume there is no reason for asymmetry and thus a zero mean
        // an initial experiment showed that this is significantly worse (36.0 vs. 37.7% frame acc) at the start, but seems to recover nearly (minor gap)
        // thought:
        //  - we could set the threshold at 0
        //  - but keep the quantization values for 0 and 1 separate
        // i.e.
        //  - do not symmetrize/pool the quantization values for 0 and 1
        //  - but hard-code the quantization threshold to be 0 instead of the mean of the two bounds
        // This should give us the best of all--fast operation yet ability to be asymmetric within a column.
        ElemType mean = 0.0f;
        if (!ZeroThresholdFor1Bit && (bits == 1))
        {
            ElemType meanacc = 0.0f;
            // (subset: compute subset sum)
            for (size_t i = subset; i < rows; i += subsets)
            {
                size_t ij = ColMIDX(i, j, M);
                meanacc += inMat[ij] + inResidual[ij];
            }
            // multi-subset (CUDA): reduce to one thread
            allReduceElem(meanacc);
            mean = meanacc / rows;
        }

        if (bits == 1)
        {
            // 1-bit case:
            // We want to minimize the (squared) reconstruction error within the two levels.
            // I.e. we should reconstruct to the respective means of each level.
            // To be able to express the range by two floats, we approximate the level threshold as the av. of the two level means.
            // compute the two level means
            ElemType meanacc0 = 0.0f, meanacc1 = 0.0f;
            unsigned int num0 = 0, num1 = 0;
            // (subset: compute subset sum)
            for (size_t i = subset; i < rows; i += subsets)
            {
                size_t ij = ColMIDX(i, j, M);
                ElemType val = inMat[ij] + inResidual[ij];
                if (val < mean)
                {
                    meanacc0 += val;
                    num0++;
                }
                else
                {
                    meanacc1 += val;
                    num1++;
                }
            }

            // multi-subset (CUDA): reduce to one thread
            allReduceElem(meanacc0);
            allReduceElem(meanacc1);
            allReduceUint(num0);
            allReduceUint(num1);

            ElemType radius;
            ElemType newmean;
            if (!ZeroThresholdFor1Bit)
            {
                // we minimize the error jointly across positive and negative numbers to make things
                // symmetrical around the mean (which may be non-zero) tying the two sides
                ElemType devacc0 = (num0 * mean) - meanacc0;
                ElemType devacc1 = meanacc1 - (num1 * mean);

                // both deviations tied, to ensure consistent mean
                ElemType dev = (devacc0 + devacc1) / rows;
                radius = 2.0f * dev;
                newmean = mean;
            }
            else
            {
                // we keep two separate reconstruction values to allow for asymmetries--but we
                // instead hard-code that the threshold is 0

                // happens for all-zero columns which do exist (mean0 is 0 in that case)
                if (num0 == 0)
                    num0 = 1;
                if (num1 == 0)
                    num1 = 1;
                ElemType mean0 = meanacc0 / num0;
                ElemType mean1 = meanacc1 / num1;

                // approximate by using their average as the threshold between 0 and 1
                // with these values, bits (0,1) which mean values (0.5,1.5) will reconstruct to mean0/1
                newmean = 0.5f * (mean0 + mean1);
                radius = 2.0f * (mean1 - newmean);
            }

            if (subset == 0)
            {
                lower = newmean - radius;
                upper = newmean + radius;
            }
        }
        else
        {
            ElemType stddevs = 4.0f; // TODO: make this a parameter
            // >1 bit:
            // We linearly quantize between 'stddevs' standard deviations.
            ElemType varacc = 0.0f;
            // (subset: compute subset sum)
            for (size_t i = subset; i < rows; i += subsets)
            {
                size_t ij = ColMIDX(i, j, M);
                ElemType val = inMat[ij] + inResidual[ij];
                varacc += (val - mean) * (val - mean);
            }
            // multi-subset (CUDA): reduce to one thread
            allReduceElem(varacc);
            ElemType stddev = sqrt(varacc / rows);
            if (subset == 0)
            {
                // stddevs = how many stddevs from the mean until outside of quantization range
                lower = mean - (stddevs * stddev);
                upper = mean + (stddevs * stddev);
            }
        }
    }

private:
    ValueQuantizer<ElemType> valQ;

    template <typename T>
    friend class QuantizedMatrix;
};

}}}
#endif