import numpy as np
import matplotlib.pyplot as plt
from numpy import vectorize
from random import shuffle

# 
# pomocné funkce
# 

# výpočet energie
def e(w, s):
    e = 0
    for i in range(0, 3):
        for j in range(0, 3):
            e += w[i,j] * s[i] * s[j]
    return -0.5 * e

# přechodová funkce (vektorizovaná)
@vectorize
def mysign(a):
    if a >= 0:
        return 1
    else:
        return -1

# 
# učení sítě (2 obrázky)
# ~ POZOR: Síť má teoretickou kapacitu pouze 1,2 obrázku!
# 

# vstupní vektor (obrázek) 1
# -1 +1 -1
# +1 -1 +1
# -1 +1 -1
print('\nUČENÍ: obrázek 1')
s_in_1 = (-1, 1, -1, 1, -1, 1, -1, 1, -1)
s_in_1 = np.array(s_in_1)

# matice vah pro 1. obrázek
w_1 = np.outer(s_in_1, s_in_1)
np.fill_diagonal(w_1, 0)
print(w_1)
print('E(s_in_1) =', e(w_1, s_in_1))

plt.matshow(s_in_1.reshape(3, 3), cmap='gray')
plt.show()
#input('?>')

# vstupní vektor (obrázek) 2
# +1 +1 -1
# -1 -1 -1
# +1 +1 -1
print('\nUČENÍ: obrázek 2')
s_in_2 = (1, 1, -1, -1, -1, -1, 1, 1, -1)
s_in_2 = np.array(s_in_2)

# matice vah pro 2. obrázek
w_2 = np.outer(s_in_2, s_in_2)
np.fill_diagonal(w_2, 0)
print(w_2)
print('E(s_in_2) =', e(w_2, s_in_2))

plt.matshow(s_in_2.reshape(3, 3), cmap='gray')
plt.show()
#input('?>')

# matice vah pro oba obrázky
w = (w_1 + w_2) / 2
print('\nUČENÍ: obrázek 1+2')
print(w)
print('E(w, s_in_1) =', e(w, s_in_1))
print('E(w, s_in_2) =', e(w, s_in_2))
#input('?>')

# 
# vybavování
# ~ nalezení naučeného obrázku nejpodobnějšího předloženému vstupu
# 
print()
print('*'*33)

# vybavování ze vstupu 's':
# +1 +1 +1
# +1 +1 -1
# +1 +1 -1
print('\nVYBAVOVÁNÍ: neznámý obrázek')
s = (1, 1, 1, 1, 1, -1, 1, 1, -1)   # OK
#s = (1, 1, 1, 1, 1, 1, 1, -1, -1)   # Cha! ^_^
s = np.array(s)
print('E(s) =', e(w, s))

plt.matshow(s.reshape(3, 3), cmap='gray')
plt.show()
#input('?>')


# iterace (10 kroků "natvrdo")
o = s   # nastavení pro 1. krok
# vlastní iterace od vstupu k zapamatovanému obrázku
i = 1
#while (q := input('?>')) not in ('q', 'Q'):
while True:
    print(f'\nVYBAVOVÁNÍ: krok {i}')
    o = mysign(w @ o)
    print('E(o) =', e(w, o))
    plt.matshow(o.reshape(3, 3), cmap='gray')
    plt.show()
    i += 1