#include "zeus/Math.hpp"
#include "ANIM.hpp"

#define DUMP_KEYS 0

namespace DataSpec::DNAANIM
{

size_t ComputeBitstreamSize(size_t keyFrameCount, const std::vector<Channel>& channels)
{
    size_t bitsPerKeyFrame = 0;
    for (const Channel& chan : channels)
    {
        switch (chan.type)
        {
        case Channel::Type::Rotation:
            bitsPerKeyFrame += 1;
        case Channel::Type::Translation:
        case Channel::Type::Scale:
            bitsPerKeyFrame += chan.q[0];
            bitsPerKeyFrame += chan.q[1];
            bitsPerKeyFrame += chan.q[2];
            break;
        case Channel::Type::KfHead:
            bitsPerKeyFrame += 1;
            break;
        case Channel::Type::RotationMP3:
            bitsPerKeyFrame += chan.q[0];
            bitsPerKeyFrame += chan.q[1];
            bitsPerKeyFrame += chan.q[2];
            bitsPerKeyFrame += chan.q[3];
            break;
        default: break;
        }
    }
    return (bitsPerKeyFrame * keyFrameCount + 31) / 32 * 4;
}

static inline QuantizedRot QuantizeRotation(const Value& quat, atUint32 div)
{
    float q = M_PIF / 2.0f / float(div);
    zeus::simd_floats f(quat.simd);
    return
    {
        {
            atInt32(std::asin(f[1]) / q),
            atInt32(std::asin(f[2]) / q),
            atInt32(std::asin(f[3]) / q),
        },
        (f[0] < 0.f)
    };
}

static inline Value DequantizeRotation(const QuantizedRot& v, atUint32 div)
{
    float q = M_PIF / 2.0f / float(div);
    athena::simd_floats f = {
        0.0f,
        std::sin(v.v[0] * q),
        std::sin(v.v[1] * q),
        std::sin(v.v[2] * q),
    };
    f[0] = std::sqrt(std::max((1.0f -
                               (f[1] * f[1] +
                                f[2] * f[2] +
                                f[3] * f[3])), 0.0f));
    f[0] = v.w ? -f[0] : f[0];
    Value retval;
    retval.simd.copy_from(f);
    return retval;
}

static inline Value DequantizeRotation_3(const QuantizedRot& v, atUint32 div)
{
    float q = 1.0f / float(div);
    athena::simd_floats f = {
        0.0f,
        v.v[0] * q,
        v.v[1] * q,
        v.v[2] * q,
    };
    f[0] = std::sqrt(std::max((1.0f -
                               (f[1] * f[1] +
                                f[2] * f[2] +
                                f[3] * f[3])), 0.0f));
    f[0] = v.w ? -f[0] : f[0];
    Value retval;
    retval.simd.copy_from(f);
    return retval;
}

bool BitstreamReader::dequantizeBit(const atUint8* data)
{
    atUint32 byteCur = (m_bitCur / 32) * 4;
    atUint32 bitRem = m_bitCur % 32;

    /* Fill 32 bit buffer with region containing bits */
    /* Make them least significant */
    atUint32 tempBuf = hecl::SBig(*reinterpret_cast<const atUint32*>(data + byteCur)) >> bitRem;

    /* That's it */
    m_bitCur += 1;
    return tempBuf & 0x1;
}

atInt32 BitstreamReader::dequantize(const atUint8* data, atUint8 q)
{
    atUint32 byteCur = (m_bitCur / 32) * 4;
    atUint32 bitRem = m_bitCur % 32;

    /* Fill 32 bit buffer with region containing bits */
    /* Make them least significant */
    atUint32 tempBuf = hecl::SBig(*reinterpret_cast<const atUint32*>(data + byteCur)) >> bitRem;

    /* If this shift underflows the value, buffer the next 32 bits */
    /* And tack onto shifted buffer */
    if ((bitRem + q) > 32)
    {
        atUint32 tempBuf2 = hecl::SBig(*reinterpret_cast<const atUint32*>(data + byteCur + 4));
        tempBuf |= (tempBuf2 << (32 - bitRem));
    }

    /* Mask it */
    atUint32 mask = (1 << q) - 1;
    tempBuf &= mask;

    /* Sign extend */
    atUint32 sign = (tempBuf >> (q - 1)) & 0x1;
    if (sign)
        tempBuf |= ~0u << q;

    /* Return delta value */
    m_bitCur += q;
    return atInt32(tempBuf);
}

std::vector<std::vector<Value>>
BitstreamReader::read(const atUint8* data,
                      size_t keyFrameCount,
                      const std::vector<Channel>& channels,
                      atUint32 rotDiv,
                      float transMult,
                      float scaleMult)
{
    m_bitCur = 0;
    std::vector<std::vector<Value>> chanKeys;
    std::vector<QuantizedValue> chanAccum;
    chanKeys.reserve(channels.size());
    chanAccum.reserve(channels.size());
    for (const Channel& chan : channels)
    {
        chanAccum.push_back(chan.i);

        chanKeys.emplace_back();
        std::vector<Value>& keys = chanKeys.back();
        keys.reserve(keyFrameCount);
        switch (chan.type)
        {
        case Channel::Type::Rotation:
        {
            QuantizedRot qr = {{chan.i[0], chan.i[1], chan.i[2]}, false};
            keys.emplace_back(DequantizeRotation(qr, rotDiv));
            break;
        }
        case Channel::Type::Translation:
        {
            keys.push_back({chan.i[0] * transMult, chan.i[1] * transMult, chan.i[2] * transMult});
            break;
        }
        case Channel::Type::Scale:
        {
            keys.push_back({chan.i[0] * scaleMult, chan.i[1] * scaleMult, chan.i[2] * scaleMult});
            break;
        }
        case Channel::Type::KfHead:
        {
            break;
        }
        case Channel::Type::RotationMP3:
        {
            QuantizedRot qr = {{chan.i[1], chan.i[2], chan.i[3]}, bool(chan.i[0] & 0x1)};
            keys.emplace_back(DequantizeRotation_3(qr, rotDiv));
            break;
        }
        default: break;
        }
    }

    for (size_t f=0 ; f<keyFrameCount ; ++f)
    {
#if DUMP_KEYS
        fprintf(stderr, "\nFRAME %" PRISize " %u %u\n", f, (m_bitCur / 32) * 4, m_bitCur % 32);
        int lastId = -1;
#endif
        auto kit = chanKeys.begin();
        auto ait = chanAccum.begin();
        for (const Channel& chan : channels)
        {
#if DUMP_KEYS
            if (chan.id != lastId)
            {
                lastId = chan.id;
                fprintf(stderr, "\n");
            }
#endif
            QuantizedValue& p = *ait;
            switch (chan.type)
            {
            case Channel::Type::Rotation:
            {
                bool wBit = dequantizeBit(data);
                p[0] += dequantize(data, chan.q[0]);
                p[1] += dequantize(data, chan.q[1]);
                p[2] += dequantize(data, chan.q[2]);
                QuantizedRot qr = {{p[0], p[1], p[2]}, wBit};
                kit->emplace_back(DequantizeRotation(qr, rotDiv));
#if DUMP_KEYS
                fprintf(stderr, "%d R: %d %d %d %d\t", chan.id, wBit, p[0], p[1], p[2]);
#endif
                break;
            }
            case Channel::Type::Translation:
            {
                atInt32 val1 = dequantize(data, chan.q[0]);
                p[0] += val1;
                atInt32 val2 = dequantize(data, chan.q[1]);
                p[1] += val2;
                atInt32 val3 = dequantize(data, chan.q[2]);
                p[2] += val3;
                kit->push_back({p[0] * transMult, p[1] * transMult, p[2] * transMult});
#if DUMP_KEYS
                fprintf(stderr, "%d T: %d %d %d\t", chan.id, p[0], p[1], p[2]);
#endif
                break;
            }
            case Channel::Type::Scale:
            {
                p[0] += dequantize(data, chan.q[0]);
                p[1] += dequantize(data, chan.q[1]);
                p[2] += dequantize(data, chan.q[2]);
                kit->push_back({p[0] * scaleMult, p[1] * scaleMult, p[2] * scaleMult});
#if DUMP_KEYS
                fprintf(stderr, "%d S: %d %d %d\t", chan.id, p[0], p[1], p[2]);
#endif
                break;
            }
            case Channel::Type::KfHead:
            {
                dequantizeBit(data);
                break;
            }
            case Channel::Type::RotationMP3:
            {
                atInt32 val1 = dequantize(data, chan.q[0]);
                p[0] += val1;
                atInt32 val2 = dequantize(data, chan.q[1]);
                p[1] += val2;
                atInt32 val3 = dequantize(data, chan.q[2]);
                p[2] += val3;
                atInt32 val4 = dequantize(data, chan.q[3]);
                p[3] += val4;
                QuantizedRot qr = {{p[1], p[2], p[3]}, bool(p[0] & 0x1)};
                kit->emplace_back(DequantizeRotation_3(qr, rotDiv));
                break;
            }
            default: break;
            }
            ++kit;
            ++ait;
        }
#if DUMP_KEYS
        fprintf(stderr, "\n");
#endif
    }

    return chanKeys;
}

void BitstreamWriter::quantizeBit(atUint8* data, bool val)
{
    atUint32 byteCur = (m_bitCur / 32) * 4;
    atUint32 bitRem = m_bitCur % 32;

    /* Fill 32 bit buffer with region containing bits */
    /* Make them least significant */
    *(atUint32*)(data + byteCur) =
    hecl::SBig(hecl::SBig(*(atUint32*)(data + byteCur)) | (val << bitRem));

    m_bitCur += 1;
}

void BitstreamWriter::quantize(atUint8* data, atUint8 q, atInt32 val)
{
    atUint32 byteCur = (m_bitCur / 32) * 4;
    atUint32 bitRem = m_bitCur % 32;

    atUint32 masked = val & ((1 << q) - 1);

    /* Fill 32 bit buffer with region containing bits */
    /* Make them least significant */
    *(atUint32*)(data + byteCur) =
    hecl::SBig(hecl::SBig(*(atUint32*)(data + byteCur)) | (masked << bitRem));

    /* If this shift underflows the value, buffer the next 32 bits */
    /* And tack onto shifted buffer */
    if ((bitRem + q) > 32)
    {
        *(atUint32*)(data + byteCur + 4) =
        hecl::SBig(hecl::SBig(*(atUint32*)(data + byteCur + 4)) | (masked >> (32 - bitRem)));
    }

    m_bitCur += q;
}

std::unique_ptr<atUint8[]>
BitstreamWriter::write(const std::vector<std::vector<Value>>& chanKeys,
                       size_t keyFrameCount, std::vector<Channel>& channels,
                       atUint32 quantRange,
                       atUint32& rotDivOut,
                       float& transMultOut,
                       float& scaleMultOut,
                       size_t& sizeOut)
{
    m_bitCur = 0;
    rotDivOut = quantRange; /* Normalized range of values */
    float quantRangeF = float(quantRange);

    /* Pre-pass to calculate translation multiplier */
    float maxTransVal = 0.0f;
    float maxScaleVal = 0.0f;
    auto kit = chanKeys.begin();
    for (Channel& chan : channels)
    {
        switch (chan.type)
        {
        case Channel::Type::Translation:
        {
            for (auto it=kit->begin();
                 it != kit->end();
                 ++it)
            {
                const Value* key = &*it;
                zeus::simd_floats f(key->simd);
                maxTransVal = std::max(maxTransVal, std::fabs(f[0]));
                maxTransVal = std::max(maxTransVal, std::fabs(f[1]));
                maxTransVal = std::max(maxTransVal, std::fabs(f[2]));
            }
            break;
        }
        case Channel::Type::Scale:
        {
            for (auto it=kit->begin();
                 it != kit->end();
                 ++it)
            {
                const Value* key = &*it;
                zeus::simd_floats f(key->simd);
                maxScaleVal = std::max(maxScaleVal, std::fabs(f[0]));
                maxScaleVal = std::max(maxScaleVal, std::fabs(f[1]));
                maxScaleVal = std::max(maxScaleVal, std::fabs(f[2]));
            }
            break;
        }
        default: break;
        }
        ++kit;
    }
    transMultOut = maxTransVal / quantRangeF;
    scaleMultOut = maxScaleVal / quantRangeF;

    /* Output channel inits */
    std::vector<QuantizedValue> initVals;
    initVals.reserve(channels.size());
    kit = chanKeys.begin();
    for (Channel& chan : channels)
    {
        chan.q[0] = 1;
        chan.q[1] = 1;
        chan.q[2] = 1;
        switch (chan.type)
        {
        case Channel::Type::Rotation:
        {
            QuantizedRot qr = QuantizeRotation((*kit)[0], rotDivOut);
            chan.i = qr.v;
            initVals.push_back(chan.i);
            break;
        }
        case Channel::Type::Translation:
        {
            zeus::simd_floats f((*kit)[0].simd);
            chan.i = {atInt32(f[0] / transMultOut),
                      atInt32(f[1] / transMultOut),
                      atInt32(f[2] / transMultOut)};
            initVals.push_back(chan.i);
            break;
        }
        case Channel::Type::Scale:
        {
            zeus::simd_floats f((*kit)[0].simd);
            chan.i = {atInt32(f[0] / scaleMultOut),
                      atInt32(f[1] / scaleMultOut),
                      atInt32(f[2] / scaleMultOut)};
            initVals.push_back(chan.i);
            break;
        }
        default: break;
        }
        ++kit;
    }

    /* Pre-pass to analyze quantization factors for channels */
    std::vector<QuantizedValue> lastVals = initVals;
    kit = chanKeys.begin();
    auto vit = lastVals.begin();
    for (Channel& chan : channels)
    {
        QuantizedValue& last = *vit++;
        switch (chan.type)
        {
        case Channel::Type::Rotation:
        {
            for (auto it=kit->begin() + 1;
                 it != kit->end();
                 ++it)
            {
                QuantizedRot qrCur = QuantizeRotation(*it, rotDivOut);
                chan.q[0] = std::max(chan.q[0], atUint8(qrCur.v.qFrom(last, 0)));
                chan.q[1] = std::max(chan.q[1], atUint8(qrCur.v.qFrom(last, 1)));
                chan.q[2] = std::max(chan.q[2], atUint8(qrCur.v.qFrom(last, 2)));
                last = qrCur.v;
            }
            break;
        }
        case Channel::Type::Translation:
        {
            for (auto it=kit->begin() + 1;
                 it != kit->end();
                 ++it)
            {
                zeus::simd_floats f(it->simd);
                QuantizedValue cur = {atInt32(f[0] / transMultOut),
                                      atInt32(f[1] / transMultOut),
                                      atInt32(f[2] / transMultOut)};
                chan.q[0] = std::max(chan.q[0], atUint8(cur.qFrom(last, 0)));
                chan.q[1] = std::max(chan.q[1], atUint8(cur.qFrom(last, 1)));
                chan.q[2] = std::max(chan.q[2], atUint8(cur.qFrom(last, 2)));
                last = cur;
            }
            break;
        }
        case Channel::Type::Scale:
        {
            for (auto it=kit->begin() + 1;
                 it != kit->end();
                 ++it)
            {
                zeus::simd_floats f(it->simd);
                QuantizedValue cur = {atInt32(f[0] / scaleMultOut),
                                      atInt32(f[1] / scaleMultOut),
                                      atInt32(f[2] / scaleMultOut)};
                chan.q[0] = std::max(chan.q[0], atUint8(cur.qFrom(last, 0)));
                chan.q[1] = std::max(chan.q[1], atUint8(cur.qFrom(last, 1)));
                chan.q[2] = std::max(chan.q[2], atUint8(cur.qFrom(last, 2)));
                last = cur;
            }
            break;
        }
        default: break;
        }
        ++kit;
    }

    /* Generate Bitstream */
    sizeOut = ComputeBitstreamSize(keyFrameCount, channels);
    std::unique_ptr<atUint8[]> newData(new atUint8[sizeOut]);
    memset(newData.get(), 0, sizeOut);

    lastVals = initVals;
    for (size_t f=0 ; f<keyFrameCount ; ++f)
    {
        kit = chanKeys.begin();
        vit = lastVals.begin();
        for (const Channel& chan : channels)
        {
            const Value& val = (*kit++)[f+1];
            QuantizedValue& last = *vit++;
            switch (chan.type)
            {
            case Channel::Type::Rotation:
            {
                QuantizedRot qrCur = QuantizeRotation(val, rotDivOut);
                quantizeBit(newData.get(), qrCur.w);
                quantize(newData.get(), chan.q[0], qrCur.v[0] - last.v[0]);
                quantize(newData.get(), chan.q[1], qrCur.v[1] - last.v[1]);
                quantize(newData.get(), chan.q[2], qrCur.v[2] - last.v[2]);
                last = qrCur.v;
                break;
            }
            case Channel::Type::Translation:
            {
                zeus::simd_floats f(val.simd);
                QuantizedValue cur = {atInt32(f[0] / transMultOut),
                                      atInt32(f[1] / transMultOut),
                                      atInt32(f[2] / transMultOut)};
                quantize(newData.get(), chan.q[0], cur[0] - last[0]);
                quantize(newData.get(), chan.q[1], cur[1] - last[1]);
                quantize(newData.get(), chan.q[2], cur[2] - last[2]);
                last = cur;
                break;
            }
            case Channel::Type::Scale:
            {
                zeus::simd_floats f(val.simd);
                QuantizedValue cur = {atInt32(f[0] / scaleMultOut),
                                      atInt32(f[1] / scaleMultOut),
                                      atInt32(f[2] / scaleMultOut)};
                quantize(newData.get(), chan.q[0], cur[0] - last[0]);
                quantize(newData.get(), chan.q[1], cur[1] - last[1]);
                quantize(newData.get(), chan.q[2], cur[2] - last[2]);
                last = cur;
                break;
            }
            default: break;
            }
        }
    }
    return newData;
}

}