#!/bin/sh
set -eu

[ "$#" -eq 4 ] || [ "$#" -eq 5 ] || {
    echo "usage: $0 BASELINE_A1 CANDIDATE_B1 CANDIDATE_B2 BASELINE_A2 [OUTPUT_JSON]" >&2
    exit 2
}
a1=$1
b1=$2
b2=$3
a2=$4
output=${5:-}
script_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
temporary=$(mktemp -d /private/tmp/gem-decode-attention-abba.XXXXXX)
trap 'rm -rf "$temporary"' EXIT INT TERM HUP

merge_pair() {
    first=$1
    second=$2
    destination=$3
    jq -s '
      def controls: {
        schema_version, engine, workload, shape, score_token_tile, value_kernel,
        value_dimension_tile, value_threadgroups, query_heads, kv_heads, head_dimension,
        tokens, live_tokens, cache_capacity, window, kv_cache_bytes, score_scratch_bytes,
        value_partition_scratch_bytes, value_partitions, iterations, warmup, dispatches
      };
      .[0] as $x | .[1] as $y |
      if ($x | controls) != ($y | controls) or
         $x.repetitions < 3 or $y.repetitions < 3 or
         ($x.gpu_seconds_samples | length) != $x.repetitions or
         ($y.gpu_seconds_samples | length) != $y.repetitions or
         ($x.stages.score.samples_seconds | length) != $x.repetitions or
         ($y.stages.score.samples_seconds | length) != $y.repetitions
      then error("ABBA arms do not preserve their tuning controls") else
        $x |
        .repetitions = ($x.repetitions + $y.repetitions) |
        .gpu_seconds_samples = ($x.gpu_seconds_samples + $y.gpu_seconds_samples) |
        .stages.score.samples_seconds =
          ($x.stages.score.samples_seconds + $y.stages.score.samples_seconds)
      end
    ' "$first" "$second" > "$destination"
}

merge_pair "$a1" "$a2" "$temporary/baseline.json"
merge_pair "$b1" "$b2" "$temporary/candidate.json"
"$script_dir/analyze-decode-attention.sh" "$temporary/baseline.json" \
    "$temporary/candidate.json" "$temporary/pooled.json"

analyze() {
    jq --slurpfile a1 "$a1" --slurpfile b1 "$b1" --slurpfile b2 "$b2" \
       --slurpfile a2 "$a2" '
      def mean($xs): $xs | add / length;
      def gain($before; $after): 100 * (1 - mean($after) / mean($before));
      def drift($first; $second): 100 * (mean($second) / mean($first) - 1);
      . + {bracket: {
        order: ["baseline-a1", "candidate-b1", "candidate-b2", "baseline-a2"],
        arm_repetitions: $a1[0].repetitions,
        total_gpu: {
          baseline_a1_mean_seconds: mean($a1[0].gpu_seconds_samples),
          candidate_b1_mean_seconds: mean($b1[0].gpu_seconds_samples),
          candidate_b2_mean_seconds: mean($b2[0].gpu_seconds_samples),
          baseline_a2_mean_seconds: mean($a2[0].gpu_seconds_samples),
          opening_improvement_pct: gain($a1[0].gpu_seconds_samples; $b1[0].gpu_seconds_samples),
          closing_improvement_pct: gain($a2[0].gpu_seconds_samples; $b2[0].gpu_seconds_samples),
          baseline_drift_pct: drift($a1[0].gpu_seconds_samples; $a2[0].gpu_seconds_samples),
          candidate_drift_pct: drift($b1[0].gpu_seconds_samples; $b2[0].gpu_seconds_samples),
          every_candidate_sample_faster:
            (([$b1[0].gpu_seconds_samples[], $b2[0].gpu_seconds_samples[]] | max) <
             ([$a1[0].gpu_seconds_samples[], $a2[0].gpu_seconds_samples[]] | min))
        },
        score_stage: {
          opening_improvement_pct:
            gain($a1[0].stages.score.samples_seconds; $b1[0].stages.score.samples_seconds),
          closing_improvement_pct:
            gain($a2[0].stages.score.samples_seconds; $b2[0].stages.score.samples_seconds)
        }
      }}
    ' "$temporary/pooled.json"
}

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