/*
 * Mild Mega r2 C++ inference engine v3 (基于魔改YOLO架构)
 * Architecture: 3-9-9-13-13-19-19-19-44-44-64-64 -> head(64-64,dw) -> out(64-9)
 */
#include "model_v10.hpp"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <new>
#include <algorithm>
#include <unordered_map>
#include <string>
#include <vector>

static inline float clamp_f(float v, float lo, float hi) { return v < lo ? lo : (v > hi ? hi : v); }
static inline float sigmoid_f(float x) {
    x = clamp_f(x, -30.0f, 30.0f);
    return 1.0f / (1.0f + std::exp(-x));
}
static inline float relu_maybe(float s, bool relu) { return (relu && s < 0.0f) ? 0.0f : s; }

struct WT { char n[128]; int nd; int s[4]; float* d; };
static int gn = 0;
static WT* gw = nullptr;

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

static void bn_precompute() {
    std::unordered_map<std::string, int> bn_weight_idx;
    for (int i = 0; i < gn; ++i) {
        const char* name = gw[i].n;
        const char* dot = std::strrchr(name, '.');
        if (!dot || std::strcmp(dot + 1, "weight") != 0) continue;
        if (gw[i].nd != 1) continue;
        bn_weight_idx[std::string(name, dot - name)] = i;
    }
    std::vector<WT> new_wts;
    for (auto& kv : bn_weight_idx) {
        const std::string& base = kv.first;
        float* w = gw[kv.second].d;
        float* b = wf((base + ".bias").c_str());
        float* m = wf((base + ".running_mean").c_str());
        float* v = wf((base + ".running_var").c_str());
        if (!m || !v) continue;
        const int C = gw[kv.second].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];
        }
        WT ta = {}; std::snprintf(ta.n, 128, "%s._alpha", base.c_str());
        ta.nd = 1; ta.s[0] = C; ta.s[1] = ta.s[2] = ta.s[3] = 1; ta.d = alpha;
        WT tb = {}; std::snprintf(tb.n, 128, "%s._beta", base.c_str());
        tb.nd = 1; tb.s[0] = C; tb.s[1] = tb.s[2] = tb.s[3] = 1; tb.d = beta;
        new_wts.push_back(ta); new_wts.push_back(tb);
    }
    WT* old = gw;
    const int old_n = gn, new_n = old_n + (int)new_wts.size();
    gw = new WT[new_n];
    std::memcpy(gw, old, old_n * sizeof(WT));
    for (size_t i = 0; i < new_wts.size(); ++i) gw[old_n + i] = new_wts[i];
    gn = new_n;
    delete[] old;
}

static constexpr int IN_W = 160, IN_H = 120;
static constexpr int BUF_FLOATS = 9 * 60 * 80;
static float*  g_preproc = nullptr;
static float*  g_A = nullptr;
static float*  g_B = nullptr;
static float*  g_T = nullptr;
static float*  g_out = nullptr;
static uint8_t g_resized[IN_W * IN_H * 3];
static float   g_se[160];
static float   g_prob[5 * 15 * 20];
static float   g_lut[256];
static char    g_dump_dir[256] = {0};

static void dump_tensor(const char* name, const float* d, int count) {
    if (!g_dump_dir[0]) return;
    char path[320];
    std::snprintf(path, sizeof(path), "%s/%s.f32", g_dump_dir, name);
    FILE* f = std::fopen(path, "wb");
    if (f) { std::fwrite(d, sizeof(float), count, f); std::fclose(f); }
}

static void conv1x1(float* o, const float* in, const float* cw, const float* cb,
                    const float* ba, const float* bb,
                    int H, int W, int iC, int oC, bool relu)
{
    const int N = H * W;
    for (int oc = 0; oc < oC; ++oc) {
        float* oo = o + oc * N;
        const float* ww = cw + oc * iC;
        const float bias = cb ? cb[oc] : 0.0f;
        for (int i = 0; i < N; ++i) oo[i] = bias;
        for (int ic = 0; ic < iC; ++ic) {
            const float w = ww[ic];
            const float* ii = in + ic * N;
            for (int i = 0; i < N; ++i) oo[i] += w * ii[i];
        }
        const float a = ba ? ba[oc] : 1.0f;
        const float b = bb ? bb[oc] : 0.0f;
        for (int i = 0; i < N; ++i)
            oo[i] = relu_maybe(oo[i] * a + b, relu);
    }
}

