#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include <omp.h>

struct TraceHeader {
    char magic[4];
    uint16_t version;
    uint16_t n_layers;
    uint16_t n_experts;
    uint16_t top_m;
    uint32_t n_tokens;
};

static inline float fp16_to_fp32(uint16_t h) {
    uint32_t sign = (uint32_t)(h & 0x8000) << 16;
    uint32_t exp  = (h >> 10) & 0x1f;
    uint32_t mant = h & 0x03ff;
    uint32_t f;
    if (exp == 0) {
        if (mant == 0) f = sign;
        else {
            exp = 1;
            while ((mant & 0x0400) == 0) { mant <<= 1; exp--; }
            mant &= 0x03ff;
            f = sign | ((exp + (127 - 15)) << 23) | (mant << 13);
        }
    } else if (exp == 31) {
        f = sign | 0x7f800000 | (mant << 13);
    } else {
        f = sign | ((exp + (127 - 15)) << 23) | (mant << 13);
    }
    float out;
    memcpy(&out, &f, 4);
    return out;
}

enum Policy { POL_LRU=0, POL_LFU=1, POL_EMA=2, POL_NEAR_PICK=3, POL_INV_ROUTER=4, POL_BELADY=5 };

struct TokenRec {
    int n_cand;
    uint16_t ids[64];
    float logits[64];
};

struct ReplayResult {
    double misses_per_tok;
    double hit_rate;
};

