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


def AutoClass2D_train():

    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 = 'Libr18Samples100Lum.h5'
    #testsetFile = 'Libr18_27SamplesLum.h5'
    trainsetFile = 'Libr18Samples3002D.h5'
    testsetFile = 'Libr18_27Samples2D.h5'
    trainset = load_stellar_spectra(trainsetFile)
    testset = load_stellar_spectra(testsetFile)
    #numNames = 31

    whichCheck = 0

    #Training set is libr18 of Gray & Corbally:
    #X_train, y_train = trainset.data, trainset.target
    X_train = trainset.data
    y1_train, y2_train = trainset.target1, trainset.target2
    #X, y = trainset.data, trainset.target
    trainnames = trainset.target_names
    spectrainnames, lumtrainnames = trainset.target1_names, trainset.target2_names
    print("trainnames ", trainnames)
    print("spectrainnames ", spectrainnames, " lumtrainnames ", lumtrainnames)
    #numNames = len(trainnames)
    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]
    #Y_train = np.eye(numNames)[y_train]
    #Y_test = np.eye(numTestNames)[y_test]
    #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)
    #print("whichCheck ", whichCheck, " y_train[whichCheck] ", y_train[whichCheck], " Y_train[whichCheck] ", Y_train[whichCheck])
    #plt.plot(X_train[1])
    #plt.plot(X_test[1])

    # --- 2. Define network architecture
    #Original M<NIST values:
    #n_inputs = 784     # 28x28 pixels
    #n_hidden = 64
    #n_outputs = 10
    
    n_inputs = len(X_train[1])     # 1600 - 30 - 1 pixels
    #n_hidden1 = 512
    #n_hidden2 = 256
    #n_hidden3 = 128
    #n_hidden1 = 768
    #n_hidden2 = 384
    #n_hidden3 = 192
    #n_outputs = numNames      #unique MK specral classes inc. subclass but NOT luminosity class!)
    #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)
    #W4 = rng.normal(0, 0.1, (n_hidden3, n_outputs)).astype(np.float32)
    #b4 = np.zeros((1, n_outputs)).astype(np.float32)
    #For double-headed 2D classification:
    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 = 200  #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            
            #Z4 = A3 @ W4 + b4         # Hidden → output
            #Y_pred = softmax(Z4)      # Probabilities for each digit
            #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            
          
            #print("W1_4.shape ", W1_4.shape, " W2_4.shape ", W2_4.shape)
            #print("Y1_pred.shape ", Y1_pred.shape, " Y1b.shape ", Y1b.shape)
            #print("Y2_pred.shape ", Y2_pred.shape, " Y2b.shape ", Y2b.shape)
            
            # ---- 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
            #dZ4 = (Y_pred - Yb) / batch_size      # derivative of loss wrt Z4
            #dW4 = A3.T @ dZ4
            #db4 = np.sum(dZ4, axis=0, keepdims=True)
            #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 
            #dA3 = dZ4 @ W4.T
            #For double-headed 2D classification:
            dA3 = dZ1_4 @ W1_4.T + dZ2_4 @ W2_4.T
            #dA3 = dZ1_4 @ W1_4.T  #Test with errors from luminosity classification only
            #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
            #W4 -= learning_rate * dW4
            #b4 -= learning_rate * db4
            #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)       
        #Z4 = A3 @ W4 + b4
        #Y_pred = softmax(Z4)
        #acc = np.mean(np.argmax(Y_pred, axis=1) == y_test)
        #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)      
        #Y2_predArgMax = np.argmax(Y2_pred, axis=1)
        #print("shapes Y2_predArgMax",  Y2_predArgMax.shape, " y2_test ", y2_test.shape)
        #print("Y2_predArgMax ", Y2_predArgMax[0:19])
        #print("y2_test ", y2_test[0:19])
        #Representative loss from last batch this epoch:
        #

        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')     


    #print("b4 ", b4)
    #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

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

#This should be kept in the same source file so consistency
#  with training above can be checked:
# - Activation functions have to be consistent with above

#def predict(x, W1, b1, W2, b2, W3, b3, W4, b4):
def predict2D(x, W1, b1, W2, b2, W3, b3, W1_4, b1_4, W2_4, b2_4):   
    
    Z1 = x @ 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)    
    #Z4 = A3 @ W4 + b4
    #Y_pred = softmax(Z4)
    #predicted_class = np.argmax(Y_pred)
    #confidence = Y_pred[0][predicted_class]
    #For double-headed @d spectral classification:
    #Spectral classification:
    Z1_4 = A3 @ W1_4 + b1_4
    Y1_pred = softmax(Z1_4)
    predicted_class = np.argmax(Y1_pred)
    confidence1 = Y1_pred[0][predicted_class]
    #Luminosity classification:
    Z2_4 = A3 @ W2_4 + b2_4
    Y2_pred = softmax(Z2_4)
    predicted_lum = np.argmax(Y2_pred)
    confidence2 = Y2_pred[0][predicted_lum]    
    
    
    return predicted_class, confidence1, predicted_lum, confidence2

"""
Benchmark comparison: MNIST with 70000 hand-drawn Arabic numeral samples:
With  learning_rate = 0.02 , n_epochs = 30 and GELU activation I get:
Epoch 1: loss=0.4473, test accuracy=0.888
Epoch 2: loss=0.4038, test accuracy=0.911
Epoch 3: loss=0.1995, test accuracy=0.924
Epoch 4: loss=0.1789, test accuracy=0.933
Epoch 5: loss=0.1804, test accuracy=0.939
Epoch 6: loss=0.1179, test accuracy=0.942
Epoch 7: loss=0.2010, test accuracy=0.945
Epoch 8: loss=0.1351, test accuracy=0.947
Epoch 9: loss=0.1399, test accuracy=0.950
Epoch 10: loss=0.0933, test accuracy=0.952
Epoch 11: loss=0.0827, test accuracy=0.953
Epoch 12: loss=0.0626, test accuracy=0.957
Epoch 13: loss=0.1589, test accuracy=0.958
Epoch 14: loss=0.0984, test accuracy=0.958
Epoch 15: loss=0.0986, test accuracy=0.960
Epoch 16: loss=0.1259, test accuracy=0.960
Epoch 17: loss=0.1175, test accuracy=0.962
Epoch 18: loss=0.0399, test accuracy=0.963
Epoch 19: loss=0.1015, test accuracy=0.964
Epoch 20: loss=0.1121, test accuracy=0.965
Epoch 21: loss=0.1196, test accuracy=0.965
Epoch 22: loss=0.1060, test accuracy=0.965
Epoch 23: loss=0.0498, test accuracy=0.964
Epoch 24: loss=0.0864, test accuracy=0.965
Epoch 25: loss=0.0351, test accuracy=0.967
Epoch 26: loss=0.0438, test accuracy=0.967
Epoch 27: loss=0.1148, test accuracy=0.968
Epoch 28: loss=0.0648, test accuracy=0.969
Epoch 29: loss=0.0374, test accuracy=0.968
Epoch 30: loss=0.0324, test accuracy=0.968

"""
    

    
    
