%matplotlib inline
import matplotlib.pyplot as plt # für plotting, wie gewohnt

import numpy as np              # für numerische Aktionen mit Arrays, wie gewohnt

import sys                      # um System-Befehle ausführen zu können

from IPython.display import Image  # um Bilder von einer URL anzuzeigen

# hier die Funktionen für den Datensatz 

from sklearn.datasets import make_moons, make_circles # zur Erzeugung von Datensets

# die Pytorch- und Lightning-spezifischen Imports kommen weiter unten

import torch                      # Die Package PyTorch selbst

import torch.nn as nn             # Ein Teil nochmal extra, als Abkürzung, für Teile von neuronalen Netzen

from torch import optim           # Das Modul für die Optimierungs-Algorithmen, auch extra nochmal

from torchmetrics.functional import accuracy    # Die Funktion zur Berechnung der Accuracy

from torch.utils.data import Dataset, DataLoader, random_split  # einige Tols für die Datenaufbereitung

import pytorch_lightning as pl    # Und die Abstraction PyTorch Lightning

# So wird unten die Zelle für Aufruf und Training aussehen. 
# ACHTUNG: NOCH NICHT AUSFÜHREN, das machen wir später, aber hier 
# wird erst einmal klar, was wir noch brauchen

# zunächst holen wir uns von PyTorch Lightning eine Instanz der Trainer-Python-Klasse
# Die maximale Epochenzahl ist wichtig, sonst läuft der Trainer erstmal unbegrenzt
trainer = pl.Trainer(max_epochs=100)

# Als nächstes erzeugen wir eine Instanz unserer eigenen Netzwerk-Python-Klasse,
# die wir noch schreiben müssen. Das klingt zunächst nach Stress, ist 
# aber sehr geradlinig, wie Sie gleich sehen werden
# Dieser Schritt ist im Prinzip analog zum Modell-Aufruf in Scikit-Learn, wie letztens
model = OurNetwork(Einpaarinputs)

# Dann rufen wir das Training auf, auch das ist analog zur vergangenen Einheit
trainer.fit(model)

# Und schließlich rufen wir den Test auf, das ist eine kleine Abkürzung im Vergleich zum
# vergangenen Setup mit Scikit-Learn, aber grundsätzlich auch das gleiche
trainer.test(model)

# das ist grundsätzlich alles.

# Zunächst eine einfache Klasse für unseren Datensatz. Das hat den Sinn, dass
# die Daten von PyTorch Lightning einfacher verarbeitet werden können, egal,
# auf welcher Hardware und in welcher Konfiguration
class OurData(Dataset):
    
    # In der Initialisierung weisen wir einfach unsere Daten den Inputs und Outputs zu
    def __init__(self):

        # Hier die Inputs. "self" bezieht sich dabei immer auf die Instanz
        self.data_X = input_data
        
        # Hier die Labels. 
        self.data_y = label_data
        
        # Schreiben wir die Länge der geladenen Daten heraus
        print("Successfully created training data of length: ", len(self.data_X))

    def __len__(self):
        # Diese Methode definiert, was ausgegeben wird, wenn nach der Länge des Datensatzes gefragt wird
        return len(self.data_X)

    def __getitem__(self, idx):
        # Diese Methode definiert, wie man das nächste Element des Datensatzes erhält
        return self.data_X[idx], self.data_y[idx]
    

