#!/usr/bin/env python3
import sys
import math

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

print("=" * 60)
print("GRNexus - 3D Image Classifier (RGB)")
print("Convolutional Neural Network Demo")
print("=" * 60)

# Generate synthetic RGB images (16x16x3)
def generate_image(pattern_type):
    image = [[[0.0 for _ in range(3)] for _ in range(16)] for _ in range(16)]
    
    if pattern_type == 'red_square':
        # Red square in center
        for i in range(6, 11):
            for j in range(6, 11):
                image[i][j][0] = 1.0  # Red channel
    elif pattern_type == 'green_circle':
        # Green circle
        center_x, center_y = 8, 8
        radius = 4
        for i in range(16):
            for j in range(16):
                dist = math.sqrt((i - center_x)**2 + (j - center_y)**2)
                if dist <= radius:
                    image[i][j][1] = 1.0  # Green channel
    elif pattern_type == 'blue_diagonal':
        # Blue diagonal line
        for i in range(16):
            image[i][i][2] = 1.0  # Blue channel
            if i != 15 - i:
                image[i][15-i][2] = 1.0
    
    return image

# Flatten 3D image to 1D vector
def flatten_image(image):
    return [pixel for row in image for pixel in row for pixel in pixel]

print("\n1. Generating synthetic RGB images...")
x_train = []
y_train = []

# Generate training data
for _ in range(5):
    x_train.append(flatten_image(generate_image('red_square')))
    y_train.append([1.0, 0.0, 0.0])  # Class 0: Red square
    
    x_train.append(flatten_image(generate_image('green_circle')))
    y_train.append([0.0, 1.0, 0.0])  # Class 1: Green circle
    
    x_train.append(flatten_image(generate_image('blue_diagonal')))
    y_train.append([0.0, 0.0, 1.0])  # Class 2: Blue diagonal

print(f"   Generated {len(x_train)} images (16x16x3 = {len(x_train[0])} features)")
print("   Classes: Red Square, Green Circle, Blue Diagonal")

# Create CNN model
print("\n2. Creating CNN model...")
model = NeuralNetwork(
    loss='cross_entropy',
    optimizer='sgd',
    learning_rate=0.01,
    name='image_classifier_3d'
)

input_dim = 16 * 16 * 3  # 768 features

# Simple CNN-like architecture (using Dense layers)
model.add(DenseLayer(units=128, input_dim=input_dim, activation=ReLU()))
model.add(BatchNormLayer())
model.add(DropoutLayer(rate=0.3))
model.add(DenseLayer(units=64, input_dim=128, activation=ReLU()))
model.add(BatchNormLayer())
model.add(DenseLayer(units=32, input_dim=64, activation=ReLU()))
model.add(DenseLayer(units=3, input_dim=32, activation=Softmax()))

# Train
print("\n3. Training model...")
model.train(x_train, y_train, epochs=100, batch_size=3, verbose=False)

# Save
model.save('image_classifier_3d.nexus')
print("   Model saved!")

# Test
print("\n4. Testing predictions:")
print("-" * 60)

test_images = [
    ('red_square', "Red Square"),
    ('green_circle', "Green Circle"),
    ('blue_diagonal', "Blue Diagonal")
]

classes = ["Red Square 🟥", "Green Circle 🟢", "Blue Diagonal 🔵"]

for pattern, name in test_images:
    image = flatten_image(generate_image(pattern))
    pred = model.predict([image])[0]
    predicted_class = pred.index(max(pred))
    confidence = round(pred[predicted_class] * 100, 1)
    
    print(f"\nInput: {name}")
    print(f"Predicted: {classes[predicted_class]} ({confidence}%)")
    print(f"Probabilities: {[round(v * 100, 1) for v in pred]}")

print("\n" + "=" * 60)
print("✅ 3D Image Classifier Demo complete!")
print("=" * 60)
print("\nNote: This demo uses flattened images with Dense layers.")
print("For better performance with real images, use Conv2DLayer!")
