#!/usr/bin/env python3
import sys

from grnexus import NeuralNetwork
from lib.grnexus_layers import DenseLayer, BatchNormLayer
from lib.grnexus_activations import ReLU
from lib.grnexus_normalization import Softmax
import random

print("Testing BatchNorm save/load...")

# Create a simple model with BatchNorm
model = NeuralNetwork(
    loss='cross_entropy',
    optimizer='sgd',
    learning_rate=0.01,
    name='test_batchnorm'
)

model.add(DenseLayer(units=10, input_dim=5, activation=ReLU()))
model.add(BatchNormLayer())
model.add(DenseLayer(units=3, input_dim=10, activation=Softmax()))

# Train with dummy data
x_train = [[random.random() for _ in range(5)] for _ in range(20)]
y_train = []
for _ in range(20):
    label = [0.0, 0.0, 0.0]
    label[random.randint(0, 2)] = 1.0
    y_train.append(label)

print("\n1. Training model...")
model.train(x_train, y_train, epochs=10, batch_size=4, verbose=False)

# Get BatchNorm layer
bn_layer = model.layers[1]
running_mean_before = bn_layer.running_mean
running_var_before = bn_layer.running_var

print("\n2. BatchNorm stats BEFORE save:")
print(f"   running_mean: {[round(v, 4) for v in running_mean_before]}")
print(f"   running_var: {[round(v, 4) for v in running_var_before]}")

# Make a prediction before save
test_input = [[random.random() for _ in range(5)]]
pred_before = model.predict(test_input)[0]
print(f"\n3. Prediction BEFORE save: {[round(v, 4) for v in pred_before]}")

# Save model
print("\n4. Saving model...")
model.save('test_bn_model.nexus')

# Load model
print("\n5. Loading model...")
loaded_model = NeuralNetwork.load('test_bn_model.nexus')

# Get BatchNorm layer from loaded model
bn_layer_loaded = loaded_model.layers[1]
running_mean_after = bn_layer_loaded.running_mean
running_var_after = bn_layer_loaded.running_var

print("\n6. BatchNorm stats AFTER load:")
print(f"   running_mean: {[round(v, 4) for v in running_mean_after]}")
print(f"   running_var: {[round(v, 4) for v in running_var_after]}")

# Make prediction with loaded model
pred_after = loaded_model.predict(test_input)[0]
print(f"\n7. Prediction AFTER load: {[round(v, 4) for v in pred_after]}")

# Compare
print("\n8. VERIFICATION:")
if running_mean_before == running_mean_after and running_var_before == running_var_after:
    print("   ✅ BatchNorm stats preserved correctly!")
else:
    print("   ❌ BatchNorm stats NOT preserved!")

if pred_before == pred_after:
    print("   ✅ Predictions are IDENTICAL!")
else:
    diff = [abs(a - b) for a, b in zip(pred_before, pred_after)]
    max_diff = max(diff)
    if max_diff < 1e-6:
        print("   ✅ Predictions are IDENTICAL (within floating point precision)")
    else:
        print(f"   ❌ Predictions are DIFFERENT! Max diff: {max_diff}")

# Cleanup
if os.path.exists('test_bn_model.nexus'):
    os.remove('test_bn_model.nexus')
print("\n✅ Test complete!")
