import random
import numpy
import argparse
import yaml
import csv

NUMBER_OF_DICE = 3
DIE_SIZE = 6
TRIALS = 10000
scoring_cache = {}

def die():
    return random.randint(1, DIE_SIZE)

def dice_set(count):
    dice = []
    for ii in range(count):
        dice.append(die())
    dice.sort()
    return dice
    
def count_pairs(dice):
    count = 0
    if len(dice) > 1:
        for ii in range(len(dice) - 1):
            for jj in range(ii + 1, len(dice)):
                if dice[ii] == dice[jj]:
                    count += 1
    return count
    
def count_runs(dice):
    count = 0
    if len(dice) > 2:
        for ii in range(len(dice) - 2):
            # for the second die, we're going to look at the next die and keep going while the value equals the first + 1
            for jj in range(ii + 1, len(dice) - 1):
                if dice[jj] > (dice[ii] + 1):
                    break
                if dice[jj] == (dice[ii] + 1):
                    for kk in range(jj + 1, len(dice)):
                        if dice[kk] > (dice[jj] + 1):
                            break
                        if dice[kk] == (dice[jj] + 1):
                            count += 1
    return count

def count_totals(dice, target_total):
    if len(dice) == 0:
        return 0
    elif len(dice) == 1:
        return 1 if dice[0] == target_total else 0
    else:
        # test for each subset of this list, subtracting each die off and recursing
        count = 0
        for ii in range(len(dice)):
            if dice[ii] < target_total:
                # recurse - new list composed of the remaining elements
                count += count_totals(dice[ii + 1:], target_total - dice[ii])
            elif dice[ii] == target_total:
                count += 1
            else:
                break
        return count
    
def count_totals_pairs(dice):
    count = 0
    if len(dice) > 1:
        for ii in range(len(dice) - 1):
            for jj in range(ii + 1, len(dice)):
                if (dice[ii] + dice[jj]) == 7:
                    count += 1
    return count

#seems to be the winner
def full_scoring_v1(dice):
    key = str(dice)
    if key in scoring_cache:
        return scoring_cache[key]
    else:
        score = count_pairs(dice) + count_totals_pairs(dice) + (count_runs(dice) * 2)
        scoring_cache[key] = score
        return score
    
def scoring_version_1(pairs, runs, totals_pairs, totals):
    return pairs + totals_pairs + (runs * 2)

def scoring_version_2(pairs, runs, totals_pairs, totals):
    return (pairs * 2) + (totals_pairs * 2) + (runs * 3)

def scoring_version_3(pairs, runs, totals_pairs, totals):
    return pairs + totals_pairs + runs

def reroll_generator(die):
    for ii in range(1,7):
        yield ii

def flip_generator(die):
    yield 7 - die

def plus_minus_generator(die):
    if die > 1:
        yield die - 1
    if die < 6:
        yield die + 1

def plus_minus_two_generator(die):
    if die > 2:
        yield die - 2
    if die > 1:
        yield die - 1
    if die < 6:
        yield die + 1
    if die < 5:
        yield die + 2

def plus_generator(die):
    if die < 6:
        yield die + 1

def plus_two_generator(die):
    if die < 6:
        yield die + 1
    if die < 5:
        yield die + 2

def minus_generator(die):
    if die > 1:
        yield die - 1

def minus_two_generator(die):
    if die > 2:
        yield die - 2
    if die > 1:
        yield die - 1
        
def choose_generator(die):
    for ii in range(1,7):
        yield ii
        
def compute_ability(dice, gen, num_steps, use_best):
    # loop through each die, re-computing the score for each possible value for the ability
    best_score = 0.0
    best_position = 0
    start_score = full_scoring_v1(dice)
    for ii in range(len(dice)):
        total_score = 0.0
        inner_best_score = 0.0
        for replacement in gen(dice[ii]):
            new_dice = dice[:]
            new_dice[ii] = replacement
            dice.sort()
            if num_steps == 1:
                score = full_scoring_v1(new_dice)
            else:
                score = compute_ability(new_dice, gen, num_steps - 1, use_best)
            if use_best:
                inner_best_score = max(inner_best_score, score)
            else:
                total_score += score
        if not use_best:
            inner_best_score = total_score/6
        if inner_best_score > best_score:
            best_position = ii
            best_score = inner_best_score
    return best_score if best_score > start_score else start_score
    
def compute_defense(current_dice, defense_count):
    modified_dice = current_dice[:]
    defense_dice = dice_set(defense_count)
    for die in defense_dice:
        if die in modified_dice:
            modified_dice.remove(die)
    return full_scoring_v1(modified_dice)
    
def compute_select_defense(current_dice, defense_count):
    # loop through each die, re-computing the score for each possible value for the ability
    worst_score = 100000.0
    for ii in range(len(current_dice)):
        new_dice = current_dice[:]
        new_dice.remove(current_dice[ii])
        if defense_count == 1:
            score = full_scoring_v1(new_dice)
        else:
            score = compute_select_defense(new_dice, defense_count - 1)
        if score < worst_score:
            worst_score = score
    return worst_score
    
def attack_trials(outwriter, trials, dice_count):
    scoring = []
    defense_count = range(1,4)
    output_arrays = {}
    for ii in defense_count:
        output_arrays[ii] = []
    for ii in range(trials):
        current_dice = dice_set(dice_count)
        scoring.append(full_scoring_v1(current_dice))
        for choice in defense_count:
            output_arrays[choice].append(compute_select_defense(current_dice, choice))
        if ii % 10 == 0:
            print('trial {0}'.format(ii))
    output_row = [str(dice_count), str(numpy.array(scoring).mean())]
    for ii in defense_count:
        output_row.append(str(numpy.array(output_arrays[ii]).mean()))
    outwriter.writerow(output_row)
    
def title_row_gen(base):
    for ii in range(3):
        yield base.format(ii + 1)
    
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--count', dest='trials', type=int, default=TRIALS, help='How many trials to run?')
    parser.add_argument('-f', '--filename', dest='filename', type=str, default='output.csv', help='Output file name')
    args = parser.parse_args()
    print("Running {0} trials.".format(args.trials))
    with open(args.filename, 'w', newline='') as csvfile:
        outwriter = csv.writer(csvfile, delimiter=',')
        output_row = ['Die count', 'Mean Scoring']
        for base in ['Mean Scoring (w/ {0} defense choices)']:
            for title in title_row_gen(base):
                output_row.append(title)
        outwriter.writerow(output_row)
        attack_trials(outwriter, args.trials, 3)
        attack_trials(outwriter, args.trials, 4)
        attack_trials(outwriter, args.trials, 5)
        attack_trials(outwriter, args.trials, 6)
        attack_trials(outwriter, args.trials, 7)
        attack_trials(outwriter, args.trials, 8)
        attack_trials(outwriter, args.trials, 9)
        attack_trials(outwriter, args.trials, 10)
    
if __name__ == '__main__':
    main()