#!/bin/sh
set -eu

[ "$#" -eq 1 ] || [ "$#" -eq 2 ] || {
    echo "usage: $0 GATE_DIRECTORY [OUTPUT_JSON]" >&2
    exit 2
}
directory=$1
output=${2:-}

[ -s "$directory/gate-config.json" ] || {
    echo "gem: missing BK128 gate input: $directory/gate-config.json" >&2
    exit 2
}
[ -s "$directory/metal-device.json" ] || {
    echo "gem: missing BK128 gate input: $directory/metal-device.json" >&2
    exit 2
}
for path in "$directory/sha256.txt" "$directory/native-linkage.txt" \
            "$directory/macos-version.txt" "$directory/run-start-utc.txt"; do
    [ -s "$path" ] || { echo "gem: missing BK128 provenance input: $path" >&2; exit 2; }
done
for tokens in 4096 65536 131072; do
    for path in "$directory/seeded-$tokens.json" \
                "$directory/host-preflight-$tokens.json" \
                "$directory/host-postflight-$tokens.json"; do
        [ -s "$path" ] || { echo "gem: missing BK128 gate input: $path" >&2; exit 2; }
    done
done

analyze()
{
    jq -n \
        --slurpfile cfg "$directory/gate-config.json" \
        --slurpfile device "$directory/metal-device.json" \
        --slurpfile r4 "$directory/seeded-4096.json" \
        --slurpfile r64 "$directory/seeded-65536.json" \
        --slurpfile r128 "$directory/seeded-131072.json" \
        --slurpfile h4 "$directory/host-preflight-4096.json" \
        --slurpfile h64 "$directory/host-preflight-65536.json" \
        --slurpfile h128 "$directory/host-preflight-131072.json" \
        --slurpfile p4 "$directory/host-postflight-4096.json" \
        --slurpfile p64 "$directory/host-postflight-65536.json" \
        --slurpfile p128 "$directory/host-postflight-131072.json" \
        --rawfile sha "$directory/sha256.txt" \
        --rawfile linkage "$directory/native-linkage.txt" \
        --rawfile macos "$directory/macos-version.txt" \
        --rawfile started "$directory/run-start-utc.txt" '
      def mean($xs): $xs | add / length;
      def sample_stddev($xs):
        mean($xs) as $m |
        (($xs | map((. - $m) * (. - $m)) | add) / (($xs | length) - 1)) | sqrt;
      def cv_pct($xs): 100 * sample_stddev($xs) / mean($xs);
      def abs: if . < 0 then -. else . end;
      def host_ok($h):
        $h.schema_version == 2 and $h.status == "pass" and
        $h.logical_cpus > 0 and $h.load_per_logical_cpu >= 0 and
        $h.load_per_logical_cpu <= $h.max_load_per_logical_cpu and
        $h.recent_aggregate_process_cpu_percent >= 0 and
        $h.short_window_busy_cpu_percent >= 0 and
        $h.short_window_busy_cpu_percent <= $h.max_short_window_busy_cpu_percent and
        $h.ac_power == true and $h.no_thermal_warning == true and
        ($h.power_state | type) == "string" and
        ($h.power_state | contains("AC Power")) and
        ($h.thermal_state | type) == "string";
      def valid_run($r; $position; $capacity):
        $r.schema_version == 1 and $r.engine == "gem-native-c-metal" and
        $r.shape == "31b-full" and $r.kv_type == "f32" and
        $r.position == $position and $r.rows == 32 and
        $r.cache_capacity == $capacity and $r.seed_prefix == true and
        $r.repetitions == 10 and $r.warmup == 10 and
        $r.bq8_pipeline_resources == {
          thread_execution_width: 32,
          max_total_threads_per_threadgroup: 1024,
          static_threadgroup_memory_bytes: 0,
          dynamic_threadgroup_memory_bytes: 18496,
          threads_per_threadgroup: 256
        } and
        $r.bq8_bk128_pipeline_resources == {
          thread_execution_width: 32,
          max_total_threads_per_threadgroup: 1024,
          static_threadgroup_memory_bytes: 0,
          dynamic_threadgroup_memory_bytes: 20544,
          threads_per_threadgroup: 256
        } and
        ($r.variants["online-bq8"].samples | length) == $r.repetitions and
        ($r.variants["online-bq8-bk128"].samples | length) == $r.repetitions and
        ($r.variants["online-bq8"].samples | all(. > 0)) and
        ($r.variants["online-bq8-bk128"].samples | all(. > 0));
      def metrics($r):
        $r.variants["online-bq8"].samples as $b |
        $r.variants["online-bq8-bk128"].samples as $c |
        [range(0; $b | length) as $i | 100 * (1 - $c[$i] / $b[$i])] as $paired |
        ($b | length / 2) as $half |
        mean($b) as $bm | mean($c) as $cm |
        cv_pct($b) as $bcv | cv_pct($c) as $ccv |
        mean($paired) as $paired_mean |
        sample_stddev($paired) as $paired_stddev |
        (100 * (1 - $cm / $bm)) as $gain |
        (100 * (mean($b[$half:]) / mean($b[0:$half]) - 1)) as $bdrift |
        (100 * (mean($c[$half:]) / mean($c[0:$half]) - 1)) as $cdrift |
        {
          baseline_mean_seconds: $bm,
          candidate_mean_seconds: $cm,
          improvement_pct: $gain,
          baseline_cv_pct: $bcv,
          candidate_cv_pct: $ccv,
          observed_noise_floor_pct: ([$bcv, $ccv] | max),
          paired_improvement_mean_pct: $paired_mean,
          paired_improvement_stddev_pct: $paired_stddev,
          paired_improvement_95ci_lower_pct:
            ($paired_mean - 2.262 * $paired_stddev / (($paired | length) | sqrt)),
          baseline_half_drift_pct: $bdrift,
          candidate_half_drift_pct: $cdrift,
          stable: ($bcv <= 3 and $ccv <= 3 and
                   ($bdrift | abs) <= 3 and ($cdrift | abs) <= 3),
          every_candidate_sample_faster: (($c | max) < ($b | min))
        };

      $r4[0] as $a | $r64[0] as $b | $r128[0] as $c |
      if (($sha | contains("gem-long-attention-profile")) and
          ($sha | contains("gem-metal-counter-info")) and
          ($sha | contains("src/metal.m")) and
          ($sha | contains("src/long_attention_profile.c")) and
          ($sha | contains("scripts/run-bq8-bk128-gate.sh")) and
          ($sha | contains("scripts/analyze-bq8-bk128-gate.sh")) and
          ($linkage | contains("gem-long-attention-profile")) and
          ($linkage | contains("gem-metal-counter-info")) and
          (($linkage | test("lib(llama|ggml)"; "i")) | not) and
          ($macos | length) > 1 and
          ($started | test("^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}Z")) and
          $device[0].schema_version == 1 and
          $device[0].workload == "metal-counter-info" and
          $device[0].device.name == "Apple M4 Max" and
          $device[0].device.has_unified_memory == true and
          $device[0].device.gpu_families.apple9 == true and
          $device[0].device.max_threadgroup_memory_bytes >= 20544 and
          $cfg[0].schema_version == 2 and $cfg[0].engine == "gem-native-c-metal" and
          $cfg[0].axis == {name: "cache_tile_width", baseline: 64, candidate: 128} and
          $cfg[0].positions == [4064, 65504, 131040] and $cfg[0].rows == 32 and
          $cfg[0].kv_type == "f32" and $cfg[0].repetitions == 10 and
          $cfg[0].warmup == 10 and $cfg[0].seed_prefix == true and
          $cfg[0].max_load_per_logical_cpu > 0 and
          $cfg[0].max_short_window_busy_cpu_percent > 0 and
          $cfg[0].requires_ac_power == true and
          $cfg[0].requires_no_thermal_warning == true and
          valid_run($a; 4064; 4096) and valid_run($b; 65504; 65536) and
          valid_run($c; 131040; 131072) and host_ok($h4[0]) and
          host_ok($h64[0]) and host_ok($h128[0]) and host_ok($p4[0]) and
          host_ok($p64[0]) and host_ok($p128[0])) | not
      then error("BK128 inputs violate the frozen one-axis gate controls")
      else
        metrics($a) as $m4 | metrics($b) as $m64 | metrics($c) as $m128 |
        ($m4.stable and $m64.stable and $m128.stable) as $stable |
        (($m4.improvement_pct >= -2 or
          (-$m4.improvement_pct) <= $m4.observed_noise_floor_pct)) as $short_guard |
        (($m64.improvement_pct >= 3 and
          $m64.improvement_pct > $m64.observed_noise_floor_pct and
          $m64.paired_improvement_95ci_lower_pct > 0) and
         ($m128.improvement_pct >= 3 and
          $m128.improvement_pct > $m128.observed_noise_floor_pct and
          $m128.paired_improvement_95ci_lower_pct > 0)) as $long_gain |
        ($stable and $short_guard and $long_gain) as $advance |
        {
          schema_version: 2,
          scope: "seeded exact-shape 31B/F32 global-prefill attention operator screen",
          axis: {name: "cache_tile_width", baseline: 64, candidate: 128},
          provenance: {
            run_started_utc: ($started | rtrimstr("\n")),
            macos_version: ($macos | rtrimstr("\n")),
            sha256_receipt: ($sha | rtrimstr("\n")),
            native_linkage_receipt: ($linkage | rtrimstr("\n"))
          },
          frozen_controls: {
            query_rows: 8, measured_rows: 32, kv_type: "f32",
            repetitions: 10, warmup: 10, seeded_prefix: true,
            maximum_cv_pct: 3, maximum_half_drift_pct: 3,
            minimum_long_context_improvement_pct: 3,
            paired_confidence_level_pct: 95,
            paired_t_critical_df9: 2.262,
            maximum_4k_regression_pct: 2
          },
          threadgroup_memory: {
            device_limit_bytes: $device[0].device.max_threadgroup_memory_bytes,
            baseline_bytes: 18496,
            candidate_bytes: 20544,
            baseline_limit_utilization_pct:
              (100 * 18496 / $device[0].device.max_threadgroup_memory_bytes),
            candidate_limit_utilization_pct:
              (100 * 20544 / $device[0].device.max_threadgroup_memory_bytes),
            candidate_additional_bytes: 2048,
            occupancy_inference: "not available from the public per-threadgroup limit"
          },
          host_preflight: {
            at_4096: $h4[0], at_65536: $h64[0], at_131072: $h128[0]
          },
          host_postflight: {
            at_4096: $p4[0], at_65536: $p64[0], at_131072: $p128[0]
          },
          contexts: {at_4096: $m4, at_65536: $m64, at_131072: $m128},
          screening: {
            stable: $stable,
            short_context_guard_passed: $short_guard,
            long_context_gain_passed: $long_gain,
            advance_to_evaluated_graph_gate: $advance,
            promotion_eligible: false,
            decision: (if $advance then "advance_to_evaluated_graph_gate"
                       else "keep_bk64_and_remove_bk128" end),
            reason: (if $advance
              then "operator screen passed; evaluated-graph/logit and correctness gates remain required"
              else "operator screen failed stability, 4K regression, or both 64K and 128K gain requirements"
              end)
          }
        }
      end'
}

if [ -n "$output" ]; then analyze > "$output"; else analyze; fi