// fh-replay: run a recording (fh-tracker --record) through the tracker offline, with the
// live scheduling, and report how well it kept the hands.
//
//   fh-replay DIR [--oracle N] [--slow F] [--timeline FILE] [--threads N] [--models DIR]
//             [--from S] [--to S] [--contrast MODE|PALM/HAND] (clahe[:CLIP], none, stretch)
//
// --oracle N: every N-th set, also search every tile of every camera (slow), to see
//             which hands were there to find. Compares that with what the tracker had.
// --slow F:   the live tracker skips the sets that arrive while it's busy; replay takes
//             each step's time here times F as the busy time (the headset is busier live).
// --cost:     instead of timing the steps, charge each round of model calls what it
//             typically costs live (10 ms landmarks, 18 ms palms): repeatable results.
// --timeline: per processed set, a line per hand (time, id, side, views, wrist) and per view
//             (hand, camera, presence, next crop, set index).
// --keep-presence P: landmark presence a tracked view needs to stay (default 0.5, as new ones).
// --contrast: how the palm search's and the landmark model's crops are equalized
//             (default clahe:2/none, as fh-tracker).
#include "record.h"
#include "tracker.h"

#include <algorithm>
#include <chrono>
#include <cstdio>
#include <cstring>
#include <map>
#include <set>
#include <string>
#include <vector>

namespace {

struct Track {
    double first = 0, last = 0;
    int sets = 0, left = 0;
    // the last two palm positions (raw, smoothed) and times, for the jitter measure
    V3 raw[2]{}, sm[2]{};
    double t[2]{};
    int line = 0;   // updates on the current unbroken run
};

double median(std::vector<double> v) {
    if (v.empty()) return 0;
    std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end());
    return v[v.size() / 2];
}

}  // namespace

