""" From
https://www.learnpytorch.io/02_pytorch_classification/
8. Putting things together by building a multi-class PyTorch mode
@author: 
 - developed into 5-layer double-headed with GELU by Ian Short

"""

from sklearn.datasets import make_circles
from sklearn.model_selection import train_test_split
import pandas as pd
import matplotlib.pyplot as plt
# Standard PyTorch imports
import torch
from torch import nn
import requests
from pathlib import Path
import torchvision
from torchmetrics import Accuracy
from torch.utils.data import DataLoader, TensorDataset

import MKClassModel

import numpy as np
import h5py
from sklearn.metrics import classification_report


from pathlib import Path

# Create device agnostic code
device = "cuda" if torch.cuda.is_available() else "cpu"
print("device = ", device)

# Calculate accuracy (a classification metric)
def accuracy_fn(y_true, y_pred):
    correct = torch.eq(y_true, y_pred).sum().item() # torch.eq() calculates where two tensors are equal
    acc = (correct / len(y_pred)) * 100 
    return acc

plt.ion()

"""
# 1. Create multi-class data
X_blob, y_blob = make_blobs(n_samples=1000,
    n_features=NUM_FEATURES, # X features
    centers=NUM_CLASSES, # y labels 
    cluster_std=1.5, # give the clusters a little shake up (try changing this to 1.0, the default)
    random_state=RANDOM_SEED
)

# 2. Turn data into tensors
X_blob = torch.from_numpy(X_blob).type(torch.float)
y_blob = torch.from_numpy(y_blob).type(torch.LongTensor)
print(X_blob[:5], y_blob[:5])

# 3. Split into train and test sets
X_blob_train, X_blob_test, y_blob_train, y_blob_test = train_test_split(X_blob,
    y_blob,
    test_size=0.2,
    random_state=RANDOM_SEED
)

# 4. Plot data
plt.figure(figsize=(10, 7))
plt.scatter(X_blob[:, 0], X_blob[:, 1], c=y_blob, cmap=plt.cm.RdYlBu);
"""

class StellarData:
    """A simple 'Bunch' object to mimic sklearn's output"""
    def __init__(self, data, target, target1, target2, target_names, target1_names, target2_names):
        self.data = data
        self.target = target
        self.target_names = target_names
        #For 2D classification:
        #Spectral classes
        self.target1 = target1
        self.target1_names = target1_names
        #Luminosity classes
        self.target2 = target2
        self.target2_names = target2_names       

def load_stellar_spectra(file_path):
    with h5py.File(file_path, 'r') as f:
        # Load everything into RAM (if it fits)
        X = f['flux'][:]
        y = f['labels'][:]
        target_names = [n.decode('utf-8') for n in f['class_names'][:]]
        #For 2D classification:
        y1 = f['spclass'][:]
        target1_names = [n.decode('utf-8') for n in f['spclass_names'][:]]
        y2 = f['lumclass'][:]
        target2_names = [n.decode('utf-8') for n in f['lclass_names'][:]]
        
    return StellarData(X, y, y1, y2, target_names, target1_names, target2_names)

trainsetFile = 'Libr18Samples3002D.h5'
testsetFile = 'Libr18_27Samples2D.h5'
trainset = load_stellar_spectra(trainsetFile)
testset = load_stellar_spectra(testsetFile)

#Training set is libr18 of Gray & Corbally:
X_train = trainset.data
y1_train, y2_train = trainset.target1, trainset.target2
trainnames = trainset.target_names
spectrainnames, lumtrainnames = trainset.target1_names, trainset.target2_names
#print("trainnames ", trainnames)
#print("spectrainnames ", spectrainnames, " lumtrainnames ", lumtrainnames)
numSpecNames = len(spectrainnames)
numLumNames = len(lumtrainnames)


#Test set is libr18_27 of Gray & Corbally:
X_test = testset.data
y1_test, y2_test = testset.target1, testset.target2
testnames = testset.target_names
spectestnames, lumtestnames = testset.target1_names, testset.target2_names
#print("testnames ", testnames)
#print("spectestnames ", spectestnames, " lumtestnames ", lumtestnames)
numSpecTestNames = len(spectestnames)
numLumTestNames = len(lumtestnames)

