#!/bin/sh
# Compare two llama-bench JSON files using gem's promotion policy.
set -eu

usage() {
    echo "usage: $0 [--primary pp|tg] [--allow-model-change] BASELINE_JSON CANDIDATE_JSON [OUTPUT_JSON]" >&2
    exit 2
}

primary=tg
allow_model_change=false

while [ "$#" -gt 0 ]; do
    case "$1" in
        --primary)
            [ "$#" -ge 2 ] || usage
            primary=$2
            shift 2
            ;;
        --allow-model-change)
            allow_model_change=true
            shift
            ;;
        --help)
            usage
            ;;
        --)
            shift
            break
            ;;
        -*)
            usage
            ;;
        *)
            break
            ;;
    esac
done

[ "$primary" = pp ] || [ "$primary" = tg ] || usage
[ "$#" -eq 2 ] || [ "$#" -eq 3 ] || usage

baseline=$1
candidate=$2
output=${3:-}

[ -r "$baseline" ] || { echo "gem: baseline is not readable: $baseline" >&2; exit 1; }
[ -r "$candidate" ] || { echo "gem: candidate is not readable: $candidate" >&2; exit 1; }
command -v jq >/dev/null 2>&1 || { echo "gem: jq is required" >&2; exit 1; }

run_analysis() {
    jq -n \
        --slurpfile baseline "$baseline" \
        --slurpfile candidate "$candidate" \
        --arg baseline_path "$baseline" \
        --arg candidate_path "$candidate" \
        --arg primary "$primary" \
        --argjson allow_model_change "$allow_model_change" '
      def require_array($name; $value):
        if ($value | type) != "array" or ($value | length) == 0
        then error("\($name) must be a non-empty llama-bench JSON array")
        else $value
        end;

      def row_for($rows; $kind):
        [ $rows[] |
          select(if $kind == "pp"
                 then (.n_prompt // 0) > 0 and (.n_gen // 0) == 0
                 else (.n_prompt // 0) == 0 and (.n_gen // 0) > 0
                 end) ] as $matches |
        if ($matches | length) != 1
        then error("expected exactly one \($kind) row, found \($matches | length)")
        else $matches[0]
        end;

      def metric($row):
        if (($row.avg_ts | type) != "number") or
           (($row.stddev_ts | type) != "number") or
           (($row.samples_ts | type) != "array") or
           (($row.samples_ts | length) < 3) or
           ($row.avg_ts <= 0) or ($row.stddev_ts < 0)
        then error("each row needs positive avg_ts, non-negative stddev_ts, and at least three samples_ts")
        else {
          tokens_per_second: $row.avg_ts,
          stddev_tokens_per_second: $row.stddev_ts,
          repetitions: ($row.samples_ts | length),
          coefficient_of_variation_pct: (100 * $row.stddev_ts / $row.avg_ts)
        }
        end;

      def difference_fields($a; $b; $fields):
        [ $fields[] as $field |
          select($a[$field] != $b[$field]) |
          {field: $field, baseline: $a[$field], candidate: $b[$field]} ];

      def comparison($base_row; $candidate_row):
        metric($base_row) as $base_metric |
        metric($candidate_row) as $candidate_metric |
        (100 * ($candidate_metric.tokens_per_second /
                $base_metric.tokens_per_second - 1)) as $delta |
        ([ $base_metric.coefficient_of_variation_pct,
           $candidate_metric.coefficient_of_variation_pct ] | max) as $noise |
        {
          baseline: $base_metric,
          candidate: $candidate_metric,
          delta_pct: $delta,
          observed_noise_floor_pct: $noise,
          outside_observed_variance: (($delta | fabs) > $noise)
        };

      require_array("baseline"; $baseline[0]) as $base_rows |
      require_array("candidate"; $candidate[0]) as $candidate_rows |
      row_for($base_rows; "pp") as $base_pp |
      row_for($base_rows; "tg") as $base_tg |
      row_for($candidate_rows; "pp") as $candidate_pp |
      row_for($candidate_rows; "tg") as $candidate_tg |

      ["backends", "cpu_info", "gpu_info", "n_threads", "n_gpu_layers",
       "n_cpu_moe", "split_mode", "main_gpu", "no_kv_offload", "devices",
       "tensor_split", "tensor_buft_overrides"] as $fixed_fields |
      ["model_filename", "model_type", "model_size", "model_n_params"] as $model_fields |
      ["build_commit", "build_number"] as $revision_fields |
      ["n_batch", "n_ubatch", "flash_attn", "attention_kernel", "command_span", "q4_kernel",
       "prefill_kernel", "rope_kernel", "kv_store_kernel", "load_mode",
       "no_op_offload", "no_host"] as $axis_fields |
      difference_fields($base_pp; $candidate_pp; $fixed_fields) as $fixed_differences |
      difference_fields($base_pp; $candidate_pp; $model_fields) as $model_differences |
      difference_fields($base_pp; $candidate_pp; $revision_fields) as $revision_differences |
      (difference_fields($base_pp; $candidate_pp; $axis_fields) +
       (if $base_pp.type_k == $base_pp.type_v and
           $candidate_pp.type_k == $candidate_pp.type_v then
          if $base_pp.type_k != $candidate_pp.type_k then
            [{field: "kv_type", baseline: $base_pp.type_k, candidate: $candidate_pp.type_k}]
          else [] end
        else
          difference_fields($base_pp; $candidate_pp; ["type_k", "type_v"])
        end)) as $axis_differences |
      ([
        if $base_pp.n_prompt != $candidate_pp.n_prompt then
          {field: "pp_tokens", baseline: $base_pp.n_prompt, candidate: $candidate_pp.n_prompt}
        else empty end,
        if $base_tg.n_gen != $candidate_tg.n_gen then
          {field: "tg_tokens", baseline: $base_tg.n_gen, candidate: $candidate_tg.n_gen}
        else empty end
      ]) as $workload_differences |
      comparison($base_pp; $candidate_pp) as $pp |
      comparison($base_tg; $candidate_tg) as $tg |
      (if $primary == "tg" then $tg else $pp end) as $primary_metric |
      (if $primary == "tg" then $pp else $tg end) as $secondary_metric |
      (($fixed_differences | length) == 0 and
       ($workload_differences | length) == 0 and
       ($allow_model_change or ($model_differences | length) == 0) and
       ($axis_differences | length) <= 1) as $comparable |
      ($primary_metric.delta_pct >= 3 and
       $primary_metric.outside_observed_variance) as $primary_pass |
      (($secondary_metric.delta_pct < -2) and
       $secondary_metric.outside_observed_variance) as $secondary_regression |
      {
        schema_version: 1,
        sources: {baseline: $baseline_path, candidate: $candidate_path},
        policy: {
          primary_kpi: $primary,
          minimum_primary_improvement_pct: 3,
          maximum_secondary_regression_pct: 2,
          repetitions_required: 3,
          variance_rule: "absolute delta must exceed the larger coefficient of variation"
        },
        comparability: {
          comparable: $comparable,
          allow_model_change: $allow_model_change,
          fixed_differences: $fixed_differences,
          workload_differences: $workload_differences,
          model_differences: $model_differences,
          runtime_revision_differences: $revision_differences,
          tuning_axis_differences: $axis_differences
        },
        metrics: {pp: $pp, tg: $tg},
        promotion: {
          eligible: ($comparable and $primary_pass and ($secondary_regression | not)),
          primary_threshold_passed: $primary_pass,
          secondary_regression_detected: $secondary_regression
        }
      }
    '
}

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