#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

namespace fs = std::filesystem;
using Clock = std::chrono::steady_clock;

uint16_t u16(const std::vector<uint8_t>& b, size_t o) {
    return uint16_t(b[o]) | uint16_t(b[o + 1]) << 8;
}
uint32_t u32(const std::vector<uint8_t>& b, size_t o) {
    uint32_t v = 0;
    for (size_t i = 0; i < 4; ++i) v |= uint32_t(b[o + i]) << (8 * i);
    return v;
}
float f32(const std::vector<uint8_t>& b, size_t o) {
    uint32_t bits = u32(b, o);
    float v = 0;
    std::memcpy(&v, &bits, 4);
    return v;
}

struct Desc {
    uint16_t semantic = 0;
    uint8_t dtype = 0;
    float scale = 1;
    uint32_t off = 0;
    uint32_t len = 0;
    uint8_t valueClass = 0;
};
struct Plane {
    Desc d;
    std::vector<int64_t> raw;
};
struct Frame {
    uint32_t w = 0, h = 0;
    std::vector<uint8_t> validity;
    std::vector<Plane> planes;
};

bool validAt(const std::vector<uint8_t>& v, uint64_t i) {
    return ((v[i >> 3] >> (i & 7u)) & 1u) != 0;
}
std::vector<uint8_t> readFile(const fs::path& p) {
    std::ifstream in(p, std::ios::binary);
    if (!in) throw std::runtime_error("open failed: " + p.string());
    return {std::istreambuf_iterator<char>(in), std::istreambuf_iterator<char>()};
}

Frame parseV2(const fs::path& p) {
    const auto b = readFile(p);
    if (b.size() < 112 || std::memcmp(b.data(), "QBV2", 4) != 0) throw std::runtime_error("not QBV2");
    Frame f;
    const uint16_t n = u16(b, 12);
    f.w = u32(b, 16);
    f.h = u32(b, 20);
    const uint64_t cells = uint64_t(f.w) * f.h;
    const size_t payload = 112 + size_t(n) * 32;
    std::vector<Desc> ds;
    for (uint16_t i = 0; i < n; ++i) {
        const size_t o = 112 + size_t(i) * 32;
        ds.push_back({u16(b,o), b[o+2], f32(b,o+4), u32(b,o+16), u32(b,o+20), b[o+28]});
    }
    const auto it = std::find_if(ds.begin(), ds.end(), [](const Desc& d){ return d.dtype == 6; });
    if (it == ds.end()) throw std::runtime_error("missing validity");
    f.validity.assign(b.begin()+payload+it->off, b.begin()+payload+it->off+it->len);

    for (const auto& d : ds) {
        if (d.dtype == 6) continue;
        Plane p0;
        p0.d = d;
        p0.raw.resize(cells);
        const uint8_t* raw = b.data() + payload + d.off;
        for (uint64_t i=0;i<cells;++i) {
            if (d.dtype == 1) {
                uint16_t bits = uint16_t(raw[i*2]) | uint16_t(raw[i*2+1]) << 8;
                int16_t x; std::memcpy(&x,&bits,2); p0.raw[i]=x;
            } else if (d.dtype == 2) {
                p0.raw[i]=uint16_t(raw[i*2]) | uint16_t(raw[i*2+1]) << 8;
            } else if (d.dtype == 3) p0.raw[i]=raw[i];
            else if (d.dtype == 4) { int8_t x; std::memcpy(&x,raw+i,1); p0.raw[i]=x; }
            else throw std::runtime_error("unsupported V2 dtype");
        }
        f.planes.push_back(std::move(p0));
    }
    return f;
}
int64_t normDir(int64_t x, uint32_t n) {
    if (x < 0) x += n;
    else if (x >= int64_t(n)) x -= n;
    return x;
}
uint64_t zigzag(int64_t r) { return r >= 0 ? uint64_t(2*r) : uint64_t(-2*r-1); }
int64_t unzigzag(uint64_t z) { return (z&1u) ? -int64_t(z/2)-1 : int64_t(z/2); }

