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

import torch
from torch import nn

# 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
