viva_tensor

//// NF4 (NormalFloat4) Quantization - QLoRA Style
////
//// DESCOBERTA VIA HUGGINGCHAT + PESQUISA:
//// NF4 usa 16 níveis derivados dos quantis da distribuição normal
//// Isso é ÓTIMO para pesos de NNs que seguem distribuição gaussiana!
////
//// Referências:
//// - QLoRA Paper: https://arxiv.org/abs/2305.14314
//// - bitsandbytes: create_normal_map function
//// - MLSys 2024: AWQ Best Paper
////
//// Vantagens sobre Q4 uniforme:
//// - Mais precisão próximo de zero (onde concentram os pesos)
//// - 8x compressão com ~0.1% erro
//// - Matematicamente ótimo para dados normais

import gleam/float
import gleam/int
import gleam/io
import gleam/list
import viva_tensor/tensor.{type Tensor, Tensor}

// ============================================================================
// CONSTANTES NF4 - Derivadas da Distribuição Normal
// ============================================================================

/// Os 16 níveis NF4 são os quantis de N(0,1) normalizados para [-1, 1]
/// Esses valores são hardcoded em bitsandbytes e usados em QLoRA
pub fn nf4_levels() -> List(Float) {
  [
    -1.0,
    // quantile(1/16)
    -0.6961928009986877,
    // quantile(2/16)
    -0.5250730514526367,
    // quantile(3/16)
    -0.39491748809814453,
    // quantile(4/16)
    -0.28444138169288635,
    // quantile(5/16)
    -0.18477343022823334,
    // quantile(6/16)
    -0.09105003625154495,
    // quantile(7/16)
    0.0,
    // zero (importante!)
    0.07958029955625534,
    // quantile(9/16)
    0.16093020141124725,
    // quantile(10/16)
    0.24611230194568634,
    // quantile(11/16)
    0.33791524171829224,
    // quantile(12/16)
    0.44070982933044434,
    // quantile(13/16)
    0.5626170039176941,
    // quantile(14/16)
    0.7229568362236023,
    // quantile(15/16)
    1.0,
    // quantile(16/16)
  ]
}

// ============================================================================
// TIPOS
// ============================================================================

/// Bloco NF4 quantizado (tipicamente 64 valores)
pub type NF4Block {
  NF4Block(
    /// Índices 4-bit (0-15) para cada valor
    indices: List(Int),
    /// Escala do bloco (abs_max original)
    abs_max: Float,
    /// Tamanho do bloco
    block_size: Int,
  )
}

/// Tensor completo quantizado em NF4
pub type NF4Tensor {
  NF4Tensor(
    /// Lista de blocos quantizados
    blocks: List(NF4Block),
    /// Shape original
    shape: List(Int),
    /// Número de elementos
    num_elements: Int,
    /// Memória em bytes
    memory_bytes: Int,
    /// Taxa de compressão
    compression_ratio: Float,
  )
}

/// Configuração de quantização
pub type NF4Config {
  NF4Config(
    /// Tamanho do bloco (64 é padrão QLoRA)
    block_size: Int,
    /// Usar Double Quantization (quantiza os scales também)
    double_quant: Bool,
  )
}

/// Configuração padrão QLoRA
pub fn default_config() -> NF4Config {
  NF4Config(block_size: 64, double_quant: False)
}

// ============================================================================
// QUANTIZAÇÃO NF4
// ============================================================================

/// Quantiza tensor para NF4
pub fn quantize(t: Tensor, config: NF4Config) -> NF4Tensor {
  let data = tensor.to_list(t)
  let shape = get_tensor_shape(t)
  let num_elements = list.length(data)

  // Divide em blocos
  let chunks = list.sized_chunk(data, config.block_size)

  // Quantiza cada bloco
  let blocks =
    list.map(chunks, fn(chunk) { quantize_block(chunk, config.block_size) })

  // Calcula memória:
  // - 4 bits por valor
  // - 16 bits (FP16) para abs_max por bloco
  let num_blocks = list.length(blocks)
  let data_bytes = num_elements / 2
  // 4 bits = 0.5 bytes
  let scale_bytes = num_blocks * 2
  // FP16 = 2 bytes
  let memory = data_bytes + scale_bytes

  // FP32 seria: num_elements * 4
  let original_memory = num_elements * 4
  let ratio = int.to_float(original_memory) /. int.to_float(memory)

  NF4Tensor(
    blocks: blocks,
    shape: shape,
    num_elements: num_elements,
    memory_bytes: memory,
    compression_ratio: ratio,
  )
}

/// Quantiza um bloco de valores para NF4
fn quantize_block(values: List(Float), block_size: Int) -> NF4Block {
  // Encontra abs_max para escala
  let abs_max =
    values
    |> list.map(float.absolute_value)
    |> list.fold(0.0, float.max)

  // Evita divisão por zero
  let safe_max = case abs_max >. 0.0 {
    True -> abs_max
    False -> 1.0
  }

  // Normaliza para [-1, 1]
  let normalized = list.map(values, fn(v) { v /. safe_max })

  // Mapeia cada valor para o índice NF4 mais próximo
  let indices = list.map(normalized, find_nearest_nf4_index)

  NF4Block(indices: indices, abs_max: safe_max, block_size: block_size)
}

