import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report
import time

# Set seeds for reproducibility
torch.manual_seed(42)
np.random.seed(42)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# Load the dataset
df = pd.read_csv('ArSenTD-LEV.tsv', sep='\t')

# Filter to Lebanese tweets only
df_lebanese = df[df['Country'] == 'lebanon'].copy()

print(f"Total samples: {len(df)}")
print(f"Lebanese samples: {len(df_lebanese)}")
print(f"\nColumns: {list(df.columns)}")
print(f"\nFirst few tweets:")
df_lebanese.head()

# Label encoding
label_map = {
    'very_negative': 0,
    'negative': 1,
    'neutral': 2,
    'positive': 3,
    'very_positive': 4
}
df_lebanese['label'] = df_lebanese['Sentiment'].map(label_map)

print("Label distribution:")
print(df_lebanese['Sentiment'].value_counts())

# Visualize dataset characteristics
fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# 1. Sentiment distribution
ax = axes[0, 0]
sentiment_counts = df_lebanese['Sentiment'].value_counts()
ax.bar(range(len(sentiment_counts)), sentiment_counts.values, color='steelblue')
ax.set_xticks(range(len(sentiment_counts)))
ax.set_xticklabels(sentiment_counts.index, rotation=45, ha='right')
ax.set_ylabel('Count')
ax.set_title('Sentiment Distribution')
ax.grid(True, alpha=0.3, axis='y')

# 2. Text length distribution
ax = axes[0, 1]
df_lebanese['text_length'] = df_lebanese['Tweet'].str.len()
mean_len = df_lebanese['text_length'].mean()
ax.hist(df_lebanese['text_length'], bins=30, color='coral', edgecolor='black')
ax.axvline(mean_len, color='red', linestyle='--', linewidth=2, label=f'Mean: {mean_len:.0f}')
ax.set_xlabel('Text Length (characters)')
ax.set_ylabel('Frequency')
ax.set_title('Text Length Distribution')
ax.legend()
ax.grid(True, alpha=0.3, axis='y')

# 3. Word count distribution
ax = axes[1, 0]
df_lebanese['word_count'] = df_lebanese['Tweet'].str.split().str.len()
mean_words = df_lebanese['word_count'].mean()
ax.hist(df_lebanese['word_count'], bins=30, color='lightgreen', edgecolor='black')
ax.axvline(mean_words, color='red', linestyle='--', linewidth=2, label=f'Mean: {mean_words:.0f}')
ax.set_xlabel('Word Count')
ax.set_ylabel('Frequency')
ax.set_title('Word Count Distribution')
ax.legend()
ax.grid(True, alpha=0.3, axis='y')

# 4. Topic distribution
ax = axes[1, 1]
top_topics = df_lebanese['Topic'].value_counts().head(10)
ax.barh(range(len(top_topics)), top_topics.values, color='plum')
ax.set_yticks(range(len(top_topics)))
ax.set_yticklabels(top_topics.index)
ax.set_xlabel('Count')
ax.set_title('Top 10 Topics')
ax.grid(True, alpha=0.3, axis='x')

plt.tight_layout()
plt.show()

print(f"\nDataset statistics:")
print(f"Average text length: {mean_len:.1f} characters")
print(f"Average word count: {mean_words:.1f} words")

# Split data
X = df_lebanese['Tweet'].values
y = df_lebanese['label'].values

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"Train size: {len(X_train)}")
print(f"Test size: {len(X_test)}")

# Create bag-of-words features
vectorizer = CountVectorizer(max_features=1000, ngram_range=(1, 2))
X_train_bow = vectorizer.fit_transform(X_train)
X_test_bow = vectorizer.transform(X_test)

print(f"BoW vocabulary size: {len(vectorizer.vocabulary_)}")
print(f"Feature matrix shape: {X_train_bow.shape}")
print(f"\nExample features (first 10 words):")
print(list(vectorizer.vocabulary_.keys())[:10])

# Train baseline model
baseline_model = LogisticRegression(max_iter=1000, random_state=42)
baseline_model.fit(X_train_bow, y_train)

y_pred_baseline = baseline_model.predict(X_test_bow)
baseline_acc = accuracy_score(y_test, y_pred_baseline)

print(f"Baseline Test Accuracy: {baseline_acc:.4f} ({baseline_acc*100:.2f}%)")
print(f"\nClassification Report:")
print(classification_report(y_test, y_pred_baseline, 
                          target_names=['very_neg', 'neg', 'neutral', 'pos', 'very_pos']))

