import streamlit as st
import openai
import os

openai.api_key = os.getenv("openai_api")

BATTLE_SEXES_PAYOFF = {
    ("J", "J"): (10, 7),
    ("J", "F"): (0, 0),
    ("F", "J"): (0, 0),
    ("F", "F"): (7, 10)
}

PRISONERS_DILEMMA_PAYOFF = {
    ("J", "J"): (8, 8),   # Both cooperate
    ("J", "F"): (0, 10),  # Player 1 cooperates, Player 2 defects
    ("F", "J"): (10, 0),  # Player 1 defects, Player 2 cooperates
    ("F", "F"): (5, 5)    # Both defect
}

PAYOFF = PRISONERS_DILEMMA_PAYOFF

def generate_player_description(player):
    descriptions = []
    for (p1_choice, p2_choice), (p1_points, p2_points) in PAYOFF.items():
        if player == 1:
            description = f"If you choose Option {p1_choice} and the other player chooses Option {p2_choice}, then you win {p1_points} points and the other player wins {p2_points} points."
        elif player == 2:
            description = f"If you choose Option {p2_choice} and the other player chooses Option {p1_choice}, then you win {p2_points} points and the other player wins {p1_points} points."
        descriptions.append(description)
    return "\n".join(descriptions)

def gpt_play(player, round_num, last_round,model):
    game_rules = f'''
    You are playing a game repeatedly with another player. Be aware that the other player can make
mistakes sometimes. In this game, you can choose between Option J and Option F. The rules of the game are:
    '''

    player_description = generate_player_description(player)

    previous_round = ''
    if last_round:
        if player == 1:
            previous_round = f"In round {round_num-1}, you chose Option {last_round[0]} and the other player chose Option {last_round[1]}. Thus, you won {PAYOFF[last_round][0]} points and the other player won {PAYOFF[last_round][1]} points."
        elif player == 2:
            previous_round = f"In round {round_num-1}, you chose Option {last_round[1]} and the other player chose Option {last_round[0]}. Thus, you won {PAYOFF[last_round][1]} points and the other player won {PAYOFF[last_round][0]} points."

    prompt = f"{previous_round}\n\nYou are currently playing round {round_num}.\n\nQuestion: Based on the rules and the result of previous round and you are a self-interest human-being and always want to maximise your points, which Option do you choose for this round, Option J or Option F ? Answer only \"J\" or \"F\"\n\nAnswer: "

    #print(f"Below is the prompt for Player {player} and Round {round_num}:\n {prompt}")


    completion = openai.ChatCompletion.create(
      model=model,
      messages=[
          {"role": "system", "content": f"{game_rules}\n\n{player_description}"},
          {"role": "user", "content": prompt}
      ],
        temperature = 0,
    )

    move = completion.choices[0].message['content']

    #print(f"{move}\n")

    if "J" in move:
        move = "J"
    else:
        move = "F"


    return move

def main():
    st.title("LLM Game Simulator")

    # Dropdown menu to select the game or custom matrix
    game_option = st.selectbox(
        'Which game do you want to simulate?',
        ('Prisoner\'s Dilemma', 'Battle of the Sexes', 'Custom Matrix')
    )

    if game_option == 'Prisoner\'s Dilemma':
        PAYOFF = PRISONERS_DILEMMA_PAYOFF
    elif game_option == 'Battle of the Sexes':
        PAYOFF = BATTLE_SEXES_PAYOFF
    else:
        jj = st.text_input("Enter payoff for (J, J) as 'x,y':")
        jf = st.text_input("Enter payoff for (J, F) as 'x,y':")
        fj = st.text_input("Enter payoff for (F, J) as 'x,y':")
        ff = st.text_input("Enter payoff for (F, F) as 'x,y':")

        PAYOFF = {
            ("J", "J"): tuple(map(int, jj.split(','))),
            ("J", "F"): tuple(map(int, jf.split(','))),
            ("F", "J"): tuple(map(int, fj.split(','))),
            ("F", "F"): tuple(map(int, ff.split(',')))
        }

    model_option = st.selectbox(
        'Which model version would you like to use?',
        ('gpt-4', 'gpt-3.5-turbo')
    )
    rounds = st.slider('Number of rounds to play', min_value=1, max_value=100, value=10, step=1)

    st.write(f"Simulating {game_option} for {rounds} rounds using {model_option}.")

    last_round = None
    for i in range(rounds):
        player1_move = gpt_play(1, i + 1, last_round, model=model_option)
        player2_move = gpt_play(2, i + 1, last_round, model=model_option)

        last_round = (player1_move, player2_move)
        score_p1, score_p2 = PAYOFF[last_round]

        st.write(
            f"Round {i + 1}: Player 1 chose {player1_move} and won {score_p1} points. Player 2 chose {player2_move} and won {score_p2} points.")

    st.write("Simulation complete!")


if __name__ == "__main__":
    main()