/*
 * NanoDetHeat v10: 3.2M MACs, 4-class
 * 优化: BN折叠 + conv_bn_relu融合 + 快速exp + 边界修正
 */
#include "model_v10.hpp"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <ctime>
#include <new>
#include <algorithm>
#include <unordered_map>
#include <string>
#include <vector>

// ============================================================
// 快速 exp — Schraudolph 方法 (IEEE 754 整数近似)
// ============================================================
static inline float fast_exp(float x) {
    // clamp to avoid overflow: exp(88) fits in float
    x = std::min(x, 88.0f);
    x = std::max(x, -87.0f);
    // exp(x) ~ 2^(x * log2(e))
    // float 按位: 2^23 * (127 + x*log2(e))
    union { float f; int i; } u;
    u.i = (int)(12102203.0f * x) + 1064866805;  // 2^23 * log2(e) = 12102203
    return u.f;
}
static inline float fast_sigmoid(float x) {
    return 1.0f / (1.0f + fast_exp(-x));
}

// ============================================================
// 权重 (新增 BN alpha/beta 预计算)
// ============================================================
struct WT { char n[128]; int nd; int s[4]; float* d; };
static int gn=0; static WT* gw=nullptr;

static float* wf(const char* k) {
    for(int i=0;i<gn;++i) if(!std::strcmp(gw[i].n,k)) return gw[i].d;
    return nullptr;
}

static void init_lut();

// 加载后扫描 BN 参数，预计算 alpha/beta 存入权重表
static void bn_precompute() {
    // 扫描所有权重，找 BN 三元组 (weight/bias/running_mean/running_var)
    std::unordered_map<std::string, int> base_idx;
    for (int i=0; i<gn; ++i) {
        const char* name = gw[i].n;
        // 匹配 xxx.weight (BN) 或 xxx.bias (BN) → base = "xxx"
        const char* dot = std::strrchr(name, '.');
        if (!dot) continue;
        std::string base(name, dot - name);
        int sz = gw[i].s[0];

        if (std::strcmp(dot+1, "weight")==0 && sz>0 && sz <= 256) {
            base_idx[base] = i;
        }
    }

    std::vector<WT> new_wts;
    for (auto& [base, wi] : base_idx) {
        // 找对应的 bias / running_mean / running_var
        std::string w_key = base + ".weight";
        std::string b_key = base + ".bias";
        std::string m_key = base + ".running_mean";
        std::string v_key = base + ".running_var";

        float* w = wf(w_key.c_str());
        float* b = wf(b_key.c_str());
        float* m = wf(m_key.c_str());
        float* v = wf(v_key.c_str());

        if (!w || !m || !v) continue;  // cls_head/size_head 没有 BN
        int C = gw[wi].s[0];
        float* alpha = new float[C];
        float* beta  = new float[C];
        const float eps = 1e-5f;
        for (int c=0; c<C; ++c) {
            alpha[c] = w[c] / std::sqrt(v[c] + eps);
            beta[c]  = (b ? b[c] : 0.0f) - m[c] * alpha[c];
        }
        // 存为新的权重条目
        new_wts.push_back({});
        std::snprintf(new_wts.back().n, 128, "%s._alpha", base.c_str());
        new_wts.back().nd=1; new_wts.back().s[0]=C; new_wts.back().s[1]=1; new_wts.back().s[2]=1; new_wts.back().s[3]=1;
        new_wts.back().d = alpha;
        new_wts.push_back({});
        std::snprintf(new_wts.back().n, 128, "%s._beta", base.c_str());
        new_wts.back().nd=1; new_wts.back().s[0]=C; new_wts.back().s[1]=1; new_wts.back().s[2]=1; new_wts.back().s[3]=1;
        new_wts.back().d = beta;
    }

    // 重新分配权重数组 (旧数组 + 新条目)
    WT* old_gw = gw;
    int old_gn = gn;
    int new_gn = old_gn + (int)new_wts.size();
    gw = new WT[new_gn];
    std::memcpy(gw, old_gw, old_gn * sizeof(WT));
    for (size_t i=0; i<new_wts.size(); ++i) {
        gw[old_gn + i] = new_wts[i];
    }
    gn = new_gn;
    delete[] old_gw;
    init_lut();
}