ReplayResult replay_one_layer(
    const std::vector<TokenRec>& trace,
    int n_experts,
    int cache_cap,
    Policy pol,
    int k_used,
    float lambda_moce,
    int protect_top_j,
    int near_pick_n = 48,
    double ema_decay = 0.99,
    int warmup_tokens = 64
) {
    int n_tok = (int)trace.size();
    int eval_tokens = std::max(1, n_tok - warmup_tokens);
    if (cache_cap >= n_experts) {
        return {0.0, 1.0};
    }

    std::vector<std::vector<int>> next_use(n_experts);
    std::vector<int> next_ptr(n_experts, 0);
    if (pol == POL_BELADY) {
        for (int t = 0; t < n_tok; ++t) {
            int lim = std::min(k_used, trace[t].n_cand);
            for (int i = 0; i < lim; ++i) {
                next_use[trace[t].ids[i]].push_back(t);
            }
        }
    }

    std::vector<uint8_t> in_cache(n_experts, 0);
    std::vector<int> cache_list;
    cache_list.reserve(cache_cap);
    std::vector<int> last_used(n_experts, 0);
    std::vector<int> freq(n_experts, 0);
    std::vector<double> ema_score(n_experts, 0.0);
    double ema_mult = 1.0;
    double inv_decay = 1.0 / ema_decay;

    int64_t misses = 0;
    int64_t total_picks = 0;

    uint16_t chosen[16];
    uint8_t in_chosen[512] = {0};

    for (int t = 0; t < n_tok; ++t) {
        const TokenRec& rec = trace[t];
        int n_chosen = 0;
        if (lambda_moce <= 0.0f || rec.n_cand <= k_used) {
            n_chosen = std::min(k_used, rec.n_cand);
            for (int i = 0; i < n_chosen; ++i) chosen[i] = rec.ids[i];
        } else {
            float delta_avg = std::max(1e-4f, rec.logits[0] - rec.logits[std::min(k_used - 1, rec.n_cand - 1)]);
            float bonus = lambda_moce * delta_avg;
            struct Cand { float score; int neg_rank; uint16_t eid; };
            Cand cands[64];
            for (int r = 0; r < rec.n_cand; ++r) {
                uint16_t eid = rec.ids[r];
                float s = rec.logits[r] + ((r < protect_top_j || in_cache[eid]) ? bonus : 0.0f);
                cands[r] = {s, -r, eid};
            }
            std::sort(cands, cands + rec.n_cand, [](const Cand& a, const Cand& b) {
                if (a.score != b.score) return a.score > b.score;
                if (a.neg_rank != b.neg_rank) return a.neg_rank > b.neg_rank;
                return a.eid > b.eid;
            });
            n_chosen = std::min(k_used, rec.n_cand);
            for (int i = 0; i < n_chosen; ++i) chosen[i] = cands[i].eid;
        }

        if (pol == POL_BELADY) {
            for (int i = 0; i < n_chosen; ++i) {
                uint16_t eid = chosen[i];
                auto& lst = next_use[eid];
                int& ptr = next_ptr[eid];
                while (ptr < (int)lst.size() && lst[ptr] <= t) ptr++;
            }
        }

        for (int i = 0; i < n_chosen; ++i) in_chosen[chosen[i]] = 1;
        if (t >= warmup_tokens) {
            int missing_slots = std::max(0, k_used - n_chosen);
            total_picks += n_chosen + missing_slots;
            misses += missing_slots;
        }

        for (int i = 0; i < n_chosen; ++i) {
            uint16_t eid = chosen[i];
            if (!in_cache[eid]) {
                if (t >= warmup_tokens) misses++;
                if ((int)cache_list.size() >= cache_cap) {
                    int victim_idx = -1;
                    for (int ci = 0; ci < (int)cache_list.size(); ++ci) {
                        int c = cache_list[ci];
                        if (in_chosen[c]) continue;
                        if (victim_idx < 0) { victim_idx = ci; continue; }
                        int best_c = cache_list[victim_idx];
                        bool better = false;
                        if (pol == POL_LRU) {
                            better = (last_used[c] < last_used[best_c]) || (last_used[c] == last_used[best_c] && c < best_c);
                        } else if (pol == POL_LFU) {
                            if (freq[c] != freq[best_c]) better = freq[c] < freq[best_c];
                            else if (last_used[c] != last_used[best_c]) better = last_used[c] < last_used[best_c];
                            else better = c < best_c;
                        } else if (pol == POL_EMA || pol == POL_NEAR_PICK || pol == POL_INV_ROUTER) {
                            if (ema_score[c] != ema_score[best_c]) better = ema_score[c] < ema_score[best_c];
                            else if (last_used[c] != last_used[best_c]) better = last_used[c] < last_used[best_c];
                            else better = c < best_c;
                        } else if (pol == POL_BELADY) {
                            int nt_c = (next_ptr[c] < (int)next_use[c].size()) ? next_use[c][next_ptr[c]] : 1000000000;
                            int nt_b = (next_ptr[best_c] < (int)next_use[best_c].size()) ? next_use[best_c][next_ptr[best_c]] : 1000000000;
                            if (nt_c != nt_b) better = nt_c > nt_b;
                            else better = c < best_c;
                        }
                        if (better) victim_idx = ci;
                    }
                    int victim = cache_list[victim_idx];
                    in_cache[victim] = 0;
                    cache_list[victim_idx] = eid;
                } else {
                    cache_list.push_back(eid);
                }
                in_cache[eid] = 1;
            }
            last_used[eid] = t + 1;
            freq[eid] += 1;
        }
        for (int i = 0; i < n_chosen; ++i) in_chosen[chosen[i]] = 0;

        if (pol == POL_EMA || pol == POL_NEAR_PICK || pol == POL_INV_ROUTER) {
            ema_mult *= inv_decay;
            if (ema_mult > 1e100) {
                double sc = 1.0 / ema_mult;
                for (int e = 0; e < n_experts; ++e) ema_score[e] *= sc;
                ema_mult = 1.0;
            }
            if (pol == POL_EMA) {
                for (int i = 0; i < n_chosen; ++i) ema_score[chosen[i]] += ema_mult;
            } else if (pol == POL_NEAR_PICK) {
                int top_n = std::min(near_pick_n, rec.n_cand);
                double denom = std::max(1, top_n - k_used);
                for (int r = 0; r < top_n; ++r) {
                    double cred = (r < k_used) ? 1.0 : 0.35 * (1.0 - (double)(r - k_used) / denom);
                    ema_score[rec.ids[r]] += cred * ema_mult;
                }
            } else if (pol == POL_INV_ROUTER && rec.n_cand > 0) {
                float z_max = rec.logits[0];
                float z_min = rec.logits[rec.n_cand - 1];
                float span = std::max(1e-4f, z_max - z_min);
                for (int r = 0; r < rec.n_cand; ++r) {
                    double util = (rec.logits[r] - z_min) / span + ((r < k_used) ? 1.0 : 0.0);
                    ema_score[rec.ids[r]] += util * ema_mult;
                }
            }
        }
    }

    double m_per_tok = (double)misses / (double)eval_tokens;
    double hit_rate = 1.0 - ((double)misses / (double)std::max<int64_t>(1, total_picks));
    return {m_per_tok, hit_rate};
}