int main(int argc, char **argv) {
    if (argc < 2 || argv[1][0] == '-') {
        std::printf("usage: %s DIR [--oracle N] [--slow F] [--timeline FILE] [--threads N] [--models DIR] [--from S] [--to S]\n", argv[0]);
        return 1;
    }
    const std::string dir = argv[1];
    int oracle = 0, threads = 2;
    double slow = 1.0, from = 0, to = 1e9;
    bool cost = false;
    Contrast palm_contrast, hand_contrast{Contrast::None};   // as fh-tracker's
    double keep_presence = 0.5;   // landmark presence a tracked view needs to stay
    std::string timeline, models = std::string(argv[0]).substr(0, std::string(argv[0]).rfind('/') + 1) + "../models/ncnn";
    for (int i = 2; i < argc; ++i) {
        const std::string a = argv[i];
        const bool more = i + 1 < argc;
        if (a == "--oracle" && more) oracle = std::atoi(argv[++i]);
        else if (a == "--slow" && more) slow = std::atof(argv[++i]);
        else if (a == "--timeline" && more) timeline = argv[++i];
        else if (a == "--threads" && more) threads = std::atoi(argv[++i]);
        else if (a == "--models" && more) models = argv[++i];
        else if (a == "--cost") cost = true;
        else if (a == "--keep-presence" && more) keep_presence = std::atof(argv[++i]);
        else if (a == "--contrast" && more) {
            if (!Contrast::parse_pair(argv[++i], palm_contrast, hand_contrast))
                return std::fprintf(stderr, "--contrast MODE or PALM/HAND, each clahe[:CLIP]|none|stretch\n"), 1;
        }
        else if (a == "--from" && more) from = std::atof(argv[++i]);
        else if (a == "--to" && more) to = std::atof(argv[++i]);
        else return std::fprintf(stderr, "unknown option %s\n", a.c_str()), 1;
    }
    std::string err;
    std::map<std::string, Camera> calib;
    Nets nets;
    SetReader in;
    if (!load_calibration(calib, err) || !nets.load(models, false, err) || !in.open(dir, err))
        return std::fprintf(stderr, "%s\n", err.c_str()), 1;
    nets.set_contrast(palm_contrast, hand_contrast);
    FILE *tl = timeline.empty() ? nullptr : std::fopen(timeline.c_str(), "w");

    std::vector<fh_set_cam_t> cams;
    std::vector<std::vector<uint8_t>> px;
    if (!in.next(cams, px)) return std::fprintf(stderr, "%s: no sets\n", dir.c_str()), 1;
    std::map<std::string, Camera> used;
    for (auto &c : cams)
        if (calib.count(c.name)) used[c.name] = calib[c.name];
    Pool pool(threads, {2, 3, 4});
    Tracker tracker(used, nets, pool);
    tracker.set_keep_presence(keep_presence);

    uint64_t t0 = 0, busy_until = 0, next_ns = 0;
    int index = -1;   // of the set in the recording
    int nsets = 0, processed = 0, left = 0, right = 0, both = 0, hist[3] = {};
    std::map<int, Track> tracks;
    std::vector<const Hand *> last_out;
    // oracle: sets where a side's hand was findable, and where the tracker had it then
    int o_sets = 0, o_left = 0, o_right = 0, o_left_hit = 0, o_right_hit = 0, o_left_extra = 0, o_right_extra = 0;
    std::map<std::string, int> o_by_cam;
    double busy_ms = 0;
    // jitter: how far each update's palm is from a straight line through the last two,
    // mm (steady motion cancels out; what's left is noise and real acceleration)
    std::vector<double> jit_raw, jit_sm;
    int near_face = 0, hand_updates = 0;   // published palms within 20 cm of the eyes
    do {
        std::map<std::string, Image> images;
        uint64_t t = UINT64_MAX;
        for (size_t i = 0; i < cams.size(); ++i) {
            if (!used.count(cams[i].name)) continue;
            images[cams[i].name] = {px[i].data(), int(cams[i].width), int(cams[i].height), int(cams[i].width)};
            t = std::min(t, cams[i].capture_ns);
        }
        if (!t0) t0 = t;
        const double ts = (t - t0) / 1e9;
        ++index;
        if (ts < from) continue;
        if (ts > to) break;
        ++nsets;

        if (t >= busy_until && t >= next_ns) {
            const auto w0 = std::chrono::steady_clock::now();
            const Stats before = tracker.stats;
            const auto out = tracker.step(images, int64_t(t));
            double ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - w0).count();
            if (cost) {   // repeatable: rounds of model calls at typical live costs, per thread
                const int hands = tracker.stats.hand_calls - before.hand_calls, palms = tracker.stats.palm_calls - before.palm_calls;
                ms = (2 + 10.0 * ((hands + threads - 1) / threads) + 18.0 * ((palms + threads - 1) / threads)) / slow;
            }
            busy_ms += ms;
            busy_until = t + uint64_t(ms * slow * 1e6) + 3'000'000;   // + the ring hand-off
            next_ns = t + uint64_t((tracker.interval() - 0.005) * 1e9);
            ++processed;
            last_out = out;
            bool l = false, r = false;
            for (const Hand *h : out) {
                (h->pts[0][0] < 0 ? l : r) = true;
                Track &tr = tracks[h->id];
                auto palm = [](const V3 *p) { return (p[0] + p[5] + p[9] + p[13] + p[17]) * 0.2; };
                const V3 raw = palm(h->pts), sm = palm(h->smooth);
                ++hand_updates, near_face += norm(sm) < 0.2;
                if (tr.line && ts - tr.t[0] >= 0.1) tr.line = 0;   // a gap: the line starts over
                if (tr.line >= 2 && tr.t[0] - tr.t[1] > 1e-3) {
                    const double k = (ts - tr.t[0]) / (tr.t[0] - tr.t[1]);
                    jit_raw.push_back(norm(raw - tr.raw[0] - (tr.raw[0] - tr.raw[1]) * k) * 1000);
                    jit_sm.push_back(norm(sm - tr.sm[0] - (tr.sm[0] - tr.sm[1]) * k) * 1000);
                }
                tr.raw[1] = tr.raw[0], tr.sm[1] = tr.sm[0], tr.t[1] = tr.t[0];
                tr.raw[0] = raw, tr.sm[0] = sm, tr.t[0] = ts;
                ++tr.line;
                if (!tr.sets) tr.first = ts;
                tr.last = ts, ++tr.sets, tr.left += h->pts[0][0] < 0;
                if (tl)
                    std::fprintf(tl, "%.3f %d %s %d %+.3f %+.3f %+.3f\n", ts, h->id, h->pts[0][0] < 0 ? "L" : "R", h->nviews,
                                 h->pts[0][0], h->pts[0][1], h->pts[0][2]);
            }
            if (tl && out.empty()) std::fprintf(tl, "%.3f -\n", ts);
            if (tl)
                for (const Seen &v : tracker.views_now())
                    std::fprintf(tl, "%.3f   view %d %s presence %.2f roi %.0f %.0f %.0f %.3f  set %d\n", ts, v.hand, v.cam.c_str(),
                                 v.lm.presence, v.roi.center[0], v.roi.center[1], v.roi.size, v.roi.rotation, index);
            left += l, right += r, both += l && r;
            ++hist[std::min<size_t>(out.size(), 2)];
        }

        if (oracle > 0 && nsets % oracle == 0) {
            const Stats keep = tracker.stats;
            const auto seen = tracker.exhaustive(images);
            tracker.stats = keep;
            bool l = false, r = false;
            for (const Seen &s : seen) {
                (s.wrist[0] < 0 ? l : r) = true;
                ++o_by_cam[s.cam + (s.wrist[0] < 0 ? " L" : " R")];
            }
            bool tl_ = false, tr_ = false;
            for (const Hand *h : last_out) (h->pts[0][0] < 0 ? tl_ : tr_) = true;
            ++o_sets;
            o_left += l, o_right += r;
            o_left_hit += l && tl_, o_right_hit += r && tr_;
            o_left_extra += !l && tl_, o_right_extra += !r && tr_;
            if (tl && ((!l && tl_) || (!r && tr_))) std::fprintf(tl, "%.3f oracle-extra %s%s set %d\n", ts, !l && tl_ ? "L" : "", !r && tr_ ? "R" : "", index);
        }
    } while (in.next(cams, px));
    if (tl) std::fclose(tl);

    const double secs = nsets > 1 ? nsets / 30.0 : 0;
    const Stats &s = tracker.stats;
    std::printf("%s: %d sets (%.0f s), processed %d (%.1f/s), %.1f ms per step\n", dir.c_str(), nsets, secs, processed,
                processed / std::max(secs, 1e-9), busy_ms / std::max(processed, 1));
    std::printf("hands per processed set: 0 %.0f%%, 1 %.0f%%, 2 %.0f%%;  a hand on the left %.0f%%, right %.0f%%, both %.0f%%\n",
                100.0 * hist[0] / processed, 100.0 * hist[1] / processed, 100.0 * hist[2] / processed,
                100.0 * left / processed, 100.0 * right / processed, 100.0 * both / processed);
    std::vector<double> lens[2];
    for (auto &[id, tr] : tracks) lens[tr.left * 2 > tr.sets ? 0 : 1].push_back(tr.last - tr.first);
    for (int k = 0; k < 2; ++k) {
        double total = 0;
        for (double d : lens[k]) total += d;
        std::printf("%s tracks: %zu, median %.1f s, total %.0f s\n", k ? "right" : "left ", lens[k].size(), median(lens[k]), total);
    }
    std::printf("views lost %d, handoff misses %d, dups %d, splits %d;  hands new %d, merged %d, forgotten %d\n", s.lost,
                s.handoff_miss, s.dups, s.splits, s.created, s.merged, s.forgotten);
    std::printf("model calls: palm %d (%.1f/s), hand %d (%.1f/s)\n", s.palm_calls, s.palm_calls / std::max(secs, 1e-9),
                s.hand_calls, s.hand_calls / std::max(secs, 1e-9));
    {
        std::vector<double> r, step;
        std::map<int, double> prev;
        for (auto &[id, x] : s.mono_ratio) {
            r.push_back(x);
            if (prev.count(id)) step.push_back(std::fabs(x - prev[id]));
            prev[id] = x;
        }
        std::sort(r.begin(), r.end());
        std::sort(step.begin(), step.end());
        if (!r.empty())
            std::printf("single-view distance / stereo: 10%% %.2f, median %.2f, 90%% %.2f; change between frames median %.3f, 90%% %.3f\n",
                        r[r.size() / 10], r[r.size() / 2], r[r.size() * 9 / 10], step[step.size() / 2], step[step.size() * 9 / 10]);
    }
    std::printf("palms within 20 cm of the eyes: %d of %d hand updates\n", near_face, hand_updates);
    std::sort(jit_raw.begin(), jit_raw.end());
    std::sort(jit_sm.begin(), jit_sm.end());
    if (!jit_raw.empty())
        std::printf("palm jitter (off a straight line through the last two updates): measured median %.1f mm, 90%% %.1f mm; "
                    "published median %.1f mm, 90%% %.1f mm\n", jit_raw[jit_raw.size() / 2], jit_raw[jit_raw.size() * 9 / 10],
                    jit_sm[jit_sm.size() / 2], jit_sm[jit_sm.size() * 9 / 10]);
    if (o_sets) {
        std::printf("oracle, %d sets: a left hand findable in %d, the tracker had it in %d (%.0f%%); right %d, had %d (%.0f%%)\n",
                    o_sets, o_left, o_left_hit, 100.0 * o_left_hit / std::max(o_left, 1), o_right, o_right_hit,
                    100.0 * o_right_hit / std::max(o_right, 1));
        std::printf("        tracker had a hand the full search didn't find: left %d, right %d\n", o_left_extra, o_right_extra);
        std::printf("        found by camera:");
        for (auto &[k, n] : o_by_cam) std::printf("  %s %d", k.c_str(), n);
        std::printf("\n");
    }
    return 0;
}