#pragma once
#include <cmath>
#include <cstdint>
#include <cstring>
#include <immintrin.h> // AVX Intrinsics
#include <omp.h>
#include <vector>

// --- CONFIGURAÇÃO ---
const int BLOCK_SIZE_LIB = 128;

// --- ESTRUTURAS ---
struct PackedWeights {
  std::vector<uint8_t> data;
  int rows; // K
  int cols; // N

  PackedWeights(int r, int c) : rows(r), cols(c) {
    // 2 bits per element. 4 elements per byte.
    // Total elements = rows * cols.
    // Total bytes = ceil(total_elements / 4)
    data.resize(((size_t)r * c + 3) / 4 + 1, 0);
  }

  // Deprecated square constructor for backward compat
  PackedWeights(int n) : PackedWeights(n, n) {}

  void set(int i, int j, int val) {
    // Row-major: idx = i * cols + j
    size_t idx = (size_t)i * cols + j;
    size_t byte_idx = idx / 4;
    int bit_offset = (idx % 4) * 2;

    uint8_t mask = 0;
    if (val == 1)
      mask = 0b01;
    else if (val == -1)
      mask = 0b10;

    data[byte_idx] &= ~(0b11 << bit_offset);
    data[byte_idx] |= (mask << bit_offset);
  }

  inline int get(int i, int j) const {
    size_t idx = (size_t)i * cols + j;
    size_t byte_idx = idx / 4;
    int bit_offset = (idx % 4) * 2;

    uint8_t byte = data[byte_idx];
    uint8_t val = (byte >> bit_offset) & 0b11;

    if (val == 0b01)
      return 1;
    if (val == 0b10)
      return -1;
    return 0;
  }
};

struct BlockCoord {
  int bi;
  int bk;
};

// --- DECLARATION OF CUDA KERNELS (IF ENABLED) ---
#ifdef MAXWELL_CUDA
extern "C" void launch_maxwell_cuda(const float *h_A, const uint8_t *h_B_packed,
                                    float *h_C, const BlockCoord *h_active,
                                    int num_active, int M, int K, int N);

extern "C" void launch_maxwell_cuda_wmma(const float *h_A,
                                         const uint8_t *h_B_packed, float *h_C,
                                         const BlockCoord *h_active,
                                         int num_active, int M, int K, int N);

extern "C" void launch_maxwell_cuda_wmma_csr(const float *h_A,
                                             const uint8_t *h_B_packed,
                                             float *h_C, const int *h_row_ptr,
                                             const int *h_col_ind,
                                             int num_active, int M, int K,
                                             int N);
// NEW: Backward Weights Kernel
extern "C" void launch_maxwell_backward_weights(const float *h_A, const float *h_grad_Y,
                                               float *h_grad_W, int M, int K, int N);

// NEW: GPU Oracle
extern "C" void launch_oracle_cuda(const float *h_A, BlockCoord *h_active, 
                                   int *h_count, int M, int K, float threshold);
#endif