n_inputs = len(X_train[1])     # 1600 - 30 - 1 pixels

"""
# One-hot encode the labels: e.g., 3 -> [0,0,0,1,0,0,0,0,0,0]
#For double-headed 2D classification:
#Spectral class lables:
Y1_train = np.eye(numSpecNames)[y1_train]
Y1_test = np.eye(numSpecTestNames)[y1_test]
#Luminosity class lables:
Y2_train = np.eye(numLumNames)[y2_train]
Y2_test = np.eye(numLumTestNames)[y2_test]   
"""

#print("len(y_train) ", len(y_train), " len(X_train[0]) ", len(X_train[0]))
numSamples = len(y1_train)
print("numSamples ", numSamples)

"""
#Truncate for testing:

X_train = X_train[:3660]
X_test = X_test[:291]
y1_train = y1_train[:3660]
y1_test = y1_test[:291]
y2_train = y2_train[:3660]
y2_test = y2_test[:291]
"""

# Turn data into tensors

#Spectra
X_train = torch.from_numpy(X_train).type(torch.float)
X_test = torch.from_numpy(X_test).type(torch.float)

"""
#One-hot encoded labels
Y1_train = torch.from_numpy(Y1_train).type(torch.LongTensor)
Y1_test = torch.from_numpy(Y1_test).type(torch.LongTensor)
Y2_train = torch.from_numpy(Y2_train).type(torch.LongTensor)
Y2_test = torch.from_numpy(Y2_test).type(torch.LongTensor)
"""
y1_train = torch.from_numpy(y1_train).type(torch.LongTensor)
y1_test = torch.from_numpy(y1_test).type(torch.LongTensor)
y2_train = torch.from_numpy(y2_train).type(torch.LongTensor)
y2_test = torch.from_numpy(y2_test).type(torch.LongTensor)

#print(f"X_train[:5], {X_train[:5]}, Y1_train[:5], {Y1_train}")
#print(f"X_test[:5], {X_test[:5]}, Y1_test[:5], {Y1_test}")
#print(f"X_train[:5], {X_train[:5]}, y1_train[:5], {y1_train}")
#print(f"X_test[:5], {X_test[:5]}, y1_test[:5], {y1_test}")

print(" ************************ ")
print("")
print(" CAUTION!!   CAUTION!!  CAUTION!!")
print("")
print("numSpecNames = ", numSpecNames, " numSpecTestNames = ", numSpecTestNames)
print("numLumNames = ", numLumNames, " numLumTestNames = ", numLumTestNames)
print("")
print(" ************************ ")

# Set the hyperparameters for data creation
#NUM_CLASSES = 31  #Number of MK spectral classes
#NUM_FEATURES = 1600 #Number of pixels in classifiction spectra
RANDOM_SEED = 42
n_hidden1 = 768
n_hidden2 = 256
n_hidden3 = 64    
n_spec_outputs = numSpecNames
n_lum_outputs = numLumNames
print("n_spec_outputs ", n_spec_outputs, " n_lum_outputs ", n_lum_outputs)

learning_rate = 0.005
#learning_rate = 0.01
n_epochs = 200  #production
#n_epochs = 10  #quick testing
batch_size = int(numSamples/122)
print("numSamples ", numSamples, " batch_size ", batch_size)

# Create a DataLoader to get the mini-batch behavior from your NumPy code
#train_loader = DataLoader(TensorDataset(X_train, y1_train), batch_size=batch_size, shuffle=True)
#Double-headed:
#print("shape, dtype, y1_train: ", y1_train.shape, y1_train.dtype, " y2_train: ", y2_train.shape, y2_train.dtype)
#print("y1_train ", y1_train[:10], " y2_train ", y2_train[:10])
train_loader = DataLoader(TensorDataset(X_train, y1_train, y2_train), batch_size=batch_size, shuffle=True, num_workers=0)

