(*
 * A bimatrix game is represented simply as a two-dimensional array of 
 * float * float pairs.  There are no names for the player's strategies -
 * they are simply the integer index into the array.
 *
 * Two examples of games are shown, a prisoners' dilemma, and the homework game.
 *)
type payoffs = float * float;;

type bimatrix_game = payoffs array array;;

let prisoners_dilemma = [| [| (-1.0, -1.0) ; (-10.0, 0.0) |] ; 
			   [| (0.0, -10.0) ; (-8.0, -8.0) |] |] ;;

let hw_game = [| [| (8.0, 13.0) ; (3.0, 3.0) ; (3.0, 3.0) ; (0.0, 5.0) |] ;
		 [| (3.0, 3.0) ; (13.0, 8.0) ; (0.0, 3.0) ; (5.0, 5.0) |] ;
		 [| (3.0, 3.0) ; (3.0, 0.0) ; (4.0, 4.0) ; (12.0, 2.0) |] ;
	         [| (5.0, 0.0) ; (5.0, 5.0) ; (2.0, 12.0) ; (10.0, 10.0) |] |] ;;

(*
 * A mixed strategy is a probability distribution over pure strategies, and is
 * represented simply as an array of floats.  If the i-th element of the array is
 * p, that means that the probability of playing the pure strategy with index i is
 * p.  No checks are made to ensure that the array represents a valid mixed 
 * strategy, i.e. that all entries are non-negative and they sum to 1.  You will 
 * have to make sure this holds yourself.
 *
 * A strategy profile consists of a mixed strategy for each player.
 *)

type mixed_strategy = float array;;

type strategy_profile = mixed_strategy * mixed_strategy;;

(*
 * Given a game G and a strategy profile (sigma_1, sigma_2), 
 * we can compute the expected payoffs by summing, 
 * over all possible pure strategy profiles (c_1, c_2),
 * sigma_1.(c_1) * sigma_1.(c_2) * G.(c_1, c_2)
 * which is the probability that the players will play (c_1, c_2) 
 * times their utilities under (c_1, c_2).  
 *)
 
let sum2 = function (x1,y1) -> function (x2,y2) -> (x1 +. y1, x2 +. y2) 

let expected_value : bimatrix_game -> strategy_profile -> payoffs = 
  fun game -> function (strategy1, strategy2) ->
    let scaled_game = 
      Array.mapi 
	(fun i row -> 
          (Array.mapi 
	     (fun j (e1, e2) -> 
	       let mul = strategy1.(i) *. strategy2.(j) in
	       (e1 *. mul, e2 *. mul))
	     row))
	game in
    Array.fold_right 
      (fun row tot -> Array.fold_right sum2 row tot) 
      game 
      (0.0, 0.0) ;;

(*
 * play_one is a function to play one realization of a bimatrix game, 
 * given a mixed strategy profile.  A pure strategy is chosen for each player,
 * according to their mixed strategy.  The outcome of the game is the pure
 * strategy profile actually played, and the resulting payoffs.
 *)

type outcome = (int * int) * payoffs ;;

let play_one : bimatrix_game -> strategy_profile -> outcome = 
  fun game -> function (strategy1, strategy2) ->
    let choose p qs = 
      let rec choose' n = function
	| q :: qs -> if p <= q then n else choose' (n+1) qs 
	| [] -> failwith "Probabilities in profile sum to less than 1" in
      choose' 0 qs in
    let action1 = choose (Random.float 1.0) (Array.to_list strategy1)
    and action2 = choose (Random.float 1.0) (Array.to_list strategy2) in
    ((action1, action2), game.(action1).(action2)) ;;

(*
 * When playing repeated games, players have a history of past games on
 * which to base future decisions.  The history is just a list of outcomes,
 * with the MOST RECENT FIRST.  
 * An empty list means there have been no previous games.
 *)

type history = outcome list ;;

(* 
 * A player for a repeated game will vary strategy based on history, so it is
 * a function from history to a mixed strategy.  It is assumed that the player
 * knows the structure of the bimatrix game being played, and that the history
 * is a history of realizations of that game.
 *)

type player = history -> mixed_strategy ;;

(*
 * play_repeated g n (p1, p2) will play the repeated game consisting of n
 * plays of the bimatrix game g, between players p1 and p2.  On each round, 
 * the players receive the history of the game up to that point, return a mixed
 * strategy, and the next round is played.  
 *
 * The final output of the function is the history of all the games.  The total
 * score can be computed using the final_score function below.
 *
 * Finally, a simple function that takes two players and
 * plays the hw_game 10 times and computes final scores has been provided.  
 *)

let rec play_repeated : bimatrix_game -> int -> player * player -> history =
  fun game rounds -> function (player1, player2) ->
    if rounds < 1 then []
    else 
      let old_history = play_repeated game (rounds-1) (player1, player2) in 
      let strategy1 = player1 old_history
      and strategy2 = player2 old_history in
      play_one game (strategy1, strategy2) :: old_history ;;

let final_score : history -> payoffs = fun hist ->
  List.fold_left sum2 (0.0, 0.0) (List.map snd hist);;

let play_hw10 : player * player -> (history * payoffs) = fun players ->
  let hist = play_repeated hw_game 10 players in
  (hist, final_score hist);;