// ============================================================
// 内存
// ============================================================
struct M10 {
    float preproc[3*120*160];  // 预处理 buffer (静态, 避免每帧 new/delete)
    float stem[6*60*80];
    float b1[8*60*80];
    float b2[12*30*40];
    float b3[16*15*20];
    float b4[32*15*20];
    float b5[48*15*20];
    float b6[48*15*20];
    float sh[24*15*20];
    float cl[5*15*20];
    float sz[2*15*20];
    float dw_out[8*60*80];
    float se_buf[80];
    uint8 peak_mask[15*20];
};
static M10* m=nullptr;

// ============================================================
// Relu
// ============================================================
static void relu_f(float* x, int N) {
    for (int i=0; i<N; ++i) if (x[i]<0) x[i]=0;
}

// ============================================================
// 融合: conv + BN + ReLU (使用预计算的 alpha/beta)
// ============================================================
static void conv_bn_relu(float* o, const float* in,
    const float* conv_w, const float* conv_b,
    const float* bn_alpha, const float* bn_beta,
    int H, int W, int iC, int oC, int K, int str, int grp, bool use_relu)
{
    if (!conv_w || !in || !o) return;
    int Ho=H/str, Wo=W/str, pad=K/2;
    int icpg=iC/grp, ocpg=oC/grp;

    if (K==1 && str==1) {
        // 1×1 fused: 指针递增避免 y*W+x
        for (int g=0; g<grp; ++g) {
            for (int oc=0; oc<ocpg; ++oc) {
                int occ = g*ocpg + oc;
                float* oo = o + occ*H*W;
                float cb = conv_b ? conv_b[occ] : 0.0f;
                const float* ww = conv_w + occ*icpg;
                float ba=bn_alpha?bn_alpha[occ]:1.0f, bb=bn_beta?bn_beta[occ]:0.0f;

                for (int y=0; y<H; ++y) {
                    float* oy = oo + y*W;
                    const float* iy = in + y*W;
                    for (int x=0; x<W; ++x) {
                        float s = cb;
                        const float* px = iy + x;
                        for (int ic=0; ic<icpg; ++ic) {
                            int icc = g*icpg + ic;
                            s += ww[ic] * px[icc*H*W];
                        }
                        s = s*ba + bb;
                        if (use_relu && s<0) s = 0;
                        oy[x] = s;
                    }
                }
            }
        }
        return;
    }

    if (grp==iC && K==3 && str==1) {
        // 3×3 dw fused + interior unrolled / border with checks
        for (int c=0; c<iC; ++c) {
            float*       const oo = o + c*H*W;
            const float* const ii = in + c*H*W;
            const float* const ww = conv_w + c*9;
            float cb = conv_b ? conv_b[c] : 0.0f;
            float ba=bn_alpha?bn_alpha[c]:1.0f, bb=bn_beta?bn_beta[c]:0.0f;
            float w0=ww[0], w1=ww[1], w2=ww[2], w3=ww[3], w4=ww[4], w5=ww[5], w6=ww[6], w7=ww[7], w8=ww[8];

            if (H>=3 && W>=3) {
                // interior: 1≤y<H-1, 1≤x<W-1 (no boundary checks)
                for (int y=1; y<H-1; ++y) {
                    float*       const oy = oo + y*W + 1;
                    const float* const i0 = ii + (y-1)*W;
                    const float* const i1 = i0 + W;
                    const float* const i2 = i1 + W;
                    for (int x=1; x<W-1; ++x) {
                        float s = cb
                            + w0*i0[x-1] + w1*i0[x] + w2*i0[x+1]
                            + w3*i1[x-1] + w4*i1[x] + w5*i1[x+1]
                            + w6*i2[x-1] + w7*i2[x] + w8*i2[x+1];
                        s = s*ba + bb;
                        if (use_relu && s<0) s = 0;
                        oy[x-1] = s;
                    }
                }
                // top row (y=0): boundary check — ky=-1 clamps to y=0
                {
                    float*       const oy = oo;
                    const float* const i0 = ii;         // y=0 (ky=-1 and ky=0)
                    const float* const i1 = ii + W;     // y=1 (ky=1)
                    for (int x=0; x<W; ++x) {
                        float s = cb;
                        // ky=-1: use i0 (clamped), ky=0: use i0, ky=1: use i1
                        int xl=(x>0)?x-1:x, xr=(x<W-1)?x+1:x;
                        s += w0*i0[xl] + w1*i0[x] + w2*i0[xr]   // ky=-1
                           + w3*i0[xl] + w4*i0[x] + w5*i0[xr]   // ky=0
                           + w6*i1[xl] + w7*i1[x] + w8*i1[xr];  // ky=1
                        s = s*ba + bb;
                        if (use_relu && s<0) s = 0;
                        oy[x] = s;
                    }
                }
                // bottom row (y=H-1): boundary check — ky=1 clamps to y=H-1
                {
                    float*       const oy = oo + (H-1)*W;
                    const float* const i0 = ii + (H-2)*W;        // y=H-2 (ky=-1)
                    const float* const i1 = ii + (H-1)*W;        // y=H-1 (ky=0 and ky=1)
                    for (int x=0; x<W; ++x) {
                        float s = cb;
                        int xl=(x>0)?x-1:x, xr=(x<W-1)?x+1:x;
                        s += w0*i0[xl] + w1*i0[x] + w2*i0[xr]   // ky=-1
                           + w3*i1[xl] + w4*i1[x] + w5*i1[xr]   // ky=0
                           + w6*i1[xl] + w7*i1[x] + w8*i1[xr];  // ky=1
                        s = s*ba + bb;
                        if (use_relu && s<0) s = 0;
                        oy[x] = s;
                    }
                }
                // middle rows left border (1≤y<H-1, x=0)
                for (int y=1; y<H-1; ++y) {
                    const float* const i0 = ii + (y-1)*W;
                    const float* const i1 = i0 + W;
                    const float* const i2 = i1 + W;
                    float s = cb
                        + w1*i0[0] + w2*i0[1]
                        + w4*i1[0] + w5*i1[1]
                        + w7*i2[0] + w8*i2[1];
                    s = s*ba + bb;
                    if (use_relu && s<0) s = 0;
                    oo[y*W] = s;
                }
                // middle rows right border (1≤y<H-1, x=W-1)
                for (int y=1; y<H-1; ++y) {
                    const float* const i0 = ii + (y-1)*W + (W-2);
                    const float* const i1 = i0 + W;
                    const float* const i2 = i1 + W;
                    float s = cb
                        + w0*i0[0] + w1*i0[1]
                        + w3*i1[0] + w4*i1[1]
                        + w6*i2[0] + w7*i2[1];
                    s = s*ba + bb;
                    if (use_relu && s<0) s = 0;
                    oo[y*W+W-1] = s;
                }
            } else {
                // degenerate (small H/W): fallback
                for (int y=0; y<H; ++y) {
                    for (int x=0; x<W; ++x) {
                        float s = cb;
                        for (int ky=0; ky<3; ++ky) {
                            int iy = y + ky - 1;
                            if (iy<0||iy>=H) continue;
                            for (int kx=0; kx<3; ++kx) {
                                int ix = x + kx - 1;
                                if (ix<0||ix>=W) continue;
                                s += ww[ky*3+kx] * ii[iy*W+ix];
                            }
                        }
                        s = s*ba + bb;
                        if (use_relu && s<0) s = 0;
                        oo[y*W+x] = s;
                    }
                }
            }
        }
        return;
    }

    // 通用路径 — 常数外提 + iy*W 预计算
    for (int g=0; g<grp; ++g) {
        int ic_base = g*icpg;
        for (int oc=0; oc<ocpg; ++oc) {
            int occ = g*ocpg + oc;
            float* oo = o + occ*Ho*Wo;
            float cb = conv_b ? conv_b[occ] : 0.0f;
            float ba=bn_alpha?bn_alpha[occ]:1.0f, bb=bn_beta?bn_beta[occ]:0.0f;

            for (int y=0; y<Ho; ++y) {
                int yi0 = y*str;
                for (int x=0; x<Wo; ++x) {
                    float s = cb;
                    int xi0 = x*str;
                    for (int ic=0; ic<icpg; ++ic) {
                        int icc = ic_base + ic;
                        const float* ww = conv_w + ((occ*icpg+ic)*K*K);
                        const float* ii = in + icc*H*W;
                        for (int ky=0; ky<K; ++ky) {
                            int iy = yi0 + ky - pad;
                            if (iy<0||iy>=H) continue;
                            int iy_off = iy*W;
                            for (int kx=0; kx<K; ++kx) {
                                int ix = xi0 + kx - pad;
                                if (ix<0||ix>=W) continue;
                                s += ww[ky*K+kx] * ii[iy_off + ix];
                            }
                        }
                    }
                    s = s*ba + bb;
                    if (use_relu && s<0) s = 0;
                    oo[y*Wo+x] = s;
                }
            }
        }
    }
}