/// Encontra o índice do nível NF4 mais próximo
fn find_nearest_nf4_index(value: Float) -> Int {
  let levels = nf4_levels()

  // Encontra o índice com menor distância
  levels
  |> list.index_map(fn(level, idx) {
    let distance = float.absolute_value(value -. level)
    #(idx, distance)
  })
  |> list.fold(#(0, 999.0), fn(best, current) {
    case current.1 <. best.1 {
      True -> current
      False -> best
    }
  })
  |> fn(result) { result.0 }
}

// ============================================================================
// DEQUANTIZAÇÃO
// ============================================================================

/// Dequantiza tensor NF4 de volta para FP32
pub fn dequantize(nf4: NF4Tensor) -> Tensor {
  let levels = nf4_levels()

  let data =
    list.flat_map(nf4.blocks, fn(block) {
      list.map(block.indices, fn(idx) {
        let level = get_at_index(levels, idx, 0.0)
        level *. block.abs_max
      })
    })

  // Trunca para tamanho original (caso último bloco seja parcial)
  let truncated = list.take(data, nf4.num_elements)

  Tensor(data: truncated, shape: nf4.shape)
}

// ============================================================================
// DOUBLE QUANTIZATION (Avançado)
// ============================================================================

/// Double Quantization: quantiza os próprios scales
/// Reduz overhead de metadados de 0.5 bits/param para 0.127 bits/param
pub type DoubleQuantNF4 {
  DoubleQuantNF4(
    /// Blocos com índices NF4
    blocks: List(NF4Block),
    /// Scales quantizados (INT8)
    quantized_scales: List(Int),
    /// Scale global para os scales
    scales_scale: Float,
    /// Shape original
    shape: List(Int),
    /// Número de elementos
    num_elements: Int,
    /// Memória em bytes
    memory_bytes: Int,
  )
}

/// Aplica Double Quantization
pub fn double_quantize(t: Tensor, config: NF4Config) -> DoubleQuantNF4 {
  // Primeiro: NF4 normal
  let nf4 = quantize(t, config)

  // Coleta todos os scales
  let scales = list.map(nf4.blocks, fn(b) { b.abs_max })

  // Quantiza os scales para INT8
  let scales_max =
    scales
    |> list.map(float.absolute_value)
    |> list.fold(0.0, float.max)

  let scales_scale = case scales_max >. 0.0 {
    True -> 127.0 /. scales_max
    False -> 1.0
  }

  let quantized_scales =
    list.map(scales, fn(s) {
      let scaled = s *. scales_scale
      float.clamp(scaled, -127.0, 127.0)
      |> float.round
    })

  // Memória: 4 bits/valor + 8 bits/bloco (scale) + 4 bytes global
  let num_blocks = list.length(nf4.blocks)
  let data_bytes = nf4.num_elements / 2
  let scale_bytes = num_blocks
  // 8 bits = 1 byte
  let memory = data_bytes + scale_bytes + 4

  DoubleQuantNF4(
    blocks: nf4.blocks,
    quantized_scales: quantized_scales,
    scales_scale: scales_scale,
    shape: nf4.shape,
    num_elements: nf4.num_elements,
    memory_bytes: memory,
  )
}

// ============================================================================
// ANÁLISE E ESTATÍSTICAS
// ============================================================================

/// Estatísticas de quantização
pub type NF4Stats {
  NF4Stats(
    original_bytes: Int,
    compressed_bytes: Int,
    compression_ratio: Float,
    mean_error: Float,
    max_error: Float,
    num_blocks: Int,
  )
}

/// Calcula estatísticas de erro
pub fn compute_stats(original: Tensor, nf4: NF4Tensor) -> NF4Stats {
  let decompressed = dequantize(nf4)

  let orig_data = tensor.to_list(original)
  let decomp_data = tensor.to_list(decompressed)

  // Calcula erros
  let errors =
    list.map2(orig_data, decomp_data, fn(o, d) { float.absolute_value(o -. d) })

  let mean_error = case errors {
    [] -> 0.0
    _ -> {
      let sum = list.fold(errors, 0.0, fn(acc, e) { acc +. e })
      sum /. int.to_float(list.length(errors))
    }
  }

  let max_error = list.fold(errors, 0.0, float.max)

  let original_bytes = list.length(orig_data) * 4

  NF4Stats(
    original_bytes: original_bytes,
    compressed_bytes: nf4.memory_bytes,
    compression_ratio: nf4.compression_ratio,
    mean_error: mean_error,
    max_error: max_error,
    num_blocks: list.length(nf4.blocks),
  )
}

// ============================================================================
// BENCHMARK
// ============================================================================

