from math import *
from random import randint
import sys


def main():
    partnerAyears = 0
    partnerByears = 0
    partnerAstrategy = 0
    partnerBstrategy = 0
    partnerAstrategy, partnerBstrategy = get_strat()
    trials = get_num_trials()
    partnerAyears, partnerByears = run_trials(partnerAstrategy, partnerBstrategy, partnerAyears, partnerByears, trials)
    print("Partner A spent ", partnerAyears / trials, " average years in jail")
    print("Partner B spent ", partnerByears / trials, " average years in jail")


def run_trials(partnerAstrat, partnerBstrat, partnerA, partnerB, count):
    if count == 0:
        return 0, 0
    return run_trials_rec(partnerAstrat, partnerBstrat, partnerA, partnerB, count)


# Recursive function
def run_trials_rec(partnerAstrat, partnerBstrat, partnerA, partnerB, count):
    if count == 0:
        return partnerA, partnerB
    # Let 1 equal coop, and 0 equal silent
    choiceA = partnerAstrat
    choiceB = partnerBstrat

    yearsAddedforA = 0
    yearsAddedforB = 0
    if choiceA == 1 and choiceB == 1:
        yearsAddedforA = 4
        yearsAddedforB = 4
    if choiceA != choiceB:
        if choiceA == 0:
            yearsAddedforA = 6
            yearsAddedforB = 0
        else:
            yearsAddedforA = 0
            yearsAddedforB = 6
    if choiceA == 0 and choiceB == 0:
        yearsAddedforA = 2
        yearsAddedforB = 2

    partnerA += yearsAddedforA
    partnerB += yearsAddedforB
    count = count - 1

    return run_trials_rec(partnerAstrat, partnerBstrat, partnerA, partnerB, count)


def get_num_trials():
    trials = int(input("Enter a number from 10-500 "))
    if trials < 10 or trials > 500:
        print("Please enter a number between 10 and 500 ")
        return get_num_trials()
    else:
        return trials


def get_strat():
    print("3 strategies available:")
    print("1: Random")
    print("2: Only-Silent")
    print("3: Only-Cooperates")
    partnerA = get_strategy('A')
    partnerB = get_strategy('B')

    # We need to change the values, because we need the value to match our recursive method
    if partnerA == 2:
        partnerA = 0
    elif partnerA == 3:
        partnerA = 1
    else:
        partnerA = randint(0, 1)

    if partnerB == 2:
        partnerB = 0
    elif partnerB == 3:
        partnerB = 1
    else:
        partnerB = randint(0, 1)

    return partnerA, partnerB


# Recursive function to get strategy
def get_strategy(letter):
    number = 0
    print("Please enter partner ", letter, "'s strategy")
    number = int(input())
    if number <= 0 or number > 3:
        print("Must enter a number between 1 and 3")
        print("Please try again")
        return get_strategy(letter)
    else:
        return number


if __name__ == "__main__":
    main()