# Build character vocabulary
all_text = ' '.join(X)
unique_chars = sorted(set(all_text))
char_to_idx = {ch: idx+1 for idx, ch in enumerate(unique_chars)}  # 0 reserved for padding
char_to_idx['<PAD>'] = 0

vocab_size = len(char_to_idx)
print(f"Vocabulary size: {vocab_size}")
print(f"\nSample characters (first 20):")
print(unique_chars[:20])
print(f"\nLast 20 characters (likely Arabic):")
print(unique_chars[-20:])

# Encode sequences
max_length = 140  # Twitter-like limit

def encode_text(text, max_len=max_length):
    """Convert text to sequence of character indices"""
    encoded = [char_to_idx.get(ch, 0) for ch in text[:max_len]]
    # Pad to max_length
    if len(encoded) < max_len:
        encoded = encoded + [0] * (max_len - len(encoded))
    return encoded

X_train_encoded = np.array([encode_text(text) for text in X_train])
X_test_encoded = np.array([encode_text(text) for text in X_test])

print(f"Encoded train shape: {X_train_encoded.shape}")
print(f"Encoded test shape: {X_test_encoded.shape}")
print(f"\nExample encoding (first tweet, first 20 chars):")
print(f"Original: {X_train[0][:20]}")
print(f"Encoded:  {X_train_encoded[0][:20]}")

# Create PyTorch dataset
class SentimentDataset(Dataset):
    def __init__(self, sequences, labels):
        self.sequences = torch.LongTensor(sequences)
        self.labels = torch.LongTensor(labels)

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        return self.sequences[idx], self.labels[idx]

train_dataset = SentimentDataset(X_train_encoded, y_train)
test_dataset = SentimentDataset(X_test_encoded, y_test)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

print(f"Train batches: {len(train_loader)}")
print(f"Test batches: {len(test_loader)}")

class SimpleRNN(nn.Module):
    def __init__(self, vocab_size, embedding_dim=64, hidden_dim=128, output_dim=5, dropout=0.5):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        self.rnn = nn.RNN(embedding_dim, hidden_dim, batch_first=True)
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        # x: (batch_size, seq_len) - character indices
        embedded = self.embedding(x)  # (batch_size, seq_len, embedding_dim)
        output, hidden = self.rnn(embedded)  # output: (batch_size, seq_len, hidden_dim)
        # Use last hidden state
        final_hidden = hidden.squeeze(0)  # (batch_size, hidden_dim)
        dropped = self.dropout(final_hidden)
        logits = self.fc(dropped)  # (batch_size, output_dim)
        return logits

class LSTMModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim=64, hidden_dim=128, output_dim=5, dropout=0.5):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        embedded = self.embedding(x)
        output, (hidden, cell) = self.lstm(embedded)
        final_hidden = hidden.squeeze(0)
        dropped = self.dropout(final_hidden)
        logits = self.fc(dropped)
        return logits

class GRUModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim=64, hidden_dim=128, output_dim=5, dropout=0.5):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        self.gru = nn.GRU(embedding_dim, hidden_dim, batch_first=True)
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        embedded = self.embedding(x)
        output, hidden = self.gru(embedded)
        final_hidden = hidden.squeeze(0)
        dropped = self.dropout(final_hidden)
        logits = self.fc(dropped)
        return logits

def train_epoch(model, loader, criterion, optimizer, device):
    """Train for one epoch"""
    model.train()
    total_loss = 0
    correct = 0
    total = 0

    for sequences, labels in loader:
        sequences, labels = sequences.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(sequences)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    return total_loss / len(loader), 100. * correct / total

def evaluate(model, loader, criterion, device):
    """Evaluate on validation/test set"""
    model.eval()
    total_loss = 0
    correct = 0
    total = 0

    with torch.no_grad():
        for sequences, labels in loader:
            sequences, labels = sequences.to(device), labels.to(device)
            outputs = model(sequences)
            loss = criterion(outputs, labels)

            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    return total_loss / len(loader), 100. * correct / total

