# frozen_string_literal: true
require 'rubygems'
gem 'rumale', '= 0.28.1'
require 'rumale/preprocessing/standard_scaler'
require 'rumale/linear_model/logistic_regression'
require 'csv'
require 'json'

module ExamExperiment
  FEATURES=['attendance_pct','practice_hours'].freeze
  def self.features(row)
    values=FEATURES.map { |key| Float(row.fetch(key)) }
    unless values.all?(&:finite?) && values[0].between?(0,100) && values[1].between?(0,40)
      raise ArgumentError,'features outside declared fixture domain'
    end
    values
  rescue KeyError,TypeError
    raise ArgumentError,'missing or invalid feature'
  end
  def self.load_rows(path)
    rows=CSV.read(path,headers:true).map(&:to_h)
    raise ArgumentError,'empty dataset' if rows.empty?
    rows.each do |r|
      features(r)
      raise ArgumentError,'binary label required' unless ['0','1'].include?(r['passed'])
      raise ArgumentError,'known cohort required' unless ['train','holdout'].include?(r['cohort'])
    end
    rows
  end
  def self.fit(rows)
    raise ArgumentError,'fit accepts training cohort only' unless rows.all? { |r| r['cohort']=='train' }
    x=Numo::DFloat.asarray(rows.map { |r| features(r) })
    y=Numo::Int32.asarray(rows.map { |r| Integer(r.fetch('passed')) })
    raise ArgumentError,'both classes required' unless y.to_a.uniq.sort == [0,1]
    scaler=Rumale::Preprocessing::StandardScaler.new
    scaler.fit(x)
    raise ArgumentError,'constant or invalid training feature' unless scaler.std_vec.to_a.all? { |v| v.finite? && v>0 }
    classifier=Rumale::LinearModel::LogisticRegression.new(reg_param: 0.1,max_iter: 1000,tol: 1e-8)
    classifier.fit(scaler.transform(x),y)
    [scaler,classifier]
  end
  def self.probabilities(fitted,rows)
    scaler,classifier=fitted
    x=Numo::DFloat.asarray(rows.map { |r| features(r) })
    column=classifier.classes.to_a.index(1)
    classifier.predict_proba(scaler.transform(x))[true,column].to_a
  end
  def self.confusion(labels,probabilities,threshold=0.5)
    raise ArgumentError,'invalid threshold' unless threshold.finite? && threshold.between?(0,1)
    raise ArgumentError,'unequal or empty inputs' if labels.empty? || labels.length != probabilities.length
    counts={tp:0,fp:0,tn:0,fn:0}
    labels.zip(probabilities).each do |y,p|
      raise ArgumentError,'invalid label or probability' unless [0,1].include?(y) && p.finite? && p.between?(0,1)
      key=p>=threshold ? (y==1 ? :tp : :fp) : (y==0 ? :tn : :fn)
      counts[key]+=1
    end
    counts
  end
end
if $PROGRAM_NAME==__FILE__
  rows=ExamExperiment.load_rows(File.join(__dir__,'students.csv'))
  training,held=rows.partition { |r| r['cohort']=='train' }
  fitted=ExamExperiment.fit(training); probs=ExamExperiment.probabilities(fitted,held)
  labels=held.map { |r| Integer(r['passed']) }
  puts JSON.generate(training_rows:training.length,holdout_rows:held.length,
                     train_mean:fitted[0].mean_vec.to_a, probabilities:probs.map { |p| p.round(4) },
                     at_05:ExamExperiment.confusion(labels,probs,0.5),
                     at_08:ExamExperiment.confusion(labels,probs,0.8),
                     always_pass:ExamExperiment.confusion(labels,Array.new(labels.length,1.0)))
end