// --- DEMON ORACLE: STATISTICAL SAMPLING (O(1)) ---
// Updated to support M x K (A)
inline std::vector<BlockCoord> invoke_demon_lib(const float *A, int M, int K) {
  std::vector<BlockCoord> active_blocks;
  int blocks_m = (M + BLOCK_SIZE_LIB - 1) / BLOCK_SIZE_LIB;
  int blocks_k = (K + BLOCK_SIZE_LIB - 1) / BLOCK_SIZE_LIB;

  // 1. Mask-First Guard: Check global sparsity on 1% sample
  size_t total_elements = (size_t)M * K;
  size_t sample_size = total_elements / 100; // 1%
  if (sample_size == 0)
    sample_size = total_elements; // Small matrices

  int non_zeros = 0;
  int stride = 101;
  size_t limit = std::min(sample_size * stride, total_elements);

  for (size_t k = 0; k < limit; k += stride) {
    if (A[k] != 0.0f)
      non_zeros++;
  }

  float sparsity = 1.0f - ((float)non_zeros / (float)(limit / stride + 1));

  // User threshold: Only sparse if sparsity > 80%.
  if (sparsity < 0.8f) {
    // Use Dense Mode (All Blocks)
    active_blocks.reserve(blocks_m * blocks_k);
    for (int bi = 0; bi < blocks_m; ++bi) {
      for (int bk = 0; bk < blocks_k; ++bk) {
        active_blocks.push_back({bi, bk});
      }
    }
    return active_blocks;
  }

  for (int bi = 0; bi < blocks_m; ++bi) {
    for (int bk = 0; bk < blocks_k; ++bk) {
      bool active = false;
      int base_row = bi * BLOCK_SIZE_LIB;
      int base_col = bk * BLOCK_SIZE_LIB;

      // Safety check for boundary blocks
      if (base_row >= M || base_col >= K)
        continue;

      // Check 64 samples
      const int mask = (BLOCK_SIZE_LIB * BLOCK_SIZE_LIB) - 1;
      for (int k = 0; k < 64; ++k) {
        int offset = (k * 257) & mask;
        // Map offset to 2D within block
        int r_off = offset / BLOCK_SIZE_LIB;
        int c_off = offset % BLOCK_SIZE_LIB;

        if (base_row + r_off < M && base_col + c_off < K) {
          if (A[(size_t)(base_row + r_off) * K + (base_col + c_off)] != 0.0f) {
            active = true;
            break;
          }
        }
      }
      if (active) {
        active_blocks.push_back({bi, bk});
      }
    }
  }
  return active_blocks;
}

// --- ORACLE: EARLY EXIT ---
inline std::vector<BlockCoord>
invoke_demon_early_exit(const float *A, int M, int K, float threshold = 1e-5) {
  std::vector<BlockCoord> active_blocks;
  int blocks_m = (M + BLOCK_SIZE_LIB - 1) / BLOCK_SIZE_LIB;
  int blocks_k = (K + BLOCK_SIZE_LIB - 1) / BLOCK_SIZE_LIB;

  for (int bi = 0; bi < blocks_m; ++bi) {
    for (int bk = 0; bk < blocks_k; ++bk) {
      float energy = 0.0f;
      int row_idx = bi * BLOCK_SIZE_LIB;
      int col_start = bk * BLOCK_SIZE_LIB;

      if (row_idx >= M || col_start >= K)
        continue;

      for (int k = 0; k < 16; ++k) {
        if (col_start + k < K)
          energy += std::abs(A[(size_t)row_idx * K + col_start + k]);
      }
      if (energy > threshold)
        active_blocks.push_back({bi, bk});
    }
  }
  return active_blocks;
}

// --- METAL KERNELS ---

