# typed: strict
# frozen_string_literal: true

module Stats
  # Immutable struct representing a single regression analysis result
  class RegressionResult < T::ImmutableStruct
    extend T::Sig

    const :equation, Stats::Equation
    const :training_r_squared, Float
    const :evaluation_r_squared, Float
    const :x_values, T::Array[Float]
    const :y_values, T::Array[Float]

    sig { returns(String) }
    def equation_string
      equation.to_s
    end

    sig { returns(String) }
    def score_summary
      "Training R² = #{Stats::Helpers.format_r_squared(training_r_squared)}, Evaluation R² = #{Stats::Helpers.format_r_squared(evaluation_r_squared)}"
    end
  end

  # Immutable struct representing the complete analysis results
  class AnalysisResults < T::ImmutableStruct
    extend T::Sig

    const :linear, RegressionResult
    const :quadratic, RegressionResult
    const :logarithmic, RegressionResult
    const :square_root, RegressionResult
  end

  # Handles regression analysis using Rumale with normalized data
  class RegressionAnalyzer
    extend T::Sig

    sig { params(records: T::Array[Stats::DataPoint]).void }
    def initialize(records)
      @records = records
    end

    # Generic regression analysis method to eliminate duplication
    sig do
      params(
        normalizer_class: T.class_of(Stats::DataNormalizer),
        equation_class: T.class_of(Stats::Equation),
        split: Stats::TrainTestSplit,
        degree: Integer,
      ).returns(Stats::RegressionResult)
    end
    def analyze_regression(normalizer_class, equation_class, split, degree: 1)
      # Create normalizers for training and evaluation data
      training_normalizer = normalizer_class.new(split.training_records)
      evaluation_normalizer = normalizer_class.new(split.evaluation_records)

      regression_x = training_normalizer.regression_x_range
      poly_features = Rumale::Preprocessing::PolynomialFeatures.new(degree:)
      regressor = Rumale::LinearModel::LinearRegression.new(fit_bias: true)
      pipeline =
        Rumale::Pipeline::Pipeline.new(
          steps: {
            transformer: poly_features,
            estimator: regressor,
          },
        )

      # Fit the pipeline on training data
      training_x_matrix = training_normalizer.normalized_x_matrix
      training_y_vector = training_normalizer.normalized_y_vector
      pipeline.fit(training_x_matrix, training_y_vector)

      # Score on training data
      training_r_squared = pipeline.score(training_x_matrix, training_y_vector)

      # Score on evaluation data
      evaluation_x_matrix = evaluation_normalizer.normalized_x_matrix
      evaluation_y_vector = evaluation_normalizer.normalized_y_vector
      evaluation_r_squared =
        pipeline.score(evaluation_x_matrix, evaluation_y_vector)

      weight_vec = pipeline.steps[:estimator].weight_vec.to_a

      # Generate regression line data in original scale
      regression_y =
        generate_regression_line(training_normalizer, regression_x, weight_vec)

      # Create equation object
      equation =
        create_equation(equation_class, training_normalizer, weight_vec)

      Stats::RegressionResult.new(
        equation:,
        training_r_squared:,
        evaluation_r_squared:,
        x_values: regression_x,
        y_values: regression_y,
      )
    end

    sig do
      returns(
        T::Array[[TrainedRegressionModel::ModelType, Stats::RegressionResult]],
      )
    end
    def analyze
      # Split data into training and evaluation sets
      split = Stats::Helpers.split_train_test(@records)

      [
        [
          TrainedRegressionModel::ModelType::Linear,
          analyze_regression(
            Stats::LinearNormalizer,
            Stats::PolynomialEquation,
            split,
            degree: 1,
          ),
        ],
        [
          TrainedRegressionModel::ModelType::Quadratic,
          analyze_regression(
            Stats::QuadraticNormalizer,
            Stats::PolynomialEquation,
            split,
            degree: 2,
          ),
        ],
        [
          TrainedRegressionModel::ModelType::Logarithmic,
          analyze_regression(
            Stats::LogarithmicNormalizer,
            Stats::LogarithmicEquation,
            split,
          ),
        ],
        [
          TrainedRegressionModel::ModelType::SquareRoot,
          analyze_regression(
            Stats::SquareRootNormalizer,
            Stats::SquareRootEquation,
            split,
          ),
        ],
      ]
    end

    # Generate regression line using appropriate denormalization method
    sig do
      params(
        normalizer: Stats::DataNormalizer,
        regression_x: T::Array[Float],
        weight_vec: T::Array[Float],
      ).returns(T::Array[Float])
    end
    def generate_regression_line(normalizer, regression_x, weight_vec)
      normalizer.denormalize_regression(regression_x, weight_vec)
    end

    # Create appropriate equation object based on type
    sig do
      params(
        equation_class: T.class_of(Stats::Equation),
        normalizer: Stats::DataNormalizer,
        weight_vec: T::Array[Float],
      ).returns(Stats::Equation)
    end
    def create_equation(equation_class, normalizer, weight_vec)
      if equation_class == Stats::PolynomialEquation
        coefficients =
          case normalizer
          when Stats::LinearNormalizer
            normalizer.denormalize_coefficients(
              T.cast(weight_vec, [Float, Float]),
            )
          when Stats::QuadraticNormalizer
            normalizer.denormalize_coefficients(
              T.cast(weight_vec, [Float, Float, Float]),
            )
          else
            raise "Unsupported normalizer for PolynomialEquation: #{normalizer.class}"
          end
        Stats::PolynomialEquation.new(normalizer.x, normalizer.y, coefficients)
      elsif equation_class == Stats::LogarithmicEquation ||
            equation_class == Stats::SquareRootEquation
        equation_class.new(
          normalizer.x,
          normalizer.y,
          T.cast(weight_vec, [Float, Float]),
        )
      else
        raise "Unsupported equation class: #{equation_class}"
      end
    end
  end
end