// ============================================================
// 纯 conv (无 BN/ReLU, cls_head / size_head / skip 的 conv 部分)
// ============================================================
static void conv2d(float* o, const float* in, const float* w, const float* b,
                   int H, int W, int iC, int oC, int K, int str, int grp) {
    conv_bn_relu(o, in, w, b, nullptr, nullptr, H, W, iC, oC, K, str, grp, false);
}

// ============================================================
// SE 通道注意力 (使用快速 sigmoid)
// ============================================================
static void se_module(float* x, int C, int N,
                      const float* w1, const float* b1,
                      const float* w2, const float* b2,
                      float* buf) {
    if (!w1||!w2) return;
    int mid=C/4;
    float* pool=buf;

    // AdaptiveAvgPool
    for (int c=0; c<C; ++c) {
        float s=0;
        float* px = x+c*N, *end = px+N;
        while (px < end) { s += *px++; }
        pool[c]=s/(float)N;
    }
    // fc1: C→mid + ReLU
    float* fc1=buf+C;
    for (int o=0; o<mid; ++o) {
        float s=b1?b1[o]:0;
        const float* ww=w1+o*C;
        for (int ic=0; ic<C; ++ic) s+=pool[ic]*ww[ic];
        fc1[o]=s>0?s:0;
    }
    // fc2: mid→C + Sigmoid (fast)
    float* fc2=buf+C+mid;
    for (int o=0; o<C; ++o) {
        float s=b2?b2[o]:0;
        const float* ww=w2+o*mid;
        for (int ic=0; ic<mid; ++ic) s+=fc1[ic]*ww[ic];
        fc2[o]=fast_sigmoid(s);
    }
    for (int c=0; c<C; ++c) {
        float* xc=x+c*N; float sc=fc2[c];
        for (int i=0; i<N; ++i) xc[i]*=sc;
    }
}

