import numpy as np
import random

letter_C = np.array([
    [1, 1, 1, 1, 1],
    [1, 0, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [1, 1, 1, 1, 1],
])
noisy_C = np.array([
    [1, 1, 1, 1, 1],
    [0, 1, 0, 0, 1],
    [1, 0, 0, 0, 0],
    [1, 0, 0, 1, 0],
    [1, 0, 1, 1, 1],
])

letter_I = np.array([
    [0, 1, 1, 1, 1],
    [0, 0, 1, 0, 0],
    [0, 0, 1, 0, 0],
    [0, 0, 1, 0, 0],
    [1, 1, 1, 1, 1],
])
noisy_I = np.array([
    [1, 1, 1, 1, 1],
    [0, 0, 1, 0, 0],
    [0, 1, 1, 0, 0],
    [0, 0, 0, 0, 0],
    [0, 1, 0, 1, 1],
])

letter_T = np.array([
    [1, 1, 1, 1, 1],
    [0, 0, 1, 0, 0],
    [0, 0, 1, 0, 0],
    [0, 0, 1, 0, 0],
    [0, 0, 1, 0, 0],
])
noisy_T = np.array([
    [1, 1, 0, 1, 0],
    [0, 0, 1, 0, 0],
    [0, 1, 1, 0, 0],
    [0, 0, 0, 0, 0],
    [0, 0, 1, 0, 0],
])


class HopfieldNet:
    def __init__(self, num_neurons, threshold=None):
        assert num_neurons <= 1000
        self.weights = np.zeros((num_neurons, num_neurons)).astype(np.int)
        self.state = np.array((1, num_neurons))
        if threshold:
            self.thresholds = np.array([threshold for _ in num_neurons])
        else:
            self.thresholds = np.zeros((num_neurons,))

    def fit(self, X):
        num_p = X.shape[0]
        num_k = X.shape[1]
        # check right number of pattern
        assert num_p < num_k * 0.138

        num_k = X.shape[1]
        for p in range(X.shape[0]):
            X_p = X[p, :].reshape((1, num_k))
            matrix_lr = np.dot(X_p.T, X_p).astype(np.int)
            np.fill_diagonal(matrix_lr, 0)
            self.weights += matrix_lr

    def predict(self, X, show_energy=False, show_char=False):
        num_k = X.shape[1]
        X_pred = X.copy()

        # loop per every pattern
        for p in range(X_pred.shape[0]):
            differ = True
            time_s = 0

            # loop until the state
            # stay the same
            while differ:
                X_prev = X_pred[p].copy()
                # print energy
                if show_energy:
                    self.print_energy(X_pred[p], p, time_s)
                # print char
                if show_char and num_k <= 100:
                    self.print_char(X_pred[p], p, time_s)

                # loop per every neuron
                for k in range(num_k):
                    val = np.dot(X_pred[p], self.weights[:, k])
                    val_thres = 1 if val > self.thresholds[k] else -1
                    X_pred[p, k] = val_thres

                # check if the new state differs from the previous one
                differ = False if np.array_equal(X_pred[p], X_prev) else True

                time_s += 1

        return X_pred

    def print_energy(self, state, num_p, time_s):
        first_term = 0
        second_term = 0

        for i in range(state.shape[0]):
            for j in range(state.shape[0]):
                first_term += self.weights[i, j] * state[i] * state[j]

        for k in range(state.shape[0]):
            second_term += self.thresholds[k] * state[k]

        energy = -0.5 * first_term + second_term
        print('Pattern: {}\t||\tTime stamp: {}\t||\tEnergy: {:7.0f}'.format(num_p, time_s, energy))
        return energy

    def print_char(self, sequence, num_p, time_s):
        sqrtK = np.sqrt(sequence.shape[0])
        # check if correct sequence
        assert sqrtK % 1 == 0

        print('Pattern: {}\t||\tTime stamp: {}'.format(num_p, time_s))

        for y in range(int(sqrtK)):
            for x in range(int(sqrtK)):
                idx = int(y * sqrtK + x)
                val = '*' if sequence[idx] > 0 else ' '
                print(val, end=' ')
            print('', sep='', end='\n')
        print('', sep='', end='\n')


def test_w_less_101():
    print('\n================')
    print('K < 101')
    print('================\n')
    X = np.array([
        letter_C.flatten(),
        letter_I.flatten(),
        letter_T.flatten(),
    ])
    X = np.where(X > 0, 1, -1)
    net = HopfieldNet(X.shape[1])
    net.fit(X)

    X_test = np.array([
        noisy_C.flatten(),
        noisy_I.flatten(),
        noisy_T.flatten(),
    ])
    X_test = np.where(X_test > 0, 1, -1)

    _ = net.predict(X_test, show_char=True)


def test_w_more_100():
    print('\n================')
    print('K > 100')
    print('================\n')
    num_k = random.randint(101, 1000)
    binary = 2
    X = np.array([
        np.random.randint(binary, size=num_k),
        np.random.randint(binary, size=num_k),
        np.random.randint(binary, size=num_k),
        np.random.randint(binary, size=num_k),
        np.random.randint(binary, size=num_k),
        np.random.randint(binary, size=num_k),
        np.random.randint(binary, size=num_k),
        np.random.randint(binary, size=num_k),
        np.random.randint(binary, size=num_k),
    ])
    X = np.where(X > 0, 1, -1)
    net = HopfieldNet(X.shape[1])
    net.fit(X)

    X_test = np.array([
        np.random.randint(binary, size=num_k),
        np.random.randint(binary, size=num_k),
        np.random.randint(binary, size=num_k),
    ])
    X_test = np.where(X_test > 0, 1, -1)

    _ = net.predict(X_test, show_energy=True)


if __name__ == '__main__':
    test_w_less_101()
    test_w_more_100()

[(4, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (11, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (19, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (26, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (34, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (41, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (146, 'numpy.where', 'np.where', True, 'import numpy as np\n'), (155, 'numpy.where', 'np.where', True, 'import numpy as np\n'), (164, 'random.randint', 'random.randint', False, 'import random\n'), (177, 'numpy.where', 'np.where', True, 'import numpy as np\n'), (186, 'numpy.where', 'np.where', True, 'import numpy as np\n'), (54, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (122, 'numpy.sqrt', 'np.sqrt', True, 'import numpy as np\n'), (56, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (58, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (70, 'numpy.fill_diagonal', 'np.fill_diagonal', True, 'import numpy as np\n'), (167, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (168, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (169, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (170, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (171, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (172, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (173, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (174, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (175, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (182, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (183, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (184, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (53, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (69, 'numpy.dot', 'np.dot', True, 'import numpy as np\n'), (95, 'numpy.dot', 'np.dot', True, 'import numpy as np\n'), (100, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n')]