outputFile = trainsetFile + "PyT_" + str(n_hidden1).strip()\
         + "_" + str(n_hidden2).strip() + "_" + str(n_hidden3).strip()\
         + "_" + str(n_epochs).strip() + "_" + str(batch_size).strip()\
                + ".out"
    
handle = open(outputFile, 'w', encoding='ascii')

handle.write("Training set: " + trainsetFile + " Test set: " + testsetFile + "\n")
handle.write("n_hidden1 " + str(n_hidden1) + " n_hidden2 " + str(n_hidden2)\
                 + " n_hidden3 " + str(n_hidden3) + " n_epochs " + str(n_epochs)\
                 + " batch_size " + str(batch_size) + "\n")
handle.write("\n")
handle.write("Epoch      loss      accuracy \n")

loss1V = [0.0 for i in range(n_epochs)]
acc1V = [0.0 for i in range(n_epochs)]
loss2V = [0.0 for i in range(n_epochs)]
acc2V = [0.0 for i in range(n_epochs)]
    
#For 2D classification:
lossTotV = [0.0 for i in range(n_epochs)]


# Gaussian Initialization
#Replace PyTorch default initialization Kaiming Uniform or Glorot.
def init_weights(m):
    # Check if the module 'm' being visited is a Linear layer
    if isinstance(m, nn.Linear):
        # Set weights to NumPy standard: rng.normal(0, 0.1)
        torch.nn.init.normal_(m.weight, mean=0.0, std=0.1)
        # ??? torch.nn-biases.
        # Set biases to your NumPy standard: np.zeros
        nn.init.constant_(m.bias, 0)
        

#8.2 Building a multi-class classification model in PyTorch

"""
# Build model
class MKClassModel(nn.Module):
    def __init__(self, input_features, output_features, hidden_units1, hidden_units2):
        #Initializes all required hyperparameters for a multi-class classification model.
        #
        #Args:
        #    input_features (int): Number of input features to the model.
        #    out_features (int): Number of output features of the model
        #      (how many classes there are).
        #    hidden_units (int): Number of hidden units between layers, default 8.
        
        super().__init__()
        self.linear_layer_stack = nn.Sequential(
            nn.Linear(in_features=input_features, out_features=hidden_units1),
            # nn.ReLU(), # <- does our dataset require non-linear layers? (try uncommenting and see if the results change)
            nn.Linear(in_features=hidden_units1, out_features=hidden_units2),
            # nn.ReLU(), # <- does our dataset require non-linear layers? (try uncommenting and see if the results change)
            nn.Linear(in_features=hidden_units2, out_features=output_features), # how many classes are there?
        )
    
    def forward(self, x):
        return self.linear_layer_stack(x)

# Create an instance of model and send it to the target device
model_4 = MKClassModel(input_features=n_inputs, 
                    output_features=n_spec_outputs, 
                    hidden_units1=n_hidden1,
                    hidden_units2=n_hidden2).to(device)
"""
"""
# Build model with non-linear activation function
class MKClassModel(nn.Module):
    def __init__(self, input_features, output1_features, output2_features, hidden1_units, hidden2_units, hidden3_units):
        super().__init__()
        self.input_features = input_features
        self.output1_features = output1_features
        self.output2_features = output2_features
        self.hidden1_units = hidden1_units
        self.hidden2_units = hidden2_units
        self.hidden3_units = hidden3_units
        #Shared layers:
        self.layer_1 = nn.Linear(in_features=self.input_features, out_features=self.hidden1_units)
        self.layer_2 = nn.Linear(in_features=self.hidden1_units, out_features=self.hidden2_units)
        self.layer_3 = nn.Linear(in_features=self.hidden2_units, out_features=self.hidden3_units)
        #self.layer_4 = nn.Linear(in_features=self.hidden_units3, out_features=self.output_features)
        #The two heads:
        self.head_1_4 = nn.Linear(in_features=self.hidden3_units, out_features=self.output1_features)
        self.head_2_4 = nn.Linear(in_features=self.hidden3_units, out_features=self.output2_features)
        #self.relu = nn.ReLU() # <- add in ReLU activation function
        self.gelu = nn.GELU(approximate='tanh')
        # Can also put sigmoid in the model 
        # This would mean you don't need to use it on the predictions
        # self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # Intersperse the ReLU activation function between layers
        #return self.layer_4(self.relu(self.layer_3(self.relu(self.layer_2(self.relu(self.layer_1(x)))))))
        #return self.layer_4(self.gelu(self.layer_3(self.gelu(self.layer_2(self.gelu(self.layer_1(x)))))))
        x = self.gelu(self.layer_3(self.gelu(self.layer_2(self.gelu(self.layer_1(x))))))
        #Double-headed 2D classification:
        logits1 = self.head_1_4(x)
        logits2 = self.head_2_4(x)
        return logits1, logits2
"""

