import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


def load_data():
    df_orig_train = pd.read_csv('uebungen/aufgabe3/mnist.csv')
    df_digits = df_orig_train.drop('label', axis=1)
    return df_digits.to_numpy()

mnist = load_data()

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))  # Sigmoidfunktion

class RBM:

    def __init__(self, visible_size: int, hidden_size: int, learnrate: float = 0.1):
        """__init__ Initialisiere der RBM.

        Args:
            visible_size (int): anzahl neuronen sichtbare schicht
            hidden_size (int): anzahl neuronen sichtbare schicht
            learnrate (float, optional): learnrate eta in [0;1]. Default als 0.1.
        """
        self.learnrate = learnrate
        self.visible_size = visible_size
        self.hidden_size = hidden_size

        # Initialisieren lernbarer Attribute
        self.weights = np.random.normal(0, 0.01, (self.visible_size, self.hidden_size))
        self.visible_bias = np.zeros(self.visible_size)
        self.hidden_bias = np.zeros(self.hidden_size)

    def activate(self, v0):
        return sigmoid(np.dot(v0, self.weights) + self.hidden_bias)

    def reactivate(self, h0):
        return sigmoid(np.dot(self.weights, h0) + self.visible_bias)

    def contrastive_divergence(self, v0, h0, v1, h1):
        # Gradient
        postive_gradient = np.outer(v0, h0)
        negative_gradient = np.outer(v1, h1)

        # Gewichte per delta anpassen
        self.weights += self.learnrate * (postive_gradient - negative_gradient)

        # Biases per delta anpassen
        self.visible_bias += self.learnrate * (v0 - v1)
        self.hidden_bias += self.learnrate * (h0 - h1)

    def train(self, v0):
        # versteckte schichten aktivieren
        h0 = self.activate(v0)

        # Reaktivieren sichtbarer Schichten
        v1 = self.reactivate(h0)

        # Aktivieren nächster versteckter Schicht
        h1 = self.activate(v1)

        # Gewichte anpassen
        self.contrastive_divergence(v0, h0, v1, h1)

    def run(self, v0):
        # Aktivieren der Schichten
        h0 = self.activate(v0)
        v1 = self.reactivate(h0)

        return h0, v1

training_epochs = 20
rbm = RBM(28 ** 2, 100, 0.02)

for epoch in range(training_epochs):
    np.random.shuffle(mnist)
    for i in range(100, 600): # 500 mnist zahlen zwischen 100 und 600 
        number = mnist[i] / 255
        rbm.train(number)
    if (epoch + 1) % 5 == 0:
        print(f"Epoch {epoch+1}/{training_epochs} completed.")

# Ergebnisse plotten
rows, columns = (9, 9)

fig = plt.figure(figsize=(10, 7))
fig.canvas.manager.set_window_title(
    "Rekonstruktion des MNIST Datensatzes mit einer Restricted Boltzmann Machine")

for i in range((rows * columns)):
    if i % 3 == 0:
        number = mnist[i] / 255
        (hidden, visible) = rbm.run(number)

        results = [hidden.reshape((10, 10)), visible.reshape(
            (28, 28)), number.reshape((28, 28))]

        for j, item in enumerate(results):
            fig.add_subplot(rows, columns, i + j + 1)
            plt.imshow(item, cmap='gray')
            plt.axis('off')

plt.show()