// ============================================================
// SE-ResDW Block v10
// ============================================================
static void se_block_v10(float* o, const float* in, int iC, int oC, int H, int W, int str, const char* pfx) {
    int Ho=H/str, Wo=W/str, No=Ho*Wo;
    char na[128], wr[128], bi[128];

    // dw conv + BN + ReLU (fused)
    std::snprintf(na,128,"%s.dw.0.weight",pfx);
    std::snprintf(wr,128,"%s.dw.1._alpha",pfx);
    std::snprintf(bi,128,"%s.dw.1._beta",pfx);
    conv_bn_relu(m->dw_out, in, wf(na), nullptr,
                 wf(wr), wf(bi), H, W, iC, iC, 3, str, iC, true);

    // pw conv + BN + ReLU (fused)
    std::snprintf(na,128,"%s.pw.0.weight",pfx);
    std::snprintf(wr,128,"%s.pw.1._alpha",pfx);
    std::snprintf(bi,128,"%s.pw.1._beta",pfx);
    conv_bn_relu(o, m->dw_out, wf(na), nullptr,
                 wf(wr), wf(bi), Ho, Wo, iC, oC, 1, 1, 1, true);

    // skip (BN fused, no ReLU)
    bool iden=(str==1 && iC==oC);
    float* skip=m->dw_out;
    if (iden) {
        std::memcpy(skip, in, iC*H*W*4);
    } else {
        std::snprintf(na,128,"%s.sk.0.weight",pfx);
        std::snprintf(wr,128,"%s.sk.1._alpha",pfx);
        std::snprintf(bi,128,"%s.sk.1._beta",pfx);
        conv_bn_relu(skip, in, wf(na), nullptr,
                     wf(wr), wf(bi), H, W, iC, oC, 1, str, 1, false);
    }
    for (int i=0; i<oC*No; ++i) o[i]+=skip[i];

    // SE
    std::snprintf(na,128,"%s.se.1.weight",pfx);
    std::snprintf(bi,128,"%s.se.1.bias",pfx);
    char se2[128], se2b[128];
    std::snprintf(se2,128,"%s.se.3.weight",pfx);
    std::snprintf(se2b,128,"%s.se.3.bias",pfx);
    se_module(o, oC, No, wf(na), wf(bi), wf(se2), wf(se2b), m->se_buf);

    relu_f(o, oC*No);
}