model_4 = MKClassModel.MKClassModel(n_inputs, n_spec_outputs, n_lum_outputs, n_hidden1, n_hidden2, n_hidden3).to(device)

# 3. Use the 'apply' method to run the function on every layer inside model_4
model_4.apply(init_weights)

#print(model_4)

# Create loss and optimizer
loss_fn = nn.CrossEntropyLoss()
#loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.SGD(model_4.parameters(), 
                            lr=learning_rate) # exercise: try changing the learning rate here and seeing what happens to the model's performance


# Perform a single forward pass on the data (we'll need to put it to the target device for it to work)
#myVarSpec, myVarLum = model_4(X_train.to(device))[:5]
#print("myVarSpec ", myVarSpec)
#print("myVarLum ", myVarLum)

# How many elements in a single prediction sample?
#myVar2 = model_4(X_train.to(device))[0].shape
#print("myVar2 ", myVar2, numSpecNames, numLumNames)

## Make prediction logits with model
#y1_logits, y2_logits = model_4(X_test.to(device))
#
## Perform softmax calculation on logits across dimension 1 to get prediction probabilities
#y1_pred_probs = torch.softmax(y1_logits, dim=1)
#y2_pred_probs = torch.softmax(y2_logits, dim=1) 
##print(f"y_logits[:5], {y_logits[:5]}")
##print(f"y_pred_probs[:5], {y_pred_probs[:5]}")

# 8.5 Creating a training and testing loop for a multi-class PyTorch model

# Fit the model
torch.manual_seed(42)

# Set number of epochs
epochs = n_epochs

# Put data to target device
#X_train, Y1_train = X_train.to(device), Y1_train.to(device)
#X_test, Y1_test = X_test.to(device), Y1_test.to(device)
X_train, y1_train, y2_train = X_train.to(device), y1_train.to(device), y2_train.to(device)
X_test, y1_test, y2_test = X_test.to(device), y1_test.to(device), y2_test.to(device)

#Or try a learning rate scheduler:
#Reduce lr by 0.5 every 50 steps:
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5)