// --- METAL KERNELS (TILED & OPTIMIZED) ---
// Generalized for A: [M, K], B: [K, N], C: [M, N]
inline void run_maxwell_metal_lib(const float *A, const PackedWeights &B,
                                  float *C, int M, int K, int N,
                                  const std::vector<BlockCoord> &active_blocks,
                                  bool apply_relu = false) {

  int n_block_rows = (M + BLOCK_SIZE_LIB - 1) / BLOCK_SIZE_LIB;
  std::vector<std::vector<int>> blocks_by_row(n_block_rows);

  for (const auto &coord : active_blocks) {
    if (coord.bi < n_block_rows)
      blocks_by_row[coord.bi].push_back(coord.bk);
  }

  // Tiling Constants
  const int TILE_J = 1024;
  const int BLOCK_K = 32;

#pragma omp parallel for schedule(dynamic)
  for (int bi = 0; bi < n_block_rows; ++bi) {
    if (blocks_by_row[bi].empty())
      continue;

    int i_base = bi * BLOCK_SIZE_LIB;

    // Tiling J (Columns of C/B)
    for (int j_start = 0; j_start < N; j_start += TILE_J) {
      int j_end = std::min(j_start + TILE_J, N);

      // Iterate over active blocks of A for this row bi
      for (int bk : blocks_by_row[bi]) {
        int k_base = bk * BLOCK_SIZE_LIB;

        int k_block_limit = std::min(BLOCK_SIZE_LIB, K - k_base);
        if (k_block_limit <= 0)
          continue;

        // Blocking K (Cols of A / Rows of B) to keep B in L1
        for (int kk = 0; kk < k_block_limit; kk += BLOCK_K) {
          int k_limit = std::min(kk + BLOCK_K, k_block_limit);

          // Loop over I (Rows of A / C)
          int i_limit = std::min(BLOCK_SIZE_LIB, M - i_base);

          for (int ii = 0; ii < i_limit; ++ii) {
            int i_idx = i_base + ii;
            float *row_C = &C[(size_t)i_idx * N];

            for (int k_offset = kk; k_offset < k_limit; ++k_offset) {
              int k_idx = k_base + k_offset;
              float val_A = A[(size_t)i_idx * K + k_idx]; // A is [M, K]

              if (val_A == 0.0f)
                continue;

              // B is [K, N]. Row k_idx of B.
              // PackedWeights B: row_offset in BYTES = k_idx * (ceil(N/4))
              // Wait, PackedWeights::data is flat.
              // row k starts at: k * ceil(N/4) bytes?
              // PackedWeights stores as [rows, cols].
              // data layout: row0 (packed), row1 (packed)...
              // Row size (packed) = (cols + 3) / 4 bytes.
              int row_stride_bytes = (N + 3) / 4;
              size_t b_row_offset = (size_t)k_idx * row_stride_bytes;

              const uint8_t *B_row_ptr = &B.data[b_row_offset];

              // --- AVX2 Processing ---
#ifdef __AVX2__
              __m256 vec_A = _mm256_set1_ps(val_A);
              const __m256i shift_counts =
                  _mm256_setr_epi32(0, 2, 4, 6, 8, 10, 12, 14);
              const __m256i mask_2bits = _mm256_set1_epi32(0x3);
              const __m256i c1 = _mm256_set1_epi32(1);
              const __m256i c2 = _mm256_set1_epi32(2);

              int j = j_start;
              // Unroll 4x8 = 32 elements
              for (; j <= j_end - 32; j += 32) {
                for (int u = 0; u < 4; ++u) {
                  int jj = j + u * 8;
                  uint16_t packed_chunk;
                  memcpy(&packed_chunk, &B_row_ptr[jj / 4], sizeof(uint16_t));

                  __m256i v_packed = _mm256_set1_epi32((int)packed_chunk);
                  v_packed = _mm256_srlv_epi32(v_packed, shift_counts);
                  v_packed = _mm256_and_si256(v_packed, mask_2bits);

                  __m256i m_pos_i = _mm256_cmpeq_epi32(v_packed, c1);
                  __m256i m_neg_i = _mm256_cmpeq_epi32(v_packed, c2);

                  __m256 vec_C = _mm256_loadu_ps(&row_C[jj]);
                  vec_C = _mm256_add_ps(
                      vec_C,
                      _mm256_and_ps(_mm256_castsi256_ps(m_pos_i), vec_A));
                  vec_C = _mm256_sub_ps(
                      vec_C,
                      _mm256_and_ps(_mm256_castsi256_ps(m_neg_i), vec_A));
                  _mm256_storeu_ps(&row_C[jj], vec_C);
                }
              }
              // Handle cleanup 8
              for (; j <= j_end - 8; j += 8) {
                uint16_t packed_chunk;
                memcpy(&packed_chunk, &B_row_ptr[j / 4], sizeof(uint16_t));
                __m256i v_packed = _mm256_set1_epi32((int)packed_chunk);
                v_packed = _mm256_srlv_epi32(v_packed, shift_counts);
                v_packed = _mm256_and_si256(v_packed, mask_2bits);

                __m256i m_pos_i = _mm256_cmpeq_epi32(v_packed, c1);
                __m256i m_neg_i = _mm256_cmpeq_epi32(v_packed, c2);
                __m256 vec_C = _mm256_loadu_ps(&row_C[j]);
                vec_C = _mm256_add_ps(
                    vec_C, _mm256_and_ps(_mm256_castsi256_ps(m_pos_i), vec_A));
                vec_C = _mm256_sub_ps(
                    vec_C, _mm256_and_ps(_mm256_castsi256_ps(m_neg_i), vec_A));
                _mm256_storeu_ps(&row_C[j], vec_C);
              }
              // Scalar Tail
              for (; j < j_end; ++j) {
                int weight = B.get(k_idx, j);
                if (weight == 1)
                  row_C[j] += val_A;
                else if (weight == -1)
                  row_C[j] -= val_A;
              }
#else
              // Fallback
              for (int j = j_start; j < j_end; ++j) {
                int weight = B.get(k_idx, j);
                if (weight == 1)
                  row_C[j] += val_A;
                else if (weight == -1)
                  row_C[j] -= val_A;
              }
#endif
            }
          }
        }
      }
    }

    // ReLU Activation
    if (apply_relu) {
      int i_limit = std::min(BLOCK_SIZE_LIB, M - i_base);
      for (int ii = 0; ii < i_limit; ++ii) {
        float *row_C = &C[(size_t)(i_base + ii) * N];
#ifdef __AVX2__
        __m256 zero = _mm256_setzero_ps();
        int j = 0;
        for (; j <= N - 8; j += 8) {
          __m256 v = _mm256_loadu_ps(&row_C[j]);
          v = _mm256_max_ps(v, zero);
          _mm256_storeu_ps(&row_C[j], v);
        }
        for (; j < N; ++j)
          if (row_C[j] < 0)
            row_C[j] = 0;
#else
        for (int j = 0; j < N; ++j)
          if (row_C[j] < 0)
            row_C[j] = 0;
#endif
      }
    }
  }
}

