module RankKing
  class OpenSkill
    # Gaussian curve where mu is the mean and sigma the stddev
    Rating = Data.define(:mu, :sigma) do
      def ordinal = self.mu - 3*self.sigma
    end

    TeamRating = Data.define(:mu, :sigma_sq, :team, :rank)

    def initialize(tau: nil)
      @z = 3.0
      @mu = 25.0
      @tau = tau || @mu / 300.0
      @sigma = @mu / @z

      @epsilon = 0.0001
      @beta = @sigma / 2.0
    end

    def rate(*teams)
      # allow for passing in single-rating teams
      teams = teams.map { Array(_1) }

      # tau keeps sigma from dropping too low so that ratings stay pliable
      # after many games
      tau_sq = @tau ** 2
      teams = teams.map {|team|
        team.map {|r| r.with(sigma: Math.sqrt(r.sigma ** 2 + tau_sq)) }
      }

      rank = (0...teams.size).to_a

      ordered_teams, tenet = self.class.unwind(teams, rank)
      new_ratings = plackett_luce(ordered_teams, rank)
      reordered_teams = self.class.unwind(new_ratings, tenet)

      # TODO prevent sigma increase?

      reordered_teams
    end

    def self.unwind(src, rank)
      fail unless src.size == rank.size

      return [[], []] if src.empty?

      src.each.with_index
        .sort_by { rank.fetch(_2) }
        .transpose
    end

    def plackett_luce(game, rank)
      team_ratings = self.team_ratings(game)
      c            = util_c(team_ratings)
      sum_q        = util_sum_q(team_ratings, c)
      a            = util_a(team_ratings)

      team_ratings.map.with_index {|x, i|
        x_mu_over_ce = Math.exp(x.mu / c) # tmp1

        omega_sum, delta_sum = team_ratings.each.with_index
          .filter {|y,_| y.rank <= x.rank }
          .inject([0, 0]) {|(omega, delta), (y, i)|
            quotient = x_mu_over_ce / sum_q.fetch(i)
            [
              omega + (x == y ? 1 - quotient : -quotient) / a.fetch(i),
              delta + (quotient * (1 - quotient)) / a.fetch(i),
            ]
          }

        x_gamma = Math.sqrt(x.sigma_sq) / c
        x_omega = omega_sum * (x.sigma_sq / c)
        x_delta = x_gamma * delta_sum * (x.sigma_sq / c ** 2)

        x.team.map {|team| Rating.new(
          mu:    team.mu + (team.sigma ** 2 / x.sigma_sq) * x_omega,
          sigma: team.sigma * Math.sqrt([1 - (team.sigma ** 2 / x.sigma_sq) * x_delta, @epsilon].max),
        )}
      }
    end

    def rating(mu: nil, sigma: nil) #:: Rating
      mu ||= @mu
      sigma ||= mu / @z
      Rating.new(mu: mu.to_f, sigma: sigma.to_f)
    end

    def team_ratings(game)
      game.map.with_index {|team, i|
        TeamRating.new(
          mu:       team.sum(&:mu),
          sigma_sq: team.sum { _1.sigma ** 2 },
          team:,
          rank:     i,
        )
      }
    end

    def util_a(team_ratings)
      team_ratings.map {|q|
        team_ratings.count {|i| i.rank == q.rank }
      }
    end

    def util_c(team_ratings)
      Math.sqrt(team_ratings.sum { _1.sigma_sq + @beta ** 2 })
    end

    def util_sum_q(team_ratings, c)
      team_ratings.map {|q|
        team_ratings
          .select {|i| i.rank >= q.rank }
          .sum {|i| Math.exp(i.mu / c) }
      }
    end
  end
end