# Nun die für unsere Zwecke recht umfangreiche Klasse für das KNN
# Hier kommen einige Methoden, die die Funktionen beim Training und weiteren Schritten
# vorbereiten und dann im Detail beschreiben und definieren
# Klassennamen verwenden üblicherweise Großbuchstaben am Anfang von Wortteilen
class OurNetwork(pl.LightningModule):
    """
    Einfaches Pytorch-Lightning-Modul, das aus einem x,y-Koordinatenpaar eine Klasse
    für einen Datensatz mit 2 Mond-Punktwolken vorhersagt
    """

    # diese Methode initialisiert wieder die Klasse
    # "self" steht dabei wieder für die Instanz, wenn diese einmal erstellt ist
    # alle mit self. referenzierten Variablen können auch problemlos innerhalb
    # der Klasse verwendet werden
    
    # wir verwenden hier zwei Variablen, die beim Instanz-Erzeugen übergeben werden
    def __init__(self, batch_size=100, hidden_dim=16):
        
        # diese Klasse basiert auf einer anderen ("LightningModule")
        # hier laden wir deren init Methode, damit erbt die Klasse auch
        # alles, was die übergeordnete (parent) Klasse kann
        super().__init__()
        
        # hier definieren wir die Batchsize aus der Eingabe-Variablen
        self.batch_size = batch_size
        
        # Hier definieren wir nun die Layers zur Verwendung innerhalb der Instanz
        # nn.Linear ist ein fully-connected Layer, das in_features Inputs und
        # out_features Outputs hat
        # Hier ist das Input-Layer bzw. die Parameter-Matrix vom Input- zum ersten hidden Layer
        self.fc_1 = nn.Linear(in_features=2, out_features=hidden_dim)

        # Hier ist die Parameter-Matrix vom ersten zum zweiten hidden Layer
        self.fc_2 = nn.Linear(in_features=hidden_dim, out_features=hidden_dim)

        # Hier ist die Parameter-Matrix vom zweiten zum dritten hidden Layer, 
        # falls wir das später einführen und verwenden wollen
        self.fc_3 = nn.Linear(in_features=hidden_dim, out_features=hidden_dim)

        # Und hier die Parameter-Matrix vom letzten hidden Layer zum Output Layer
        # Das Output Layer hat zwei outputs, die den Wahrscheinlichkeiten für
        # die beiden Klassen entsprechen
        self.fc_4 = nn.Linear(in_features=hidden_dim, out_features=2)

        # Hier die Loss-Funktion. Wir wählen Cross-Entropy, eine gute 
        # Möglichkeit für ein Klassifikationsproblem
        self.loss = nn.CrossEntropyLoss()

        # Noch ein Demo-Input (das muss nicht unbedingt sein, aber es ist praktisch, 
        # falls man etwas testen will)
        self.example_input_array = torch.zeros(self.batch_size, 2)
        
        # Hier starten wir noch Listen fürs Plotten hinterher
        # Einmal die Loss-Funktion auf dem Training-Set
        self.train_loss = []
        
        # Und die Loss-Funktion auf dem Validierungs-Set
        self.val_loss = []
        
        # Hier endet die Init-Funktion der Klasse
        return

        
    # Jetzt kommt die zentrale Methode der Klasse, in der definiert wird,
    # wie ein Datenpunkt (genauer gesagt, ein Batch von Datenpunkten)
    # durch das Netz von vorne bis hinten durchgereicht wird
    
    # Hier können Sie beliebige Netzwerke konstruieren, mit allem, was PyTorch
    # so an Layern und anderen Dingen bietet. Wir starten einmal simpel: 
    # Nur ein paar fully-connected Layers und Aktivierungsfunktionen dazwischen, sonst nichts
    
    # Input-Vairable hier ist x, das für einen Batch (de facto eine Liste) von Inputs steht
    def forward(self, x):
        """
        Hier wird ein Dateninput durch das Netz geschleust
        """
        # Am Anfang müssen wir (leider) dafür sorgen, dass das Datenformat passt
        # Der Standard-Container für Daten in PyTorch ist ein sogenannter "Tensor"
        # Dieser speichert Werte an verschiedenen Stellen im KNN, aber gleichzeitig
        # auch noch die Gradienten an diesen Stellen, sodass beim Optimieren einfach
        # und schnell darauf zugegriffen werden kann
        x = torch.tensor(x).float()
        
        # x geht von Input zu hidden Layer 1
        x = self.fc_1(x)
        
        # x geht durch eine Sigmoid-Funktion
        x = torch.sigmoid(x)

        # x geht von hidden Layer 1 zu hidden Layer 2
        x = self.fc_2(x)
        
        # x geht durch eine Sigmoid-Funktion
        # x = torch.sigmoid(x)

        # x geht von hidden Layer 2 zu hidden Layer 3 (für später, falls gewünscht)
        # x = self.fc_3(x)
        
        # x geht durch eine Sigmoid-Funktion
        x = torch.sigmoid(x)

        # x geht ins Output Layer
        x = self.fc_4(x)
        
        # Die outputs werden zurückgegeben
        return x
    

    # Als nächstes wird definiert, was in einem Trainingsschritt alles passieren soll
    # Ein Trainingsschritt bedeutet hier, dass ein Batch von Daten innerhalb einer
    # Epoche zum Training verwendet wird
    
    # Der Batch und ein zugehöriger Index sind daher auch Inputs dieser Methode, die
    # man verwenden kann (aber nicht muss)
    def training_step(self, batch, batch_idx):
        
        # zunächst nehmen wir den Batch in die Inputs und Labels auseinander wie gewohnt
        X, y = batch
        
        # Hier kommt der sogenannte "forward pass", d.h. wir rufen "self" auf, was im
        # konkreten Fall eines Pytorch-Lightning oder PyTorch Moduls de facto die 
        # "forward"-Methode aufruft
        y_hat = self(X)
        
        # Wir berechnen die Loss-Funktion für diese Vorhersagen im Vergleich zu
        # den echten Labels
        loss = self.loss(y_hat, y)
 
        # und hier kommt ein Befehl, der die laufende Ausgabe während des Trainings
        # und das "Logging", also die Aufzeichnungen der Trainingsfortschritts
        # steuert. Logging ist sehr mächtig und kann vielfältig eingesetzt werden.
        # Z.B. kann es die Plots ersetzen, die wir hier nachher mit der Hand aus 
        # unseren Listen erzeugen werden.
        # Wir begnügen uns hier aber damit, einfach die laufende Anzeige zu verfolgen.
        self.log('train_loss', loss, on_step=False, on_epoch=True, prog_bar=True) 
    
        # hier hängen wir den Wert noch an eine unserer Listen an. Dazu müssen wir 
        # zunächst den Zahlenwert mit "detach" aus dem Tensor holen und danach in
        # ein Numpy-Array verwandeln, damit wir es nachher einfach verwenden können.
        self.train_loss.append(loss.detach().numpy())
    
        # Zurückgegeben wird hier der Wert des Losses, denn daran orientiert sich
        # der Optimierungsprozess, der hier fast komplett "unter der Motorhaube" läuft.
        return loss

    
    # Als nächstes wird definiert, was in einem Validierungsschritt alles passieren soll
    # Ein Validierungsschritt bedeutet hier, dass ein Batch von Daten innerhalb einer
    # Epoche aus dem Validierungsset verwendet wird
    
    # Diese Methode ist analog zu verwenden und zu schreiben wie der Trainingsschritt
    def validation_step(self, batch, batch_idx):

        X, y = batch
        
        y_hat = self(X)
        
        loss = self.loss(y_hat, y)
 
        # hier hängen wir den Loss-Wert an die entsprechende Liste für die Validierung an
        self.val_loss.append(loss.detach().numpy())
    
        self.log('val_loss', loss, on_step=False, on_epoch=True, prog_bar=True) 
        
        return loss
    

    # Als nächstes wird definiert, was in einem Testschritt alles passieren soll
    # Ein Testschritt bedeutet hier, dass ein Batch von Daten nach dem fertigen 
    # Training aus dem Testset verwendet wird
    
    # Diese Methode ist wieder analog zu verwenden und zu schreiben wie der Trainingsschritt
    def test_step(self, batch, batch_idx):

        X, y = batch
        
        y_hat = self(X)
        
        loss = self.loss(y_hat, y)
        
        # hier berechnen wir zusätzlich noch die Accuracy, so wie wir das auch in der
        # vergangenen Einheit bereits gemacht haben. Diesmal verwenden wir die Funktion,
        # wie Sie von PyTorch bereit gestellt wird
        acc = accuracy(y_hat, y)
 
        # Für den Log stellen wir diesmal zwei Variablen in einem Dictionary bereit
        metrics = {"test_acc": acc, "test_loss": loss}
    
        # und wir loggen dieses Dictionary
        self.log_dict(metrics)
        
        # außerdem wird das Dictionary zurückgegeben
        return metrics   
    

    # Jetzt kommt die (in PyTorch Lightning sehr kurze) Definition des zu verwendenden
    # Optimierungs-Algorithmus
    def configure_optimizers(self):

        # ein üblicher Algorithmus für Gradient Descent beim Deep Learning ist "Adam"
        # Hier wird übrigens die Schrittweite (der oben bereits erwähnte Hyperparameter
        # der "learning rate") definiert
        # Das erste Argument im Adam-Aufruf sind die zu optimierenden Parameter, hier
        # alle Parameter im Netz
        optimizer = optim.Adam(self.parameters(), lr=0.001)

        # Der eingerichtete Optimierungs-Algorithmus wird zurückgegeben
        return optimizer


    # Jetzt kommt die (optionale) Einrichtung der Daten nach bestimmten Ideen oder
    # Kriterien am Anfang des Trainings.
    # Wir nutzen dies, um die Aufteilung in Training, Validation und Test vorzunehmen
    
    # als Argument gibt es hier "stage", das wir aber nicht verwenden
    def setup(self, stage):
        
        # Zunächst erzeugen wir eine Instanz unserer Datensatz-Klasse von oben
        # Hier befinden sich nun unsere vorher erzeugten Daten, aber in leicht
        # zugänglichem Format 
        all_data = OurData()
        
        # Wir fragen die Länge des Datensatzes ab und schreiben sie raus
        data_length = len(all_data)
        print("Länge der gesamten Daten: ", data_length)
        
        # Als nächstes berechnen wir die Längen von Validierungs- und Test-Set
        # Für die Validierung nehmen wir 20% der Daten
        self.validation_number = int(data_length * 0.2)
        
        # Und auch zum Testen nehmen wir 20% der Daten
        self.test_number = int(data_length * 0.2)
        
        # Den Rest nehmen wir für das Training
        self.training_number = data_length - self.test_number - self.validation_number
        
        # die Ausgabe dieser Zahlen, zur Kontrolle
        print("Daten-Partitionierung:", self.training_number, self.validation_number, self.test_number, data_length)

        # Und hier wird der Datensatz dann aufgesplittet
        train_part, val_part, test_part = random_split(all_data, [self.training_number, self.validation_number, self.test_number])

        # Zum Schluss machen wir noch diese Teile des Datensatzes für die Instanz verfügbar
        self.train_dataset = train_part
        self.val_dataset = val_part
        self.test_dataset = test_part
        
        return
    
     
    # Hier definieren wir noch drei sogenannte "Dataloader". Das hat nur den Zweck, die
    # einzelnen Teile des Datensatzes den richtigen Prozessen beim Training, Validierung
    # und Testen zuzuordnen. Außerdem wird hier die Batchsize für die verschiedenen 
    # Prozesse eingestellt (die könnten auch verschieden sein).
    # Der Parameter "num_workers" ist etwas problematisch. Ich muss ihn auf dem Mac auf 0
    # setzen, dann läuft das Training gut und auch schnell. Auf anderer Hardware kann
    # damit eingestellt werden, wie viele Prozesse mit der Datenvorbereitung z.B. für die 
    # GPU beschäftigt sein sollen.
    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, num_workers=0)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=self.batch_size, num_workers=0)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size, num_workers=0)