// ============================================================
// 解码 — softmax + NMS (预计算 peak mask)
// ============================================================
static constexpr int OH=15, OW=20, STRIDE=8, NC=4;
static constexpr float RS[NC][2] = {
    {56.85f,45.39f}, {101.44f,66.53f}, {96.75f,62.86f}, {66.01f,35.49f},
};

static int decode(DetectBoxV10* boxes, int max, const float* th) {
    int N=OH*OW, cnt=0;

    // 预计算 peak mask: 每个类别做一次 maxpool
    uint8* mask = m->peak_mask;
    for (int c=0; c<NC; ++c) {
        const float* cp = m->cl + c*N;
        uint8* mp = mask + c*N;
        for (int gy=0; gy<OH; ++gy) {
            for (int gx=0; gx<OW; ++gx) {
                if (cp[gy*OW+gx] < 0) { mp[gy*OW+gx]=0; continue; }
                bool peak=true;
                float val=cp[gy*OW+gx];
                for (int dy=-1; dy<=1&&peak; ++dy)
                for (int dx=-1; dx<=1&&peak; ++dx) {
                    int ny=gy+dy, nx=gx+dx;
                    if (ny<0||ny>=OH||nx<0||nx>=OW) continue;
                    if (cp[ny*OW+nx] > val) peak=false;
                }
                mp[gy*OW+gx] = peak ? 1 : 0;
            }
        }
    }

    // softmax + check per cell
    for (int gy=0; gy<OH && cnt<max; ++gy) {
        for (int gx=0; gx<OW && cnt<max; ++gx) {
            int idx=gy*OW+gx;

            // softmax over NC+1 classes (fast exp)
            float sm[5], mx=-1e9f, sum=0;
            for (int c=0; c<=NC; ++c) {
                float v=m->cl[c*N+idx];
                sm[c]=v; if(v>mx) mx=v;
            }
            for (int c=0; c<=NC; ++c) { sm[c]=fast_exp(sm[c]-mx); sum+=sm[c]; }

            // best class (excluding bg=NC), per-class threshold
            int bc=-1; float bs=0;
            for (int c=0; c<NC; ++c) {
                if (!mask[c*N+idx]) continue;
                float p=sm[c]/sum;
                if (p>bs && p>=th[c]) { bs=p; bc=c; }
            }
            if (bc<0) continue;

            float pw=std::max(1.0f, m->sz[0*N+idx]*RS[bc][0]);
            float ph=std::max(1.0f, m->sz[1*N+idx]*RS[bc][1]);
            boxes[cnt].cls=bc; boxes[cnt].conf=bs;
            boxes[cnt].cx=((float)gx+0.5f)*STRIDE;
            boxes[cnt].cy=((float)gy+0.5f)*STRIDE;
            boxes[cnt].w=pw; boxes[cnt].h=ph;
            cnt++;
        }
    }
    return cnt;
}

// ============================================================
// 预处理 LUT: uint8→float normalized
// ============================================================
static float g_lut[256];
static void init_lut() {
    for (int i=0; i<256; ++i) g_lut[i] = (float)i / 255.0f;
}

