redux-scraper/spec/models/trained_regression_model_spec.rb

# frozen_string_literal: true

require "rails_helper"

RSpec.describe TrainedRegressionModel, type: :model do
  describe "#predict" do
    it "predicts a value for a linear model" do
      model =
        described_class.new(
          name: "test-linear",
          model_type: "linear",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 0.0,
          x_max: 10.0,
          y_min: 0.0,
          y_max: 20.0,
          coefficients: [0.5, 1.0], # norm_intercept, norm_slope
          training_r_squared: 1.0,
          evaluation_r_squared: 1.0,
          equation_string: "y = x",
        )
      expect(model.predict(1.0)).to be_a(Float)
    end

    it "predicts a value for a quadratic model" do
      model =
        described_class.new(
          name: "test-quadratic",
          model_type: "quadratic",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 0.0,
          x_max: 10.0,
          y_min: 0.0,
          y_max: 100.0,
          coefficients: [0.1, 0.2, 0.3], # c, b, a for ax² + bx + c
          training_r_squared: 0.95,
          evaluation_r_squared: 0.90,
          equation_string: "y = 0.3x² + 0.2x + 0.1",
        )
      result = model.predict(5.0)
      expect(result).to be_a(Float)
      expect(result).to be_finite
    end

    it "predicts a value for a logarithmic model" do
      model =
        described_class.new(
          name: "test-log",
          model_type: "logarithmic",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 1.0,
          x_max: 100.0,
          y_min: 0.0,
          y_max: 50.0,
          coefficients: [0.1, 0.5], # intercept, slope
          training_r_squared: 0.85,
          evaluation_r_squared: 0.80,
          equation_string: "y = 0.5 * ln(x) + 0.1",
        )
      result = model.predict(10.0)
      expect(result).to be_a(Float)
      expect(result).to be_finite
      expect(result).to be > 0
    end

    it "predicts a value for a square root model" do
      model =
        described_class.new(
          name: "test-sqrt",
          model_type: "square_root",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 0.0,
          x_max: 100.0,
          y_min: 0.0,
          y_max: 50.0,
          coefficients: [0.2, 0.8], # intercept, slope
          training_r_squared: 0.88,
          evaluation_r_squared: 0.83,
          equation_string: "y = 0.8 * √x + 0.2",
        )
      result = model.predict(25.0)
      expect(result).to be_a(Float)
      expect(result).to be_finite
      expect(result).to be > 0
    end
  end

  describe "#equation" do
    it "constructs and caches a Stats::PolynomialEquation for linear models" do
      model =
        described_class.new(
          name: "test-linear",
          model_type: "linear",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 0.0,
          x_max: 10.0,
          y_min: 0.0,
          y_max: 20.0,
          coefficients: [0.5, 1.0],
          training_r_squared: 1.0,
          evaluation_r_squared: 1.0,
          equation_string: "y = x",
        )

      equation = model.equation
      expect(equation).to be_a(Stats::PolynomialEquation)
      expect(equation.coefficients).to eq([0.5, 1.0])

      # Verify caching
      expect(model.equation).to be(equation)
    end

    it "constructs and caches a Stats::PolynomialEquation for quadratic models" do
      model =
        described_class.new(
          name: "test-quadratic",
          model_type: "quadratic",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 0.0,
          x_max: 10.0,
          y_min: 0.0,
          y_max: 100.0,
          coefficients: [0.1, 0.2, 0.3],
          training_r_squared: 0.95,
          evaluation_r_squared: 0.90,
          equation_string: "y = 0.3x² + 0.2x + 0.1",
        )

      equation = model.equation
      expect(equation).to be_a(Stats::PolynomialEquation)
      expect(equation.coefficients).to eq([0.1, 0.2, 0.3])
    end

    it "constructs and caches a Stats::LogarithmicEquation for logarithmic models" do
      model =
        described_class.new(
          name: "test-log",
          model_type: "logarithmic",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 1.0,
          x_max: 100.0,
          y_min: 0.0,
          y_max: 50.0,
          coefficients: [0.1, 0.5],
          training_r_squared: 0.85,
          evaluation_r_squared: 0.80,
          equation_string: "y = 0.5 * ln(x) + 0.1",
        )

      equation = model.equation
      expect(equation).to be_a(Stats::LogarithmicEquation)
      expect(equation.coefficients).to eq([0.1, 0.5])
    end

    it "constructs and caches a Stats::SquareRootEquation for square root models" do
      model =
        described_class.new(
          name: "test-sqrt",
          model_type: "square_root",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 0.0,
          x_max: 100.0,
          y_min: 0.0,
          y_max: 50.0,
          coefficients: [0.2, 0.8],
          training_r_squared: 0.88,
          evaluation_r_squared: 0.83,
          equation_string: "y = 0.8 * √x + 0.2",
        )

      equation = model.equation
      expect(equation).to be_a(Stats::SquareRootEquation)
      expect(equation.coefficients).to eq([0.2, 0.8])
    end
  end

  describe "consistency with math" do
    context "linear model" do
      let(:model) do
        described_class.new(
          name: "test-linear-consistency",
          model_type: "linear",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 0.0,
          x_max: 10.0,
          y_min: 0.0,
          y_max: 20.0,
          coefficients: [0.5, 1.0], # intercept: 0.5, slope: 1.0
          training_r_squared: 1.0,
          evaluation_r_squared: 1.0,
          equation_string: "y = x + 10",
        )
      end

      it "maintains linear relationship" do
        x1, x2, x3 = 2.0, 5.0, 8.0
        y1 = model.predict(x1)
        y2 = model.predict(x2)
        y3 = model.predict(x3)

        # For linear: (y2-y1)/(x2-x1) = (y3-y2)/(x3-x2)
        slope1 = (y2 - y1) / (x2 - x1)
        slope2 = (y3 - y2) / (x3 - x2)

        expect(slope1).to be_within(0.001).of(slope2)
      end
    end

    context "quadratic model" do
      let(:model) do
        described_class.new(
          name: "test-quadratic-consistency",
          model_type: "quadratic",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 0.0,
          x_max: 10.0,
          y_min: 0.0,
          y_max: 100.0,
          coefficients: [0.1, 0.2, 0.3], # c, b, a
          training_r_squared: 0.95,
          evaluation_r_squared: 0.90,
          equation_string: "y = 0.3x² + 0.2x + 0.1",
        )
      end

      it "maintains quadratic relationship through second differences" do
        x_values = [1.0, 2.0, 3.0, 4.0, 5.0]
        y_values = x_values.map { |x| model.predict(x) }

        # Calculate first differences
        first_diffs =
          (1...y_values.length).map { |i| y_values[i] - y_values[i - 1] }

        # Calculate second differences
        second_diffs =
          (1...first_diffs.length).map do |i|
            first_diffs[i] - first_diffs[i - 1]
          end

        # For quadratic, second differences should be approximately constant
        second_diffs.each_with_index do |diff, i|
          next if i == 0
          expect(diff).to be_within(0.01).of(second_diffs[0])
        end
      end
    end
  end

  describe "error handling" do
    it "raises error for unsupported model type during initialization" do
      expect {
        described_class.new(
          name: "test-unsupported",
          model_type: "unsupported",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 0.0,
          x_max: 10.0,
          y_min: 0.0,
          y_max: 20.0,
          coefficients: [0.5, 1.0],
          training_r_squared: 1.0,
          evaluation_r_squared: 1.0,
          equation_string: "y = x",
        )
      }.to raise_error(ArgumentError, /'unsupported' is not a valid model_type/)
    end

    it "produces consistent results across multiple calls" do
      model =
        described_class.new(
          name: "test-consistency",
          model_type: "linear",
          total_records_count: 10,
          training_records_count: 8,
          evaluation_records_count: 2,
          train_test_split_ratio: 0.8,
          random_seed: 42,
          x_min: 0.0,
          x_max: 10.0,
          y_min: 0.0,
          y_max: 20.0,
          coefficients: [0.5, 1.0],
          training_r_squared: 1.0,
          evaluation_r_squared: 1.0,
          equation_string: "y = x",
        )

      x_value = 5.0
      result1 = model.predict(x_value)
      result2 = model.predict(x_value)

      expect(result1).to eq(result2)
    end
  end
end