#!/bin/sh
# Derive minimum logical traffic for Gemma 4 31B split-score decode attention.
set -eu

if [ "$#" -lt 1 ] || [ "$#" -gt 2 ]; then
    echo "usage: $0 PROFILE.json [OUTPUT.json]" >&2
    exit 2
fi
profile=$1
output=${2:-}
[ -f "$profile" ] || { echo "gem: decode profile is missing: $profile" >&2; exit 2; }
if [ -n "$output" ] && { [ -e "$output" ] || [ -L "$output" ]; }; then
    echo "gem: traffic analysis output already exists: $output" >&2
    exit 2
fi

jq -e '
  .schema_version == 1 and .engine == "gem-native-c-metal" and
  .workload == "synthetic-exact-shape-decode-attention-stage-profile" and
  .shape == "31b-full" and .query_heads == 32 and .kv_heads == 4 and
  .head_dimension == 512 and .window == 0 and
  (.live_tokens | type) == "number" and .live_tokens >= 8192 and
  .live_tokens == (.live_tokens | floor) and
  (.value_dimension_tile | type) == "number" and .value_dimension_tile > 0 and
  .value_dimension_tile == (.value_dimension_tile | floor) and
  (.head_dimension % .value_dimension_tile) == 0 and
  (.score_scratch_bytes == .query_heads * .live_tokens * 4) and
  ((.value_partitions // 1) | type) == "number" and
  (.value_partitions // 1) >= 1 and
  (.value_partitions // 1) == ((.value_partitions // 1) | floor) and
  (if (.value_partitions // 1) > 1 then
     .value_partition_scratch_bytes ==
       .query_heads * .head_dimension * (.value_partitions // 1) * 4
   else true end)
' "$profile" >/dev/null || {
    echo "gem: profile does not satisfy the 31B F32 decode traffic contract" >&2
    exit 1
}

profile_hash=$(shasum -a 256 "$profile" | awk '{print $1}')
temporary=$(mktemp /private/tmp/gem-decode-traffic.XXXXXX)
trap 'rm -f "$temporary"' EXIT INT TERM HUP
jq --arg source "$profile" --arg profile_hash "$profile_hash" '
  (.query_heads * .live_tokens * 4) as $score |
  (.kv_heads * .head_dimension * .live_tokens * 4) as $cache_side |
  (.head_dimension / .value_dimension_tile) as $dimension_groups |
  (.value_partitions // 1) as $partitions |
  (if $partitions > 1 then .value_partition_scratch_bytes else 0 end) as $partials |
  (.query_heads * .head_dimension * 4) as $output |
  ($score * (1 + 5 + $dimension_groups)) as $score_intermediate |
  ($cache_side * 2 + $score_intermediate + $partials * 2 + $output) as $total |
  {schema_version: 1, engine: "gem-native-c-metal",
   workload: "decode-attention-minimum-logical-traffic",
   source_profile: $source, source_profile_sha256: $profile_hash,
   shape: {name: .shape, query_heads, kv_heads, head_dimension,
           live_tokens, value_dimension_tile,
           value_dimension_groups: $dimension_groups,
           value_partitions: $partitions},
   bytes: {key_cache_read: $cache_side, value_cache_read: $cache_side,
           raw_score_write: $score,
           normalize_score_io: (5 * $score),
           value_score_read: ($dimension_groups * $score),
           avoidable_score_intermediate: $score_intermediate,
           value_partial_write: $partials, value_partial_reduce_read: $partials,
           output_write: $output, modeled_total: $total},
   ratios: {
     score_intermediate_over_kv_cache_read_pct:
       (100 * $score_intermediate / (2 * $cache_side)),
     bounded_partial_over_kv_cache_read_pct:
       (100 * (2 * $partials + $output) / (2 * $cache_side))},
   interpretation: {
     accounting: "minimum-logical-bytes-not-device-counter-traffic",
     normalization_passes: ["max-read", "exp-read-write", "scale-read-write"],
     value_score_reuse: "one full score-row read per output-dimension group",
     promotion_eligible: false,
     required_next_evidence: "Metal bandwidth/cache counters and authentic whole-token A/B"}}
' "$profile" > "$temporary"

if [ -n "$output" ]; then
    mv "$temporary" "$output"
    trap - EXIT INT TERM HUP
else
    cat "$temporary"
fi