Masons/Masons.qnt

module mafia {
    import basicSpells.* from "basicspells"
    const players : Set[str]
    type Alignment = Scum | Town
    type Role = Mafia | Banana | Doctor | Mover | Gunsmith | Vigilante(int) | Mason | Villager
    type Phase = Day | Night
    type LifeState = Alive | Dead
    type Status = Pending | Done(Alignment)
    type PlayerFeatures = {
         name: str,
         alignment: Alignment,
         role: Role,
         status: LifeState,
         voted: bool,
         mascot: bool,
         nominated: bool,
    }

    var players_to_features: str -> PlayerFeatures
    var votes_by_player: str -> int
    var game_phase: Phase
    var game_status: Status
    var last_saved: str

    pure def all_voted(players: str -> PlayerFeatures): bool =
         players.values().filter(p => p.status == Alive).forall(p => p.voted == true)

    pure def alive_mafia(players: str -> PlayerFeatures): int =
        size(players.values().filter((p) => p.status == Alive and p.alignment == Scum))

    pure def alive_town(players: str -> PlayerFeatures): int =
         size(players.values().filter((p) => p.status == Alive and p.alignment == Town))

    pure def mascots_alive(players: str -> PlayerFeatures): int =
        size(players.values().filter((p) => p.status == Alive and p.mascot))

    pure def doc_alive(players: str -> PlayerFeatures): int =
        size(players.values().filter((p) => p.status == Alive and p.role == Doctor))

    pure def vig0_alive(players: str -> PlayerFeatures): int =
        size(players.values().filter((p) => p.status == Alive and p.role == Vigilante(0)))

    pure def vig1_alive(players: str -> PlayerFeatures): int =
        size(players.values().filter((p) => p.status == Alive and p.role == Vigilante(1)))

    pure def mafia_lose_condition_active(players: str -> PlayerFeatures):bool = {
        val mafias = alive_mafia(players)
        val mafia_aligned = if (size(keys(players)) < 11)
            (mafias + mascots_alive(players))
            else mafias
         2 > mafia_aligned
    }

    pure def mafia_win_condition_active(players: str -> PlayerFeatures): bool = alive_mafia(players) == alive_town(players)

    pure def gen_list_of_roles(n: int): List[Role] = {
     val special: List[Role] =
             if (Set(7).contains(n))           [Banana]
        else if (Set(8).contains(n))           [Doctor]
        else if (Set(9).contains(n))           [Mover]
        else if (Set(10).contains(n))          [Doctor]
        else if (Set(11,12).contains(n))       [Doctor, Gunsmith]
        else if (Set(13,14).contains(n))       [Doctor, Banana]
        else if (Set(15,16).contains(n))       [Mason, Mason, Gunsmith]
        else if (Set(17,18).contains(n))       [Mover, Gunsmith]
        else if (Set(19,20).contains(n))       [Mason,Mason,Doctor,Vigilante(0)]
        else if (Set(21,22).contains(n))       [Vigilante(0), Vigilante(0), Doctor, Gunsmith]
        else if (Set(23,24,25,26).contains(n)) [Vigilante(0), Vigilante(0), Vigilante(1), Vigilante(1), Banana]
        else                                   []
      val mafias: List [Role] = repeat((n+1)/4, Mafia)
      if (n < 7 or n > 26) [] else
        special.concat(mafias)
            .concat(repeat(n-length(special)-length(mafias),Villager))
    }

    val get_most_voted_players = {
        if (all_voted(players_to_features)) {
            val max_votes = players.fold(-1, (acc, p) => {
                val votes = votes_by_player.get(p)
                max(votes,acc)
            })
            players.filter(p => votes_by_player.get(p) == max_votes)
        } else Set()  // Return an empty set if not all players have voted
    }

    pure def update_status(players: str -> PlayerFeatures): Status =
        if (mafia_lose_condition_active(players)) Done(Town) else
        if (mafia_win_condition_active(players))  Done(Scum) else
        Pending

    pure def update_after_kill(victims: List[str], players: str -> PlayerFeatures): str -> PlayerFeatures =
        victims.foldl(players, (curr, victim) => curr.setBy(victim, p => {...p, status: Dead}))

    pure def update_after_hanging(victim: str, players: str -> PlayerFeatures): str -> PlayerFeatures = {
        val playersNew = update_after_kill([victim], players)
        playersNew.transformValues(p => {... p, voted: false})
    }

    action nightPhase = all {
        nondet victimScum = players_to_features.values().filter(p => p.status == Alive and p.alignment == Town).oneOf()
        nondet doctorSave = players_to_features.values().filter(p => p.status == Alive and p.role != Doctor and p.name != last_saved).oneOf()
        nondet victimVig0 = players_to_features.values().filter(p => p.status == Alive and p.role != Vigilante(0)).oneOf()
        nondet victimVig1 = players_to_features.values().filter(p => p.status == Alive and p.role != Vigilante(1)).oneOf()
        val victim = if (victimScum == doctorSave and doc_alive(players_to_features) > 0) [] else [victimScum]
        val victims: List[PlayerFeatures] = {
            val temp = if (vig0_alive(players_to_features) > 0) victim.append(victimVig0) else victim
            if (vig1_alive(players_to_features) > 0) temp.append(victimVig1) else temp
        }
        val updated_features = update_after_kill(victims.listMap(p => p.name),players_to_features)
        val new_game_status = update_status(updated_features)
        all {
            players_to_features.values().exists(p => p.status == Alive and p.role == Mafia),
            game_phase == Night,
            players_to_features' = updated_features,
            game_status' = new_game_status,
            game_phase' = Day,
            votes_by_player' = votes_by_player
        }
    }

    // TODO: Nomination, Voting, Execution, in that order


    // TODO: invariants
    // TODO: Win rates
}
module play_mafia {
    import mafia(players = Set("Pingu", "T-Dor", "Prince", "Benjamin", "Barwe", "Draken", "Ikka", "Krig", "Izza", "Underhill", "Geting", "Igelkott", "Ugglan", "Joppe", "Människan", "Wisdom")).*
}