pub fn main() {
  benchmark_nf4()
}

pub fn benchmark_nf4() {
  io.println(
    "╔══════════════════════════════════════════════════════════════════╗",
  )
  io.println(
    "║  NF4 (NormalFloat4) QUANTIZATION - QLoRA Style                   ║",
  )
  io.println(
    "║  8x compressão com distribuição gaussiana otimizada!             ║",
  )
  io.println(
    "╚══════════════════════════════════════════════════════════════════╝\n",
  )

  io.println("NÍVEIS NF4 (16 quantis da distribuição normal):")
  nf4_levels()
  |> list.index_map(fn(level, idx) {
    io.println("  [" <> int.to_string(idx) <> "]: " <> float_to_string(level))
  })

  io.println("\n━━━ BENCHMARK: Tensor 1024x512 ━━━")
  let t = tensor.random_uniform([1024, 512])
  let config = default_config()

  // NF4 normal
  let #(time_nf4, nf4) = timer_tc(fn() { quantize(t, config) })

  let stats = compute_stats(t, nf4)

  io.println("\nNF4 Quantization:")
  io.println("  Tempo:       " <> int.to_string(time_nf4 / 1000) <> "ms")
  io.println(
    "  Original:    " <> int.to_string(stats.original_bytes / 1024) <> " KB",
  )
  io.println(
    "  Comprimido:  " <> int.to_string(stats.compressed_bytes / 1024) <> " KB",
  )
  io.println(
    "  Compressão:  " <> float_to_string(stats.compression_ratio) <> "x",
  )
  io.println("  Erro médio:  " <> float_to_string(stats.mean_error))
  io.println("  Erro máx:    " <> float_to_string(stats.max_error))
  io.println("  Blocos:      " <> int.to_string(stats.num_blocks))

  // Double Quantization
  io.println("\n━━━ DOUBLE QUANTIZATION (reduz overhead de metadados) ━━━")
  let #(time_dq, dq) = timer_tc(fn() { double_quantize(t, config) })

  let dq_ratio =
    int.to_float(stats.original_bytes) /. int.to_float(dq.memory_bytes)

  io.println("  Tempo:       " <> int.to_string(time_dq / 1000) <> "ms")
  io.println(
    "  Comprimido:  " <> int.to_string(dq.memory_bytes / 1024) <> " KB",
  )
  io.println("  Compressão:  " <> float_to_string(dq_ratio) <> "x")

  // Comparação com outros formatos
  io.println("\n━━━ COMPARAÇÃO DE FORMATOS ━━━")
  io.println(
    "  FP32:  " <> int.to_string(stats.original_bytes / 1024) <> " KB (1x)",
  )
  io.println(
    "  FP16:  " <> int.to_string(stats.original_bytes / 2 / 1024) <> " KB (2x)",
  )
  io.println(
    "  INT8:  " <> int.to_string(stats.original_bytes / 4 / 1024) <> " KB (4x)",
  )
  io.println(
    "  NF4:   "
    <> int.to_string(stats.compressed_bytes / 1024)
    <> " KB ("
    <> float_to_string(stats.compression_ratio)
    <> "x)",
  )
  io.println(
    "  NF4+DQ: "
    <> int.to_string(dq.memory_bytes / 1024)
    <> " KB ("
    <> float_to_string(dq_ratio)
    <> "x)",
  )

  io.println(
    "\n╔══════════════════════════════════════════════════════════════════╗",
  )
  io.println(
    "║  POR QUE NF4 É MELHOR QUE Q4 UNIFORME:                           ║",
  )
  io.println(
    "║                                                                  ║",
  )
  io.println(
    "║  1. Níveis derivados da distribuição normal                      ║",
  )
  io.println(
    "║  2. Mais precisão próximo de zero (onde concentram pesos)        ║",
  )
  io.println(
    "║  3. Matematicamente ótimo para dados gaussianos                  ║",
  )
  io.println(
    "║  4. Usado em QLoRA, bitsandbytes, Hugging Face                   ║",
  )
  io.println(
    "║                                                                  ║",
  )
  io.println(
    "║  Com NF4: 24GB VRAM → ~192GB efetivo (8x)!                       ║",
  )
  io.println(
    "╚══════════════════════════════════════════════════════════════════╝",
  )
}

// ============================================================================
// HELPERS
// ============================================================================

fn get_tensor_shape(t: Tensor) -> List(Int) {
  case t {
    Tensor(_, shape) -> shape
    tensor.StridedTensor(_, shape, _, _) -> shape
  }
}

fn get_at_index(lst: List(Float), idx: Int, default: Float) -> Float {
  case list.drop(lst, idx) {
    [first, ..] -> first
    [] -> default
  }
}

fn float_to_string(f: Float) -> String {
  let rounded = int.to_float(float.round(f *. 10_000.0)) /. 10_000.0
  float.to_string(rounded)
}

@external(erlang, "timer", "tc")
fn timer_tc(f: fn() -> a) -> #(Int, a)