// Parse simple keep_map file exported as 48 lines of space-separated expert ints
std::vector<std::vector<int>> load_keep_lines(const std::string& path) {
    std::vector<std::vector<int>> out(48);
    if (path == "NONE") return out;
    std::ifstream f(path);
    for (int l = 0; l < 48; ++l) {
        int k;
        f >> k;
        out[l].resize(k);
        for (int i = 0; i < k; ++i) f >> out[l][i];
    }
    return out;
}

int main(int argc, char** argv) {
    if (argc < 7) {
        fprintf(stderr, "Usage: %s <trace.q4rt> <keep.txt|NONE> <cache_cap> <pol 0..5> <lambda> <top_j>\n", argv[0]);
        return 1;
    }
    FILE* fp = fopen(argv[1], "rb");
    if (!fp) { fprintf(stderr, "Failed to open trace file: %s\n", argv[1]); return 1; }
    TraceHeader hdr;
    if (fread(&hdr, sizeof(hdr), 1, fp) != 1) return 1;
    fseek(fp, 0, SEEK_END);
    long fsz = ftell(fp);
    fseek(fp, sizeof(hdr), SEEK_SET);
    int rec_bytes = hdr.top_m * 4;
    int actual_tokens = (fsz - 16) / (hdr.n_layers * rec_bytes);
    int n_tok = (hdr.n_tokens == 0 || actual_tokens < (int)hdr.n_tokens) ? actual_tokens : (int)hdr.n_tokens;

    std::vector<uint8_t> raw((size_t)n_tok * hdr.n_layers * rec_bytes);
    if (fread(raw.data(), 1, raw.size(), fp) != raw.size()) return 1;
    fclose(fp);

    auto keep = load_keep_lines(argv[2]);
    int cache_cap = atoi(argv[3]);
    if (cache_cap < 10) { fprintf(stderr, "cache_cap (%d) must be >= k_used (10)\n", cache_cap); return 1; }
    Policy pol = (Policy)atoi(argv[4]);
    float lam = atof(argv[5]);
    int top_j = atoi(argv[6]);

    double tot_misses = 0.0;
    double tot_hit = 0.0;

    #pragma omp parallel for reduction(+:tot_misses,tot_hit) schedule(dynamic)
    for (int l = 0; l < hdr.n_layers; ++l) {
        int remap[512];
        for (int e = 0; e < 512; ++e) remap[e] = -1;
        int n_exp = hdr.n_experts;
        if (!keep[l].empty()) {
            n_exp = (int)keep[l].size();
            for (int i = 0; i < n_exp; ++i) remap[keep[l][i]] = i;
        } else {
            for (int i = 0; i < 512; ++i) remap[i] = i;
        }

        std::vector<TokenRec> l_trace(n_tok);
        for (int t = 0; t < n_tok; ++t) {
            const uint8_t* p = raw.data() + ((size_t)t * hdr.n_layers + l) * rec_bytes;
            const uint16_t* ids = (const uint16_t*)p;
            const uint16_t* lgs = (const uint16_t*)(p + hdr.top_m * 2);
            TokenRec& tr = l_trace[t];
            tr.n_cand = 0;
            for (int i = 0; i < hdr.top_m; ++i) {
                int nid = remap[ids[i]];
                if (nid >= 0) {
                    tr.ids[tr.n_cand] = (uint16_t)nid;
                    tr.logits[tr.n_cand] = fp16_to_fp32(lgs[i]);
                    tr.n_cand++;
                }
            }
        }
        ReplayResult r = replay_one_layer(l_trace, n_exp, cache_cap, pol, 10, lam, top_j);
        tot_misses += r.misses_per_tok;
        tot_hit += r.hit_rate;
    }
    printf("%.6f %.6f\n", tot_misses, tot_hit / hdr.n_layers);
    return 0;
}