import random
import csv

NUM_ROUND =72
NUM_BLOCKS = 2
NUM_GROUPS = 10
NUM_ROUND_IN_BLOCK = 36
NUM_SESSIONS = 2
COOP_SUCCESS = ['7-7-7', '8-8-8', '9-9-9', '10-10-10', '6-7-8', '7-8-9', '8-9-10', '9-10-11', '5-7-9', '6-8-10', '7-9-11', '8-10-12']
NONCOOP_SUCCESS = ["10-10-10"]
COOP_FAIL_SAME = ["3-3-3"]
NONCOOP_FAIL = ["5-5-5"]
BLOCKS = ["known_full"]
THRESHOLDS = [1,2,3]
csv_name = 'scenarios_blocks_1124.csv'

# generate the column names
column_types = ['scenario', 'threshold', 'noncoop', 'block']
iterations = 10
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)

# generate case of fixed contribution cost or fixed failure payoff
half_block = NUM_ROUND_IN_BLOCK/2
cases_in_block =["same_difference" for i in range(int(half_block))]  + ["same_payoff" for i in range(int(half_block))] 
random.shuffle(cases_in_block)
print(cases_in_block)

def generate_coop_success(cases_in_block): # generate all the combination of the cooperate and success payoff
    coop_success = []
    coop_success_payoff = COOP_SUCCESS
    thresholds = THRESHOLDS
    for i in range(len(COOP_SUCCESS)): 
        for j in thresholds:
            if cases_in_block[i] == "same_difference":
                coop_success_tem = decode_payoff_list(coop_success_payoff[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_success.append([coop_success_payoff[i],j,coop_fail_tem])
            elif cases_in_block[i] == "same_payoff":
                coop_success.append([coop_success_payoff[i],j,COOP_FAIL_SAME[0]])
            
    return coop_success
    

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

def encode_payoff_list(scenario):
    return str('-'.join(map(str, scenario)))


# create the csv file for scenarios

coop_success = generate_coop_success(cases_in_block)
print(coop_success)
scenarios_list = []
block_group = NUM_BLOCKS*NUM_GROUPS
for i in range(block_group): # generate sequence of scenarios for number of group
    random.shuffle(coop_success)
    scenarios_list.append(list(coop_success))

blocks = BLOCKS
block_list = []
for j in range(NUM_GROUPS): 
    tem_block = []
    for z in range(NUM_SESSIONS): 
        random.shuffle(blocks)
        tem_block += list(blocks)
    block_list.append(tem_block)

with open(csv_name, 'w', newline='') as csvfile:
    scenarios_writer = csv.writer(csvfile)
    scenarios_writer.writerow(column_names)
    for b in range(NUM_BLOCKS):
        for r in range(NUM_ROUND_IN_BLOCK):
            row = []
            for s in range(10):
                row.append(scenarios_list[s+10*b][r][0])
                row.append(scenarios_list[s+10*b][r][1])
                row.append(scenarios_list[s+10*b][r][2])
                row.append(block_list[s][b])             
            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]])"