for epoch in range(epochs):
    
    ### Training
    model_4.train()

    #for batch_X, batch_y in train_loader:
    #Now double-headed
    #print("epoch ", epoch)
    for batch_X, batch_y1, batch_y2 in train_loader:

        #print("Inner loop started")

        batch_X = batch_X.to(device)
        batch_y1 = batch_y1.to(device)
        batch_y2 = batch_y2.to(device)        

        # 1. Forward pass
        #y_logits = model_4(X_train) # model outputs raw logits
        #y_logits = model_4(batch_X) # model outputs raw logits 
        #y_pred = torch.softmax(y_logits, dim=1).argmax(dim=1) # go from logits -> prediction probabilities -> prediction labels
        #Double-headed:
        y1_logits, y2_logits = model_4(batch_X) # model outputs raw logits 
        y1_pred = torch.softmax(y1_logits, dim=1).argmax(dim=1) # go from logits -> prediction probabilities -> prediction labels
        y2_pred = torch.softmax(y2_logits, dim=1).argmax(dim=1)
        # print(y_logits)
        # 2. Calculate loss and accuracy
        #print("check loss args: ", y_logits.shape, y_logits.dtype, y1_train.shape, y1_train.dtype)
        #loss = loss_fn(y_logits, y1_train)
        #loss = loss_fn(y_logits, batch_y)
        #Now double-headed:
        loss1 = loss_fn(y1_logits, batch_y1)
        loss2 = loss_fn(y2_logits, batch_y2)
        #total_loss = loss1 + loss2
        #Or Weighted - luminosity branch converges more quickly:
        total_loss = ((2.0*loss1) + (1.0*loss2)) / 3.0
        #print("check acc args: ", y_pred.shape, y_pred.dtype)
        #acc = accuracy_fn(y_true=y1_train,
        #                  y_pred=y_pred)
        #acc = accuracy_fn(y_true=batch_y,
        #                  y_pred=y_pred)
        acc1 = accuracy_fn(y_true=batch_y1,
                          y_pred=y1_pred)        
        acc2 = accuracy_fn(y_true=batch_y2,
                          y_pred=y2_pred)        


        # 3. Optimizer zero grad
        optimizer.zero_grad()
        #loss on the number of input neurons ($1/\sqrt{n}$)

        # 4. Loss backwards
        #loss.backward()
        #Double-headed:
        total_loss.backward()

        # 5. Optimizer step
        optimizer.step()

    ## after the training steps:
    #Slows down convergence with these schedule parameters
    #scheduler.step()


    ### Testing
    model_4.eval()
    with torch.inference_mode():
      # 1. Forward pass
      test1_logits, test2_logits = model_4(X_test)
      test1_pred = torch.softmax(test1_logits, dim=1).argmax(dim=1)
      test2_pred = torch.softmax(test2_logits, dim=1).argmax(dim=1)
      # 2. Calculate test loss and accuracy
      test1_loss = loss_fn(test1_logits, y1_test)
      test2_loss = loss_fn(test2_logits, y2_test)
      test1_acc = accuracy_fn(y_true=y1_test,
                             y_pred=test1_pred)
      test2_acc = accuracy_fn(y_true=y2_test,
                             y_pred=test2_pred)

    # Print out what's happening
    if epoch % 10 == 0:
        print(f"Epoch: {epoch} | Loss1: {loss1:.5f}, Acc1: {acc1:.2f}% | Test Loss1: {test1_loss:.5f}, Test1 Acc: {test1_acc:.2f}%\
             | Loss2: {loss2:.5f}, Acc2: {acc2:.2f}% | Test Loss2: {test2_loss:.5f}, Test2 Acc: {test2_acc:.2f}%")

    loss1V[epoch] = loss1.detach().cpu().numpy()
    loss2V[epoch] = loss2.detach().cpu().numpy()

    acc1V[epoch] = test1_acc
    acc2V[epoch] = test2_acc
    #For 2D classification:
    lossTotV = total_loss.cpu()
    #print(f"Epoch {epoch+1}: loss={loss:.4f}, test accuracy={acc:.3f}")
    handle.write(f"{epoch+1} {loss1:.4f} {loss2:.4f} {acc1:.3f} {acc2:.3f}\n")
    
#Report from final epoch:

"""Not working with PyTorch tensors
handle.write("\n")
handle.write("From sklearn.metrics.classification_report:\n")
handle.write("Spectral classification:\n")
handle.write(classification_report(y1_test, np.argmax(y1_pred, axis=1), target_names=spectestnames, zero_division=0.0))
handle.write("Luminosity classification:\n")
handle.write(classification_report(y2_test, np.argmax(y2_pred, axis=1), target_names=lumtestnames, zero_division=0.0))
"""    

plt.figure()
plt.subplot(1, 1, 1)    
plt.title("")
plt.ylabel('Accuracy')
plt.xlabel('Epoch')
#plt.ylim([0.0, 1.1])
    
plt.plot(acc1V)
plt.plot(acc2V)

plt.draw()
plt.pause(0.1)

fileStem = "AutoClassPyTAcc"
epsName = fileStem + '.eps'
plt.savefig(epsName, format='eps', dpi=1000)
#Save as jpg for PowerPoint
jpgName = fileStem + '.jpg'
plt.savefig(jpgName, format='jpg')

