#include <algorithm>
#include <chrono>
#include <cstdint>
#include <format>
#include <iostream>
#include <print>
#include <random>
#include <string>
#include <vector>

#include <arm_neon.h>

#include "counters/bench.h"

double pretty_print(const std::string &name, size_t num_values,
                    counters::event_aggregate agg) {
  std::print("{:<50} : ", name);
  std::print(" {:5.3f} ns ", agg.fastest_elapsed_ns() / double(num_values));
  std::print(" {:5.2f} Gv/s ", double(num_values) / agg.fastest_elapsed_ns());
  if (counters::has_performance_counters()) {
    std::print(" {:5.2f} GHz ", agg.cycles() / double(agg.elapsed_ns()));
    std::print(" {:5.2f} c ", agg.fastest_cycles() / double(num_values));
    std::print(" {:5.2f} i ", agg.fastest_instructions() / double(num_values));
    std::print(" {:5.2f} i/c ",
               agg.fastest_instructions() / double(agg.fastest_cycles()));
  }
  std::print("\n");
  return double(num_values) / agg.fastest_elapsed_ns();
}

void collect_benchmark_results(size_t input_size, size_t number_strings) {
  // Generate many strings of varying content, including leading/trailing spaces
  std::vector<std::string> strings;
  strings.reserve(number_strings);
  // Create strings with varying lengths up to input_size
  for (size_t i = 0; i < number_strings; ++i) {
    size_t len = (i % input_size) + 1;
    std::string s;
    s.reserve(len + 10);
    // Fill with random printable characters and occasional spaces
    for (size_t k = 0; k < len; ++k) {
      char c = char('!' + (k + i) % 94);
      s.push_back(c);
    }
    strings.push_back(std::move(s));
  }

  volatile uint64_t counter_classic = 0;

  auto count_classic = [&strings, &counter_classic]() {
    size_t c = 0;
    for (const auto &str : strings) {
      c += std::count(str.begin(), str.end(), '!');
    }
    counter_classic += c;
  };

  volatile uint64_t counter_ranges = 0;

  auto count_ranges = [&strings, &counter_ranges]() {
    size_t c = 0;
    for (const auto &str : strings) {
      c += std::ranges::count(str, '!');
    }
    counter_ranges += c;
  };

  volatile uint64_t counter_assembly_claude = 0;
  auto count_assembly_claude = [&strings, &counter_assembly_claude]() {
    size_t c = 0;
    for (const auto &str : strings) {
      size_t local_c = 0;
      const char* data = str.data();
      size_t len = str.size();
      __asm__ volatile (
        // NEON: process 16 bytes at a time
        "dup v1.16b, %w[ch]\n"       // v1 = splat '!'
        "movi v2.16b, #0\n"          // v2 = byte accumulator
        "mov x0, %[data]\n"
        "mov x1, %[len]\n"
        "cmp x1, #16\n"
        "b.lt 2f\n"
        // Main NEON loop
        "1:\n"
        "ld1 {v0.16b}, [x0], #16\n"
        "cmeq v0.16b, v0.16b, v1.16b\n" // 0xFF for match, 0x00 otherwise
        "sub v2.16b, v2.16b, v0.16b\n"  // sub -1 = add 1 per match
        "sub x1, x1, #16\n"
        "cmp x1, #16\n"
        "b.ge 1b\n"
        // Horizontal sum of byte accumulator
        "2:\n"
        "uaddlv h3, v2.16b\n"        // sum all 16 bytes into h3
        "umov w2, v3.h[0]\n"         // extract integer to GPR
        // Scalar tail
        "3:\n"
        "cbz x1, 4f\n"
        "ldrb w3, [x0], #1\n"
        "cmp w3, %w[ch]\n"
        "cinc w2, w2, eq\n"          // branchless increment
        "sub x1, x1, #1\n"
        "b 3b\n"
        "4:\n"
        "mov %w[out], w2\n"
        : [out] "=r"(local_c)
        : [data] "r"(data), [len] "r"(len), [ch] "r"((int)'!')
        : "x0", "x1", "x2", "x3", "v0", "v1", "v2", "v3"
      );
      c += local_c;
    }
    counter_assembly_claude += c;
  };
  volatile uint64_t counter_assembly_grok = 0;
  auto count_assembly_grok = [&strings, &counter_assembly_grok]() {
    size_t c = 0;
    for (const auto &str : strings) {
      size_t local_c = 0;
      const char* data = str.data();
      size_t len = str.size();
      __asm__ volatile (
        "dup v1.16b, %w[ch]\n"       // v1 = splat '!'
        "movi v2.16b, #0\n"          // v2 = byte accumulator
        "mov x0, %[data]\n"
        "mov x1, %[len]\n"
        "cmp x1, #32\n"
        "b.lt 2f\n"
        // Main NEON loop: process 32 bytes per iteration
        "1:\n"
        "ld1 {v0.16b}, [x0], #16\n"
        "cmeq v0.16b, v0.16b, v1.16b\n"
        "sub v2.16b, v2.16b, v0.16b\n"
        "ld1 {v0.16b}, [x0], #16\n"
        "cmeq v0.16b, v0.16b, v1.16b\n"
        "sub v2.16b, v2.16b, v0.16b\n"
        "sub x1, x1, #32\n"
        "cmp x1, #32\n"
        "b.ge 1b\n"
        // Process remaining 16 bytes if any
        "2:\n"
        "cmp x1, #16\n"
        "b.lt 3f\n"
        "ld1 {v0.16b}, [x0], #16\n"
        "cmeq v0.16b, v0.16b, v1.16b\n"
        "sub v2.16b, v2.16b, v0.16b\n"
        "sub x1, x1, #16\n"
        // Horizontal sum
        "3:\n"
        "uaddlv h3, v2.16b\n"
        "umov w2, v3.h[0]\n"
        // Scalar tail
        "4:\n"
        "cbz x1, 5f\n"
        "ldrb w3, [x0], #1\n"
        "cmp w3, %w[ch]\n"
        "cinc w2, w2, eq\n"
        "sub x1, x1, #1\n"
        "b 4b\n"
        "5:\n"
        "mov %w[out], w2\n"
        : [out] "=r"(local_c)
        : [data] "r"(data), [len] "r"(len), [ch] "r"((int)'!')
        : "x0", "x1", "x2", "x3", "v0", "v1", "v2", "v3"
      );
      c += local_c;
    }
    counter_assembly_grok += c;
  };

  volatile uint64_t counter_claude2 = 0;
  auto count_assembly_claude_2 = [&strings, &counter_claude2]() {
    size_t c = 0;
    for (const auto &str : strings) {
      size_t local_c = 0;
      const char* data = str.data();
      size_t len = str.size();
      __asm__ volatile (
        "dup v1.16b, %w[ch]\n"
        "movi v2.16b, #0\n"          // accumulator 0
        "movi v4.16b, #0\n"          // accumulator 1
        "movi v7.16b, #0\n"          // accumulator 2
        "movi v16.16b, #0\n"         // accumulator 3
        "mov x0, %[data]\n"
        "mov x1, %[len]\n"
        // 64-byte loop
        "cmp x1, #64\n"
        "b.lt 20f\n"
        "10:\n"
        "ldp q0, q3, [x0]\n"        // load 32 bytes
        "ldp q5, q6, [x0, #32]\n"   // load next 32 bytes
        "add x0, x0, #64\n"
        "cmeq v0.16b, v0.16b, v1.16b\n"
        "cmeq v3.16b, v3.16b, v1.16b\n"
        "cmeq v5.16b, v5.16b, v1.16b\n"
        "cmeq v6.16b, v6.16b, v1.16b\n"
        "sub v2.16b, v2.16b, v0.16b\n"
        "sub v4.16b, v4.16b, v3.16b\n"
        "sub v7.16b, v7.16b, v5.16b\n"
        "sub v16.16b, v16.16b, v6.16b\n"
        "sub x1, x1, #64\n"
        "cmp x1, #64\n"
        "b.ge 10b\n"
        // 16-byte loop for remainder
        "20:\n"
        "cmp x1, #16\n"
        "b.lt 30f\n"
        "21:\n"
        "ld1 {v0.16b}, [x0], #16\n"
        "cmeq v0.16b, v0.16b, v1.16b\n"
        "sub v2.16b, v2.16b, v0.16b\n"
        "sub x1, x1, #16\n"
        "cmp x1, #16\n"
        "b.ge 21b\n"
        // Merge 4 accumulators and horizontal sum
        "30:\n"
        "add v2.16b, v2.16b, v4.16b\n"
        "add v7.16b, v7.16b, v16.16b\n"
        "add v2.16b, v2.16b, v7.16b\n"
        "uaddlv h3, v2.16b\n"
        "umov w2, v3.h[0]\n"
        // Scalar tail
        "31:\n"
        "cbz x1, 40f\n"
        "ldrb w3, [x0], #1\n"
        "cmp w3, %w[ch]\n"
        "cinc w2, w2, eq\n"
        "sub x1, x1, #1\n"
        "b 31b\n"
        "40:\n"
        "mov %w[out], w2\n"
        : [out] "=r"(local_c)
        : [data] "r"(data), [len] "r"(len), [ch] "r"((int)'!')
        : "x0", "x1", "x2", "x3",
          "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v16"
      );
      c += local_c;
    }
    counter_claude2 += c;
  };
  volatile uint64_t counter_assembly_grok2 = 0;
  auto count_assembly_grok_2 = [&strings, &counter_assembly_grok2]() {
    size_t c = 0;
    for (const auto &str : strings) {
      size_t local_c = 0;
      const char* data = str.data();
      size_t len = str.size();
      __asm__ volatile (
        "dup v1.16b, %w[ch]\n"
        "movi v2.16b, #0\n"          // accumulator 0
        "movi v4.16b, #0\n"          // accumulator 1
        "movi v7.16b, #0\n"          // accumulator 2
        "movi v16.16b, #0\n"         // accumulator 3
        "mov x0, %[data]\n"
        "mov x1, %[len]\n"
        // 64-byte loop with consecutive loads
        "cmp x1, #64\n"
        "b.lt 20f\n"
        "10:\n"
        "ldp q0, q3, [x0], #32\n"    // load 32 bytes into q0, q3
        "ldp q5, q6, [x0], #32\n"    // load next 32 bytes into q5, q6
        "cmeq v0.16b, v0.16b, v1.16b\n"
        "cmeq v3.16b, v3.16b, v1.16b\n"
        "cmeq v5.16b, v5.16b, v1.16b\n"
        "cmeq v6.16b, v6.16b, v1.16b\n"
        "sub v2.16b, v2.16b, v0.16b\n"
        "sub v4.16b, v4.16b, v3.16b\n"
        "sub v7.16b, v7.16b, v5.16b\n"
        "sub v16.16b, v16.16b, v6.16b\n"
        "sub x1, x1, #64\n"
        "cmp x1, #64\n"
        "b.ge 10b\n"
        // 16-byte loop for remainder
        "20:\n"
        "cmp x1, #16\n"
        "b.lt 30f\n"
        "21:\n"
        "ld1 {v0.16b}, [x0], #16\n"
        "cmeq v0.16b, v0.16b, v1.16b\n"
        "sub v2.16b, v2.16b, v0.16b\n"
        "sub x1, x1, #16\n"
        "cmp x1, #16\n"
        "b.ge 21b\n"
        // Merge 4 accumulators and horizontal sum
        "30:\n"
        "add v2.16b, v2.16b, v4.16b\n"
        "add v7.16b, v7.16b, v16.16b\n"
        "add v2.16b, v2.16b, v7.16b\n"
        "uaddlv h3, v2.16b\n"
        "umov w2, v3.h[0]\n"
        // Scalar tail
        "31:\n"
        "cbz x1, 40f\n"
        "ldrb w3, [x0], #1\n"
        "cmp w3, %w[ch]\n"
        "cinc w2, w2, eq\n"
        "sub x1, x1, #1\n"
        "b 31b\n"
        "40:\n"
        "mov %w[out], w2\n"
        : [out] "=r"(local_c)
        : [data] "r"(data), [len] "r"(len), [ch] "r"((int)'!')
        : "x0", "x1", "x2", "x3",
          "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v16"
      );
      c += local_c;
    }
    counter_assembly_grok2 += c;
  };

  volatile uint64_t counter_claude3 = 0;
  auto count_assembly_claude_3 = [&strings, &counter_claude3]() {
    size_t c = 0;
    for (const auto &str : strings) {
      size_t local_c = 0;
      const char* data = str.data();
      size_t len = str.size();
      __asm__ volatile (
        "dup v1.16b, %w[ch]\n"
        "movi v2.16b, #0\n"
        "movi v4.16b, #0\n"
        "movi v7.16b, #0\n"
        "movi v16.16b, #0\n"
        "mov x0, %[data]\n"
        "mov x1, %[len]\n"
        // 64-byte loop using single ld1 for 4 regs
        "subs x1, x1, #64\n"
        "b.lt 20f\n"
        "10:\n"
        "ld1 {v17.16b, v18.16b, v19.16b, v20.16b}, [x0], #64\n"
        "cmeq v17.16b, v17.16b, v1.16b\n"
        "cmeq v18.16b, v18.16b, v1.16b\n"
        "cmeq v19.16b, v19.16b, v1.16b\n"
        "cmeq v20.16b, v20.16b, v1.16b\n"
        "sub v2.16b, v2.16b, v17.16b\n"
        "sub v4.16b, v4.16b, v18.16b\n"
        "sub v7.16b, v7.16b, v19.16b\n"
        "sub v16.16b, v16.16b, v20.16b\n"
        "subs x1, x1, #64\n"
        "b.ge 10b\n"
        // Restore remainder
        "20:\n"
        "adds x1, x1, #64\n"
        "b.eq 30f\n"
        // 16-byte remainder
        "subs x1, x1, #16\n"
        "b.lt 25f\n"
        "21:\n"
        "ld1 {v0.16b}, [x0], #16\n"
        "cmeq v0.16b, v0.16b, v1.16b\n"
        "sub v2.16b, v2.16b, v0.16b\n"
        "subs x1, x1, #16\n"
        "b.ge 21b\n"
        "25:\n"
        "adds x1, x1, #16\n"
        // Merge 4 accumulators
        "30:\n"
        "add v2.16b, v2.16b, v4.16b\n"
        "add v7.16b, v7.16b, v16.16b\n"
        "add v2.16b, v2.16b, v7.16b\n"
        "uaddlv h3, v2.16b\n"
        "umov w2, v3.h[0]\n"
        // Scalar tail (x1 has 0-15 remaining)
        "b.eq 40f\n"
        "31:\n"
        "ldrb w3, [x0], #1\n"
        "cmp w3, %w[ch]\n"
        "cinc w2, w2, eq\n"
        "subs x1, x1, #1\n"
        "b.gt 31b\n"
        "40:\n"
        "mov %w[out], w2\n"
        : [out] "=r"(local_c)
        : [data] "r"(data), [len] "r"(len), [ch] "r"((int)'!')
        : "x0", "x1", "x2", "x3",
          "v0", "v1", "v2", "v3", "v4", "v7",
          "v16", "v17", "v18", "v19", "v20"
      );
      c += local_c;
    }
    counter_claude3 += c;
  };

  volatile uint64_t counter_neon_intrinsics = 0;
  auto count_neon_intrinsics_claude_3_grok = [&strings, &counter_neon_intrinsics]() {
    size_t c = 0;
    for (const auto &str : strings) {
      size_t local_c = 0;
      const char* data = str.data();
      size_t len = str.size();
      uint8x16_t splat = vdupq_n_u8('!');
      uint8x16_t acc0 = vdupq_n_u8(0);
      uint8x16_t acc1 = vdupq_n_u8(0);
      uint8x16_t acc2 = vdupq_n_u8(0);
      uint8x16_t acc3 = vdupq_n_u8(0);
      const uint8_t* ptr = reinterpret_cast<const uint8_t*>(data);
      while (len >= 64) {
        uint8x16_t v0 = vld1q_u8(ptr); ptr += 16;
        uint8x16_t v1 = vld1q_u8(ptr); ptr += 16;
        uint8x16_t v2 = vld1q_u8(ptr); ptr += 16;
        uint8x16_t v3 = vld1q_u8(ptr); ptr += 16;
        uint8x16_t cmp0 = vceqq_u8(v0, splat);
        uint8x16_t cmp1 = vceqq_u8(v1, splat);
        uint8x16_t cmp2 = vceqq_u8(v2, splat);
        uint8x16_t cmp3 = vceqq_u8(v3, splat);
        acc0 = vsubq_u8(acc0, cmp0);
        acc1 = vsubq_u8(acc1, cmp1);
        acc2 = vsubq_u8(acc2, cmp2);
        acc3 = vsubq_u8(acc3, cmp3);
        len -= 64;
      }
      while (len >= 16) {
        uint8x16_t v = vld1q_u8(ptr); ptr += 16;
        uint8x16_t cmp = vceqq_u8(v, splat);
        acc0 = vsubq_u8(acc0, cmp);
        len -= 16;
      }
      // Merge accumulators
      acc0 = vaddq_u8(acc0, acc1);
      acc2 = vaddq_u8(acc2, acc3);
      acc0 = vaddq_u8(acc0, acc2);
      // Horizontal sum
      uint16x8_t sum16 = vpaddlq_u8(acc0);
      uint32x4_t sum32 = vpaddlq_u16(sum16);
      uint64x2_t sum64 = vpaddlq_u32(sum32);
      uint64_t total = vgetq_lane_u64(sum64, 0) + vgetq_lane_u64(sum64, 1);
      // Scalar tail
      for (size_t i = 0; i < len; ++i) {
        if (ptr[i] == '!') ++total;
      }
      local_c = total;
      c += local_c;
    }
    counter_neon_intrinsics += c;
  };

  volatile uint64_t counter_neon_intrinsics_claude = 0;
  auto count_neon_intrinsics_claude_3_claude = [&strings, &counter_neon_intrinsics_claude]() {
    size_t c = 0;
    for (const auto &str : strings) {
      const uint8_t* ptr = reinterpret_cast<const uint8_t*>(str.data());
      size_t len = str.size();
      uint8x16_t splat = vdupq_n_u8('!');
      uint8x16_t acc0 = vdupq_n_u8(0);
      uint8x16_t acc1 = vdupq_n_u8(0);
      uint8x16_t acc2 = vdupq_n_u8(0);
      uint8x16_t acc3 = vdupq_n_u8(0);
      // 64-byte loop: single 4-register ld1, mirroring the asm
      while (len >= 64) {
        uint8x16x4_t q = vld1q_u8_x4(ptr);
        ptr += 64;
        acc0 = vsubq_u8(acc0, vceqq_u8(q.val[0], splat));
        acc1 = vsubq_u8(acc1, vceqq_u8(q.val[1], splat));
        acc2 = vsubq_u8(acc2, vceqq_u8(q.val[2], splat));
        acc3 = vsubq_u8(acc3, vceqq_u8(q.val[3], splat));
        len -= 64;
      }
      // 16-byte remainder
      while (len >= 16) {
        uint8x16_t v = vld1q_u8(ptr);
        ptr += 16;
        acc0 = vsubq_u8(acc0, vceqq_u8(v, splat));
        len -= 16;
      }
      // Merge 4 accumulators, then a single uaddlv horizontal sum
      acc0 = vaddq_u8(acc0, acc1);
      acc2 = vaddq_u8(acc2, acc3);
      acc0 = vaddq_u8(acc0, acc2);
      uint32_t total = vaddlvq_u8(acc0);
      // Scalar tail (0-15 bytes)
      for (size_t i = 0; i < len; ++i) {
        total += (ptr[i] == '!');
      }
      c += total;
    }
    counter_neon_intrinsics_claude += c;
  };

  pretty_print("count_classic", number_strings, counters::bench(count_classic));
  pretty_print("count_ranges", number_strings, counters::bench(count_ranges));
  pretty_print("count_assembly_claude", number_strings, counters::bench(count_assembly_claude));
  pretty_print("count_assembly_grok", number_strings, counters::bench(count_assembly_grok));
  pretty_print("count_assembly_claude_2", number_strings, counters::bench(count_assembly_claude_2));
  pretty_print("count_assembly_grok_2", number_strings, counters::bench(count_assembly_grok_2));
  pretty_print("count_assembly_claude_3", number_strings, counters::bench(count_assembly_claude_3));
  pretty_print("count_neon_intrinsics_claude_3_grok", number_strings, counters::bench(count_neon_intrinsics_claude_3_grok));
  pretty_print("count_neon_intrinsics_claude_3_claude", number_strings, counters::bench(count_neon_intrinsics_claude_3_claude));

  // Validate correctness with a single run of each
  counter_classic = 0;
  counter_ranges = 0;
  counter_assembly_claude = 0;
  counter_assembly_grok = 0;
  counter_claude2 = 0;
  counter_assembly_grok2 = 0;
  counter_claude3 = 0;
  counter_neon_intrinsics = 0;
  counter_neon_intrinsics_claude = 0;
  count_classic();
  count_ranges();
  count_assembly_claude();
  count_assembly_grok();
  count_assembly_claude_2();
  count_assembly_grok_2();
  count_assembly_claude_3();
  count_neon_intrinsics_claude_3_grok();
  count_neon_intrinsics_claude_3_claude();
  if (counter_classic != counter_assembly_claude) {
    std::cout << "Error: counts differ (classic vs claude): " << counter_classic << " vs " << counter_assembly_claude << std::endl;
  }
  if (counter_classic != counter_ranges) {
    std::cout << "Error: counts differ (classic vs ranges): " << counter_classic << " vs " << counter_ranges << std::endl;
  }
  if (counter_classic != counter_assembly_grok) {
    std::cout << "Error: counts differ (classic vs grok): " << counter_classic << " vs " << counter_assembly_grok << std::endl;
  }
  if (counter_classic != counter_claude2) {
    std::cout << "Error: counts differ (classic vs claude2): " << counter_classic << " vs " << counter_claude2 << std::endl;
  }
  if (counter_classic != counter_assembly_grok2) {
    std::cout << "Error: counts differ (classic vs grok2): " << counter_classic << " vs " << counter_assembly_grok2 << std::endl;
  }
  if (counter_classic != counter_claude3) {
    std::cout << "Error: counts differ (classic vs claude3): " << counter_classic << " vs " << counter_claude3 << std::endl;
  }
  if (counter_classic != counter_neon_intrinsics) {
    std::cout << "Error: counts differ (classic vs neon_intrinsics): " << counter_classic << " vs " << counter_neon_intrinsics << std::endl;
  }
  if (counter_classic != counter_neon_intrinsics_claude) {
    std::cout << "Error: counts differ (classic vs neon_intrinsics_claude): " << counter_classic << " vs " << counter_neon_intrinsics_claude << std::endl;
  }
}

int main(int argc, char **argv) { collect_benchmark_results(1024, 100000); }
