#pragma once

#include <cassert>
#include <cmath>
#include <cstddef>
#include <memory>
#include <vector>

#include <athena/Types.hpp>

namespace DataSpec::DNAANIM {

struct Value {
  athena::simd<float> simd;
  Value() = default;
  Value(const athena::simd<float>& s) : simd(s) {}
  Value(const atVec3f& v) : simd(v.simd) {}
  Value(const atVec4f& v) : simd(v.simd) {}
  Value(float x, float y, float z) : simd(x, y, z, 0.f) {}
  Value(float w, float x, float y, float z) : simd(w, x, y, z) {}
};
struct QuantizedValue {
  atInt32 v[4];
  atInt32& operator[](size_t idx) { return v[idx]; }
  atInt32 operator[](size_t idx) const { return v[idx]; }

  int qFrom(const QuantizedValue& other, size_t idx) const {
    atInt32 delta = v[idx] - other.v[idx];
    atInt32 absDelta = std::abs(delta);
    if (absDelta == 0)
      return 1;
    int ret = int(std::ceil(std::log2(absDelta))) + 1;
    if (delta > 0 && (delta >> (ret - 1)))
      ++ret;
    assert(ret <= 24 && "Bad q value");
    return ret;
  }
};
struct QuantizedRot {
  QuantizedValue v;
  bool w;
};
struct Channel {
  enum class Type { Rotation, Translation, Scale, KfHead, RotationMP3 } type;
  atInt32 id = -1;
  QuantizedValue i = {};
  atUint8 q[4] = {};
};

size_t ComputeBitstreamSize(size_t keyFrameCount, const std::vector<Channel>& channels);

class BitstreamReader {
  size_t m_bitCur;
  atInt32 dequantize(const atUint8* data, atUint8 q);
  bool dequantizeBit(const atUint8* data);

public:
  std::vector<std::vector<Value>> read(const atUint8* data, size_t keyFrameCount, const std::vector<Channel>& channels,
                                       atUint32 rotDiv, float transMult, float scaleMult);
};

class BitstreamWriter {
  size_t m_bitCur;
  void quantize(atUint8* data, atUint8 q, atInt32 val);
  void quantizeBit(atUint8* data, bool val);

public:
  std::unique_ptr<atUint8[]> 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);
};

} // namespace DataSpec::DNAANIM