plt.figure()
plt.subplot(1, 1, 1)    
plt.title("")
plt.ylabel('Cross-entropy loss')
plt.xlabel('Epoch')
#plt.ylim([0.0, 1.1])

print("loss1V  ", loss1V[:10])
plt.plot(loss1V)
plt.plot(loss2V)

plt.draw()
plt.pause(0.1)

fileStem = "AutoClassPyTLoss"
epsName = fileStem + '.eps'
plt.savefig(epsName, format='eps', dpi=1000)
#Save as jpg for PowerPoint
jpgName = fileStem + '.jpg'
plt.savefig(jpgName, format='jpg')     


# 8.6 Making and evaluating predictions with a PyTorch multi-class model

# Make predictions
model_4.eval()
with torch.inference_mode():
    y1_logits, y2_logits = model_4(X_test)

# View the first 10 predictions
print(y1_logits[:10]), print(y2_logits[:10])

# Turn predicted logits in prediction probabilities
y1_pred_probs = torch.softmax(y1_logits, dim=1)
y2_pred_probs = torch.softmax(y2_logits, dim=1)

# Turn prediction probabilities into prediction labels
y1_preds = y1_pred_probs.argmax(dim=1)
y2_preds = y2_pred_probs.argmax(dim=1)

# Compare first 10 model preds and test labels
print(f"Predictions: {y1_preds[:10]}\nLabels: {y1_test[:10]}, {y2_preds[:10]}\nLabels: {y2_test[:10]}")
print(f"Test accuracy: {accuracy_fn(y_true=y1_test, y_pred=y1_preds)}%, {accuracy_fn(y_true=y2_test, y_pred=y2_preds)}%")

"""
# Download helper functions from Learn PyTorch repo (if not already downloaded)
if Path("helper_functions.py").is_file():
  print("helper_functions.py already exists, skipping download")
else:
  print("Downloading helper_functions.py")
  request = requests.get("https://raw.githubusercontent.com/mrdbourke/pytorch-deep-learning/main/helper_functions.py")
  with open("helper_functions.py", "wb") as f:
    f.write(request.content)

from helper_functions import plot_predictions, plot_decision_boundary
"""


# Setup metric and make sure it's on the target device
torchmetrics_accuracy = Accuracy(task='multiclass', num_classes=numSpecNames).to(device)

# Calculate accuracy
print(f"Check metrics acc1, {y1_preds.shape}, {y1_preds.dtype}, {y1_test.shape}, {y1_test.dtype}")
print(torchmetrics_accuracy(y1_preds, y1_test))
print(f"Check metrics acc2, {y2_preds.shape}, {y2_preds.dtype}, {y2_test.shape}, {y2_test.dtype}")
print(torchmetrics_accuracy(y2_preds, y2_test))

#Save model:

# 1. Create models directory 
MODEL_PATH = Path("models")
MODEL_PATH.mkdir(parents=True, exist_ok=True)

# 2. Create model save path 
MODEL_NAME = "FCN2Head.pth"
MODEL_SAVE_PATH = MODEL_PATH / MODEL_NAME

# 3. Save the model state dict 
print(f"Saving model to: {MODEL_SAVE_PATH}")
torch.save(obj=model_4.state_dict(), # only saving the state_dict() only saves the models learned parameters
           f=MODEL_SAVE_PATH)

"""
How to load:

# Instantiate a fresh instance of LinearRegressionModelV2
#loaded_model_1 = LinearRegressionModelV2()
loaded_model_4 = MKClassModel(n_inputs, n_spec_outputs, n_lum_outputs, n_hidden1, n_hidden2, n_hidden3).to(device)

# Load model state dict 
loaded_model_4.load_state_dict(torch.load(MODEL_SAVE_PATH))

# Put model to target device (if your data is on GPU, model will have to be on GPU to make predictions)
loaded_model_4.to(device)

print(f"Loaded model:\n{loaded_model_4}")
print(f"Model on device:\n{next(loaded_model_4.parameters()).device}")
"""

plt.draw()
plt.pause(0.1)


