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,
         hasGun: bool,
    }

    var players_to_features: str -> PlayerFeatures
    var votes_by_player: str -> int
    var game_phase: Phase
    var game_status: Status
    var last_saved: Option[str]
    var last_moved1: Option[str]
    var last_moved2: Option[str]
    var gun_done: bool

    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 mover_alive(players: str -> PlayerFeatures): int =
        size(players.values().filter((p) => p.status == Alive and p.role == Mover))

    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 init = all {
        pure val list_of_roles = gen_list_of_roles(size(players))
        pure val set_of_roles = list_of_roles.listToSet()
        pure val role_by_player_powerset:Set[str -> Role] = players.setOfMaps(set_of_roles)
        pure val role_by_player_subset:Set[str -> Role] =
            role_by_player_powerset.filter(m => [] == (m.keys().fold(list_of_roles, (acc, n) => dropFirst(acc,x => x == m.get(n)))))
        nondet role_by_player:str -> Role = role_by_player_subset.oneOf()
        pure val mascot_potential:Set[str] = role_by_player.mapToSet().filter(x => x._2 != Mafia).map(x => x._1)
        nondet mascot_choice:str = mascot_potential.oneOf()
        players_to_features' = players.mapBy(p =>
          { name: p,
            alignment: if (role_by_player.get(p) == Mafia) Scum else Town,
            role: role_by_player.get(p),
            status: Alive,
            voted: false,
            nominated: false,
            mascot: p == mascot_choice,
            hasGun: false,
          }),
        last_saved' = None,
        last_moved1' = None,
        last_moved2' = None,
        game_phase' = Day,  // Start with the Day phase
        game_status' = Pending,  // Game is in Pending status
        votes_by_player' = players.mapBy(p => 0)  // Initialize vote counts to 0
    }
    

    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 Some(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()
        nondet move1 = players_to_features.values().filter(p => p.status == Alive and Some(p.name) != last_moved1 and Some(p.name) != last_moved2 ).oneOf()
        nondet move2 = players_to_features.values().filter(p => p.status == Alive and Some(p.name) != last_moved1 and Some(p.name) != last_moved2 and p.name != move1.name).oneOf()
        nondet gunChoice = players_to_features.values().filter(p => p.status == Alive and p.role != Gunsmith).oneOf()
        nondet handoutGun = Set(true,false).oneOf()
        val mover = mover_alive(players_to_features) > 0
        val doc = doc_alive(players_to_features) > 0
        val vig0 = vig0_alive(players_to_features) > 0
        val vig1 = vig1_alive(players_to_features) > 0
        val moved = players.mapBy(x => if (not(mover)) x else
            if (x == move1.name) move2.name else if (x == move2.name) move1.name else x)
        val victim = if (victimScum == doctorSave and doc) [] else [{...victimScum, name: moved.get(victimScum.name)}]
        val victims: List[PlayerFeatures] = {
            val temp = if (vig0) victim.append({...victimVig0, name: moved.get(victimVig0.name)}) else victim
            if (vig1) temp.append({...victimVig1, name: moved.get(victimVig1.name)}) else temp
        }
        val updated_features = update_after_kill(victims.listMap(p => p.name),players_to_features).setBy(gunChoice.name, p => {...p, hasGun: handoutGun and not(gun_done)})
        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,
            last_saved' = if (doc and victimScum == doctorSave) None else if (doc) Some(doctorSave.name) else None,
            last_moved1' = if (mover and (move1 == victimScum or move2 == victimScum)) None else if (mover) Some(move1.name) else None,
            last_moved2' = if (mover and (move1 == victimScum or move2 == victimScum)) None else if (mover) Some(move2.name) else None,
            gun_done' = (handoutGun or gun_done),
        }
    }

    // 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")).*
}