#NEWMARCH
#from data_fetcher import fetch_personalized_data 
#json_data = fetch_personalized_data(session_index)
#print(json_data)

from data_fetcher import fetch_data_as_json, fetch_data_core
json_data = fetch_data_as_json()
print(json_data)
core_data = fetch_data_core()
print(core_data)

# From own code

BOS, EOS = "<s>", "</s>" 
BINST, EINST = "[INST]", "[/INST]"
BSYS, ESYS = "<<SYS>>\n", "\n<</SYS>>\n\n"

def get_default_system_prompt():
    DEFAULT_SYSTEM_PROMPT = """\  You are an intelligent and fair game guide in a 2-player trust game, assisting players in making decisions to win. 
    Answer in a consistent style. Answer per question should be maximum 2 sentences long. The players are called The Investor and The Dealer and keep their role throughout the whole game. 
    Both start with 10€ in round 1. The game consists of 3 rounds. In round 1, The Investor invests between 0€ and 10€. 
    This amount is tripled automatically, and The Dealer can then distribute the tripled amount. After that, round 1 is over. 
    Both go into the next round with their current asset: The Investor with 10€ minus what he invested plus what he received back from The Dealer.
    The Dealer with 10€ plus what he kept from the tripled amount. 
    You will receive a JSON with information on who trusted whom with how much money after each round as context.
    Your goal is to guide players through the game, providing clear instructions and explanations. 
    If any question or action seems unclear, explain it rather than providing inaccurate information. 
    If you're unsure about an answer, it's better not to guess.

    Example JSON context after a round: {json_data}

    You are currently assisting player: check the groupNr variable from {core_data}
    You are currently in round: check the onPage variable from {core_data}
    There are x rounds left.
    Please keep that in mind when giving advice. 

    Few-shot training examples
    {BSYS} Give an overview of the trust game. {ESYS}
    {BSYS} Explain how trust amounts are calculated. {ESYS}
    {BSYS} What happens if a player doesn't trust in a round? {ESYS}
    """.format(json_data=json_data,BSYS=BSYS, ESYS=ESYS, core_data=core_data)

    print(DEFAULT_SYSTEM_PROMPT)
    return DEFAULT_SYSTEM_PROMPT


def construct_input_prompt(chat_history, message):
    input_prompt = f"<s>[INST] <<SYS>>\n{get_default_system_prompt()}\n<</SYS>>\n\n "

    for user, assistant in chat_history:
        input_prompt += f"{user} [/INST] {assistant} <s>[INST] "

    input_prompt += f"{message} [/INST] "

    return input_prompt



'''
# From Llama 2 Local 
# General code from here: https://www.youtube.com/watch?v=WzCS8z9GqHw 
# https://github.com/thisserand/llama2_local 
# These variables represent the beginning-of-sequence (BOS) and end-of-sequence (EOS) tokens. 
BOS, EOS = "<s>", "</s>" 

# These variables represent the beginning and end markers for an instruction. 
BINST, EINST = "[INST]", "[/INST]"

# These variables define markers to denote the beginning and end of a system message. 
# In the context of a chatbot, a system message is information or instructions provided by the system rather than the user. 
# The markers here use <<SYS>> to indicate the start and <</SYS>> to indicate the end, with newline characters \n for formatting.
BSYS, ESYS = "<<SYS>>\n", "\n<</SYS>>\n\n"

# This line initializes a multiline string (""") for the default system prompt. 
# It's the set of instructions for the chatbot to follow during a conversation in the trust game scenario that will be contained in the system message. 
DEFAULT_SYSTEM_PROMPT = """\ Your task is to answer in a consistent style. You answer per question is maximum 2 sentences long. 
You are an intelligent and fair game guide in a 2-player trust game. 
Your role is to assist players in making decisions during the game. 
The game consists of 3 rounds, and each player starts with an initial asset of 10€.
In each round, both players can trust each other with an amount between 0€ and 10€. 
The trusted amounts are added, multiplied by 3, divided by 2, and then evenly distributed among the participants. 
This sum, along with what's left of their initial assets, becomes their new asset for the next round.
For example, if player A trusts player B with 5€, and player B trusts player A with 3€, the combined trust is (5 + 3) = 8€. 
After the multiplier and division, both players receive (8 * 3 / 2) = 12€. 
Adding this to what's left from their initial 10€ forms their asset for the next round.
After 3 rounds, the final earnings are calculated using the same process. 
You will receive a JSON with information on who trusted whom with how much money after each round as context.
Your goal is to guide players through the game, providing clear instructions and explanations. 
If any question or action seems unclear, explain it rather than providing inaccurate information. 
If you're unsure about an answer, it's better not to guess.

Example JSON context after a round: {json_data}

You are currently assisting player: check the groupNr variable from {core_data}
You are currently in round: check the onPage variable from {core_data}
There are x rounds left.
Please keep that in mind when giving advice. 

Few-shot training examples
{BSYS} Give an overview of the trust game. {ESYS}
{BSYS} Explain how trust amounts are calculated. {ESYS}
{BSYS} What happens if a player doesn't trust in a round? {ESYS}
""".format(json_data=json_data,BSYS=BSYS, ESYS=ESYS, core_data=core_data)

print(DEFAULT_SYSTEM_PROMPT)

# Original: DEFAULT_SYSTEM_PROMPT = """\ You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information."""

def format_to_llama_chat_style(history) -> str:
    # history has the following structure:
    # - dialogs
    # --- instruction
    # --- response (None for the most recent dialog)
    prompt = ""
    for i, dialog in enumerate(history[:-1]):
      instruction, response = dialog[0], dialog[1]
      # prepend system instruction before first instruction
      if i == 0:
        instruction = f"{BSYS}{DEFAULT_SYSTEM_PROMPT}{ESYS}" + instruction
      else:
        # the tokenizer automatically adds a bos_token during encoding,
        # for this reason the bos_token is not added for the first instruction
        prompt += BOS
      prompt += f"{BINST} {instruction.strip()} {EINST} {response.strip()} " + EOS

    # new instruction from the user
    new_instruction = history[-1][0].strip()

    # the tokenizer automatically adds a bos_token during encoding,
    # for this reason the bos_token is not added for the first instruction
    if len(history) > 1:
      prompt += BOS
    else:
      # prepend system instruction before first instruction
      new_instruction = f"{BSYS}{DEFAULT_SYSTEM_PROMPT}{ESYS}" + new_instruction

    prompt += f"{BINST} {new_instruction} {EINST}"
    return prompt
'''