static void dwconv3x3(float* o, const float* in, const float* cw,
                      const float* ba, const float* bb,
                      int H, int W, int C, bool relu)
{
    for (int c = 0; c < C; ++c) {
        const float* ii = in + c * H * W;
        float* oo = o + c * H * W;
        const float* ww = cw + c * 9;
        const 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];
        const float a = ba ? ba[c] : 1.0f;
        const float b = bb ? bb[c] : 0.0f;

        for (int y = 1; y < H - 1; ++y) {
            const float *i0 = ii + (y - 1) * W, *i1 = i0 + W, *i2 = i1 + W;
            float* oy = oo + y * W;
            for (int x = 1; x < W - 1; ++x) {
                float s = 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];
                oy[x] = relu_maybe(s * a + b, relu);
            }
        }
        auto edge = [&](int y, int x) {
            float s = 0.0f;
            for (int ky = 0; ky < 3; ++ky) {
                const int iy = y + ky - 1;
                if (iy < 0 || iy >= H) continue;
                for (int kx = 0; kx < 3; ++kx) {
                    const int ix = x + kx - 1;
                    if (ix < 0 || ix >= W) continue;
                    s += ww[ky * 3 + kx] * ii[iy * W + ix];
                }
            }
            oo[y * W + x] = relu_maybe(s * a + b, relu);
        };
        for (int x = 0; x < W; ++x) { edge(0, x); edge(H - 1, x); }
        for (int y = 1; y < H - 1; ++y) { edge(y, 0); edge(y, W - 1); }
    }
}

static void conv_generic(float* o, const float* in, const float* cw, const float* cb,
                         const float* ba, const float* bb,
                         int H, int W, int iC, int oC, int K, int str, int grp, bool relu)
{
    const int Ho = H / str, Wo = W / str, pad = K / 2;
    const int icpg = iC / grp, ocpg = oC / grp;
    for (int g = 0; g < grp; ++g) {
        const int ic_base = g * icpg;
        for (int oc = 0; oc < ocpg; ++oc) {
            const int occ = g * ocpg + oc;
            float* oo = o + occ * Ho * Wo;
            const float bias = cb ? cb[occ] : 0.0f;
            const float a = ba ? ba[occ] : 1.0f;
            const float b = bb ? bb[occ] : 0.0f;
            for (int y = 0; y < Ho; ++y) {
                const int yi0 = y * str - pad;
                const int ky0 = std::max(0, -yi0);
                const int ky1 = std::min(K, H - yi0);
                for (int x = 0; x < Wo; ++x) {
                    const int xi0 = x * str - pad;
                    const int kx0 = std::max(0, -xi0);
                    const int kx1 = std::min(K, W - xi0);
                    float s = bias;
                    for (int ic = 0; ic < icpg; ++ic) {
                        const float* ww = cw + ((occ * icpg + ic) * K * K);
                        const float* ii = in + (ic_base + ic) * H * W;
                        for (int ky = ky0; ky < ky1; ++ky) {
                            const float* irow = ii + (yi0 + ky) * W + xi0;
                            const float* wrow = ww + ky * K;
                            for (int kx = kx0; kx < kx1; ++kx)
                                s += wrow[kx] * irow[kx];
                        }
                    }
                    oo[y * Wo + x] = relu_maybe(s * a + b, relu);
                }
            }
        }
    }
}

static void conv_bn_relu(float* o, const float* in,
                         const float* cw, const float* cb,
                         const float* ba, const float* bb,
                         int H, int W, int iC, int oC, int K, int str, int grp, bool relu)
{
    if (!cw || !in || !o) return;
    if (K == 1 && str == 1 && grp == 1)
        conv1x1(o, in, cw, cb, ba, bb, H, W, iC, oC, relu);
    else if (K == 3 && str == 1 && grp == iC && iC == oC)
        dwconv3x3(o, in, cw, ba, bb, H, W, iC, relu);
    else
        conv_generic(o, in, cw, cb, ba, bb, H, W, iC, oC, K, str, grp, relu);
}

