"""
Pentagonal numbers are generated by the formula, Pn=n(3n−1)/2. 
The first ten pentagonal numbers are:

1, 5, 12, 22, 35, 51, 70, 92, 117, 145, ...

It can be seen that P4 + P7 = 22 + 70 = 92 = P8. 
However, their difference, 70 − 22 = 48, is not pentagonal.

Find the pair of pentagonal numbers, Pj and Pk, 
for which their sum and difference are pentagonal and D = |Pk − Pj| is minimised; 
what is the value of D?
"""

import numpy as np
import matplotlib.pyplot as plt

from garageofcode.common.utils import Heap

pents = [1]

def is_pent(i):
    while i > pents[-1]:
        n = len(pents)
        p_new = pent(n)
        pents.append(p_new)

    return i in pents

def pent(n):
    return int(n * (3*n - 1) / 2)

def pent_diff(nm):
    n, m = nm
    return pent(n) - pent(m)

def main():
    pairs = Heap([(2, 1)], key=lambda x: x[0])

    popped_pairs = np.zeros([100, 100])
    popped_keys = []
    while True:
        key, (i, j) = pairs.pop()
        popped_keys.append(key)
        #if len(popped_keys) > 100000:
        #    print(popped_keys)
        #    break
        '''
        if i >= 100 or j >= 100:
            print(pents)
            print(popped_keys)

            plt.imshow(popped_pairs)
            plt.show()
            break
        elif i >= 100 or j >= 100:
            pass
        else:
            popped_pairs[i, j] = 1
        '''
        #print(key)
        #print(i, j)
        pi, pj = pent(i), pent(j)
        if is_pent(pi + pj) and is_pent(pi - pj):
            print("pi, pj:", pi, pj)
            print("S:", pi + pj)
            print("D:", pi - pj)        
            break

        pairs.push((i + 1, j))
        if i - j == 1:
            pairs.push((i + 1, j + 1))


if __name__ == '__main__':
    #main()

    # The stupid brute fore way worked
    for i in range(1, 1000):
        for j in range(1, i):
            pi, pj = pent(i), pent(j)
            if is_pent((pi + pj) / 2)  and is_pent((pi - pj) / 2):
                print("pi, pj:", pi, pj)
                print("S:", pi + pj)
                print("D:", pi - pj)


    #print(is_pent(92))
    #print(pents)