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

@author: original 3-layer, single-dead network with RELU by ChatGPT (OpenAI)
 - developed into 5-layer double-headed with GELU by Ian Short
"""

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

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)


# --- 1. Load the MNIST dataset (from sklearn for simplicity)
#from sklearn.datasets import fetch_openml
#from sklearn.model_selection import train_test_split

def AutoClass2DQuery(x_sample, y1_true, y2_true, W1, b1, W2, b2, W3, b3, W1_4, b1_4, W2_4, b2_4):
    
    #Query the model:
    prediction1, confidence1, prediction2, confidence2 = AutoClass2D.predict2D(x_sample, W1, b1, W2, b2, W3, b3, W1_4, b1_4, W2_4, b2_4)

    #print(f"Predicted MK class: {prediction} (confidence {confidence:.2f}) True label: {y_true}")
    #print(f"True label: {y_true}")
    
    return prediction1, confidence1, prediction2, confidence2

# ---------------------------------

plt.ion()

#
##Train:
#
#W1, b1, W2, b2, W3, b3, W4, b4 = AutoClass.AutoClass_train()

Wb = np.load("test.npz")


trainnames = Wb['tn']
spectrainnames = Wb['tn1']

lumtrainnames = Wb['tn2']
W1 = Wb['w1']
b1 = Wb['b1']
W2 = Wb['w2']
b2 = Wb['b2']
W3 = Wb['w3']
b3 = Wb['b3']
W1_4 = Wb['w14']
b1_4 = Wb['b14']
W2_4 = Wb['w24']
b2_4 = Wb['b24']

#print("b4 ", b4)

numTrainNames = len(trainnames)
numSpecTrainNames = len(spectrainnames)
numLumTrainNames = len(lumtrainnames)
print("trainnames ", trainnames)
print("spectrainnames ", spectrainnames, " lumtrainnames ", lumtrainnames)

# ----------------------------------
#
## Query:
#

#testfile = 'Libr18_27Samples.h5'  #Used for training - libr18_27 of Gray & Corbally:
testfile = 'Libr18_225Samples2D.h5'  #Truly independent of training
#Expects H5:
testset = load_stellar_spectra(testfile)

X_test, y_test = testset.data, testset.target
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)

X_test = X_test.astype(np.float32)

# One-hot encode the labels: e.g., 3 -> [0,0,0,1,0,0,0,0,0,0]
#Y_test = np.eye(numTestNames)[y_test]
#Spectral class lables:
Y1_test = np.eye(numSpecTestNames)[y1_test]
#Luminosity class lables:
Y2_test = np.eye(numLumTestNames)[y2_test]   

#idx = 23  # any test sample index

outputFile = "AutoClass2D_Query"+testfile + "250.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")


y1_trueV = [0.0 for i in range(numTestNames)]
prediction1V = [0.0 for i in range(numTestNames)]
y2_trueV = [0.0 for i in range(numTestNames)]
prediction2V = [0.0 for i in range(numTestNames)]

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

for idx in range(numTestNames):

    #print("idx ", idx)
    x_sample = X_test[idx]  
    y1_true = np.argmax(Y1_test[idx])  # correct spectral class label
    y2_true = np.argmax(Y2_test[idx])  # correct luminosity class label
    #plt.plot(x_sample)
    
    prediction1, confidence1, prediction2, confidence2 = AutoClass2DQuery(x_sample, y1_true, y2_true, W1, b1, W2, b2, W3, b3, W1_4, b1_4, W2_4, b2_4)
    #print(f"Predicted spectral class: {prediction1} (confidence {confidence1:.2f}) True label: {y1_true} Predicted luminosity class: {prediction2} (confidence {confidence2:.2f}) True label: {y2_true}")
    #print(spectrainnames[prediction1], " ", spectestnames[y1_true], " ", lumtrainnames[prediction2], " ", lumtestnames[y2_true])
    #print(classification_report(y_true, prediction, target_names=testnames, zero_division=0.0))

    handle.write(f"{prediction1} {confidence1} {y1_true} {spectrainnames[prediction1]} {spectestnames[y1_true]} {prediction2} {confidence2} {y2_true} {lumtrainnames[prediction2]} {lumtestnames[y2_true]}\n")

    y1_trueV[idx] = y1_true
    prediction1V[idx] = prediction1
    y2_trueV[idx] = y2_true
    prediction2V[idx] = prediction2

    #Luminosity class truth table:
    if y2_true == 0:
        if prediction2 == 0:
            lumTruth[0][0]+=1
        if prediction2 == 1:
            lumTruth[0][1]+=1        
        if prediction2 == 2:
            lumTruth[0][2]+=1
    if y2_true == 1:
        if prediction2 == 0:
            lumTruth[1][0]+=1
        if prediction2 == 1:
            lumTruth[1][1]+=1        
        if prediction2 == 2:
            lumTruth[1][2]+=1
    if y2_true == 2:
        if prediction2 == 0:
            lumTruth[2][0]+=1
        if prediction2 == 1:
            lumTruth[2][1]+=1        
        if prediction2 == 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)

allSpecNames = ['O9','B0','B1','B2','B3','B5','B7','B8','B9','A0','A1','A2','A3','A5','A7','F0','F2','F3','F5','F6','F8','F9','G0','G2','G5','G8','K0','K2','K3','K5','K7']
#allSpecNames = ['B2','B3','B5','B7','B8','B9','A0','A1','A2','A3','A5','A7','F0','F2','F3','F5','F6','F8','F9','G0','G2','G5','G8','K0','K2','K3','K5','K7']
numAll = len(allSpecNames)
prediction1V2 = [0 for i in range(numSpecTrainNames)]
y1_trueV2 = [0 for i in range(numSpecTestNames)]

for i in range(numSpecTrainNames):
    #print('i ', i, ' spectrainnames[i] ', spectrainnames[i])
    prediction1V2[i] = allSpecNames.index(spectrainnames[prediction1V[i]])
    #print("spectrainnames[i] ", spectrainnames[i], " allSpecNames ", allSpecNames[y1_trueV2[i]])

for i in range(numSpecTestNames):
    y1_trueV2[i] = allSpecNames.index(spectestnames[y1_trueV[i]])
    #print("spectestnames[i] ", spectestnames[i], " allSpecNames ", allSpecNames[prediction1V2[i]])
    
       
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.yticks([i for i in range(numAll)], allSpecNames, fontsize=7)
plt.xticks([i for i in range(numAll)], allSpecNames, fontsize=7)
    
#plt.plot(y1_trueV, prediction1V, 'o', color='black')
##Debugging:
#for i in range(numTestNames):
#    print("y1_trueV ", y1_trueV[i], " prediction1V ", prediction1V[i])
    
plt.plot([0, numAll], [0, numAll], 'gray')

for i in range(numSpecTestNames):
    #print("i ", i, " y1_trueV2[i] ", y1_trueV2[i], " prediction1V2[i] ", prediction1V2[i])
    if y2_trueV[i] == 0:   #LC I
        col = 'red'
        sz = 4
    if y2_trueV[i] == 1:   #LC III
        col = 'green'
        sz = 4
    if y2_trueV[i] == 2:   #LC V
        col = 'blue'
        sz = 4
    plt.plot(y1_trueV2[i], prediction1V2[i], marker='o', color=col, markersize=sz)


plt.text(4, 24, 'V', color='blue')
plt.text(4, 22, 'III', color='green')
plt.text(4, 20, 'I', color='red')

plt.show()
plt.pause(0.1)

fileStem = "AutoClass_Query_Spec250"
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_trueV, prediction2V, '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 = "AutoClass_Query_Lum250"
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')



