import pandas as pd
from random import randint, sample

# Generate data
list_list = []
for i in range(20):
    for j in range(20):
        if j == i:
            continue
        temp_list = [20] * 20
        temp_list[i] = 19
        temp_list[j] = 21
        list_list = [temp_list] + list_list

df_dict = {}
for i in range(20):
    df_dict['col_' + str(i)] = [temp_list[i] for temp_list in list_list]

df = pd.DataFrame(df_dict)

# Loop through data finding solutions
rand_rounds = 50
max_on_one_side = 5
min_on_one_side = 2
index_sample = ['col_' + str(i) for i in list(range(20))]
junk = ['left_sum', 'right_sum', 'result']
one_best = df.shape[0]

def split_list(a_list):
    """returns first half and second half of a list"""
    half = len(a_list)//2
    return a_list[:half], a_list[half:]

def get_score(solutions_df, index_left, index_right):
    """Lower scores are better. Score returned is the maximum solutions remaining
    of the three outcomes of a particular weighing."""
    solutions_df['left_sum'] = solutions_df[index_left].sum(axis=1)
    solutions_df['right_sum'] = solutions_df[index_right].sum(axis=1)
    solutions_df['result'] = '=='
    solutions_df.loc[solutions_df['left_sum'] > solutions_df['right_sum'], 'result'] = '>'
    solutions_df.loc[solutions_df['left_sum'] < solutions_df['right_sum'], 'result'] = '<'
    return solutions_df.result.value_counts().max()

# Notes
# After round 1 we should have 125-130
# After round 2 we should have < 44
# After round 3 we should have < 15
# After round 4 we should have < 7
# After round 5 we should have < 3
# After round 6 we should be done

result = {
        'step_one_left':[],
        'step_one_right':[],
        'step_one_path':[],
        'step_one_count':[],
        'step_two_left':[],
        'step_two_right':[],
        'step_two_path':[],
        'step_two_count':[],
        'step_three_left':[],
        'step_three_right':[],
        'step_three_path':[],
        'step_three_count':[],
        'step_four_left':[],
        'step_four_right':[],
        'step_four_path':[],
        'step_four_count':[],
        'step_five_left':[],
        'step_five_right':[],
        'step_five_path':[],
        'step_five_count':[],
        'step_six_left':[],
        'step_six_right':[],
        'step_six_path':[],
        'step_six_count':[],
        }
results_df = pd.DataFrame(result)

for rand in range(rand_rounds):
    n = randint(min_on_one_side, max_on_one_side)
    index_left, index_right = split_list(sample(index_sample, n*2))
    new_score = get_score(df, index_left, index_right)
    if new_score < one_best:
        one_left, one_right = index_left, index_right
        print("Round 1: "+str(one_best)+" to "+str(new_score)+" after "+str(rand)+" rounds.")
        one_best = new_score
        one_df_final = df.copy()
