# -*- coding: utf-8 -*-
"""
Created on Thu Nov 13 13:52:27 2025

 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
"""

# Standard PyTorch imports
import torch
from torch import nn

import numpy as np
from matplotlib import pyplot as plt
import h5py
from sklearn.metrics import classification_report
from pathlib import Path

import MKClassModel

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)

# 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

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

trainfile = 'Libr18Samples3002D.h5'   # This is the training set 
testfile = 'Libr18_225Samples2D.h5'  #Truly independent of training
#Expects H5:
trainset = load_stellar_spectra(trainfile)
testset = load_stellar_spectra(testfile)

#Train 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("testnames ", testnames)
#print("spectestnames ", spectestnames, " lumtestnames ", lumtestnames)
numTrainNames = len(X_train)
numSpecTrainNames = len(spectrainnames)
numLumTrainNames = len(lumtrainnames)

#Test set is libr18_225 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)
numTestNames = len(X_test)
numSpecTestNames = len(spectestnames)
numLumTestNames = len(lumtestnames)

# Turn test data into tensors
X_testT = torch.from_numpy(X_test).type(torch.float)
y1_testT = torch.from_numpy(y1_test).type(torch.LongTensor)
y2_testT = torch.from_numpy(y2_test).type(torch.LongTensor)

n_inputs = len(X_train[1])     # 1600 - 30 - 1 pixels
n_hidden1 = 768
n_hidden2 = 256
n_hidden3 = 64
n_spec_outputs = numSpecTrainNames
n_lum_outputs = numLumTrainNames

# 1. Create models directory 
MODEL_PATH = Path("models")

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

# 3. Read the model state dict 
print(f"REading model from: {MODEL_SAVE_PATH}")

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

# Load model state dict 
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)
model_4.to(device)

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

# 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_testT)

# 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_testT, y_pred=y1_preds)}%, {accuracy_fn(y_true=y2_testT, y_pred=y2_preds)}%")

plt.ion()

outputFile = "AutoClassPyT_Query"+testfile + ".out"
    
handle = open(outputFile, 'w', encoding='ascii')

handle.write("AutoClass2D Query: Test file: " + testfile + "\n")
handle.write("\n")
handle.write("Prediction Confidence Truth Prediction Truth \n")

#Luminosity class truth table:
lumTruth = [[0,0,0], [0,0,0], [0,0,0]]
lumTruthTxt = [['','',''],['','',''],['','','']]

for idx in range(numTestNames):

    #Luminosity class truth table:
    if y2_test[idx] == 0:
        if y2_preds[idx] == 0:
            lumTruth[0][0]+=1
        if y2_preds[idx] == 1:
            lumTruth[0][1]+=1        
        if y2_preds[idx] == 2:
            lumTruth[0][2]+=1
    if y2_test[idx] == 1:
        if y2_preds[idx] == 0:
            lumTruth[1][0]+=1
        if y2_preds[idx] == 1:
            lumTruth[1][1]+=1        
        if y2_preds[idx] == 2:
            lumTruth[1][2]+=1
    if y2_test[idx] == 2:
        if y2_preds[idx] == 0:
            lumTruth[2][0]+=1
        if y2_preds[idx] == 1:
            lumTruth[2][1]+=1        
        if y2_preds[idx] == 2:
            lumTruth[2][2]+=1

       
handle.close()

lumTruthTxt = [['','',''],['','',''],['','','']]
for i in range(3):
    for j in range(3):
        lumTruthTxt[i][j] = str(lumTruth[i][j])

print("Luminosity truth table:")
print(lumTruth)
                
plt.figure()
plt.subplot(1, 1, 1)    
plt.title("libr18_225: Predicted vs Actual Spectral Class")
plt.ylabel('Predicted MK spectral class')
plt.xlabel('True MK spectral class')
#plt.ylim([0.0, 1.1])

plt.yticks([i for i in range(numSpecTrainNames)], spectrainnames, fontsize=7)
plt.xticks([i for i in range(numSpecTestNames)], spectestnames, fontsize=7)
    
plt.plot(y1_test, y1_preds, 'o', 'black')

plt.show()
plt.pause(0.1)

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

plt.figure()
plt.subplot(1, 1, 1)    
plt.title("libr18_225: Predicted vs Actual Luminosity Class")
plt.ylabel('Predicted MK luminosity class')
plt.xlabel('True MK luminosity class')
#plt.ylim([0.0, 1.1])

plt.yticks([i for i in range(numLumTrainNames)], lumtrainnames, fontsize=7)
plt.xticks([i for i in range(numLumTestNames)], lumtestnames, fontsize=7)
    
plt.plot(y2_test, y2_preds, 'o', color='gray')

plt.plot([0,2],[0,2],'--', color='gray')
for i in range(3):
    for j in range(3):
        if i == j:
            plt.plot(i, j, marker='o', color='green', markersize=4)
            plt.annotate(lumTruthTxt[i][j], xy=(i,j), xytext=(i+0.05, j-0.05))
        if i != j:
            plt.plot(i, j, marker='o', color='red', markersize=1)
            plt.annotate(lumTruthTxt[i][j], xy=(i,j), xytext=(i+0.05, j-0.05))

plt.show()
plt.pause(0.1)
                
fileStem = "AutoClassPyT_Query_Lum"
epsName = fileStem + testfile + '.eps'
plt.savefig(epsName, format='eps', dpi=1000)
#Save as jpg for PowerPoint
jpgName = fileStem + testfile + '.jpg'
plt.savefig(jpgName, format='jpg')