static void se_module(float* x, int C, int N,
                      const WT* w1t, const float* b1,
                      const WT* w2t, const float* b2)
{
    if (!w1t || !w2t) return;
    const float* w1 = w1t->d;
    const float* w2 = w2t->d;
    const int mid = w1t->s[0];

    float* pool = g_se;
    float* fc1  = g_se + C;
    float* fc2  = g_se + C + mid;

    const float inv_n = 1.0f / (float)N;
    for (int c = 0; c < C; ++c) {
        float s = 0.0f;
        const float* px = x + c * N;
        for (int i = 0; i < N; ++i) s += px[i];
        pool[c] = s * inv_n;
    }
    for (int o = 0; o < mid; ++o) {
        float s = b1 ? b1[o] : 0.0f;
        const float* ww = w1 + o * C;
        for (int ic = 0; ic < C; ++ic) s += pool[ic] * ww[ic];
        fc1[o] = s > 0.0f ? s : 0.0f;
    }
    for (int o = 0; o < C; ++o) {
        float s = b2 ? b2[o] : 0.0f;
        const float* ww = w2 + o * mid;
        for (int ic = 0; ic < mid; ++ic) s += fc1[ic] * ww[ic];
        fc2[o] = sigmoid_f(s);
    }
    for (int c = 0; c < C; ++c) {
        float* xc = x + c * N;
        const float sc = fc2[c];
        for (int i = 0; i < N; ++i) xc[i] *= sc;
    }
}

static constexpr int STRIDE = 8, OH = 15, OW = 20, NC = 4;

static void model_compute() {
    struct Lay { int iC, oC, H, W, s; const char* bn; const char* sn; };
    static const Lay lays[] = {
        {3,  9,  120, 160, 2, "b1",  "se_0"},
        {9,  9,  60,  80,  1, "b2",  "se_1"},
        {9,  13, 60,  80,  2, "b3",  "se_2"},
        {13, 13, 30,  40,  1, "b4",  "se_3"},
        {13, 19, 30,  40,  2, "b5",  "se_4"},
        {19, 19, 15,  20,  1, "b6",  "se_5"},
        {19, 19, 15,  20,  1, "b7",  "se_6"},
        {19, 44, 15,  20,  1, "b8",  "se_7"},
        {44, 44, 15,  20,  1, "b9",  "se_8"},
        {44, 64, 15,  20,  1, "b10", "se_9"},
        {64, 64, 15,  20,  1, "b11", "se_10"},
    };

    const float* in = g_preproc;
    float* bufs[2] = { g_A, g_B };
    int cur = 0;

    for (const Lay& l : lays) {
        int K, grp;
        if (l.s == 2)          { K = 3; grp = 1;    }
        else if (l.iC == l.oC) { K = 3; grp = l.iC; }
        else                   { K = 1; grp = 1;    }

        const int Ho = l.H / l.s, Wo = l.W / l.s;
        float* out = bufs[cur];

        char wkey[128], akey[128], bkey[128];
        std::snprintf(wkey, 128, "%s.0.weight", l.bn);
        std::snprintf(akey, 128, "%s.1._alpha", l.bn);
        std::snprintf(bkey, 128, "%s.1._beta",  l.bn);
        conv_bn_relu(out, in, wf(wkey), nullptr, wf(akey), wf(bkey),
                     l.H, l.W, l.iC, l.oC, K, l.s, grp, true);

        char s1[128], s1b[128], s3[128], s3b[128];
        std::snprintf(s1,  128, "%s.1.weight", l.sn);
        std::snprintf(s1b, 128, "%s.1.bias",   l.sn);
        std::snprintf(s3,  128, "%s.3.weight", l.sn);
        std::snprintf(s3b, 128, "%s.3.bias",   l.sn);
        se_module(out, l.oC, Ho * Wo, wt(s1), wf(s1b), wt(s3), wf(s3b));

        in = out;
        cur ^= 1;
    }

    conv_bn_relu(g_T, in, wf("head.0.weight"), nullptr,
                 wf("head.1._alpha"), wf("head.1._beta"),
                 15, 20, 64, 64, 3, 1, 64, true);

    conv_bn_relu(g_out, g_T, wf("out.weight"), wf("out.bias"),
                 nullptr, nullptr, 15, 20, 64, 9, 1, 1, 1, false);
}

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

    for (int idx = 0; idx < N; ++idx) {
        float mx = -1e30f;
        for (int c = 0; c <= NC; ++c) mx = std::max(mx, g_out[c * N + idx]);
        float sum = 0.0f;
        for (int c = 0; c <= NC; ++c) {
            const float e = std::exp(g_out[c * N + idx] - mx);
            g_prob[c * N + idx] = e;
            sum += e;
        }
        const float inv = 1.0f / sum;
        for (int c = 0; c <= NC; ++c) g_prob[c * N + idx] *= inv;
    }

    for (int c = 0; c < NC && cnt < max_boxes; ++c) {
        float thr = thresh ? thresh[c] : 0.6f;
        for (int gy = 0; gy < OH; ++gy) {
            for (int gx = 0; gx < OW; ++gx) {
                const int idx = gy * OW + gx;
                const float pc = g_prob[c * N + idx];
                if (pc <= thr || pc <= g_prob[NC * N + idx]) continue;

                bool peak = true;
                for (int dy = -1; dy <= 1 && peak; ++dy)
                    for (int dx = -1; dx <= 1 && peak; ++dx) {
                        const int ny = gy + dy, nx = gx + dx;
                        if (ny < 0 || ny >= OH || nx < 0 || nx >= OW) continue;
                        if (g_prob[c * N + ny * OW + nx] > pc) peak = false;
                    }
                if (!peak) continue;

                const float tx = clamp_f(g_out[(NC + 1) * N + idx], 0.0f, 1.0f);
                const float ty = clamp_f(g_out[(NC + 2) * N + idx], 0.0f, 1.0f);
                const float tw = std::max(g_out[(NC + 3) * N + idx], 0.5f);
                const float th = std::max(g_out[(NC + 4) * N + idx], 0.5f);
                const float cx = (gx + tx) * STRIDE, cy = (gy + ty) * STRIDE;
                const float w = tw * STRIDE, h = th * STRIDE;

                boxes[cnt].cls  = c;
                boxes[cnt].conf = pc;
                boxes[cnt].cx = cx;
                boxes[cnt].cy = cy;
                boxes[cnt].w  = w;
                boxes[cnt].h  = h;
                if (++cnt >= max_boxes) return cnt;
            }
        }
    }
    return cnt;
}