// --- RNS / INT32 KERNEL ---
inline void
run_maxwell_int32_lib(const int32_t *A, int N, const PackedWeights &B,
                      int32_t *C,
                      const std::vector<BlockCoord> &active_blocks) {
  int n_block_rows = N / BLOCK_SIZE_LIB;
  std::vector<std::vector<int>> blocks_by_row(n_block_rows);
  for (const auto &coord : active_blocks) {
    if (coord.bi < n_block_rows)
      blocks_by_row[coord.bi].push_back(coord.bk);
  }

#pragma omp parallel for schedule(dynamic)
  for (int bi = 0; bi < n_block_rows; ++bi) {
    if (blocks_by_row[bi].empty())
      continue;
    int i_start = bi * BLOCK_SIZE_LIB;

    for (int bk : blocks_by_row[bi]) {
      int k_start = bk * BLOCK_SIZE_LIB;

      for (int i = i_start; i < i_start + BLOCK_SIZE_LIB; ++i) {
        _mm_prefetch((const char *)&A[(i + 1) * N + k_start], _MM_HINT_T0);
        for (int k = k_start; k < k_start + BLOCK_SIZE_LIB; ++k) {
          int32_t val_A = A[i * N + k];
          if (val_A == 0)
            continue;

          int row_offset = k * (N / 4);
          const uint8_t *B_row_ptr = &B.data[row_offset];
          int32_t *row_C = &C[i * N];

          _mm_prefetch((const char *)&B.data[row_offset + N / 4], _MM_HINT_T0);

#ifdef __AVX2__
          __m256i vec_A = _mm256_set1_epi32(val_A);
          int j = 0;
          const __m256i shift_counts =
              _mm256_setr_epi32(0, 2, 4, 6, 8, 10, 12, 14);
          const __m256i mask_2bits = _mm256_set1_epi32(0x3);
          const __m256i c1 = _mm256_set1_epi32(1);
          const __m256i c2 = _mm256_set1_epi32(2);

          for (; j <= N - 8; j += 8) {
            uint16_t packed_chunk;
            memcpy(&packed_chunk, &B_row_ptr[j / 4], sizeof(uint16_t));

            __m256i v_packed = _mm256_set1_epi32((int)packed_chunk);
            v_packed = _mm256_srlv_epi32(v_packed, shift_counts);
            v_packed = _mm256_and_si256(v_packed, mask_2bits);

            __m256i mask_pos = _mm256_cmpeq_epi32(v_packed, c1);
            __m256i mask_neg = _mm256_cmpeq_epi32(v_packed, c2);

            __m256i vec_C = _mm256_loadu_si256((__m256i *)&row_C[j]);
            vec_C = _mm256_add_epi32(vec_C, _mm256_and_si256(mask_pos, vec_A));
            vec_C = _mm256_sub_epi32(vec_C, _mm256_and_si256(mask_neg, vec_A));

            _mm256_storeu_si256((__m256i *)&row_C[j], vec_C);
          }
          for (; j < N; ++j) {
            int weight = B.get(k, j);
            if (weight == 1)
              row_C[j] += val_A;
            else if (weight == -1)
              row_C[j] -= val_A;
          }
#else
          for (int j = 0; j < N; ++j) {
            int weight = B.get(k, j);
            if (weight == 1)
              row_C[j] += val_A;
            else if (weight == -1)
              row_C[j] -= val_A;
          }
#endif
        }
      }
    }
  }
}