def train_and_evaluate(model_name, model, epochs=30, lr=0.001, weight_decay=1e-5):
    """Train a model and track metrics"""
    print(f"\n{'='*70}")
    print(f"Training {model_name}")
    print(f"{'='*70}")

    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)

    results = {
        'model_name': model_name,
        'train_acc': [],
        'test_acc': [],
        'train_loss': [],
        'test_loss': []
    }

    start_time = time.time()

    for epoch in range(epochs):
        train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
        test_loss, test_acc = evaluate(model, test_loader, criterion, device)

        results['train_acc'].append(train_acc)
        results['test_acc'].append(test_acc)
        results['train_loss'].append(train_loss)
        results['test_loss'].append(test_loss)

        if (epoch + 1) % 5 == 0:
            print(f"Epoch {epoch+1}/{epochs} | "
                  f"Train: {train_acc:.2f}% (loss: {train_loss:.3f}) | "
                  f"Test: {test_acc:.2f}% (loss: {test_loss:.3f})")

    training_time = time.time() - start_time
    results['training_time'] = training_time
    results['final_test_acc'] = test_acc

    print(f"\nTraining time: {training_time:.2f}s")
    print(f"Final test accuracy: {test_acc:.2f}%")

    return results, model

# Train all models
all_results = []

# Simple RNN
rnn_model = SimpleRNN(vocab_size)
results_rnn, trained_rnn = train_and_evaluate("Simple RNN", rnn_model, epochs=30)
all_results.append(results_rnn)

# LSTM
lstm_model = LSTMModel(vocab_size)
results_lstm, trained_lstm = train_and_evaluate("LSTM", lstm_model, epochs=30)
all_results.append(results_lstm)

# GRU
gru_model = GRUModel(vocab_size)
results_gru, trained_gru = train_and_evaluate("GRU", gru_model, epochs=30)
all_results.append(results_gru)

# Plot training curves
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# 1. Training curves
ax = axes[0]
epochs_range = range(1, 31)
for model in all_results:
    ax.plot(epochs_range, model['train_acc'], marker='o', linestyle='--', 
            alpha=0.6, label=f"{model['model_name']} (train)")
    ax.plot(epochs_range, model['test_acc'], marker='s', linewidth=2, 
            label=f"{model['model_name']} (test)")

ax.axhline(y=baseline_acc*100, color='red', linestyle=':', linewidth=2, label='Baseline (BoW)')
ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Accuracy (%)', fontsize=12)
ax.set_title('Training Progress: RNN vs Baseline', fontsize=14, fontweight='bold')
ax.legend(loc='best', fontsize=9)
ax.grid(True, alpha=0.3)

# 2. Final comparison
ax = axes[1]
model_names = ['Baseline\n(BoW+LR)'] + [m['model_name'] for m in all_results]
final_accs = [baseline_acc*100] + [m['final_test_acc'] for m in all_results]
colors = ['red', 'steelblue', 'coral', 'lightgreen']

bars = ax.bar(range(len(model_names)), final_accs, color=colors, alpha=0.7, edgecolor='black')
ax.set_xticks(range(len(model_names)))
ax.set_xticklabels(model_names, fontsize=10)
ax.set_ylabel('Test Accuracy (%)', fontsize=12)
ax.set_title('Final Test Accuracy Comparison', fontsize=14, fontweight='bold')
ax.grid(True, alpha=0.3, axis='y')

for bar, acc in zip(bars, final_accs):
    height = bar.get_height()
    ax.text(bar.get_x() + bar.get_width()/2., height,
            f'{acc:.1f}%', ha='center', va='bottom', fontweight='bold')

# 3. Overfitting analysis
ax = axes[2]
for model in all_results:
    train_acc = model['train_acc']
    test_acc = model['test_acc']
    gap = [train_acc[i] - test_acc[i] for i in range(len(train_acc))]
    ax.plot(epochs_range, gap, marker='o', linewidth=2, label=model['model_name'])

ax.axhline(y=0, color='black', linestyle='--', alpha=0.5)
ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Train-Test Gap (%)', fontsize=12)
ax.set_title('Overfitting Analysis', fontsize=14, fontweight='bold')
ax.legend()
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# Print final comparison
print("\n" + "="*70)
print("FINAL COMPARISON")
print("="*70)
print(f"{'Model':<25} {'Test Accuracy':<15} {'Training Time':<15}")
print("-"*70)
print(f"{'Baseline (BoW + LR)':<25} {baseline_acc*100:>6.2f}%         {'N/A':<15}")
for r in all_results:
    print(f"{r['model_name']:<25} {r['final_test_acc']:>6.2f}%         {r['training_time']:>6.2f}s")

print("\n" + "="*70)
print("OVERFITTING ANALYSIS (Final Train-Test Gap)")
print("="*70)
for r in all_results:
    final_gap = r['train_acc'][-1] - r['test_acc'][-1]
    print(f"{r['model_name']:<25} {final_gap:>6.2f}% gap")

# Tokenize text into words
# Build word vocabulary
# Use word embedding layer