for i in ['==','<','>']:
    print("Beginning round two optimization for case: Left " + i + " Right.")
    two_df = one_df_final.copy()
    two_df = two_df.loc[two_df.result==i].reset_index(drop=True)
    two_best = two_df.shape[0]
    for rand in range(rand_rounds):
        if i != '==':
            break
        n = randint(min_on_one_side, max_on_one_side)
        index_left, index_right = split_list(sample(index_sample, n*2))
        new_score = get_score(two_df, index_left, index_right)
        if new_score < two_best:
            two_left, two_right = index_left, index_right
            print("Round 2: "+str(two_best)+" to "+str(new_score)+" after "+str(rand)+" rounds.")
            two_best = new_score
            two_df_final = two_df.copy()
    for j in ['==','<','>']:
        print("Beginning round three optimization for case: Left " + j + " Right.")
        three_df = two_df_final.copy()
        three_df = three_df.loc[three_df.result==j].reset_index(drop=True)
        three_best = three_df.shape[0]
        for rand in range(rand_rounds):
            if j != '==':
                    break
            n = randint(min_on_one_side, max_on_one_side)
            index_left, index_right = split_list(sample(index_sample, n*2))
            new_score = get_score(three_df, index_left, index_right)
            if new_score < three_best:
                three_left, three_right = index_left, index_right
                print("Round 3: "+str(three_best)+" to "+str(new_score)+" after "+str(rand)+" rounds.")
                three_best = new_score
                three_df_final = three_df.copy()
        for k in ['==','<','>']:
            print("Beginning round four optimization for case: Left " + k + " Right.")
            four_df = three_df_final.copy()
            four_df = four_df.loc[four_df.result==k].reset_index(drop=True)
            four_best = four_df.shape[0]
            for rand in range(rand_rounds):
                if k != '==':
                    break
                n = randint(min_on_one_side, max_on_one_side)
                index_left, index_right = split_list(sample(index_sample, n*2))
                new_score = get_score(four_df, index_left, index_right)
                if new_score < four_best:
                    four_left, four_right = index_left, index_right
                    print("Round 4: "+str(four_best)+" to "+str(new_score)+" after "+str(rand)+" rounds.")
                    four_best = new_score
                    four_df_final = four_df.copy()
            for l in ['==','<','>']:
                print("Beginning round five optimization for case: Left " + l + " Right.")
                five_df = four_df_final.copy()
                five_df = five_df.loc[five_df.result==l].reset_index(drop=True)
                five_best = five_df.shape[0]
                for rand in range(rand_rounds):
                    if l != '==':
                        break
                    n = randint(min_on_one_side, max_on_one_side)
                    index_left, index_right = split_list(sample(index_sample, n*2))
                    new_score = get_score(five_df, index_left, index_right)
                    if new_score < five_best:
                        five_left, five_right = index_left, index_right
                        print("Round 5: "+str(five_best)+" to "+str(new_score)+" after "+str(rand)+" rounds.")
                        five_best = new_score
                        five_df_final = five_df.copy()
                for m in ['==','<','>']:
                    print("Beginning round six optimization for case: Left " + m + " Right.")
                    six_df = five_df_final.copy()
                    six_df = six_df.loc[six_df.result==m].reset_index(drop=True)
                    six_best = six_df.shape[0]
                    for rand in range(rand_rounds):
                        if m != '==':
                            break
                        n = randint(min_on_one_side, max_on_one_side)
                        index_left, index_right = split_list(sample(index_sample, n*2))
                        new_score = get_score(six_df, index_left, index_right)
                        if new_score < six_best:
                            six_left, six_right = index_left, index_right
                            print("Round 6: "+str(six_best)+" to "+str(new_score)+" after "+str(rand)+" rounds.")
                            six_best = new_score
                            six_df_final = six_df.copy()
                    for n in ['==','<','>']:
                        final_df = six_df_final.copy()
                        final_df = final_df.loc[final_df.result==n].reset_index(drop=True)
                        result = {
                                'step_one_left':one_left,
                                'step_one_right':one_right,
                                'step_one_path':"Left " + i + " Right",
                                'step_one_count':two_df.shape[0],
                                'step_two_left':two_left,
                                'step_two_right':two_right,
                                'step_two_path':"Left " + j + " Right",
                                'step_two_count':three_df.shape[0],
                                'step_three_left':three_left,
                                'step_three_right':three_right,
                                'step_three_path':"Left " + k + " Right",
                                'step_three_count':four_df.shape[0],
                                'step_four_left':four_left,
                                'step_four_right':four_right,
                                'step_four_path':"Left " + l + " Right",
                                'step_four_count':five_df.shape[0],
                                'step_five_left':five_left,
                                'step_five_right':five_right,
                                'step_five_path':"Left " + m + " Right",
                                'step_five_count':six_df.shape[0],
                                'step_six_left':six_left,
                                'step_six_right':six_right,
                                'step_six_path':"Left " + n + " Right",
                                'step_six_count':final_df.shape[0],
                                }
                        results_df = results_df.append(result, ignore_index=True)

results_df.to_csv('brute_force_results.csv', index=False)