// ============================================================
// 前向推理 — 模型计算体 (preproc 已由调用方完成)
// ============================================================
static void model_compute() {
    float* in = m->preproc;

    // stem: 3→6, s2 (fused)
    conv_bn_relu(m->stem, in, wf("s.0.weight"), nullptr,
                 wf("s.1._alpha"), wf("s.1._beta"), 120, 160, 3, 6, 3, 2, 1, true);

    // blocks
    se_block_v10(m->b1, m->stem, 6, 8,   60, 80, 1, "b1");
    se_block_v10(m->b2, m->b1,   8, 12,  60, 80, 2, "b2");
    se_block_v10(m->b3, m->b2,   12, 16, 30, 40, 2, "b3");
    se_block_v10(m->b4, m->b3,   16, 32, 15, 20, 1, "b4");
    se_block_v10(m->b5, m->b4,   32, 48, 15, 20, 1, "b5");
    se_block_v10(m->b6, m->b5,   48, 48, 15, 20, 1, "b6");

    // shared: 48→24, 1×1 (fused)
    conv_bn_relu(m->sh, m->b6, wf("sh.0.weight"), nullptr,
                 wf("sh.1._alpha"), wf("sh.1._beta"), 15, 20, 48, 24, 1, 1, 1, true);

    // heads (pure conv, no BN)
    conv2d(m->cl, m->sh, wf("ch.weight"), wf("ch.bias"), 15, 20, 24, 5, 1, 1, 1);
    conv2d(m->sz, m->sh, wf("sz.weight"), wf("sz.bias"), 15, 20, 24, 2, 1, 1, 1);
}

// ── 自写 640×480 → 3×120×160 float planar (4×4 box avg + BGR→RGB + /255) ──
static void preproc_640(const uint8* bgr) {
    float* in = m->preproc;
    const int sw=640, sh=480;
    for(int c=0;c<3;++c){
        int sc=2-c;
        float* ch = in + c*120*160;
        for(int y=0;y<120;++y){
            int iy=y*4;
            for(int x=0;x<160;++x){
                int ix=x*4, sum=0;
                for(int dy=0;dy<4;++dy)
                    for(int dx=0;dx<4;++dx)
                        sum += bgr[(iy+dy)*sw*3 + (ix+dx)*3 + sc];
                ch[y*160+x] = (float)sum * (1.0f/(16.0f*255.0f));
            }
        }
    }
}

// ── 原版: 160×120 BGR → 3×120×160 float planar ──
static void preproc_160(const uint8* bgr) {
    float* in = m->preproc;
    for (int c=0; c<3; ++c) {
        int sc=2-c; float* ch=in+c*120*160;
        for (int y=0; y<120; ++y) {
            const uint8* row=bgr+y*160*3;
            for (int x=0; x<160; ++x) ch[y*160+x]=g_lut[row[x*3+sc]];
        }
    }
}

// ── 内部: preproc + compute ──
static void forward(const uint8* bgr) { preproc_160(bgr); model_compute(); }

// ============================================================
// 接口
// ============================================================
static bool g_rdy=false;

bool model_v10_init(const char* path) {
    FILE* f=fopen(path,"rb");
    if(!f) return false;
    fread(&gn,sizeof(int),1,f);
    gw=new WT[gn];
    for (int i=0; i<gn; ++i) {
        WT& t=gw[i]; int nl; fread(&nl,4,1,f); fread(t.n,1,nl,f); t.n[nl]=0;
        fread(&t.nd,4,1,f); int tot=1;
        for (int d=0; d<t.nd; ++d) { fread(&t.s[d],4,1,f); tot*=t.s[d]; }
        for (int d=t.nd; d<4; ++d) t.s[d]=1;
        t.d=new float[tot]; fread(t.d,4,tot,f);
    }
    fclose(f);
    bn_precompute();
    m=new(std::nothrow) M10();
    if(!m) { delete[] gw; gw=nullptr; return false; }
    std::memset(m,0,sizeof(M10));
    g_rdy=true;
    return true;
}

int model_v10_detect(const uint8* bgr, int w, int h, DetectBoxV10* boxes, int max, const float* thresh) {
    if(!g_rdy) return 0;
    if(w==640 && h==480) {
        preproc_640(bgr);
        model_compute();
        return decode(boxes,max,thresh);
    }
    if(w!=160||h!=120) return 0;
    forward(bgr);
    return decode(boxes,max,thresh);
}

void model_v10_deinit() {
    if(gw) { for(int i=0; i<gn; ++i) delete[] gw[i].d; delete[] gw; gw=nullptr; }
    delete m; m=nullptr; g_rdy=false;
}

bool model_v10_ready() { return g_rdy; }