int64_t predictor(const Frame& f, const Plane& p, const std::vector<int64_t>& raw,
                  uint32_t row, uint32_t col, uint32_t n) {
    const uint64_t i=uint64_t(row)*f.w+col;
    const bool l=col>0 && validAt(f.validity,i-1);
    const bool u=row>0 && validAt(f.validity,i-f.w);
    const bool ul=row>0 && col>0 && validAt(f.validity,i-f.w-1);
    int64_t pred=0;
    if (l&&u&&ul) pred=raw[i-1]+raw[i-f.w]-raw[i-f.w-1];
    else if(l) pred=raw[i-1];
    else if(u) pred=raw[i-f.w];
    if (n) pred=normDir(pred,n);
    return pred;
}
int64_t residualFor(const Frame& f, const Plane& p, uint32_t row, uint32_t col, uint32_t n) {
    const uint64_t i=uint64_t(row)*f.w+col;
    const int64_t pred=predictor(f,p,p.raw,row,col,n);
    int64_t r=p.raw[i]-pred;
    if (!n) return r;
    int64_t m=r;
    if(m<0)m+=n;
    if(2*m>n)return m-n;
    if(2*m<int64_t(n))return m;
    return -int64_t(n/2);
}

struct BitWriter {
    std::vector<uint8_t> b;
    size_t pos=0;
    void bit(uint8_t v) {
        if ((pos & 7u)==0) b.push_back(0);
        if(v) b.back() |= uint8_t(1u << (pos & 7u));
        ++pos;
    }
    void bits(uint32_t v,uint8_t n){ for(uint8_t i=0;i<n;++i) bit((v>>i)&1u); }
};
struct EncodedPlane { std::vector<uint8_t> bytes; uint64_t blocks=0, nonempty=0; };

EncodedPlane encodePlane(const Frame& f,const Plane& p,uint32_t block) {
    EncodedPlane out;
    const uint32_t n = p.d.valueClass==2 ? uint32_t(std::floor(360.0/double(p.d.scale)+0.5)) : 0;
    for(uint32_t rb=0;rb<f.h;rb+=block) for(uint32_t cb=0;cb<f.w;cb+=block) {
        ++out.blocks;
        std::vector<uint64_t> z;
        for(uint32_t r=rb;r<std::min(rb+block,f.h);++r)
            for(uint32_t c=cb;c<std::min(cb+block,f.w);++c) {
                const uint64_t i=uint64_t(r)*f.w+c;
                if(validAt(f.validity,i)) z.push_back(zigzag(residualFor(f,p,r,c,n)));
            }
        if(z.empty()) continue;
        ++out.nonempty;
        uint64_t bestBits=~uint64_t(0); uint8_t bestK=0;
        for(uint8_t k=0;k<=31;++k) {
            uint64_t bits=0;
            for(auto v:z) bits += (v>>k)+1u+k;
            if(bits<bestBits){bestBits=bits;bestK=k;}
        }
        out.bytes.push_back(bestK);
        BitWriter bw;
        for(auto v:z) {
            const uint64_t q=v>>bestK;
            for(uint64_t j=0;j<q;++j) bw.bit(0);
            bw.bit(1);
            if(bestK) bw.bits(uint32_t(v & ((uint64_t(1)<<bestK)-1)),bestK);
        }
        out.bytes.insert(out.bytes.end(),bw.b.begin(),bw.b.end());
    }
    return out;
}

struct BitReader {
    const uint8_t* p=nullptr; size_t bytes=0,pos=0;
    uint8_t bit(){ if(pos>=bytes*8) throw std::runtime_error("truncated"); return (p[pos>>3]>>(pos++&7u))&1u; }
    uint32_t bitsN(uint8_t n){ uint32_t v=0; for(uint8_t i=0;i<n;++i)v|=uint32_t(bit())<<i; return v; }
};

