#!/usr/bin/env python3
import sys

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

print("=" * 60)
print("GRNexus - Sentiment Analysis with Embeddings")
print("Using Embedding Layer for better text representation")
print("=" * 60)

# Training data - More samples for better learning
train_texts = [
    # Positive
    "I love this amazing product",
    "Best purchase ever made",
    "Absolutely wonderful quality",
    "Great experience highly recommend",
    "Very happy with results",
    "Fantastic service and quality",
    "Excellent product love it",
    "Amazing quality highly satisfied",
    # Negative
    "Terrible waste of money",
    "Worst product ever bought",
    "Horrible quality disappointed",
    "Very bad experience",
    "Complete disaster never again",
    "Poor quality not recommended",
    "Awful product terrible service",
    "Disappointed with purchase"
]

train_labels = [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]

# Sequence Vectorizer
class SequenceVectorizer:
    def __init__(self, max_features=100, max_length=10):
        self.max_features = max_features
        self.max_length = max_length
        self.vocab = {'<PAD>': 0, '<UNK>': 1}
        self.next_idx = 2
    
    def fit(self, texts):
        word_count = {}
        for text in texts:
            for word in self.tokenize(text):
                word_count[word] = word_count.get(word, 0) + 1
        
        sorted_words = sorted(word_count.items(), key=lambda x: -x[1])[:self.max_features - 2]
        for word, _ in sorted_words:
            self.vocab[word] = self.next_idx
            self.next_idx += 1
    
    def transform(self, texts):
        return [self.text_to_sequence(text) for text in texts]
    
    def tokenize(self, text):
        import re
        return re.sub(r'[^a-z\s]', '', text.lower()).split()
    
    def text_to_sequence(self, text):
        tokens = self.tokenize(text)
        sequence = [self.vocab.get(word, self.vocab['<UNK>']) for word in tokens]
        
        # Pad or truncate
        if len(sequence) < self.max_length:
            sequence += [self.vocab['<PAD>']] * (self.max_length - len(sequence))
        else:
            sequence = sequence[:self.max_length]
        
        return [float(idx) for idx in sequence]

# Create vectorizer
print("\n1. Creating vocabulary...")
vectorizer = SequenceVectorizer(max_features=100, max_length=10)
vectorizer.fit(train_texts)
print(f"   Vocabulary size: {len(vectorizer.vocab)}")
print(f"   Max sequence length: {vectorizer.max_length}")

# Transform
print("\n2. Creating sequences...")
x_train = vectorizer.transform(train_texts)
y_train = [[0.0, 1.0] if label == 1 else [1.0, 0.0] for label in train_labels]

# Create model
print("\n3. Creating model with Embedding layer...")
model = NeuralNetwork(
    loss='cross_entropy',
    optimizer='sgd',
    learning_rate=0.1,
    name='sentiment_embedding'
)

vocab_size = len(vectorizer.vocab)
embedding_dim = 16

# Embedding → Flatten → Dense architecture
model.add(EmbeddingLayer(vocab_size=vocab_size, embedding_dim=embedding_dim))
model.add(FlattenLayer())
model.add(DenseLayer(units=64, input_dim=vectorizer.max_length * embedding_dim, activation=ReLU()))
model.add(BatchNormLayer())
model.add(DropoutLayer(rate=0.3))
model.add(DenseLayer(units=32, input_dim=64, activation=ReLU()))
model.add(DenseLayer(units=2, input_dim=32, activation=Softmax()))

print(f"   Architecture: Embedding({vocab_size}, {embedding_dim}) → Flatten → Dense(64) → Dense(32) → Dense(2)")

# Train
print("\n4. Training model...")
model.train(x_train, y_train, epochs=200, batch_size=4, verbose=False)

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

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

test_texts = [
    "I love this product",
    "Terrible experience",
    "Amazing quality",
    "Very disappointed",
    "Excellent purchase",
    "Horrible waste"
]

for text in test_texts:
    sequence = vectorizer.transform([text])[0]
    pred = model.predict([sequence])[0]
    sentiment = "POSITIVE 😊" if pred[1] > pred[0] else "NEGATIVE 😞"
    confidence = round(max(pred) * 100, 1)
    
    print(f"\nText: \"{text}\"")
    print(f"Sentiment: {sentiment} ({confidence}%)")
    print(f"Scores: Neg={round(pred[0]*100, 1)}% Pos={round(pred[1]*100, 1)}%")

print("\n" + "=" * 60)
print("✅ Embedding-based Sentiment Analysis complete!")
print("=" * 60)
print("\nNote: Embedding layer learns word representations during training.")
print("This captures semantic meaning better than one-hot encoding.")
