"""
A script for finding equal or near-equal partitions in a group.
Do parts a, b, and g
"""
from itertools import combinations
import random
import numpy as np
from matplotlib import pyplot as plt
from pathlib import Path
from progressbar import progressbar as pbar

DIR = Path(__file__).parent

group1 = [10, 13, 23, 6, 20]
group2 = [6, 4, 9, 14, 12, 3, 15, 15]
group3 = [93, 58, 141, 209, 179, 48, 225, 228]
group4 = [2474, 1129, 1388, 3752, 821, 2082, 201, 739]


def random_partition_problem(num_vars, max_power):
    group_vals = [random.randint(1, 2**max_power) for _ in range(num_vars)]
    if sum(group_vals) % 2:
        group_vals[0] += 1
    return group_vals


def find_partition(group_vals):
    num_vars = len(group_vals)
    fullset = {i for i in range(num_vars)}
    subsets = {frozenset([i]) for i in fullset}
    for L in range(1, (num_vars//2)+1):
        subsets = subsets | set(map(frozenset, combinations(fullset, L)))

    min_cost = sum(group_vals)
    for subset in subsets:
        part1 = sum([group_vals[i] for i in subset])
        part2 = sum([group_vals[i] for i in fullset - subset])
        cost = abs(part1 - part2)
        if cost < min_cost:
            min_cost = cost
        if min_cost == 0:
            break
    return min_cost


def frac_perfect(sample_size=100):
    for num_vars in [3, 5, 7, 9, 11]:
        powers = np.arange(1, 2*num_vars)
        ratios = powers / num_vars
        perfect = np.zeros_like(ratios)
        for i, power in enumerate(pbar(powers)):
            perfect[i] = np.mean([find_partition(random_partition_problem(num_vars, power)) == 0
                                  for _ in range(sample_size)]
                                 )
        np.save(f'{DIR}/perfect_{num_vars}vars.npy', np.row_stack((ratios, perfect)))
    return


def p_scaling(x):
    return 1 - np.exp(-np.sqrt(3/(2*np.pi)) * (2**-x))


def p_perfect(ratio, num_vars):
    r_crit = 1
    return 1 - np.exp(-np.sqrt(3/(2*np.pi*num_vars)) * 2**(-num_vars * (ratio - r_crit)))


def sub_x(ratio, num_vars):
    r_crit = 1
    return num_vars * (ratio - r_crit) + 0.5*np.log2(num_vars)


def plot_perfect():
    _, ax1 = plt.subplots()
    ax1.set_xlabel('M/N ratio')
    ax1.set_ylabel('Fraction of perfect partitions')

    _, ax2 = plt.subplots()
    ax2.set_xlabel('x')
    ax2.set_ylabel('p(x)')

    for num_vars in [3, 5, 7, 9, 11]:
        data = np.load(f'{DIR}/perfect_{num_vars}vars.npy')
        ax1.plot(data[0], data[1], label=f'N={num_vars}')
        ax2.plot(sub_x(data[0], num_vars), data[1], label=f'N={num_vars}')
    x = np.linspace(-7.5, 11, 100)
    ax2.plot(x, p_scaling(x), c='k', lw=2, label='scaling')
    ax1.legend()
    ax2.legend()
    plt.show()


if __name__ == '__main__':
    # frac_perfect(1000)
    plot_perfect()