std::vector<int64_t> decodePlane(const Frame& f,const Plane& p,const EncodedPlane& enc,uint32_t block) {
    std::vector<int64_t> raw(uint64_t(f.w)*f.h,0);
    size_t cursor=0;
    const uint32_t n=p.d.valueClass==2 ? uint32_t(std::floor(360.0/double(p.d.scale)+0.5)) : 0;
    for(uint32_t rb=0;rb<f.h;rb+=block) for(uint32_t cb=0;cb<f.w;cb+=block) {
        uint32_t count=0;
        for(uint32_t r=rb;r<std::min(rb+block,f.h);++r)
            for(uint32_t c=cb;c<std::min(cb+block,f.w);++c)
                if(validAt(f.validity,uint64_t(r)*f.w+c)) ++count;
        if(!count) continue;
        if(cursor>=enc.bytes.size()) throw std::runtime_error("missing k");
        const uint8_t k=enc.bytes[cursor++];
        if(k>31) throw std::runtime_error("bad k");
        BitReader br{enc.bytes.data()+cursor,enc.bytes.size()-cursor,0};
        for(uint32_t r=rb;r<std::min(rb+block,f.h);++r)
            for(uint32_t c=cb;c<std::min(cb+block,f.w);++c) {
                const uint64_t i=uint64_t(r)*f.w+c;
                if(!validAt(f.validity,i)) continue;
                uint64_t q=0; while(br.bit()==0) ++q;
                const uint32_t rem=k?br.bitsN(k):0;
                const uint64_t z=(q<<k)|rem;
                const int64_t res=unzigzag(z);
                int64_t value=predictor(f,p,raw,r,c,n)+res;
                if(n)value=normDir(value,n);
                raw[i]=value;
            }
        cursor += (br.pos+7u)/8u;
    }
    if(cursor!=enc.bytes.size()) throw std::runtime_error("extra bytes");
    return raw;
}

struct Totals {
    uint64_t bytes=0,blocks=0,nonempty=0;
    double encodeMs=0,decodeMs=0;
};

void runOne(const fs::path& file,uint32_t block,Totals& totals) {
    const Frame frame=parseV2(file);
    std::vector<EncodedPlane> encoded(frame.planes.size());
    constexpr int EncodeRepeats=3, DecodeRepeats=8;

    auto start=Clock::now();
    for(int rep=0;rep<EncodeRepeats;++rep)
        for(size_t i=0;i<frame.planes.size();++i) encoded[i]=encodePlane(frame,frame.planes[i],block);
    auto mid=Clock::now();

    for(size_t i=0;i<encoded.size();++i) {
        totals.bytes += encoded[i].bytes.size();
        totals.blocks += encoded[i].blocks;
        totals.nonempty += encoded[i].nonempty;
    }

    for(int rep=0;rep<DecodeRepeats;++rep) {
        for(size_t i=0;i<frame.planes.size();++i) {
            const auto raw=decodePlane(frame,frame.planes[i],encoded[i],block);
            if(raw!=frame.planes[i].raw) throw std::runtime_error("round-trip mismatch: "+file.string());
        }
    }
    auto end=Clock::now();
    totals.encodeMs += std::chrono::duration<double,std::milli>(mid-start).count()/EncodeRepeats;
    totals.decodeMs += std::chrono::duration<double,std::milli>(end-mid).count()/DecodeRepeats;
}
int main(int argc,char** argv) {
    try {
        const fs::path root=argc>1 ? fs::path(argv[1])
                                   : fs::path("integration/react-native/MapLibreParticleRNTestHost/assets/qbin-golden");
        std::vector<fs::path> files;
        for(const auto& e:fs::recursive_directory_iterator(root))
            if(e.is_regular_file() && e.path().extension()==".qbin") files.push_back(e.path());
        std::sort(files.begin(),files.end());
        if(files.empty()) throw std::runtime_error("no qbin fixtures");

        for(uint32_t block:{16u,32u}) {
            Totals total;
            std::cout<<"BLOCK="<<block<<"\n";
            for(const auto& file:files) {
                Totals one;
                runOne(file,block,one);
                std::cout<<file.parent_path().filename().string()<<"/"<<file.filename().string()
                         <<",bytes="<<one.bytes
                         <<",blocks="<<one.blocks
                         <<",nonempty="<<one.nonempty
                         <<",encodeMs="<<one.encodeMs
                         <<",decodeMs="<<one.decodeMs<<"\n";
                total.bytes+=one.bytes; total.blocks+=one.blocks; total.nonempty+=one.nonempty;
                total.encodeMs+=one.encodeMs; total.decodeMs+=one.decodeMs;
            }
            std::cout<<"TOTAL,block="<<block
                     <<",bytes="<<total.bytes
                     <<",blocks="<<total.blocks
                     <<",nonempty="<<total.nonempty
                     <<",encodeMs="<<total.encodeMs
                     <<",decodeMs="<<total.decodeMs<<"\n";
        }
        return 0;
    } catch(const std::exception& e) {
        std::cerr<<"FAIL: "<<e.what()<<"\n";
        return 1;
    }
}