#!/usr/bin/env python3
"""
Local Inference SCript (NumPy Version)
=====================================
Carica i pesi .json scaricati ed esegue inferenza usando SOLO NumPy (no PyTorch necessario).
"""

import json
import numpy as np
import os
import sys

def relu(x):
    return np.maximum(0, x)

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def linear_forward(x, W, b):
    # W shape: (out_features, in_features)
    # x shape: (in_features) or (batch, in_features)
    return x @ W.T + b

def predict(model_weights, X):
    # Layer 1
    w1 = np.array(model_weights['net.0.weight'])
    b1 = np.array(model_weights['net.0.bias'])
    z1 = linear_forward(X, w1, b1)
    a1 = relu(z1)
    
    # Layer 2
    w2 = np.array(model_weights['net.2.weight'])
    b2 = np.array(model_weights['net.2.bias'])
    z2 = linear_forward(a1, w2, b2)
    a2 = relu(z2)
    
    # Layer 3
    w3 = np.array(model_weights['net.4.weight'])
    b3 = np.array(model_weights['net.4.bias'])
    z3 = linear_forward(a2, w3, b3)
    a3 = sigmoid(z3)
    
    return a3

def main():
    model_path = "simple_nn_weights.json"
    
    # Check current dir or ../output/
    if not os.path.exists(model_path):
        alt_path = os.path.join("..", "output", "simple_nn_weights.json")
        if os.path.exists(alt_path):
            model_path = alt_path
            
    print(f"📂 Loading weights from {model_path}...")
    
    if not os.path.exists(model_path):
        print("❌ Model weights not found! Run 'python examples/train_nn.py' first.")
        sys.exit(1)
        
    try:
        with open(model_path) as f:
            weights = json.load(f)
        print("✅ Weights loaded successfully!")
    except Exception as e:
        print(f"❌ Error loading JSON: {e}")
        sys.exit(1)
        
    # Generate Test Data
    print("\n🔮 Running Inference (NumPy Engine)...")
    test_points = np.array([
        [0.0, 0.0],   # Center (Should be 1)
        [0.8, 0.8],   # Outside (Should be 0)
        [-0.5, 0.5],  # Inside (Should be 1)
        [2.0, 2.0]    # Far outside (Should be 0)
    ], dtype=np.float32)
    
    predictions = predict(weights, test_points)
        
    print("-" * 40)
    print(f"{'X':<10} {'Y':<10} | {'Prob':<10} {'Class':<10}")
    print("-" * 40)
    
    for i, point in enumerate(test_points):
        prob = predictions[i][0]
        cls = 1 if prob > 0.5 else 0
        print(f"{point[0]:<10.2f} {point[1]:<10.2f} | {prob:<10.4f} {cls:<10}")
        
    print("-" * 40)
    print("✅ Inference Pipeline Verified (No PyTorch needed locally).")

if __name__ == "__main__":
    main()