# author: TestRunner
# MIT License: Free use, do with it what you will etc
#
# Usage:
#   first enter the number of "colors", which is the
#   number of different types of items that can be offered.
#   In each round ender in a 5 digit number representing
#   the game's response. 0 = Red, 1 = Yellow, 2 = Green.
#
#   - For 5 or fewer colors, if you don't get an expected
#     response of 50% in the first round, then you should
#     just restart and try again.
#   - For 6 colors, it will be hard to get good odds in 3
#     rounds. But you can generally reset until you get
#     a lucky first round that give 33% or better in 3 rounds.
#   - For 7 colors, getting it done in 3 rounds is besically
#     impossible, but can almost always get a 4 with 95%+.
#     As before, reset until the 4th round can be done in high
#     odds.
#
#  Note 1: this code is not efficient at all. When the odds
#  is low, it may take a really long time to generate a result
#  (especially on 7+ colors where the odds are poor). If the
#  process is taking too long to run, you might want to just
#  terminate it because the odds are likely so bad that it's
#  not worth doing.
#
#  Note 2: The reported odds may not actually be correct, and
#  it's possible that the suggested guess is not the best.
#  However it typically quite and generally similar to the best.

import math
import random
import itertools
import time

def partitions(n, depth=-1, I=1):
    if depth != 0:
        yield (n,)
        for i in range(I, n//2 + 1):
            for p in partitions(n-i, depth-1, i):
                yield p + (i,)


def inits():
    for p in partitions(holes, colors):
        yield tuple(
            i + 1
            for i,c in enumerate(p)
            for j in range(c)
        )


def space(holes, min, max):
    return itertools.product(range(min,max+1), repeat=holes)


def sub_game(game, response, guess, check_color):
    if all(map(lambda hole: hole == 1, response)):
        print('no changes in sub game')
        return game
    if all(map(lambda hole: hole == 2, response)):
        print('empty sub game')
        return set()

    if check_color:
        dead_colors = set(
            c
            for i,c in enumerate(guess)
            if response[i] == 0
        )
    else:
        dead_colors = set()

    return set([
        tuple(
            p[i]
            for i,hole in enumerate(response)
            if hole != 2
        )
        for p in game
        if not any(map(lambda h: h in dead_colors, p))
    ])


def get_guess(key, guess):
    key = list(key)
    guess = list(guess)
    holes = len(key)
    assert len(key) == len(guess)

    response = list(blank)

    reds = 0
    for i in range(holes):
        if key[i] == guess[i]:
            response[i] = 2
            key[i] = -1
            guess[i] = -1

    whites = 0
    for i in range(holes):
        if guess[i] < 0:
            continue
        if guess[i] in key:
            response[i] = 1

    return tuple(response)


def enter_guess(guess, response, candidates):
    new_candidates = []
    for p in candidates:
        if response == get_guess(p, guess):
            new_candidates.append(p)
    return new_candidates


def filter_responses(response, responses):
    for r in responses:
        if all(map(lambda i: response[i[0]] != 2 or i[1] == 2, enumerate(r))):
            yield r


def elim_count(guess, response, candidates, sample):
    elim = 0
    for p in random.sample(candidates, min(sample, len(candidates))):
        if response != get_guess(p, guess):
            elim += 1
    return elim


def solve_game(key, init, guesses, candidates, responses, output):
    guess = init
    while guesses > 0:
        guesses -= 1
        if output:
            print(f'guess - {guess}')

        guess_r = get_guess(key, guess)
        if guess_r == correct:
            if output:
                print(f'Match found')
            return (guesses, guess)

        candidates = enter_guess(guess, guess_r, candidates)
        responses = list(filter_responses(guess_r, responses))

        guess = get_best_guess(candidates, responses)[0]

    if output:
        print(f'Match not found')
    return (-1, None)


def get_best_guess(candidates, responses):
    best = []
    best_min = 0

    for p in random.sample(candidates, min(s_size, len(candidates))):
        min_e = 100000
        for r in random.sample(responses, min(s_size, len(responses))):
            min_e = min(min_e, elim_count(p, r, candidates, s_size))

        if min_e == best_min:
            best.append(p)
        elif min_e > best_min:
            best = [p]
            best_min = min_e

    return sorted(best)





holes = 5
colors = int(input("Enter Color Count: "))

if colors > 10:
    raise Exception("color count invalid")

correct = tuple([2] * holes)
blank = [0] * holes
S = list(space(holes, 1, colors))
R = list(space(holes, 0, 2))
init_list = {
    1 : (1,1,1,1,1),
    2 : (1,1,1,1,1),
    3 : (1,1,2,2,3),
    4 : (1,1,2,3,4),
    5 : (1,2,3,4,5),
    6 : (1,2,3,4,5),
    7 : (1,2,3,4,5),
}
init = init_list[colors]

k_size = 200
s_size = 1000000

guess = init
remaining_guesses = 4
candidates = S
responses = R


print(f" * Suggested guess: {guess}")
while True:
    remaining_guesses -= 1
    if remaining_guesses == 0:
        print("Game finished.")
        exit(0)

    while True:
        try:
            text = input("Enter response (x to quit): ").strip()
            if text == 'x':
                exit(0)
            response = tuple(map(int, [c for c in text]))
            if len(response) != holes:
                raise Exception
            break
        except Exception as e:
            print("Invalid input")

    print(f'Response {response}')

    print(f'candidates {len(candidates)}, responses {len(responses)}')

    candidates = enter_guess(guess, response, candidates)
    responses = list(filter_responses(response, responses))

    print(f'candidates {len(candidates)}, responses {len(responses)}')

    guess_options = get_best_guess(candidates, responses)
    guess = guess_options[0]
    print(f' * Suggested guess {guess}')

    success_4 = 0
    success_3 = 0
    avg_guesses = 0
    sample_size = min(k_size, len(candidates))
    for key in random.sample(candidates, sample_size):
        #print(f'key - {key}')
        (guesses, final_guess) = solve_game(key, guess, remaining_guesses-1, candidates, responses, False)
        if guesses >= 0:
            success_3 += 1
    for key in random.sample(candidates, sample_size):
        #print(f'key - {key}')
        (guesses, final_guess) = solve_game(key, guess, remaining_guesses, candidates, responses, False)
        if guesses >= 0:
            success_4 += 1
            avg_guesses += guesses
    print(f' - Success rate (3 turns): {(success_3 / sample_size) * 100}%')
    print(f' - Success rate (4 turns): {(success_4 / sample_size) * 100}%')
    print(f' - Expected remaining gueses: {(avg_guesses / sample_size)}')






#for init in inits():
print(f'init - {init}')
print(f'depth 3')
success = 0
avg_guesses = 0
for key in random.sample(S, k_size):
    #print(f'key - {key}')
    (guesses, final_guess) = solve_game(key, init, 3, S, R, False)
    if guesses >= 0:
        success += 1
        avg_guesses += guesses
print(f' - Success rate: {(success / k_size) * 100}%')
print(f' - Expected gueses: {3 - (avg_guesses / k_size)}')

print(f'init - {init}')
print(f'depth 4')
success = 0
avg_guesses = 0
for key in random.sample(S, k_size):
    #print(f'key - {key}')
    (guesses, final_guess) = solve_game(key, init, 4, S, R, False)
    if guesses >= 0:
        success += 1
        avg_guesses += guesses
print(f' - Success rate: {(success / k_size) * 100}%')
print(f' - Expected gueses: {3 - (avg_guesses / k_size)}')


exit(0)







guesses = 0

#init = (1, 1, 1, 2, 2)
#key = tuple(random.randint(1, colors) for i in range(holes))

min_init = []
min_guesses = 1000000

count = len(S)
c = 0

start = time.time()
print(f'ETA: Unknown - {c}/{count}', end='\r')

for init in S:
    c += 1

    guesses = 0
    for key in S:
        guesses_1 = solve_game(key, init)
        guesses += guesses_1
    guesses = guesses / count

    if guesses == min_guesses:
        min_init.append(init)
    elif guesses < min_guesses:
        min_init = [init]
        min_guesses = guesses

    elapsed = time.time() - start
    if c == count:
        ETA = 0
    else:
        ETA = ((count / c) - 1) * elapsed
    ftime = time.strftime("%H:%M:%S", time.gmtime(ETA))
    print(f'ETA: {ftime} - {c}/{count}', end='\r')


print(f'done')

print(f'avg guess: {min_guesses}')
for i in min_init:
    print(f'init - {i}')