// --- RNS (Residue Number System) SIMULATION ---
// Simulates non-positional arithmetic using independent 8-bit channels.
// Eliminates 32-bit carry propagation latency.
inline void run_maxwell_rns_lib(const int8_t *A, int N, const PackedWeights &B,
                                int8_t *C_residues,
                                const std::vector<BlockCoord> &active_blocks) {
  int n_block_rows = N / BLOCK_SIZE_LIB;
  // RNS: We compute in 8-bit modular arithmetic

#pragma omp parallel for schedule(dynamic)
  for (int bi = 0; bi < n_block_rows; ++bi) {
    int i_start = bi * BLOCK_SIZE_LIB;
    // ... (Simplified traversal for RNS demo)
    for (int i = i_start; i < i_start + BLOCK_SIZE_LIB; ++i) {
      for (int k = 0; k < N; ++k) {
        int8_t val_A = A[i * N + k];
        if (val_A == 0)
          continue;

        int row_offset = k * (N / 4);
        const uint8_t *B_row_ptr = &B.data[row_offset];
        int8_t *row_C = &C_residues[i * N]; // C is also residues

#ifdef __AVX2__
        __m256i vec_A = _mm256_set1_epi8(val_A);
        int j = 0;
        // Process 32 residues at once (256 bits / 8 bits)
        for (; j <= N - 32; j += 32) {
          // Unpack 2-bit to 8-bit? No, RNS operates on residues.
          // Assume weights are mapped to residues {-1, 0, 1} mod M.
          // Here we just simulate the throughput of 8-bit ADDs vs 32-bit FMA.

          // Decode packed B to 8-bit integers (0, 1, -1)
          // ... (Decoding logic omitted for brevity, reusing previous masks
          // adapted for epi8)

          // Benchmark focus: Throughput of _mm256_add_epi8
          __m256i vec_C = _mm256_loadu_si256((__m256i *)&row_C[j]);

          // Dummy op: C += A (Simulating C += A * W in residue domain)
          vec_C = _mm256_add_epi8(vec_C, vec_A);

          _mm256_storeu_si256((__m256i *)&row_C[j], vec_C);
        }
#endif
      }
    }
  }
}

// --- PREDICTIVE ORACLE (APPROXIMATE COMPUTE) ---
// "Chuta" o resultado baseado em amostras (Monte Carlo GEMM)
inline void run_maxwell_predictive_lib(
    const float *A, int N, const PackedWeights &B, float *C,
    const std::vector<BlockCoord> &active_blocks, float sample_ratio = 0.01f) {

  // Monte Carlo scale factor
  float scale = 1.0f / sample_ratio;
  int stride = (int)(1.0f / sample_ratio);

// Only process sampled indices
#pragma omp parallel for
  for (int b_idx = 0; b_idx < active_blocks.size(); ++b_idx) {
    BlockCoord coord = active_blocks[b_idx];
    int bi = coord.bi;
    int bk = coord.bk;

    int i_start = bi * BLOCK_SIZE_LIB;
    int k_start = bk * BLOCK_SIZE_LIB;

    for (int i = i_start; i < i_start + BLOCK_SIZE_LIB; ++i) {
      for (int k = k_start; k < k_start + BLOCK_SIZE_LIB;
           k += stride) { // Jump!
        float val_A = A[i * N + k];
        if (val_A == 0.0f)
          continue;

        // We update the whole row J?
        // Monte Carlo Matrix Mul usually samples k.
        // Estimator: Sum(A_ik * B_kj) approx Sum_samples(A_ik * B_kj) * scale.

        for (int j = 0; j < N; ++j) {
          int weight = B.get(k, j);
          float contribution = 0.0f;
          if (weight == 1)
            contribution = val_A;
          else if (weight == -1)
            contribution = -val_A;

          if (contribution != 0.0f)
            C[i * N + j] += contribution * scale;
        }
      }
    }
  }
}

