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, 10):
        for j in range(0, 10):
            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ě (5 obrázků)
# ~ POZOR: Síť má teoretickou kapacitu pouze 13 obrázků!
# 

# vstupní vektor (obrázek) 1
print('\nUČENÍ: obrázek 1')
s_in_1 = (
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, +1, -1, -1, -1, +1, -1, -1, -1,
    -1, -1, -1, +1, -1, +1, -1, -1, -1, -1,
    -1, -1, -1, -1, +1, -1, -1, -1, -1, -1,
    -1, -1, -1, +1, -1, +1, -1, -1, -1, -1,
    -1, -1, +1, -1, -1, -1, +1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -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(10, 10), cmap='gray')
plt.show()

# vstupní vektor (obrázek) 2
print('\nUČENÍ: obrázek 2')
s_in_2 = (
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    +1, +1, +1, +1, +1, +1, +1, +1, +1, +1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -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(10, 10), cmap='gray')
plt.show()

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

# matice vah pro 3. obrázek
w_3 = np.outer(s_in_3, s_in_3)
np.fill_diagonal(w_3, 0)
print(w_3)
print('E(s_in_3) =', e(w_3, s_in_3))

plt.matshow(s_in_3.reshape(10, 10), cmap='gray')
plt.show()

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

# matice vah pro 4. obrázek
w_4 = np.outer(s_in_4, s_in_4)
np.fill_diagonal(w_4, 0)
print(w_4)
print('E(s_in_4) =', e(w_4, s_in_4))

plt.matshow(s_in_4.reshape(10, 10), cmap='gray')
plt.show()

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

# matice vah pro 5. obrázek
w_5 = np.outer(s_in_5, s_in_5)
np.fill_diagonal(w_5, 0)
print(w_5)
print('E(s_in_5) =', e(w_5, s_in_5))

plt.matshow(s_in_5.reshape(10, 10), cmap='gray')
plt.show()

# matice vah pro všechny obrázky
w = (w_1 + w_2 + w_3 + w_4 + w_5) / 5
print('\nUČENÍ: obrázek 1+2+3+4+5')
print(w)
print('E(w, s_in_1) =', e(w, s_in_1))
print('E(w, s_in_2) =', e(w, s_in_2))
print('E(w, s_in_3) =', e(w, s_in_3))
print('E(w, s_in_4) =', e(w, s_in_4))
print('E(w, s_in_5) =', e(w, s_in_5))
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':
print('\nVYBAVOVÁNÍ: neznámý obrázek')
s = (
    +1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, +1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, +1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, +1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, +1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, +1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, +1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, +1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, +1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, +1,
)
s = (
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, +1, +1, +1, +1, +1, +1, +1, +1, -1,
    -1, +1, -1, -1, -1, -1, -1, -1, +1, -1,
    -1, +1, -1, -1, -1, -1, -1, -1, +1, -1,
    -1, +1, -1, -1, -1, -1, -1, -1, +1, -1,
    -1, +1, -1, -1, -1, -1, -1, -1, +1, -1,
    -1, +1, -1, -1, -1, -1, -1, -1, +1, -1, 
    -1, +1, -1, -1, -1, -1, -1, -1, +1, -1, 
    -1, +1, +1, +1, +1, +1, +1, +1, +1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
)
s = np.array(s)
print('E(s) =', e(w, s))

plt.matshow(s.reshape(10, 10), 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(10, 10), cmap='gray')
    plt.show()
    i += 1