# Loading the relevant libraries 
import numpy as np  # numpy
import pandas as pd  # manipulation of tabular data
import matplotlib.pylab as plt  # data plotting

# Loading data
df = pd.read_csv('../DL-Assisted-NHA-Inverse-Design-/Dataset 6655.csv')  #load csv in dataframe
df  # display dataframe in notebook

# Loading data
filename = '../DL-Assisted-NHA-Inverse-Design-/Dataset 6655.csv'  # define filename
df = pd.read_csv(filename)  #load csv in dataframe
df  # display dataframe in notebook

# Formatting the data
df['Spectra'] = df.values[:,5:][:,::-1].tolist()  # gather 200 spectra column in one single list
df['Spectra'] = df['Spectra'].apply(np.array)  # list to numpy array
df.drop(df.columns[5:-1], axis=1, inplace=True)  # drop the gathered columns
df.columns = ['Lattice','Material','Thickness','Radius','Pitch','Spectra']  # name all columns
df  # display dataframe in notebook

# Inspecting the data for unique values
df['Lattice'].sort_values().unique()	
df['Material'].sort_values().unique()
df['Thickness'].sort_values().unique()
df['Radius'].sort_values().unique()
df['Pitch'].sort_values().unique()

# Unique values
array([0, 1])  # Lattice
array([0, 1, 2])  # Material
array([100, 105, 110, 115, 120, 125, 130, 135, 140, 145, 150])  # Thickness
array([ 50,  55,  60,  65,  70,  75,  80,  85,  90,  95, 100, 105, 110, 115, 120, 125, 130, 135, 140, 145, 150])  # Radius
array([475, 480, 485, 490, 495, 500, 505, 510, 515, 520, 525])  # Pitch

# Importing auxiliary libraries
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from datetime import datetime

# import pytorch
import torch
import torch.nn as nn
from torch.utils.data import DataLoader


# import tensorboard for logging
from torch.utils.tensorboard import SummaryWriter

# importing local code
import sys
sys.path.append('../')
from datasets import SpectraDataset
from models import MLP, ResidualMLP
from loops import train,val

# filename
filename = './DL-Assisted-NHA-Inverse-Design-/Dataset 6655.csv'

# load in dataframe
df = pd.read_csv(filename)

# clean and rearrange data
df['Spectra'] = df.values[:,5:][:,::-1].tolist()
df['Spectra'] = df['Spectra'].apply(np.array)
df.drop(df.columns[5:-1], axis=1, inplace=True)
df.columns = ['Lattice','Material','Thickness','Radius','Pitch','Spectra']

# select and normalize input features (X)
X_df = df[['Lattice','Material','Thickness','Radius','Pitch']]
X_df = X_df/X_df.max()

# select output labels (y)
y_df = df['Spectra']

# split in training and validation set
test_val_split = 0.1  # portion of data assigned to validation set
X_train, X_val, y_train, y_val = train_test_split(X_df, y_df, test_size=test_val_split, random_state=42)

# batch size
batch_size = 64

# Create data loaders
train_dataloader = DataLoader(training_dataset, batch_size=batch_size, shuffle= True)
val_dataloader = DataLoader(val_dataset, batch_size=batch_size)

for X, y in val_dataloader:
    print(f"Shape of X [N, C]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

Shape of X [N, C]: torch.Size([64, 5])
Shape of y: torch.Size([64, 200]) torch.float32

# pytorch building blocks
import torch.nn as nn

# Simple example of multilayer perceptron
class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()

        # first layer
        self.linear1 = nn.Linear(5, 1000, bias=False)
        self.activation1 = nn.ReLU()

        # second layer
        self.linear2 = nn.Linear(1000, 1000, bias=False)
        self.activation2 = nn.ReLU()

        # third layer
        self.linear3 = nn.Linear(1000, 1000, bias=False)
        self.activation3 = nn.ReLU()

        # fourth layer
        self.linear4 = nn.Linear(1000, 1000, bias=False)
        self.activation4 = nn.ReLU()

        # fifth layer
        self.linear5 = nn.Linear(1000, 1000, bias=False)
        self.activation5 = nn.ReLU()

        # sixth layer
        self.linear6 = nn.Linear(1000, 200, bias=False)

    def forward(self, x):

        # forward pass of the nn
        x = self.linear1(x)
        x = self.activation1(x)
        x = self.linear2(x)
        x = self.activation2(x)
        x = self.linear3(x)
        x = self.activation3(x)
        x = self.linear4(x)
        x = self.activation4(x)
        x = self.linear4(x)
        x = self.activation4(x)
        x = self.linear5(x)
        x = self.activation5(x)
        x = self.linear6(x)

        return x

# Get cpu or gpu device for training.
device = (
    "cuda"
    if torch.cuda.is_available()
    else "cpu"
)
print(f"Using {device} device")

# base learning rate
lr = 1.1e-4

# defining loss and optimizer
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(),lr=lr)

