(**************************************************************************)
(*                                BELENIOS                                *)
(*                                                                        *)
(*  Copyright © 2012-2021 Inria                                           *)
(*                                                                        *)
(*  This program is free software: you can redistribute it and/or modify  *)
(*  it under the terms of the GNU Affero General Public License as        *)
(*  published by the Free Software Foundation, either version 3 of the    *)
(*  License, or (at your option) any later version, with the additional   *)
(*  exemption that compiling, linking, and/or using OpenSSL is allowed.   *)
(*                                                                        *)
(*  This program is distributed in the hope that it will be useful, but   *)
(*  WITHOUT ANY WARRANTY; without even the implied warranty of            *)
(*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU     *)
(*  Affero General Public License for more details.                       *)
(*                                                                        *)
(*  You should have received a copy of the GNU Affero General Public      *)
(*  License along with this program.  If not, see                         *)
(*  <http://www.gnu.org/licenses/>.                                       *)
(**************************************************************************)

open Serializable_t
open Common

(** Transform a ballot in belenios format (e.g. [4,1,2,5,3]) into a list of
    choices (represented as their index in the vector) in preference order (e.g.
    [1,2,4,0,3]) *)
let process ballot =
  let nchoices = Array.length ballot in
  let used = Array.make nchoices false in
  let rec lookup rank i =
    if i < nchoices then
      if ballot.(i) = rank then (
        assert (not used.(i));
        used.(i) <- true;
        Some i)
      else lookup rank (i + 1)
    else None
  in
  let rec build_preference_list rank accu =
    match lookup rank 0 with
    | Some i -> build_preference_list (rank + 1) (i :: accu)
    | None -> List.rev accu
  in
  let preference_list = build_preference_list 1 [] in
  let rec check i =
    if i < nchoices then
      if used.(i) || ballot.(i) = 0 then check (i + 1) else false
    else true
  in
  if check 0 then Some preference_list else None

(** Here, [choices] is a map from choices to ballots that has them as first
    choice. This function updates [choices] with a new ballot: it assigns the
    ballot to its first choice, and makes sure all other choices also exist in
    [choices]. *)
let assign choices ((_, ballot) as b) =
  let prepend ballots choices x =
    match IMap.find_opt x choices with
    | None -> IMap.add x ballots choices
    | Some bs -> IMap.add x (List.rev_append ballots bs) choices
  in
  match ballot with
  | [] -> choices
  | x :: xs -> List.fold_left (prepend []) (prepend [ b ] choices x) xs

(** Here, [scores] is an association list mapping from choices to ballots and
    total score. This function collects all the ballots, filters out [i] from
    them, and multiplies [i]'s ballots by [coef]. *)
let transfer coef i scores =
  List.fold_left
    (fun accu (ai, (ab, _)) ->
      List.fold_left
        (fun accu (w, b) ->
          let w = if ai = i then w *. coef else w in
          let b = List.filter (fun x -> x <> i) b in
          (w, b) :: accu)
        accu ab)
    [] scores

(** This function performs a round of the STV algorithm. It tail-recursively
    calls itself until [nseats] is [0] or there is not enough remaining choices,
    and returns the list of events (Win|Lose|TieWin|TieLose) that occured during
    the process. *)
let rec run quota ballots events nseats =
  if nseats > 0 then
    let choices = List.fold_left assign IMap.empty ballots in
    if IMap.cardinal choices <= nseats then
      (* there is not enough choices: they all win *)
      choices |> IMap.bindings |> List.map fst
      |> (fun x -> `Win x :: events)
      |> List.rev
    else
      let scores =
        (* for each choice, compute the sum of scores of its assigned
           ballots *)
        choices
        |> IMap.map (fun bs ->
               (bs, List.fold_left (fun accu (w, _) -> accu +. w) 0. bs))
        |> IMap.bindings
        |> List.sort
             (* we sort the choices, with greater total score first,
                then in question order (this is our "arbitrary" tie
                breaking, chosen by the election administrator) *)
             (fun (ai, (_, aw)) (bi, (_, bw)) -> compare (bw, ai) (aw, bi))
      in
      match scores with
      | (ai, (_, aw)) :: xs when aw >= quota ->
          (* the first choice is above the quota *)
          let events =
            match xs with
            | (bi, (_, bw)) :: _ when aw = bw ->
                (* the second choice has the same total score, we chose
                   the first one, but log the tie *)
                `TieWin [ ai; bi ] :: events
            | _ -> events
          in
          (* note that we select a single winner, even if there are
             several choices above quota *)
          let c = (aw -. quota) /. aw in
          run quota (transfer c ai scores) (`Win [ ai ] :: events) (nseats - 1)
      | scores -> (
          match List.rev scores with
          | (ai, (_, aw)) :: xs ->
              (* we select the last choice *)
              let events =
                match xs with
                | (bi, (_, bw)) :: _ when aw = bw ->
                    (* the second last choice has the same total score,
                       we chose the last one, but log the tie *)
                    `TieLose [ ai; bi ] :: events
                | _ -> events
              in
              run quota (transfer 1. ai scores) (`Lose ai :: events) nseats
          | [] ->
              (* should not happen, because if there is no choices left,
                 the condition "there is not enough choices" above
                 should have been triggered *)
              assert false)
  else List.rev events

let compute ~nseats ballots =
  let nballots = Array.length ballots in
  let rec partition accu invalid i =
    if i < nballots then
      let ballot = ballots.(i) in
      match process ballot with
      | None -> partition accu (ballot :: invalid) (i + 1)
      | Some x -> partition (x :: accu) invalid (i + 1)
    else (List.sort compare accu, Array.of_list (List.sort compare invalid))
  in
  let stv_ballots, stv_invalid = partition [] [] 0 in
  let n = List.length stv_ballots in
  let quota = floor (float n /. float (nseats + 1)) +. 1. in
  let wballots = List.map (fun b -> (1., b)) stv_ballots in
  let stv_events = run quota wballots [] nseats in
  let stv_winners =
    stv_events |> List.map (function `Win x -> x | _ -> []) |> List.flatten
  in
  { stv_ballots; stv_invalid; stv_events; stv_winners }
