# typed: strict
# frozen_string_literal: true

module Stats
  class QuadraticNormalizer < DataNormalizer
    # Denormalize quadratic regression results back to original scale
    sig do
      override
        .params(regression_x: T::Array[Float], weight_vec: T::Array[Float])
        .returns(T::Array[Float])
    end
    def denormalize_regression(regression_x, weight_vec)
      norm_c = T.must(weight_vec[0])
      norm_b = T.must(weight_vec[1])
      norm_a = T.must(weight_vec[2])
      regression_x.map do |x|
        x_norm = @x.normalize(x)
        y_norm = norm_a * x_norm * x_norm + norm_b * x_norm + norm_c
        @y.denormalize(y_norm)
      end
    end

    # Denormalize quadratic regression coefficients back to original scale
    sig do
      params(coefficients: [Float, Float, Float]).returns([Float, Float, Float])
    end
    def denormalize_coefficients(coefficients)
      norm_c = coefficients[0]
      norm_b = coefficients[1]
      norm_a = coefficients[2]
      a_orig = norm_a * @y.scale / (@x.scale * @x.scale)
      b_orig = norm_b * @y.scale / @x.scale - 2 * a_orig * @x.min
      c_orig =
        (norm_c * @y.scale + @y.min) - b_orig * @x.min -
          a_orig * @x.min * @x.min

      [c_orig, b_orig, a_orig]
    end

    protected

    sig { override.returns(T.class_of(Stats::Axis)) }
    def axis_class
      Stats::LinearAxis
    end
  end
end
