import sys

from grnexus import NeuralNetwork
from lib.grnexus_layers import *
from lib.grnexus_activations import *
import random

print("=" * 80)
print("GRNexus Neural Network Framework - Comprehensive Tests (Python)")
print("=" * 80)

# Test 1: Simple XOR Problem (Classification)
print("\n[Test 1] XOR Problem - Binary Classification")
print("-" * 80)

# XOR dataset
x_train = [
    [0.0, 0.0],
    [0.0, 1.0],
    [1.0, 0.0],
    [1.0, 1.0]
]

y_train = [
    [1.0, 0.0],  # 0 XOR 0 = 0
    [0.0, 1.0],  # 0 XOR 1 = 1
    [0.0, 1.0],  # 1 XOR 0 = 1
    [1.0, 0.0]   # 1 XOR 1 = 0
]

# Create model
model = NeuralNetwork(
    loss='cross_entropy',
    optimizer='sgd',
    learning_rate=0.5
)

# Add layers
model.add(DenseLayer(
    units=4,
    input_dim=2,
    activation=Tanh()
))

model.add(DenseLayer(
    units=2,
    input_dim=4,
    activation=Sigmoid()
))

# Show model summary
model.summary()

# Train model
print("\nTraining XOR model...")
model.train(x_train, y_train, epochs=1000, batch_size=4, verbose=False)

# Test predictions
print("\nPredictions:")
for i, input_data in enumerate(x_train):
    prediction = model.predict([input_data])[0]
    expected = y_train[i]
    pred_class = prediction.index(max(prediction))
    true_class = expected.index(max(expected))
    
    print(f"Input: {input_data} => Predicted: {pred_class} ({[round(v, 3) for v in prediction]}) | Expected: {true_class}")

# Evaluate
results = model.evaluate(x_train, y_train)
print("\nEvaluation Results:")
print(f"Loss: {results['loss']:.6f}")
print(f"Accuracy: {results['accuracy']:.2f}%")

# Test 2: Model Saving and Loading
print("\n" + "=" * 80)
print("[Test 2] Model Persistence - Save and Load")
print("-" * 80)

model_path = 'xor_model.nexus'
print(f"Saving model to {model_path}...")
model.save(model_path)

print(f"\nLoading model from {model_path}...")
loaded_model = NeuralNetwork.load(model_path)

print("\nTesting loaded model...")
loaded_predictions = loaded_model.predict(x_train)
original_predictions = model.predict(x_train)
predictions_match = all(
    all(abs(a - b) < 1e-6 for a, b in zip(pred1, pred2))
    for pred1, pred2 in zip(loaded_predictions, original_predictions)
)
print(f"Loaded model predictions match original: {predictions_match}")

# Clean up
if os.path.exists(model_path):
    os.remove(model_path)

# Test 3: Regression Problem
print("\n" + "=" * 80)
print("[Test 3] Simple Regression - Function Approximation")
print("-" * 80)

# Generate synthetic data: y = 2x + 1
x_reg = [[i / 10.0] for i in range(20)]
y_reg = [[2.0 * x[0] + 1.0] for x in x_reg]

# Create regression model
reg_model = NeuralNetwork(
    loss='mse',
    optimizer='sgd',
    learning_rate=0.01
)

reg_model.add(DenseLayer(
    units=8,
    input_dim=1,
    activation=ReLU()
))

reg_model.add(DenseLayer(
    units=1,
    input_dim=8,
    activation=None
))

print("\nTraining regression model...")
reg_model.train(x_reg, y_reg, epochs=100, batch_size=10, verbose=False)

# Test predictions
print("\nSample Predictions:")
for x in [0.0, 0.5, 1.0, 1.5, 2.0]:
    pred = reg_model.predict([[x]])[0][0]
    expected = 2.0 * x + 1.0
    error = abs(pred - expected)
    print(f"x={x} => Predicted: {pred:.3f}, Expected: {expected:.3f}, Error: {error:.3f}")

# Test 4: Multi-layer Network with Dropout and Batch Normalization
print("\n" + "=" * 80)
print("[Test 4] Advanced Architecture - Dropout & Batch Normalization")
print("-" * 80)

# Create more complex dataset
x_complex = []
y_complex = []

