# frozen_string_literal: true
require "set"
gem "minitest", "6.0.6"
require "minitest/autorun"

class NaiveBayes
  def initialize(alpha: 1.0)
    raise ArgumentError, "alpha must be a real number" unless alpha.is_a?(Numeric) && alpha.real?
    @alpha = alpha.to_f
    raise ArgumentError, "alpha must convert to a finite positive Float" unless @alpha.finite? && @alpha.positive?
    @documents = Hash.new(0)
    @words = {}
    @totals = Hash.new(0)
    @vocabulary = Set.new
  end

  def tokens(text)
    raise ArgumentError, "text must be a String" unless text.is_a?(String)
    text.downcase.scan(/[a-z]+/)
  end

  def train(label, text)
    raise ArgumentError, "label must be a nonblank String" unless label.is_a?(String) && !label.strip.empty?
    terms = tokens(text)
    raise ArgumentError, "training document has no tokens" if terms.empty?
    label = label.dup.freeze
    @documents[label] += 1
    counts = (@words[label] ||= Hash.new(0))
    terms.each do |term|
      counts[term] += 1
      @totals[label] += 1
      @vocabulary << term
    end
    self
  end

  def scores(text)
    raise ArgumentError, "train before predicting" if @documents.empty?
    counts = tokens(text).select { |term| @vocabulary.include?(term) }.tally
    total_documents = @documents.values.sum.to_f
    @documents.keys.sort.to_h do |label|
      score = Math.log(@documents.fetch(label) / total_documents)
      denominator = @totals.fetch(label) + @alpha * @vocabulary.length
      counts.each do |term, frequency|
        probability = (@words.fetch(label)[term] + @alpha) / denominator
        score += frequency * Math.log(probability)
      end
      [label, score]
    end
  end

  def predict(text)
    scores(text).max_by { |_label, score| score }.first
  end
end

class NaiveBayesTest < Minitest::Test
  def small_model
    NaiveBayes.new.train("ham", "meet team").train("spam", "win win")
  end
  def test_hand_calculated_log_scores
    scores = small_model.scores("win")
    assert_in_delta Math.log(0.5) + Math.log(3.0 / 5), scores.fetch("spam"), 1e-12
    assert_in_delta Math.log(0.5) + Math.log(1.0 / 5), scores.fetch("ham"), 1e-12
    assert_equal "spam", small_model.predict("win")
  end
  def test_repeated_words_count_as_repeated_evidence
    model = small_model
    once, twice = model.scores("win"), model.scores("win win")
    assert_in_delta Math.log(3.0), once["spam"] - once["ham"], 1e-12
    assert_in_delta 2 * Math.log(3.0), twice["spam"] - twice["ham"], 1e-12
  end
  def test_unknown_and_empty_inputs_use_priors
    model = small_model.train("ham", "meeting notes")
    assert_equal "ham", model.predict("unknownword")
    assert_equal "ham", model.predict("")
    assert_equal model.scores(""), model.scores("unknownword")
  end
  def test_deterministic_tie_breaking
    assert_equal "ham", small_model.predict("")
  end
  def test_invalid_training_does_not_mutate_model
    model = small_model
    before = model.scores("win")
    assert_raises(ArgumentError) { model.train("spam", "123") }
    assert_raises(ArgumentError) { model.train("", "win") }
    assert_equal before, model.scores("win")
    assert_raises(ArgumentError) { NaiveBayes.new.predict("hello") }
    [0, -1, Float::NAN, Float::INFINITY].each { |a| assert_raises(ArgumentError) { NaiveBayes.new(alpha: a) } }
  end
  def test_alpha_must_be_real_numeric
    [Complex(1, 1), Complex(1, 0), "1", nil].each do |alpha|
      assert_raises(ArgumentError) { NaiveBayes.new(alpha: alpha) }
    end
  end
  def test_alpha_is_validated_after_float_conversion
    [10**400, Rational(1, 10**400)].each do |alpha|
      assert_raises(ArgumentError) { NaiveBayes.new(alpha: alpha) }
    end
  end
  def test_real_alpha_uses_float_arithmetic
    model = NaiveBayes.new(alpha: Rational(1, 2)).train("ham", "meet team").train("spam", "win win")
    scores = model.scores("win")
    assert_in_delta Math.log(0.5) + Math.log(2.5 / 3.5), scores.fetch("spam"), 1e-12
    assert_in_delta Math.log(0.5) + Math.log(0.5 / 3.5), scores.fetch("ham"), 1e-12
  end
  def test_long_document_uses_finite_log_scores
    scores = small_model.scores(("win " * 20_000).strip)
    assert scores.values.all?(&:finite?)
    assert_equal "spam", scores.max_by { |_k, v| v }.first
  end
  def test_fixed_held_out_fixture
    model = NaiveBayes.new
    [["ham", "team meeting schedule"], ["ham", "project notes meeting"],
     ["spam", "win cash prize"], ["spam", "cash prize offer"]].each { |label, text| model.train(label, text) }
    held_out = [["ham", "team notes"], ["ham", "project schedule"],
                ["spam", "win prize"], ["spam", "cash offer"]]
    predictions = held_out.map { |label, text| [label, model.predict(text)] }
    assert_equal held_out.map(&:first), predictions.map(&:last)
  end
  def test_tokenization_contract
    assert_equal %w[caf win], small_model.tokens("Café WIN 123")
    assert_raises(ArgumentError) { small_model.tokens(nil) }
  end
end