# Code: Generate Scenarios
# Data: 01/02/2024
# Creator: Lien Lin
# Design: 3 subjects, two sessions, 2 blocks, fixed contribution cost and fixed payoff cost separated in each session
# Done: Fix the error of randomization of blocks

#//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////



import random
import csv

NUM_ROUND = 66
NUM_BLOCKS = 2
NUM_GROUPS = 10
NUM_ROUND_IN_BLOCK = 33
NUM_SESSIONS = 2
COOP_SUCCESS = ["5_5_5", "6_6_6", '7_7_7', '8_8_8', '9_9_9', '5_6_7', '6_7_8', '7_8_9', '8_9_10', '5_7_9', '6_8_10']
NONCOOP_SUCCESS = ["10_10_10"]
COOP_FAIL_FIX = ["2_2_2"]
NONCOOP_FAIL = ["5_5_5"]
BLOCKS = ["fix_payoff", "fix_contribution", ]
THRESHOLDS = [1,2,3]
csv_name = 'scenarios_blocks_0105_fix_payoff.csv'

# generate the column names
column_types = ['coop_success', 'coop_fail', 'threshold', 'block']
iterations = NUM_GROUPS
column_names = []
for i in range(1, iterations + 1):
    for column_type in column_types:
        column_name = f"{column_type}{i}"
        column_names.append(column_name)


def generate_scenarios(blocks, coop_success, noncoop_success, noncoop_fail, coop_fail_fix, thresholds): # generate all the combination of the cooperate payoff (including success and fail) and the threshold (including 1,2,3
    coop_fix_contribution = []
    coop_fix_payoff =[]
    for b in range(len(blocks)):
        for i in range(len(coop_success)): 
            for j in thresholds:
                if blocks[b] == "fix_contribution":
                    coop_success_tem = decode_payoff_list(coop_success[i])
                    noncoop_success_tem = decode_payoff_list(noncoop_success[0])
                    noncoop_fail_tem = decode_payoff_list(noncoop_fail[0])
                    coop_fail_tem = [coop_success_tem[k] - noncoop_success_tem[k] + noncoop_fail_tem[k] for k in range(len(coop_success_tem))]
                    coop_fail_tem = encode_payoff_list(coop_fail_tem)
                    coop_fix_contribution.append([coop_success[i],coop_fail_tem, j, "fix_contribution"])
                elif blocks[b] == "fix_payoff":
                    coop_fix_payoff.append([coop_success[i],coop_fail_fix[0], j, "fix_payoff"])
        
    random.shuffle(coop_fix_payoff)
    random.shuffle(coop_fix_contribution)
    if blocks[0] == "fix_payoff":
        coop = coop_fix_payoff + coop_fix_contribution
    elif blocks[0] == "fix_contribution":
        coop = coop_fix_contribution + coop_fix_payoff
    
    return coop
    

def decode_payoff_list(scenario): # decode the string payoff list into list
    payoff_list = scenario.split('_')
    payoff_list = [int(s.strip()) for s in payoff_list]
    return payoff_list

def encode_payoff_list(scenario): # encode the payoff list into string
    return str('_'.join(map(str, scenario)))


# create the csv file for scenarios

blocks = BLOCKS # when blocks is equal to session
block_list = []
for j in range(NUM_GROUPS): 
    # random.shuffle(blocks)
    block_list.append(blocks)


scenarios_list = []
for g in range(NUM_GROUPS): # generate sequence of scenarios for number of group
    scenarios_tem = generate_scenarios(block_list[g], COOP_SUCCESS, NONCOOP_SUCCESS, NONCOOP_FAIL, COOP_FAIL_FIX, THRESHOLDS)
    scenarios_list.append(scenarios_tem)



with open(csv_name, 'w', newline='') as csvfile:
    scenarios_writer = csv.writer(csvfile)
    scenarios_writer.writerow(column_names) # writing each row one by one
    for r in range(NUM_ROUND):
        row = []
        for g in range(NUM_GROUPS):
            row.append(scenarios_list[g][r][0])
            row.append(scenarios_list[g][r][1])
            row.append(scenarios_list[g][r][2])
            row.append(scenarios_list[g][r][3])             
        scenarios_writer.writerow(row)

"[scenarios_list[0+10*b][r][0], scenarios_list[0+10*b][r][1],block_list[0][b], scenarios_list[1+10*b][r][0], scenarios_list[1+10*b][r][1], block_list[1][b], scenarios_list[2+10*b][r][0], scenarios_list[2+10*b][r][1], block_list[2][b], scenarios_list[3+10*b][r][0], scenarios_list[3+10*b][r][1], block_list[3][b], scenarios_list[4+10*b][r][0], scenarios_list[4+10*b][r][1], block_list[4][b], scenarios_list[5+10*b][r][0], scenarios_list[5+10*b][r][1], block_list[5][b], scenarios_list[6+10*b][r][0], scenarios_list[6+10*b][r][1], block_list[6][b], scenarios_list[7+10*b][r][0], scenarios_list[7+10*b][r][1], block_list[7][b], scenarios_list[8+10*b][r][0], scenarios_list[8+10*b][r][1], block_list[8][b], scenarios_list[9+10*b][r][0], scenarios_list[9+10*b][r][1], block_list[9][b]])"