// cv2.resize INTER_LINEAR compatible
static void resize_bilinear_hwc3(const uint8_t* src, int sw, int sh, uint8_t* dst) {
    const float fx = (float)sw / (float)IN_W;
    const float fy = (float)sh / (float)IN_H;
    static int   x0i[IN_W], x1i[IN_W];
    static float xw[IN_W];
    for (int x = 0; x < IN_W; ++x) {
        float sx = (x + 0.5f) * fx - 0.5f;
        if (sx < 0) sx = 0;
        int x0 = (int)sx;
        if (x0 > sw - 1) x0 = sw - 1;
        int x1 = std::min(x0 + 1, sw - 1);
        x0i[x] = x0 * 3; x1i[x] = x1 * 3; xw[x] = sx - (float)x0;
    }
    for (int y = 0; y < IN_H; ++y) {
        float sy = (y + 0.5f) * fy - 0.5f;
        if (sy < 0) sy = 0;
        int y0 = (int)sy;
        if (y0 > sh - 1) y0 = sh - 1;
        const int y1 = std::min(y0 + 1, sh - 1);
        const float wy = sy - (float)y0;
        const uint8_t* r0 = src + (size_t)y0 * sw * 3;
        const uint8_t* r1 = src + (size_t)y1 * sw * 3;
        uint8_t* d = dst + y * IN_W * 3;
        for (int x = 0; x < IN_W; ++x) {
            const float wx = xw[x];
            const int a = x0i[x], b = x1i[x];
            for (int c = 0; c < 3; ++c) {
                const float top = r0[a + c] + (r0[b + c] - r0[a + c]) * wx;
                const float bot = r1[a + c] + (r1[b + c] - r1[a + c]) * wx;
                const float v = top + (bot - top) * wy;
                d[x * 3 + c] = (uint8_t)(v + 0.5f);
            }
        }
    }
}