# the training loop
def train(dataloader, model, loss_fn, optimizer, device):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"Train loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

    return loss

# the validation loop
def val(dataloader, model, loss_fn, device):
    num_batches = len(dataloader)
    model.eval()
   val_loss = 0

    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            val_loss += loss_fn(pred, y).item()
    val_loss /= num_batches
    print("\n")
    print(f"Validation loss: {val_loss:>8f} \n")

    return val_loss

# create timestamp
now = datetime.now()  # current date and time
date_time = now.strftime("%d%m%y_%H%M%S")

# create summary writer for tensorboard
writer_path = '../tb_logs/' + date_time + '/'
writer = SummaryWriter(writer_path)

# loop over epochs
epochs = 2500
epoch_threshold = 500
save_checkpoint = './best_model_' + date_time + '.ckpt'
best_loss = 1.0
for epoch in range(epochs):

    # log epoch to console
    print(f"Epoch {epoch+1}\n-------------------------------")

    # performe training and validation loops
    train_loss = train(train_dataloader, model, loss_fn, optimizer, device)
    val_loss = val(val_dataloader, model, loss_fn, device)

    # log training and validation loss to console
    writer.add_scalar("Loss/train", train_loss, epoch)
    writer.add_scalar("Loss/val", val_loss, epoch)

    # save checkpoint
    if (val_loss < best_loss) and (epoch>epoch_threshold):
        
        # save checkpoint
        model.train()
        torch.save(model.state_dict(), save_checkpoint)
        best_loss = val_loss

# close connection to tensorboard
writer.flush()
writer.close()

# finished
print("Done!")

# 1pytorch building blocks
import torch.nn as nn

# a flexible, fully connected, multi headed network
class FlexibleInverseMLP(nn.Module):
    def __init__(
        self,
        hidden_layers=[200, 1000, 1000, 1000, 1000, 1000, 2, 3, 3],
        p=0.2,
        activation=nn.GELU(),
    ):
        super().__init__()

        self.layers = nn.ModuleList()
        for i in range(1, len(hidden_layers) - 3):
            self.layers.append(
                BaseBlock(
                    hidden_layers[i - 1], hidden_layers[i], p=p, activation=activation
                )
            )

        # define 2 classification heads and 1 regression head
        self.lattice_head = LinearBlock(hidden_layers[-4], hidden_layers[-3])
        self.material_head = LinearBlock(hidden_layers[-4], hidden_layers[-2])
        self.geometry_head = LinearBlock(hidden_layers[-4], hidden_layers[-1])

    def forward(self, x):

        # common path
        for layer in self.layers:
            x = layer(x)

        # three heads
        x_lattice = self.lattice_head(x)
        x_material = self.material_head(x)
        x_geometry = self.geometry_head(x)

        return x_lattice,x_material,x_geometry

# the inverse training loop
def train_inverse(dataloader, model, loss_reg, loss_ce, optimizer, device):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred_l, pred_m, pred_g = model(X)
        true_l = y[:,0].type(torch.LongTensor).to(device)
        true_m = y[:,1].type(torch.LongTensor).to(device)
        true_g = y[:,2:].to(device)

        loss = loss_ce(pred_l, true_l) + loss_ce(pred_m, true_m) + 5.0*loss_reg(pred_g, true_g) 

        # Backpropagation
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 20 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"Train loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

    return loss

# the inverse validation loop
def val_inverse(dataloader, model, loss_reg, loss_ce, device):
    num_batches = len(dataloader)
    model.eval()
    val_loss = 0

    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred_l, pred_m, pred_g = model(X)
            true_l = y[:,0].type(torch.LongTensor).to(device)
            true_m = y[:,1].type(torch.LongTensor).to(device)
            true_g = y[:,2:].to(device)
            loss = loss_ce(pred_l, true_l) + loss_ce(pred_m, true_m) + loss_reg(pred_g, true_g) 
            val_loss += loss.item()

    val_loss /= num_batches
    print("\n")
    print(f"Validation loss: {val_loss:>8f} \n")

    return val_loss