'''
BOS, EOS = "<s>", "</s>" 
BINST, EINST = "[INST]", "[/INST]"
BSYS, ESYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
DEFAULT_SYSTEM_PROMPT = """\ You are an intelligent and fair game guide in a 2-player trust game, assisting players in making decisions to win. 
Answer in a consistent style. Answer per question should be maximum 2 sentences long. The players are called The Investor and The Dealer and keep their role throughout the whole game. 
Both start with 10€ in round 1. The game consists of 3 rounds. In round 1, The Investor invests between 0€ and 10€. 
This amount is tripled automatically, and The Dealer can then distribute the tripled amount. After that, round 1 is over. 
Both go into the next round with their current asset: The Investor with 10€ minus what he invested plus what he received back from The Dealer.
The Dealer with 10€ plus what he kept from the tripled amount. 
You will receive a JSON with information on who trusted whom with how much money after each round as context.
Your goal is to guide players through the game, providing clear instructions and explanations. 
If any question or action seems unclear, explain it rather than providing inaccurate information. 
If you're unsure about an answer, it's better not to guess.

Example JSON context after a round: {json_data}

You are currently assisting player: check the groupNr variable from {core_data}
You are currently in round: check the onPage variable from {core_data}
There are x rounds left.
Please keep that in mind when giving advice. 

Few-shot training examples
{BSYS} Give an overview of the trust game. {ESYS}
{BSYS} Explain how trust amounts are calculated. {ESYS}
{BSYS} What happens if a player doesn't trust in a round? {ESYS}
""".format(json_data=json_data,BSYS=BSYS, ESYS=ESYS, core_data=core_data)

print(DEFAULT_SYSTEM_PROMPT)

def construct_input_prompt(chat_history, message) -> str:
    # chat_history has the following structure:
    # - dialogs
    # --- instruction
    # --- response (None for the most recent dialog)
    prompt = ""
    for i, dialog in enumerate(chat_history[:-1]):
      instruction, response = dialog[0], dialog[1]
      # prepend system instruction before first instruction
      if i == 0:
        instruction = f"{BSYS}{DEFAULT_SYSTEM_PROMPT}{ESYS}" + instruction
      else:
        # the tokenizer automatically adds a bos_token during encoding,
        # for this reason the bos_token is not added for the first instruction
        prompt += BOS
      prompt += f"{BINST} {instruction.strip()} {EINST} {response.strip()} " + EOS

    # new instruction from the user
    new_instruction = chat_history[-1][0].strip()

    # the tokenizer automatically adds a bos_token during encoding,
    # for this reason the bos_token is not added for the first instruction
    if len(chat_history) > 1:
      prompt += BOS
    else:
      # prepend system instruction before first instruction
      new_instruction = f"{BSYS}{DEFAULT_SYSTEM_PROMPT}{ESYS}" + new_instruction

    prompt += f"{BINST} {new_instruction} {message} {EINST}"
    return prompt
'''