# %%
# Define the neural network
class TwistedMNISTModel(nn.Module):
    def __init__(self):
        super(TwistedMNISTModel, self).__init__()
        
        # First path (main)
        self.fc1 = nn.Linear(784, 64)
        self.fc2 = nn.Linear(64, 32)

        # Second path (skip connection)
        self.skip_fc1 = nn.Linear(784, 32)

        # Processing after concatenation
        self.concat_fc = nn.Linear(64, 128)
        self.skip_fc2 = nn.Linear(32, 128)

        # Final output layer
        self.final_fc = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)  # Flatten the input (28x28 -> 784)

        # First path (main)
        z = torch.relu(self.fc1(x))
        h = torch.relu(self.fc2(z))

        # Second path (skip connection)
        u = torch.relu(self.skip_fc1(x))

        # Concatenation
        q = torch.cat((h, u), dim=1)

        # Processing concatenated output
        v = torch.relu(self.concat_fc(q))

        # Further processing of skip connection
        k = torch.relu(self.skip_fc2(u))

        # Add the processed outputs
        d = v + k

        # Final output layer
        hat_y = self.final_fc(d)
        
        return hat_y