# -*- coding: utf-8 -*-
"""
Created on Mon Nov 10 13:38:21 2025

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

"""
1. Forward pass:
   Z1​=XW1​+b1​,A1​=ReLU(Z1​)
   Z2​=A1​W2​+b2​,Y^=softmax(Z2​)
   
   The network computes a predicted probability distribution over digits.
   
2. Loss:
    Cross-entropy measures how far ^Y^ is from the true one-hot Y.
    
3. Backpropagation:
    
    Chain rule of calculus applied to compute gradients of the loss with respect to each parameter.
    
4. Gradient descent step:
    W ← W − η∂W/∂L​

"""

import numpy as np
import matplotlib.pyplot as plt
import h5py

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

#Activation functions:
    
# --- 3. Define activation functions and helpers


#Rectified (Re) linear unit activation:
        
def relu(x):
    return np.maximum(0, x)
    
def relu_deriv(x):
    return (x > 0).astype(float)


#Gaussian error (GE) linear unit activation:
  # This is the faster tanh approximation - from ChatGPT
    
def gelu(x):
    return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0/np.pi) * (x + 0.044715 * x**3)))

def gelu_deriv(x):
    sqrt_2_over_pi = np.sqrt(2.0 / np.pi)
    c = 0.044715
    x3 = x**3
    t = np.tanh(sqrt_2_over_pi * (x + c * x3))    
    # derivative of tanh(u) = sech^2(u) = 1 - tanh^2(u)
    dtanh = 1.0 - t**2
    term = sqrt_2_over_pi * (1.0 + 3.0 * c * x**2)    
    return 0.5 * (1.0 + t) + 0.5 * x * dtanh * term

def cross_entropy(y_pred, y_true):
    # Average negative log-likelihood
    m = y_true.shape[0]
    # Here * is a matrix multiplication - dot product? 
    return -np.sum(y_true * np.log(y_pred + 1e-12)) / m

def softmax(x):
    exps = np.exp(x - np.max(x, axis=1, keepdims=True))
    return exps / np.sum(exps, axis=1, keepdims=True)


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  #Needed in 2D???
        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)


plt.ion()

# Download stellar data
# y is a numClasses-length vector of numerically encoded true MK class labels

"""
2) Train/test split — critical  warning
generated noisy variants of the same template spectrum:
MUST split by star identity, not by sample.
Otherwise the network sees the same star in both train and test and accuracy becomes meaningless
split groups → then expand noise realizations
"""

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)

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

X_train = X_train.astype(np.float32)
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]
#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)
    
n_inputs = len(X_train[1])     # 1600 - 30 - 1 pixels
#For double-headed 2D classification:
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)

# Initialize weights with small random values
#rng is an object ??
#42 is an arbitrary seed value for the random number generator
rng = np.random.default_rng(42)  
#rng.normal chooses a random number weighted by the normal distribution
W1 = rng.normal(0, 0.1, (n_inputs, n_hidden1)).astype(np.float32)
b1 = np.zeros((1, n_hidden1)).astype(np.float32)
W2 = rng.normal(0, 0.1, (n_hidden1, n_hidden2)).astype(np.float32)
b2 = np.zeros((1, n_hidden2)).astype(np.float32)
W3 = rng.normal(0, 0.1, (n_hidden2, n_hidden3)).astype(np.float32)
b3 = np.zeros((1, n_hidden3)).astype(np.float32)

W1_4 = rng.normal(0, 0.1, (n_hidden3, n_spec_outputs)).astype(np.float32)
b1_4 = np.zeros((1, n_spec_outputs)).astype(np.float32)
W2_4 = rng.normal(0, 0.1, (n_hidden3, n_lum_outputs)).astype(np.float32)
b2_4 = np.zeros((1, n_lum_outputs)).astype(np.float32)


# --- 4. Training loop

learning_rate = 0.005
#Or try a learning rate scheduler:  0.5 every 50 steps?
n_epochs = 600  #production
#n_epochs = 3  #quick testing
#    batch_size = int(numSamples/122)  # 300 samples
batch_size = int(numSamples/366)  # 100 samples
print("numSamples ", numSamples, " batch_size ", batch_size)

outputFile = trainsetFile + "_" + 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)]

#print("Y1_train.shape ", Y1_train.shape, " Y2_train.shape ", Y2_train.shape)

for epoch in range(n_epochs):
        
    # Shuffle data - stochastic gradient descent
    indices = np.random.permutation(X_train.shape[0])
    X_train = X_train[indices]
    #Y_train = Y_train[indices]
    Y1_train, Y2_train = Y1_train[indices], Y2_train[indices]

    for i in range(0, X_train.shape[0], batch_size):
        Xb = X_train[i:i+batch_size]
        #Yb = Y_train[i:i+batch_size]
        Y1b = Y1_train[i:i+batch_size]
        Y2b = Y2_train[i:i+batch_size]            

        # ---- FORWARD PASS ----
        # '@' means matrix multiplication of - takes np.matrix() operands'
        Z1 = Xb @ W1 + b1         # Linear combination (inputs → hidden)
        #A1 = relu(Z1)             # Activation
        A1 = gelu(Z1)             # Activation
        Z2 = A1 @ W2 + b2         # Hidden → Hidden
        #A2 = relu(Z2)             # Activation
        A2 = gelu(Z2)             # Activation
        Z3 = A2 @ W3 + b3         # Hidden → Hidden
        #A3 = relu(Z3)             # Activation
        A3 = gelu(Z3)             # Activation            
        #For double-headed 2D classification:
        Z1_4 = A3 @ W1_4 + b1_4         # Hidden → output
        Y1_pred = softmax(Z1_4)      # Probabilities for each digit
        Z2_4 = A3 @ W2_4 + b2_4         # Hidden → output
        Y2_pred = softmax(Z2_4)      # Probabilities for each digit            
          

            
        # ---- LOSS ----
        #loss = cross_entropy(Y_pred, Yb)
        #For double-headed 2D classification:
        loss1 = cross_entropy(Y1_pred, Y1b)  #Spectral class dimension
        loss2 = cross_entropy(Y2_pred, Y2b)  #Luminosity class dimension
        #lossTot = loss1 + loss2
        #Or Weighted - luminosity branch converges more quickly:
        lossTot = ((2.0*loss1) + (1.0*loss2)) / 3.0

        loss1V[epoch] = loss1
        loss2V[epoch] = loss2

        # ---- BACKPROP ----
        # Output layer gradients
        #For double-headed 2D classification:
        #Spectral class training:
        dZ1_4 = (Y1_pred - Y1b) / batch_size      # derivative of loss wrt Z4
        dW1_4 = A3.T @ dZ1_4
        db1_4 = np.sum(dZ1_4, axis=0, keepdims=True)
        #Luminosity class training:
        dZ2_4 = (Y2_pred - Y2b) / batch_size      # derivative of loss wrt Z4
        dW2_4 = A3.T @ dZ2_4
        db2_4 = np.sum(dZ2_4, axis=0, keepdims=True)
        #print("dZ1_4.shape ", dZ1_4.shape, " dZ2_4.shape ", dZ2_4.shape)

        # 3rd Hidden layer gradients 
        #For double-headed 2D classification:
        dA3 = dZ1_4 @ W1_4.T + dZ2_4 @ W2_4.T
        #dZ3 = dA3 * relu_deriv(Z3)
        dZ3 = dA3 * gelu_deriv(Z3)
        dW3 = A2.T @ dZ3
        db3 = np.sum(dZ3, axis=0, keepdims=True)
        
        # 2nd Hidden layer gradients 
        dA2 = dZ3 @ W3.T
        #dZ2 = dA2 * relu_deriv(Z2)
        dZ2 = dA2 * gelu_deriv(Z2)
        dW2 = A1.T @ dZ2
        db2 = np.sum(dZ2, axis=0, keepdims=True)
        
        # 1st Hidden layer gradients 
        dA1 = dZ2 @ W2.T
        #dZ1 = dA1 * relu_deriv(Z1)
        dZ1 = dA1 * gelu_deriv(Z1)
        dW1 = Xb.T @ dZ1
        db1 = np.sum(dZ1, axis=0, keepdims=True)
        

        # ---- UPDATE ----
        W1 -= learning_rate * dW1
        b1 -= learning_rate * db1
        W2 -= learning_rate * dW2
        b2 -= learning_rate * db2
        W3 -= learning_rate * dW3
        b3 -= learning_rate * db3
        #For double-headed 2D classification:
        #Spectral class training:
        W1_4 -= learning_rate * dW1_4
        b1_4 -= learning_rate * db1_4
        #Luminosity class training:
        W2_4 -= learning_rate * dW2_4
        b2_4 -= learning_rate * db2_4

        #Try a learning rate schedule:
        if (epoch > 0) and (epoch%50) == 0:
            learning_rate = learning_rate/2.0
            

    # ---- Epoch summary ----
    Z1 = X_test @ W1 + b1
    #A1 = relu(Z1)
    A1 = gelu(Z1)
    Z2 = A1 @ W2 + b2
    #A2 = relu(Z2)
    A2 = gelu(Z2)
    Z3 = A2 @ W3 + b3
    #A3 = relu(Z3)
    A3 = gelu(Z3)       
    #For double-headed 2D classification:
    #Spectral class:
    Z1_4 = A3 @ W1_4 + b1_4
    Y1_pred = softmax(Z1_4)
    acc1 = np.mean(np.argmax(Y1_pred, axis=1) == y1_test)
    #Luminosity class:
    Z2_4 = A3 @ W2_4 + b2_4
    Y2_pred = softmax(Z2_4)
    acc2 = np.mean(np.argmax(Y2_pred, axis=1) == y2_test)      

    acc1V[epoch] = acc1
    acc2V[epoch] = acc2
    #For 2D classification:
    lossTotV = lossTot
    #print(f"Epoch {epoch+1}: loss={loss:.4f}, test accuracy={acc:.3f}")
    print(f"Epoch {epoch+1}: loss1={loss1:.4f}, loss2={loss2:.4f}, lossTot={lossTot:.4f}, accuracy1={acc1:.3f}, accuracy2={acc2:.3f}")
    handle.write(f"{epoch+1} {loss1:.4f} {loss2:.4f} {acc1:.3f} {acc2:.3f}\n")

#Report from final epoch:
print("From sklearn.metrics.classification_report:")
print("Spectral classification:")
print(classification_report(y1_test, np.argmax(Y1_pred, axis=1), target_names=spectestnames, zero_division=0.0))
print("Luminosity classification:")
print(classification_report(y2_test, np.argmax(Y2_pred, axis=1), target_names=lumtestnames, zero_division=0.0))
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 = "AutoClass2DAcc"
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])
    
plt.plot(loss1V)
plt.plot(loss2V)

plt.draw()
plt.pause(0.1)

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


#Save weights:
#weightsFile = "AutoClassWeights3002D.npz"
weightsFile = "test.npz"
np.savez(weightsFile, tn=trainnames, tn1=spectrainnames, tn2=lumtrainnames, w1=W1, b1=b1, w2=W2, b2=b2, w3=W3, b3=b3, w14=W1_4, b14=b1_4, w24=W2_4, b24=b2_4)
        
##return trainnames, spectrainnames, lumtrainnames, W1, b1, W2, b2, W3, b3, W1_4, b1_4, W2_4, b2_4

    

    
    