static void preproc_160(const uint8_t* img, MildPixFmt fmt) {
    for (int c = 0; c < 3; ++c) {
        const int sc = (fmt == MILD_FMT_BGR) ? (2 - c) : c;
        float* ch = g_preproc + c * IN_H * IN_W;
        for (int y = 0; y < IN_H; ++y) {
            const uint8_t* row = img + y * IN_W * 3;
            for (int x = 0; x < IN_W; ++x) ch[y * IN_W + x] = g_lut[row[x * 3 + sc]];
        }
    }
}

static bool g_rdy = false;

bool model_v10_init(const char* path) {
    FILE* f = std::fopen(path, "rb");
    if (!f) return false;
    if (std::fread(&gn, sizeof(int), 1, f) != 1 || gn <= 0 || gn > 4096) { std::fclose(f); return false; }
    gw = new WT[gn];
    for (int i = 0; i < gn; ++i) {
        WT& t = gw[i];
        int nl = 0;
        bool ok = std::fread(&nl, 4, 1, f) == 1 && nl > 0 && nl < 127
               && std::fread(t.n, 1, nl, f) == (size_t)nl;
        if (ok) {
            t.n[nl] = 0;
            ok = std::fread(&t.nd, 4, 1, f) == 1 && t.nd >= 0 && t.nd <= 4;
        }
        if (ok) {
            int tot = 1;
            for (int d = 0; d < t.nd && ok; ++d) {
                ok = std::fread(&t.s[d], 4, 1, f) == 1 && t.s[d] > 0;
                tot *= t.s[d];
            }
            for (int d = t.nd; d < 4; ++d) t.s[d] = 1;
            if (ok) {
                t.d = new float[tot];
                ok = std::fread(t.d, 4, tot, f) == (size_t)tot;
            }
        }
        if (!ok) { gn = (t.d ? i + 1 : i); std::fclose(f); model_v10_deinit(); return false; }
    }
    std::fclose(f);

    bn_precompute();
    for (int i = 0; i < 256; ++i) g_lut[i] = (float)i / 255.0f;

    g_preproc = new(std::nothrow) float[3 * IN_H * IN_W];
    g_A   = new(std::nothrow) float[BUF_FLOATS];
    g_B   = new(std::nothrow) float[BUF_FLOATS];
    g_T   = new(std::nothrow) float[BUF_FLOATS];
    g_out = new(std::nothrow) float[9 * 15 * 20];
    if (!g_preproc || !g_A || !g_B || !g_T || !g_out) { model_v10_deinit(); return false; }

    g_rdy = true;
    return true;
}

int model_v10_detect_fmt(const uint8_t* img, int w, int h, MildPixFmt fmt,
                    DetectBoxV10* boxes, int max_boxes, const float* thresh) {
    if (!g_rdy || !img || !boxes || max_boxes <= 0 || w <= 1 || h <= 1) return 0;

    const uint8_t* net_in = img;
    if (w != IN_W || h != IN_H) {
        resize_bilinear_hwc3(img, w, h, g_resized);
        net_in = g_resized;
    }
    preproc_160(net_in, fmt);
    model_compute();
    const int n = decode(boxes, max_boxes, thresh);

    const float sx = (float)w / (float)IN_W;
    const float sy = (float)h / (float)IN_H;
    for (int i = 0; i < n; ++i) {
        boxes[i].cx *= sx; boxes[i].w  *= sx;
        boxes[i].cy *= sy; boxes[i].h  *= sy;
    }
    return n;
}

int model_v10_detect(const uint8_t* img, int w, int h, DetectBoxV10* boxes, int max_boxes, const float* thresh) {
    return model_v10_detect_fmt(img, w, h, MILD_FMT_BGR, boxes, max_boxes, thresh);
}

const float* model_v10_forward_debug(const float* chw) {
    if (!g_rdy || !chw) return nullptr;
    std::memcpy(g_preproc, chw, 3 * IN_H * IN_W * sizeof(float));
    model_compute();
    return g_out;
}

void model_v10_set_dump_dir(const char* dir) {
    if (dir) { std::strncpy(g_dump_dir, dir, sizeof(g_dump_dir) - 1); g_dump_dir[sizeof(g_dump_dir)-1] = 0; }
    else g_dump_dir[0] = 0;
}

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

bool model_v10_ready() { return g_rdy; }