# rbs_inline: enabled

# Author::    Lucas Carlson  (mailto:lucas@rufy.com)
# Copyright:: Copyright (c) 2024 Lucas Carlson
# License::   LGPL

require 'json'
require 'mutex_m'

module Classifier
  # Logistic Regression (MaxEnt) classifier using Stochastic Gradient Descent.
  # Often provides better accuracy than Naive Bayes while remaining fast and interpretable.
  #
  # Example:
  #   classifier = Classifier::LogisticRegression.new(:spam, :ham)
  #   classifier.train(spam: ["Buy now!", "Free money!!!"])
  #   classifier.train(ham: ["Meeting tomorrow", "Project update"])
  #   classifier.classify("Claim your prize!") # => "Spam"
  #   classifier.probabilities("Claim your prize!") # => {"Spam" => 0.92, "Ham" => 0.08}
  #
  class LogisticRegression # rubocop:disable Metrics/ClassLength
    include Mutex_m
    include Streaming

    # @rbs @categories: Array[Symbol]
    # @rbs @weights: Hash[Symbol, Hash[Symbol, Float]]
    # @rbs @bias: Hash[Symbol, Float]
    # @rbs @vocabulary: Hash[Symbol, bool]
    # @rbs @training_data: Array[{category: Symbol, features: Hash[Symbol, Integer]}]
    # @rbs @learning_rate: Float
    # @rbs @regularization: Float
    # @rbs @max_iterations: Integer
    # @rbs @tolerance: Float
    # @rbs @fitted: bool
    # @rbs @dirty: bool
    # @rbs @storage: Storage::Base?
    # @rbs @min_word_length: Integer

    attr_accessor :storage

    DEFAULT_LEARNING_RATE = 0.1
    DEFAULT_REGULARIZATION = 0.01
    DEFAULT_MAX_ITERATIONS = 100
    DEFAULT_TOLERANCE = 1e-4

    # Creates a new Logistic Regression classifier with the specified categories.
    #
    #   classifier = Classifier::LogisticRegression.new(:spam, :ham)
    #   classifier = Classifier::LogisticRegression.new('Positive', 'Negative', 'Neutral')
    #   classifier = Classifier::LogisticRegression.new(['Positive', 'Negative', 'Neutral'])
    #
    # Options:
    # - learning_rate: Step size for gradient descent (default: 0.1)
    # - regularization: L2 regularization strength (default: 0.01)
    # - max_iterations: Maximum training iterations (default: 100)
    # - tolerance: Convergence threshold (default: 1e-4)
    # - min_word_length: Minimum word length filter in tokenization
    #
    # @rbs (*String | Symbol | Array[String | Symbol], ?learning_rate: Float, ?regularization: Float,
    #       ?max_iterations: Integer, ?tolerance: Float, ?min_word_length: Integer) -> void
    # rubocop:disable Metrics/ParameterLists
    def initialize(*categories, learning_rate: DEFAULT_LEARNING_RATE,
                   regularization: DEFAULT_REGULARIZATION,
                   max_iterations: DEFAULT_MAX_ITERATIONS,
                   tolerance: DEFAULT_TOLERANCE,
                   min_word_length: Classifier.config.min_word_length)
      super()
      categories = categories.flatten
      @categories = categories.map { |c| c.to_s.prepare_category_name }
      @weights = @categories.to_h { |c| [c, {}] }
      @bias = @categories.to_h { |c| [c, 0.0] }
      @vocabulary = {}
      @training_data = []
      @learning_rate = learning_rate
      @regularization = regularization
      @max_iterations = max_iterations
      @tolerance = tolerance
      @fitted = false
      @dirty = false
      @storage = nil
      @min_word_length = min_word_length
    end
    # rubocop:enable Metrics/ParameterLists

    # Trains the classifier with text for a category.
    #
    #   classifier.train(spam: "Buy now!", ham: ["Hello", "Meeting tomorrow"])
    #   classifier.train(:spam, "legacy positional API")
    #
    # @rbs (?(String | Symbol)?, ?String?, **(String | Array[String])) -> void
    def train(category = nil, text = nil, **categories)
      return train_single(category, text) if category && text

      categories.each do |cat, texts|
        (texts.is_a?(Array) ? texts : [texts]).each { |t| train_single(cat, t) }
      end
    end

    # Fits the model to all accumulated training data.
    # Called automatically during classify/probabilities if not already fitted.
    #
    # @rbs () -> self
    def fit
      synchronize do
        return self if @training_data.empty?
        raise ArgumentError, 'At least two categories required for fitting' if @categories.size < 2

        optimize_weights
        @fitted = true
        @dirty = false
      end
      self
    end

    # Returns the best matching category for the provided text.
    #
    #   classifier.classify("Buy now!") # => "Spam"
    #
    # @rbs (String) -> String
    def classify(text)
      probs = probabilities(text)
      best = probs.max_by { |_, v| v }
      raise StandardError, 'No classifications available' unless best

      best.first
    end

    # Returns probability distribution across all categories.
    # Probabilities are well-calibrated (unlike Naive Bayes).
    # Raises NotFittedError if model has not been fitted.
    #
    #   classifier.probabilities("Buy now!")
    #   # => {"Spam" => 0.92, "Ham" => 0.08}
    #
    # @rbs (String) -> Hash[String, Float]
    def probabilities(text)
      raise NotFittedError, 'Model not fitted. Call fit() after training.' unless @fitted

      features = text.word_hash(@min_word_length)
      synchronize do
        softmax(compute_scores(features))
      end
    end

    # Returns log-odds scores for each category (before softmax).
    # Raises NotFittedError if model has not been fitted.
    #
    # @rbs (String) -> Hash[String, Float]
    def classifications(text)
      raise NotFittedError, 'Model not fitted. Call fit() after training.' unless @fitted

      features = text.word_hash(@min_word_length)
      synchronize do
        compute_scores(features).transform_keys(&:to_s)
      end
    end

    # Returns feature weights for a category, sorted by importance.
    # Positive weights indicate the feature supports the category.
    #
    #   classifier.weights(:spam)
    #   # => {:free => 2.3, :buy => 1.8, :money => 1.5, ...}
    #
    # @rbs (String | Symbol, ?limit: Integer?) -> Hash[Symbol, Float]
    def weights(category, limit: nil)
      fit unless @fitted

      cat = category.to_s.prepare_category_name
      raise StandardError, "No such category: #{cat}" unless @weights.key?(cat)

      sorted = @weights[cat].sort_by { |_, v| -v.abs }
      sorted = sorted.first(limit) if limit
      sorted.to_h
    end

    # Returns the list of categories.
    #
    # @rbs () -> Array[String]
    def categories
      synchronize { @categories.map(&:to_s) }
    end

    # Adds a new category to the classifier.
    # Allows dynamic category creation for CLI and incremental training.
    #
    # @rbs (String | Symbol) -> void
    def add_category(category)
      cat = category.to_s.prepare_category_name
      synchronize do
        return if @categories.include?(cat)

        @categories << cat
        @weights[cat] = {}
        @bias[cat] = 0.0
        @fitted = false
        @dirty = true
      end
    end

    # Returns true if the model has been fitted.
    #
    # @rbs () -> bool
    def fitted?
      @fitted
    end

    # Returns true if there are unsaved changes.
    #
    # @rbs () -> bool
    def dirty?
      @dirty
    end

    # Provides training methods for the categories.
    #   classifier.train_spam "Buy now!"
    def method_missing(name, *args)
      category_match = name.to_s.match(/train_(\w+)/)
      return super unless category_match

      category = category_match[1].to_s.prepare_category_name
      raise StandardError, "No such category: #{category}" unless @categories.include?(category)

      args.each { |text| train(category, text) }
    end

    # @rbs (Symbol, ?bool) -> bool
    def respond_to_missing?(name, include_private = false)
      !!(name.to_s =~ /train_(\w+)/) || super
    end

    # Returns a hash representation of the classifier state.
    # Does NOT auto-fit; saves current state including unfitted models.
    #
    # @rbs (?untyped) -> Hash[Symbol, untyped]
    def as_json(_options = nil)
      {
        version: 1,
        type: 'logistic_regression',
        categories: @categories.map(&:to_s),
        weights: @weights.transform_keys(&:to_s).transform_values { |v| v.transform_keys(&:to_s) },
        bias: @bias.transform_keys(&:to_s),
        vocabulary: @vocabulary.keys.map(&:to_s),
        training_data: @training_data.map { |d| { category: d[:category].to_s, features: d[:features].transform_keys(&:to_s) } },
        learning_rate: @learning_rate,
        regularization: @regularization,
        max_iterations: @max_iterations,
        tolerance: @tolerance,
        fitted: @fitted,
        min_word_length: @min_word_length
      }
    end

    # Serializes the classifier state to a JSON string.
    #
    # @rbs (?untyped) -> String
    def to_json(_options = nil)
      JSON.generate(as_json)
    end

    # Loads a classifier from a JSON string or Hash.
    #
    # @rbs (String | Hash[String, untyped]) -> LogisticRegression
    def self.from_json(json)
      data = json.is_a?(String) ? JSON.parse(json) : json
      raise ArgumentError, "Invalid classifier type: #{data['type']}" unless data['type'] == 'logistic_regression'

      categories = data['categories'].map(&:to_sym)
      instance = allocate
      instance.send(:restore_state, data, categories)
      instance
    end

    # Saves the classifier to the configured storage.
    #
    # @rbs () -> void
    def save
      raise ArgumentError, 'No storage configured' unless storage

      storage.write(to_json)
      @dirty = false
    end

    # Saves the classifier state to a file.
    #
    # @rbs (String) -> Integer
    def save_to_file(path)
      result = File.write(path, to_json)
      @dirty = false
      result
    end

    # Loads a classifier from the configured storage.
    #
    # @rbs (storage: Storage::Base) -> LogisticRegression
    def self.load(storage:)
      data = storage.read
      raise StorageError, 'No saved state found' unless data

      instance = from_json(data)
      instance.storage = storage
      instance
    end

    # Loads a classifier from a file.
    #
    # @rbs (String) -> LogisticRegression
    def self.load_from_file(path)
      from_json(File.read(path))
    end

    # Reloads the classifier from storage, raising if there are unsaved changes.
    #
    # @rbs () -> self
    def reload
      raise ArgumentError, 'No storage configured' unless storage
      raise UnsavedChangesError, 'Unsaved changes would be lost. Call save first or use reload!' if @dirty

      data = storage.read
      raise StorageError, 'No saved state found' unless data

      restore_from_json(data)
      @dirty = false
      self
    end

    # Force reloads the classifier from storage, discarding any unsaved changes.
    #
    # @rbs () -> self
    def reload!
      raise ArgumentError, 'No storage configured' unless storage

      data = storage.read
      raise StorageError, 'No saved state found' unless data

      restore_from_json(data)
      @dirty = false
      self
    end

    # Custom marshal serialization to exclude mutex state.
    #
    # @rbs () -> Array[untyped]
    def marshal_dump
      fit unless @fitted
      [@categories, @weights, @bias, @vocabulary, @learning_rate, @regularization,
       @max_iterations, @tolerance, @fitted, @min_word_length]
    end

    # Custom marshal deserialization to recreate mutex.
    #
    # @rbs (Array[untyped]) -> void
    def marshal_load(data)
      mu_initialize
      @categories, @weights, @bias, @vocabulary, @learning_rate, @regularization,
        @max_iterations, @tolerance, @fitted, @min_word_length = data
      @training_data = []
      @dirty = false
      @storage = nil
    end

    # Loads a classifier from a checkpoint.
    #
    # @rbs (storage: Storage::Base, checkpoint_id: String) -> LogisticRegression
    def self.load_checkpoint(storage:, checkpoint_id:)
      raise ArgumentError, 'Storage must be File storage for checkpoints' unless storage.is_a?(Storage::File)

      dir = File.dirname(storage.path)
      base = File.basename(storage.path, '.*')
      ext = File.extname(storage.path)
      checkpoint_path = File.join(dir, "#{base}_checkpoint_#{checkpoint_id}#{ext}")

      checkpoint_storage = Storage::File.new(path: checkpoint_path)
      instance = load(storage: checkpoint_storage)
      instance.storage = storage
      instance
    end

    # Trains the classifier from an IO stream.
    # Each line in the stream is treated as a separate document.
    # Note: The model is NOT automatically fitted after streaming.
    # Call #fit to train the model after adding all data.
    #
    # @example Train from a file
    #   classifier.train_from_stream(:spam, File.open('spam_corpus.txt'))
    #   classifier.fit  # Required to train the model
    #
    # @example With progress tracking
    #   classifier.train_from_stream(:spam, io, batch_size: 500) do |progress|
    #     puts "#{progress.completed} documents processed"
    #   end
    #   classifier.fit
    #
    # @rbs (?(String | Symbol | nil), ?IO?, ?batch_size: Integer, **IO) { (Streaming::Progress) -> void } -> void
    def train_from_stream(category = nil, io = nil, batch_size: Streaming::DEFAULT_BATCH_SIZE, **categories, &)
      raise ArgumentError, 'Provide either (category, io) or keyword category: io pairs' if category.nil? && io.nil? && categories.empty?
      raise ArgumentError, 'Provide both category and io, or use keyword arguments' if [category, io].one?(&:nil?)

      pairs = category && io ? { category => io } : categories
      pairs.each do |cat, stream|
        stream_train_category(cat, stream, batch_size:, &)
      end
    end

    # Trains the classifier with an array of documents in batches.
    # Note: The model is NOT automatically fitted after batch training.
    # Call #fit to train the model after adding all data.
    #
    # @example Positional style
    #   classifier.train_batch(:spam, documents, batch_size: 100)
    #   classifier.fit
    #
    # @example Keyword style
    #   classifier.train_batch(spam: documents, ham: other_docs)
    #   classifier.fit
    #
    # @rbs (?(String | Symbol)?, ?Array[String]?, ?batch_size: Integer, **Array[String]) { (Streaming::Progress) -> void } -> void
    def train_batch(category = nil, documents = nil, batch_size: Streaming::DEFAULT_BATCH_SIZE, **categories, &block)
      if category && documents
        train_batch_for_category(category, documents, batch_size: batch_size, &block)
      else
        categories.each do |cat, docs|
          train_batch_for_category(cat, Array(docs), batch_size: batch_size, &block)
        end
      end
    end

    private

    # Trains from an IO stream with a single category.
    # @rbs (String | Symbol, IO, batch_size: Integer) { (Streaming::Progress) -> void } -> void
    def stream_train_category(category, io, batch_size:)
      category = category.to_s.prepare_category_name
      raise ArgumentError, "No such category: #{category}" unless @categories.include?(category)
      raise ArgumentError, 'Stream must respond to #each_line' unless io.respond_to?(:each_line)

      reader = Streaming::LineReader.new(io, batch_size: batch_size)
      total = reader.estimate_line_count
      progress = Streaming::Progress.new(total: total)

      reader.each_batch do |batch|
        synchronize do
          batch.each do |text|
            features = text.word_hash(@min_word_length)
            features.each_key { |word| @vocabulary[word] = true }
            @training_data << { category: category, features: features }
          end
          @fitted = false
          @dirty = true
        end
        progress.completed += batch.size
        progress.current_batch += 1
        yield progress if block_given?
      end
    end

    # Trains a batch of documents for a single category.
    # @rbs (String | Symbol, Array[String], ?batch_size: Integer) { (Streaming::Progress) -> void } -> void
    def train_batch_for_category(category, documents, batch_size: Streaming::DEFAULT_BATCH_SIZE)
      category = category.to_s.prepare_category_name
      raise StandardError, "No such category: #{category}" unless @categories.include?(category)

      progress = Streaming::Progress.new(total: documents.size)

      documents.each_slice(batch_size) do |batch|
        synchronize do
          batch.each do |text|
            features = text.word_hash(@min_word_length)
            features.each_key { |word| @vocabulary[word] = true }
            @training_data << { category: category, features: features }
          end
          @fitted = false
          @dirty = true
        end
        progress.completed += batch.size
        progress.current_batch += 1
        yield progress if block_given?
      end
    end

    # Core training logic for a single category and text.
    # @rbs (String | Symbol, String) -> void
    def train_single(category, text)
      category = category.to_s.prepare_category_name
      raise StandardError, "No such category: #{category}" unless @categories.include?(category)

      features = text.word_hash(@min_word_length)
      synchronize do
        features.each_key { |word| @vocabulary[word] = true }
        @training_data << { category: category, features: features }
        @fitted = false
        @dirty = true
      end
    end

    # Optimizes weights using mini-batch SGD with L2 regularization.
    # @rbs () -> void
    def optimize_weights
      return if @training_data.empty?

      initialize_weights
      prev_loss = Float::INFINITY

      @max_iterations.times do
        total_loss = run_training_epoch
        break if (prev_loss - total_loss).abs < @tolerance

        prev_loss = total_loss
      end

      @training_data = []
    end

    # @rbs () -> void
    def initialize_weights
      @vocabulary.each_key do |word|
        @categories.each { |cat| @weights[cat][word] ||= 0.0 }
      end
    end

    # @rbs () -> Float
    def run_training_epoch
      total_loss = 0.0

      @training_data.shuffle.each do |sample|
        probs = softmax(compute_scores(sample[:features]))
        update_weights(sample[:features], sample[:category], probs)
        total_loss -= Math.log([probs[sample[:category].to_s], 1e-15].max)
      end

      total_loss + l2_penalty
    end

    # @rbs (Hash[Symbol, Integer], Symbol, Hash[String, Float]) -> void
    def update_weights(features, true_category, probs)
      @categories.each do |cat|
        error = probs[cat.to_s] - (cat == true_category ? 1.0 : 0.0)
        @bias[cat] -= @learning_rate * error

        features.each do |word, count|
          gradient = (error * count) + (@regularization * (@weights[cat][word] || 0.0))
          @weights[cat][word] ||= 0.0
          @weights[cat][word] -= @learning_rate * gradient
        end
      end
    end

    # @rbs () -> Float
    def l2_penalty
      penalty = 0.0
      @weights.each_value do |cat_weights|
        cat_weights.each_value { |w| penalty += 0.5 * @regularization * w * w }
      end
      penalty
    end

    # Computes raw scores (logits) for each category.
    # @rbs (Hash[Symbol, Integer]) -> Hash[Symbol, Float]
    def compute_scores(features)
      @categories.to_h do |cat|
        score = @bias[cat]
        features.each { |word, count| score += (@weights[cat][word] || 0.0) * count }
        [cat, score]
      end
    end

    # Applies softmax to convert scores to probabilities.
    # @rbs (Hash[Symbol, Float]) -> Hash[String, Float]
    def softmax(scores)
      max_score = scores.values.max || 0.0
      exp_scores = scores.transform_values { |s| Math.exp(s - max_score) }
      sum = exp_scores.values.sum.to_f
      exp_scores.transform_keys(&:to_s).transform_values { |e| (e / sum).to_f }
    end

    # Restores classifier state from JSON string.
    # @rbs (String) -> void
    def restore_from_json(json)
      data = JSON.parse(json)
      categories = data['categories'].map(&:to_sym)
      restore_state(data, categories)
    end

    # Restores classifier state from parsed JSON data.
    # @rbs (Hash[String, untyped], Array[Symbol]) -> void
    def restore_state(data, categories)
      mu_initialize
      @categories = categories
      restore_weights_and_bias(data)
      restore_hyperparameters(data)
      @fitted = data.fetch('fitted', true)
      @dirty = false
      @storage = nil
      @min_word_length = data['min_word_length'] || Classifier.config.min_word_length
    end

    def restore_weights_and_bias(data)
      @weights = {}
      @bias = {}
      data['weights'].each { |cat, words| @weights[cat.to_sym] = words.transform_keys(&:to_sym).transform_values(&:to_f) }
      data['bias'].each { |cat, value| @bias[cat.to_sym] = value.to_f }
      @vocabulary = data['vocabulary'].to_h { |v| [v.to_sym, true] }
      @training_data = (data['training_data'] || []).map do |d|
        { category: d['category'].to_sym, features: d['features'].transform_keys(&:to_sym).transform_values(&:to_i) }
      end
    end

    def restore_hyperparameters(data)
      @learning_rate = data['learning_rate']
      @regularization = data['regularization']
      @max_iterations = data['max_iterations']
      @tolerance = data['tolerance']
    end
  end
end
