require 'test/unit'

require 'roc'

class ROCTest < Test::Unit::TestCase

  def initialize(arg)
    super(arg)
  end

  def test_area_under_curve
    x = [1,2,3]
    y = [2,3,4]
    _test_auc(6,x,y)

    x = [1,2,3]
    y = [-2,-3,-4]
    _test_auc(-6, x, y)

    x = [1,2,3]
    y = [4,3,2]
    _test_auc(6, x, y)

    x = [1,2,3]
    y = [-4,-3,-2]
    _test_auc(-6, x, y)

    x = [4,5,6]
    y = [2,1,2]
    _test_auc(3, x, y)

    x = [4,5,6]
    y = [-2,-1,-2]
    _test_auc(-3, x, y)
  end

  def _test_auc(expected_area,x,y)
    area = ROC.new.area_under_curve(x,y)
    assert_equal(expected_area, area)
  end

  def test_prep_list
    t = true
    f = false
    x,y = ROC.new.prep_list([[0,f],[1,f],[2,f],[3,t],[3,f],[0,f],[4,f],[1,t],[2,t]])
    assert_equal([1,2,3], x)
    assert_equal([0,0,1,2,3,4], y)
  end

  def test_tps_and_ppv
    tp = %w(1 2 3 4 5 6 6 6 7 8 9 10 10 10 10 11 12 ).collect {|c| c.to_f } # 17 total
    fp = %w(3.5 4 5 5 5 6 6 6.5 7 8 9 9.5 10 15).collect {|c| c.to_f } # 14 total
    xe = [1, 2, 3, 4, 5, 8, 9, 10, 11, 15, 16, 17]
    #       1, 2, 3, 4        5,   6,             7,             8,   9
    #       10              11                12
    ye = [1, 1, 1, 4.0/6.0, 0.5, 8.0/(7.0+8.0), 9.0/(9.0+9.0), 0.5, 11.0/(11.0+ 11.0),  15.0/(15.0+13.0), 16.0/(16.0+13.0), 17.0/(17.0+13.0)]
    _test_tps_and_ppv_method(tp,fp,xe,ye,"complex real-life-like scenario")

    ## leading fp's
    tp = [1,2,3]
    fp = [0,0,1,2,3,4]
    xe = [1,2,3]
    ye = [1.0/(1+3), 2.0/(2+4), 3.0/(3+5)]
    _test_tps_and_ppv_method(tp,fp,xe,ye, "leading fps")

    ## leading tp's
    tp = [-1,2,3]
    fp = [0,4]
    xe = [1,2,3]
    ye = [1.0/(1+0), 2.0/(2+1), 3.0/(3+1)]
    _test_tps_and_ppv_method(tp,fp,xe,ye, "leading tps")

    ## equal tp's leading
    tp = [0.0001,0.0001,0.0001,2]
    fp = [0.01,4.0]
    xe = [3,4]
    ye = [3.0/(3+0), 4.0/(4+1)]
    _test_tps_and_ppv_method(tp,fp,xe,ye, "equal tps leading")

    ## equal arrays with some repeated values
    tp = %w(1 2 3 3 4 5 6 ).collect {|x| x.to_f } # 17 total
    fp = %w(1 2 3 3 4 5 6 ).collect {|x| x.to_f } # 17 total
    xe = [1,2,4,5,6,7]
    ye = [0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
    _test_tps_and_ppv_method(tp,fp,xe,ye, "equal arrays with some repeated values")

  end

  def _test_tps_and_ppv_method(tp,fp,xe,ye,message='')
    (x,y) = ROC.new.tps_and_ppv(tp,fp)
    assert_equal(x.size, y.size)
    assert_equal(xe, x)
    assert_equal(ye, y, "tps_and_ppv: #{message}")
  end

end

class DecoyROCTest < ROCTest

  def test_pred_tps_ppv__leading_fps
    ## leading fp's
    hits = [1,2,3]
    decoys = [0,0,1,2,3,4]
    num_hits_e = [1,2,3]
    num_fps = [3,4,5]
    tps_e = make_tps_e(num_fps, num_hits_e)
    ppv_e = make_ppv_e(tps_e, num_hits_e)
    _test_pred_and_tps_and_ppv(hits, decoys, num_hits_e, tps_e, ppv_e)
  end

  def test_pred_tps_ppv__leading_tps
    ## leading tp's
    hits = [-1,2,3]
    decoys = [0,4]
    num_hits_e = [1,2,3]
    num_fps = [0,1,1]
    tps_e = make_tps_e(num_fps, num_hits_e)
    ppv_e = make_ppv_e(tps_e, num_hits_e)
    _test_pred_and_tps_and_ppv(hits, decoys, num_hits_e, tps_e, ppv_e)
  end

  def test_pred_tps_ppv__equal_tps_leading
    hits = [0.0001,0.0001,0.0001,2]
    decoys = [0.01,4.0]
    num_hits_e = [3,4]
    num_fps = [0,1]
    tps_e = make_tps_e(num_fps, num_hits_e)
    ppv_e = make_ppv_e(tps_e, num_hits_e)
    _test_pred_and_tps_and_ppv(hits, decoys, num_hits_e, tps_e, ppv_e)
  end

  def test_pred_tps_ppv__equal_arrays_with_some_repeated_values
    hits = %w(1 2 3 3 4 5 6 ).collect {|x| x.to_f } # 17 total
    decoys = %w(1 2 3 3 4 5 6 ).collect {|x| x.to_f } # 17 total
    num_hits_e = [1,2,4,5,6,7]
    num_fps = [1,2,4,5,6,7]
    tps_e = make_tps_e(num_fps, num_hits_e)
    ppv_e = make_ppv_e(tps_e, num_hits_e)
    _test_pred_and_tps_and_ppv(hits, decoys, num_hits_e, tps_e, ppv_e)
  end

  def _test_pred_and_tps_and_ppv(hits, decoys, num_hits_e, tps_e, ppv_e)
    answer = DecoyROC.new.pred_and_tps_and_ppv(hits, decoys)
    expected = [num_hits_e, tps_e, ppv_e]
    %w(num_hits num_tps ppv).each_with_index do |cat, i|
      assert_equal(expected[i], answer[i], cat)
    end
  end

  def make_tps_e(num_fps, num_hits_e)
    tps_e = []
    num_hits_e.each_with_index do |v,i|
      tps_e[i] = v - num_fps[i]
    end
    tps_e
  end

  def make_ppv_e(tps_e, num_hits_e)
    ppv_e = []
    tps_e.each_with_index {|v,i| ppv_e[i] = v.to_f/num_hits_e[i] }
    ppv_e
  end

end