#!/bin/sh
set -eu

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

analyze() {
    jq '
      def sum_indices($xs; $indices): [$indices[] as $i | $xs[$i]] | add;
      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);

      . as $p |
      ["attention_norm", "qkv_projections", "qkv_transform", "kv_attention",
       "attention_output", "ffn_norm", "ffn_gate_up", "ffn_activation",
       "ffn_down_output"] as $expected_names |
      ($p.stages.all.mean_seconds | add) as $layer_stage_total |
      sum_indices($p.stages.all.mean_seconds; [1, 4, 6, 8]) as $projection_total |
      $p.stages.all.mean_seconds[3] as $attention_total |
      sum_indices($p.stages.all.mean_seconds; [0, 2, 5, 7]) as $elementwise_total |
      $p.stages.sliding.mean_seconds[3] as $sliding_attention |
      $p.stages.full.mean_seconds[3] as $full_attention |
      if ($p.schema_version != 1 or $p.engine != "gem-native-c-metal" or
          $p.workload != "long-prefill-stage-profile" or
          $p.stage_names != $expected_names or
          $p.layer_count != 60 or $p.sliding_layer_count != 50 or
          $p.full_layer_count != 10 or $p.tile_tokens != 32 or
          $p.repetitions < 3 or
          ($p.wall_seconds_samples | length) != $p.repetitions or
          ($p.gpu_seconds_samples | length) != $p.repetitions or
          ($p.sampled_gpu_seconds_samples | length) != $p.repetitions or
          ($p.stages.all.mean_seconds | length) != 9 or
          ($p.stages.sliding.mean_seconds | length) != 9 or
          ($p.stages.full.mean_seconds | length) != 9 or
          ($p.stages.all.samples_seconds | length) != $p.repetitions or
          ($p.stages.sliding.samples_seconds | length) != $p.repetitions or
          ($p.stages.full.samples_seconds | length) != $p.repetitions or
          ($p.stages.all.samples_seconds | any(length != 9)) or
          ($p.stages.sliding.samples_seconds | any(length != 9)) or
          ($p.stages.full.samples_seconds | any(length != 9)))
      then error("invalid production prefill-stage profile")
      else
        {
          schema_version: 1,
          scope: "production-graph prefill stage attribution",
          source: {
            model: $p.model,
            attention_kernel: $p.attention_kernel,
            kv_type: $p.kv_type,
            prefix_mode: ($p.prefix_mode // "evaluated-prompt"),
            prefix_tokens: $p.prefix_tokens,
            tile_tokens: $p.tile_tokens,
            repetitions: $p.repetitions,
            layer_count: $p.layer_count,
            sliding_layer_count: $p.sliding_layer_count,
            full_layer_count: $p.full_layer_count
          },
          timing: {
            mean_wall_seconds: mean($p.wall_seconds_samples),
            mean_gpu_seconds: mean($p.gpu_seconds_samples),
            mean_sampled_gpu_seconds: mean($p.sampled_gpu_seconds_samples),
            sampled_gpu_fraction: (mean($p.sampled_gpu_seconds_samples) /
                                   mean($p.gpu_seconds_samples)),
            wall_cv_pct: cv_pct($p.wall_seconds_samples),
            gpu_cv_pct: cv_pct($p.gpu_seconds_samples)
          },
          ranked_layer_stages: ([range(0; 9) as $i | {
            name: $p.stage_names[$i],
            mean_seconds: $p.stages.all.mean_seconds[$i],
            layer_stage_fraction: ($p.stages.all.mean_seconds[$i] / $layer_stage_total),
            sliding_seconds: $p.stages.sliding.mean_seconds[$i],
            full_seconds: $p.stages.full.mean_seconds[$i]
          }] | sort_by(-.mean_seconds)),
          families: {
            projections: {
              mean_seconds: $projection_total,
              sampled_gpu_fraction: ($projection_total / $p.mean_sampled_gpu_seconds)
            },
            kv_attention: {
              mean_seconds: $attention_total,
              sampled_gpu_fraction: ($attention_total / $p.mean_sampled_gpu_seconds),
              sliding_seconds: $sliding_attention,
              full_seconds: $full_attention
            },
            elementwise: {
              mean_seconds: $elementwise_total,
              sampled_gpu_fraction: ($elementwise_total / $p.mean_sampled_gpu_seconds)
            },
            embedding_and_logits: {
              mean_seconds: ($p.mean_embedding_seconds + $p.mean_logits_seconds),
              sampled_gpu_fraction: (($p.mean_embedding_seconds + $p.mean_logits_seconds) /
                                     $p.mean_sampled_gpu_seconds)
            }
          },
          selected_next_axis: (
            if $projection_total > $attention_total then "q4-projection-kernels"
            elif $sliding_attention > $full_attention then
              "sliding-attention-preserving-rowwise-arithmetic"
            else "full-attention-head-dim512" end),
          screening: {
            promotion_eligible: false,
            reason: "stage attribution only; implement and benchmark one selected axis in the full graph"
          }
        }
      end
    ' "$profile"
}

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