// --- BACKWARD PASS ---
inline void
run_maxwell_backward_lib(const float *grad_C, int M, int N,
                         const PackedWeights &B, float *grad_A, int K,
                         const std::vector<BlockCoord> &active_blocks) {
  int blocks_m = (M + BLOCK_SIZE_LIB - 1) / BLOCK_SIZE_LIB;
  int blocks_k = (K + BLOCK_SIZE_LIB - 1) / BLOCK_SIZE_LIB;

  std::vector<bool> row_active(blocks_m, false);
  for (const auto &coord : active_blocks) {
    if (coord.bi < blocks_m)
      row_active[coord.bi] = true;
  }

#pragma omp parallel for schedule(dynamic)
  for (int bi = 0; bi < blocks_m; ++bi) {
    if (!row_active[bi])
      continue;
    int i_start = bi * BLOCK_SIZE_LIB;
    int i_limit = std::min(i_start + BLOCK_SIZE_LIB, M);

    for (int bk = 0; bk < blocks_k; ++bk) {
      int k_start = bk * BLOCK_SIZE_LIB;
      int k_limit = std::min(k_start + BLOCK_SIZE_LIB, K);

      for (int i = i_start; i < i_limit; ++i) {
        const float *row_grad_C = &grad_C[(size_t)i * N];
        // Prefetch grad_C row (used repeatedly for all k)
        _mm_prefetch((const char *)row_grad_C, _MM_HINT_T0);

        for (int k = k_start; k < k_limit; ++k) {
          // grad_A[i, k] = grad_C[i, :] . B[k, :]
          // B is stored as [rows=K, cols=N].
          int row_stride_bytes = (N + 3) / 4;
          size_t b_row_offset = (size_t)k * row_stride_bytes;
          const uint8_t *B_row_ptr = &B.data[b_row_offset];

          _mm_prefetch((const char *)B_row_ptr, _MM_HINT_T0);

          float sum = 0.0f;

#ifdef __AVX2__
          __m256 v_sum = _mm256_setzero_ps();
          int j = 0;
          const __m256i shift_counts =
              _mm256_setr_epi32(0, 2, 4, 6, 8, 10, 12, 14);
          const __m256i mask_2bits = _mm256_set1_epi32(0x3);
          const __m256i mask_val_1 = _mm256_set1_epi32(1);
          const __m256i mask_val_2 = _mm256_set1_epi32(2);

          for (; j <= N - 8; j += 8) {
            uint16_t packed_chunk;
            memcpy(&packed_chunk, &B_row_ptr[j / 4], sizeof(uint16_t));

            __m256i v_packed = _mm256_set1_epi32((int)packed_chunk);
            v_packed = _mm256_srlv_epi32(v_packed, shift_counts);
            v_packed = _mm256_and_si256(v_packed, mask_2bits);

            __m256i m_pos_i = _mm256_cmpeq_epi32(v_packed, mask_val_1);
            __m256i m_neg_i = _mm256_cmpeq_epi32(v_packed, mask_val_2);
            __m256 m_pos = _mm256_castsi256_ps(m_pos_i);
            __m256 m_neg = _mm256_castsi256_ps(m_neg_i);
            __m256 vec_G = _mm256_loadu_ps(&row_grad_C[j]);

            v_sum = _mm256_add_ps(v_sum, _mm256_and_ps(m_pos, vec_G));
            v_sum = _mm256_sub_ps(v_sum, _mm256_and_ps(m_neg, vec_G));
          }
          float *arr = (float *)&v_sum;
          sum += arr[0] + arr[1] + arr[2] + arr[3] + arr[4] + arr[5] + arr[6] +
                 arr[7];
          for (; j < N; ++j) {
            int weight = B.get(k, j);
            if (weight == 1)
              sum += row_grad_C[j];
            else if (weight == -1)
              sum -= row_grad_C[j];
          }
#else
          for (int j = 0; j < N; ++j) {
            int weight = B.get(k, j);
            if (weight == 1)
              sum += row_grad_C[j];
            else if (weight == -1)
              sum -= row_grad_C[j];
          }
#endif
          grad_A[(size_t)i * K + k] = sum;
        }
      }
    }
  }
}