random.seed(42)
for _ in range(100):
    x1 = random.random() * 2 - 1
    x2 = random.random() * 2 - 1
    
    # Non-linear decision boundary: inside circle
    if x1**2 + x2**2 < 0.5:
        x_complex.append([x1, x2])
        y_complex.append([1.0, 0.0])
    else:
        x_complex.append([x1, x2])
        y_complex.append([0.0, 1.0])

# Create advanced model
adv_model = NeuralNetwork(
    loss='cross_entropy',
    optimizer='sgd',
    learning_rate=0.1
)

adv_model.add(DenseLayer(
    units=16,
    input_dim=2,
    activation=ReLU()
))

adv_model.add(BatchNormLayer(
    epsilon=1e-5,
    momentum=0.1
))

adv_model.add(DropoutLayer(rate=0.3))

adv_model.add(DenseLayer(
    units=8,
    input_dim=16,
    activation=ReLU()
))

adv_model.add(DenseLayer(
    units=2,
    input_dim=8,
    activation=Sigmoid()
))

adv_model.summary()

print("\nTraining advanced model...")
adv_model.train(x_complex, y_complex, epochs=50, batch_size=16, verbose=False)

# Evaluate
results = adv_model.evaluate(x_complex, y_complex)
print("\nFinal Results:")
print(f"Loss: {results['loss']:.6f}")
print(f"Accuracy: {results['accuracy']:.2f}%")

# Test 5: Different Activation Functions
print("\n" + "=" * 80)
print("[Test 5] Testing Various Activation Functions")
print("-" * 80)

activations = [
    ('Sigmoid', Sigmoid()),
    ('Tanh', Tanh()),
    ('ReLU', ReLU()),
    ('LeakyReLU', LeakyReLU(alpha=0.01)),
    ('ELU', ELU(alpha=1.0)),
    ('Swish', Swish()),
    ('GELU', GELU())
]

for name, activation in activations:
    test_model = NeuralNetwork(
        loss='cross_entropy',
        optimizer='sgd',
        learning_rate=0.3
    )
    
    test_model.add(DenseLayer(
        units=4,
        input_dim=2,
        activation=activation
    ))
    
    test_model.add(DenseLayer(
        units=2,
        input_dim=4,
        activation=Sigmoid()
    ))
    
    # Train on XOR
    test_model.train(x_train, y_train, epochs=500, batch_size=4, verbose=False)
    
    # Evaluate
    results = test_model.evaluate(x_train, y_train)
    print(f"{name.ljust(15)} - Accuracy: {results['accuracy']:.2f}%")

# Test 6: Testing Layer Types
print("\n" + "=" * 80)
print("[Test 6] Testing Different Layer Types")
print("-" * 80)

# Test LSTM Layer
print("\nTesting LSTM Layer...")
lstm_layer = LSTMLayer(units=4, input_size=3)
sequence_input = [[[1.0, 0.5, 0.2], [0.8, 0.3, 0.1], [0.6, 0.4, 0.3]]]  # batch_size=1, seq_len=3, input_size=3
lstm_output = lstm_layer.forward(sequence_input)
print(f"LSTM input shape: (1, 3, 3), output shape: {len(lstm_output)}x{len(lstm_output[0])}x{len(lstm_output[0][0])}")

# Test Conv2D Layer
print("\nTesting Conv2D Layer...")
conv_layer = Conv2DLayer(filters=2, kernel_size=3, stride=1, padding=0)
image_input = [[[random.random() for _ in range(5)] for _ in range(5)]]  # 5x5 image
conv_output = conv_layer.forward(image_input)
print(f"Conv2D input shape: (1, 5, 5), output shape: {len(conv_output)}x{len(conv_output[0])}x{len(conv_output[0][0])}")

# Test MaxPooling Layer
print("\nTesting MaxPooling Layer...")
pool_layer = MaxPoolingLayer(pool_size=2, stride=2)
pool_input = [[[random.random() for _ in range(4)] for _ in range(4)]]  # 4x4 input
pool_output = pool_layer.forward(pool_input)
print(f"MaxPooling input shape: (1, 4, 4), output: {type(pool_output).__name__} with {len(pool_output)} elements")

print("\n" + "=" * 80)
print("All tests completed successfully!")
print("=" * 80)
