# -*- coding: utf-8 -*-
"""Sentiment_Analysis_IMDB.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1uL_ftP6OruQtWT9MjSvesQl2go2sh8-x
"""

# ============================================================
# Cell 1: Install required packages
# ============================================================
!pip install datasets transformers accelerate -q
!pip install shap wordcloud nltk -q
print("All packages installed successfully!")

# ============================================================
# Cell 2: Import all libraries
# ============================================================
import os, re, time, warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# NLP and text processing
import nltk
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
from nltk.tokenize import word_tokenize

# Scikit-learn: traditional ML pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    classification_report, confusion_matrix, ConfusionMatrixDisplay
)

# Hugging Face: Transformer pipeline
from datasets import load_dataset
from transformers import (
    DistilBertTokenizerFast, DistilBertForSequenceClassification,
    Trainer, TrainingArguments, EarlyStoppingCallback
)
import torch
from torch.utils.data import Dataset

# Interpretability
import shap
from wordcloud import WordCloud

# Download NLTK data
nltk.download('punkt', quiet=True)
nltk.download('punkt_tab', quiet=True)
nltk.download('stopwords', quiet=True)
nltk.download('wordnet', quiet=True)

warnings.filterwarnings('ignore')
sns.set_theme(style='whitegrid', palette='muted')

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

# Check GPU availability
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
if device.type == 'cuda':
    print(f"GPU name: {torch.cuda.get_device_name(0)}")
    print(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
print(f"PyTorch version: {torch.__version__}")

# ============================================================
# Cell 3: Load the IMDB Large Movie Review Dataset
# ============================================================
dataset = load_dataset('imdb')

# Convert to pandas DataFrames for easier exploration
train_df = pd.DataFrame(dataset['train'])
test_df  = pd.DataFrame(dataset['test'])

print(f"Training set size: {len(train_df)}")
print(f"Test set size:     {len(test_df)}")
print(f"Columns: {list(train_df.columns)}")
print(f"\nLabel mapping: 0 = Negative, 1 = Positive")
print(f"\n--- First 3 rows ---")
train_df.head(3)

# ============================================================
# Cell 4: EDA - Class Distribution and Review Lengths
# ============================================================
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# (a) Class distribution
class_counts = train_df['label'].value_counts().sort_index()
labels = ['Negative (0)', 'Positive (1)']
colours = ['#e74c3c', '#27ae60']
axes[0].bar(labels, class_counts.values, color=colours, edgecolor='black')
for i, v in enumerate(class_counts.values):
    axes[0].text(i, v + 200, str(v), ha='center', fontweight='bold')
axes[0].set_title('Class Distribution (Training Set)', fontsize=13)
axes[0].set_ylabel('Count')

# (b) Review length distribution (word count)
train_df['word_count'] = train_df['text'].apply(lambda x: len(x.split()))
axes[1].hist(train_df[train_df['label']==0]['word_count'], bins=50,
             alpha=0.6, label='Negative', color='#e74c3c', edgecolor='black')
axes[1].hist(train_df[train_df['label']==1]['word_count'], bins=50,
             alpha=0.6, label='Positive', color='#27ae60', edgecolor='black')
axes[1].set_title('Review Length Distribution (Words)', fontsize=13)
axes[1].set_xlabel('Word Count')
axes[1].set_ylabel('Frequency')
axes[1].legend()

# (c) Average review length by class
avg_len = train_df.groupby('label')['word_count'].mean()
axes[2].bar(labels, avg_len.values, color=colours, edgecolor='black')
for i, v in enumerate(avg_len.values):
    axes[2].text(i, v + 2, f'{v:.1f}', ha='center', fontweight='bold')
axes[2].set_title('Mean Review Length by Sentiment', fontsize=13)
axes[2].set_ylabel('Mean Word Count')

plt.tight_layout()
plt.show()

print(f"\nDataset balance: {class_counts[0]/len(train_df)*100:.1f}% negative, "
      f"{class_counts[1]/len(train_df)*100:.1f}% positive")
print(f"Mean review length: {train_df['word_count'].mean():.0f} words")
print(f"Median review length: {train_df['word_count'].median():.0f} words")

# ============================================================
# Cell 5: Word Clouds for Positive and Negative Reviews
# ============================================================
def simple_clean(text):
    """Quick clean for word frequency analysis."""
    text = re.sub(r'<[^>]+>', '', text)       # remove HTML tags
    text = re.sub(r'[^a-zA-Z\s]', '', text)  # keep only letters
    return text.lower()

stop_words = set(stopwords.words('english'))

pos_text = ' '.join(train_df[train_df['label']==1]['text'].apply(simple_clean))
neg_text = ' '.join(train_df[train_df['label']==0]['text'].apply(simple_clean))

fig, axes = plt.subplots(1, 2, figsize=(16, 6))

# Positive word cloud
wc_pos = WordCloud(width=800, height=400, background_color='white',
                   stopwords=stop_words, colormap='Greens', max_words=100)
wc_pos.generate(pos_text)
axes[0].imshow(wc_pos, interpolation='bilinear')
axes[0].set_title('Positive Reviews - Most Frequent Words', fontsize=14)
axes[0].axis('off')

# Negative word cloud
wc_neg = WordCloud(width=800, height=400, background_color='white',
                   stopwords=stop_words, colormap='Reds', max_words=100)
wc_neg.generate(neg_text)
axes[1].imshow(wc_neg, interpolation='bilinear')
axes[1].set_title('Negative Reviews - Most Frequent Words', fontsize=14)
axes[1].axis('off')

plt.tight_layout()
plt.show()

# ============================================================
# Cell 6: Text Preprocessing for TF-IDF Pipeline
# ============================================================
lemmatizer = WordNetLemmatizer()

def preprocess_text_tfidf(text):
    """
    Full preprocessing pipeline for traditional ML:
    1. Remove HTML tags (common in IMDB data)
    2. Lowercase
    3. Remove non-alphabetic characters
    4. Tokenise
    5. Remove stopwords
    6. Lemmatise
    """
    text = re.sub(r'<[^>]+>', '', text)         # strip HTML
    text = text.lower()                          # lowercase
    text = re.sub(r'[^a-zA-Z\s]', '', text)    # remove punctuation/numbers
    tokens = word_tokenize(text)                 # tokenise
    tokens = [lemmatizer.lemmatize(t) for t in tokens
              if t not in stop_words and len(t) > 2]
    return ' '.join(tokens)

print("Preprocessing training data for TF-IDF (this takes ~2-3 minutes)...")
start = time.time()
train_df['clean_text'] = train_df['text'].apply(preprocess_text_tfidf)
test_df['clean_text']  = test_df['text'].apply(preprocess_text_tfidf)
elapsed = time.time() - start
print(f"Done in {elapsed:.1f}s")

# Show a before/after example
print("\n" + "="*60)
print("BEFORE preprocessing:")
print("="*60)
print(train_df['text'].iloc[0][:300])
print("\n" + "="*60)
print("AFTER preprocessing:")
print("="*60)
print(train_df['clean_text'].iloc[0][:300])

# ============================================================
# Cell 7: Minimal Preprocessing for DistilBERT
# ============================================================
def preprocess_text_bert(text):
    """
    Minimal preprocessing for BERT/DistilBERT.
    Only strip HTML - the tokeniser handles everything else.
    Preserving punctuation, casing, and stopwords is important
    because the model was pretrained with these features.
    """
    text = re.sub(r'<[^>]+>', '', text)
    return text.strip()

train_df['bert_text'] = train_df['text'].apply(preprocess_text_bert)
test_df['bert_text']  = test_df['text'].apply(preprocess_text_bert)

print("BERT preprocessing complete (minimal - HTML removal only)")
print(f"\nExample BERT-preprocessed text (first 200 chars):")
print(train_df['bert_text'].iloc[0][:200])

# ============================================================
# Cell 8: TF-IDF Vectorisation
# ============================================================
# Configuration:
# - max_features=50000: top 50k most important terms
# - ngram_range=(1,2): unigrams + bigrams (e.g., "not good")
# - sublinear_tf=True: log scaling to prevent frequent terms dominating
# - min_df=5: ignore very rare terms (< 5 documents)
# - max_df=0.95: ignore very common terms (> 95% of documents)

tfidf = TfidfVectorizer(
    max_features=50000,
    ngram_range=(1, 2),
    sublinear_tf=True,
    min_df=5,
    max_df=0.95
)

print("Fitting TF-IDF vectoriser on training data...")
start = time.time()
X_train_tfidf = tfidf.fit_transform(train_df['clean_text'])
X_test_tfidf  = tfidf.transform(test_df['clean_text'])
print(f"Done in {time.time()-start:.1f}s")

print(f"\nVocabulary size: {len(tfidf.vocabulary_):,}")
print(f"Training feature matrix shape: {X_train_tfidf.shape}")
print(f"Test feature matrix shape: {X_test_tfidf.shape}")
print(f"Matrix density: {X_train_tfidf.nnz / (X_train_tfidf.shape[0] * X_train_tfidf.shape[1]) * 100:.2f}%")

# ============================================================
# Cell 9: Hyperparameter Tuning - GridSearchCV
# ============================================================
param_grid = {
    'C': [0.1, 1.0, 10.0],        # regularisation strength
    'penalty': ['l1', 'l2'],        # L1 (sparse) vs L2 (dense)
    'solver': ['liblinear'],        # supports both L1 and L2
    'max_iter': [1000]
}

print("Running GridSearchCV (5-fold CV) for Logistic Regression...")
print("Testing 6 parameter combinations x 5 folds = 30 fits\n")
start = time.time()

grid_search = GridSearchCV(
    LogisticRegression(random_state=SEED),
    param_grid,
    cv=5,
    scoring='f1',
    n_jobs=-1,
    verbose=1
)
grid_search.fit(X_train_tfidf, train_df['label'])

lr_train_time = time.time() - start

print(f"\nBest parameters: {grid_search.best_params_}")
print(f"Best CV F1-score: {grid_search.best_score_:.4f}")
print(f"Total training time (incl. CV): {lr_train_time:.1f}s")

# Show all results
results_df = pd.DataFrame(grid_search.cv_results_)[
    ['param_C', 'param_penalty', 'mean_test_score', 'std_test_score', 'rank_test_score']
].sort_values('rank_test_score')
print("\nAll CV Results:")
results_df

# ============================================================
# Cell 10: Evaluate TF-IDF + Logistic Regression on Test Set
# ============================================================
lr_model = grid_search.best_estimator_

start = time.time()
y_pred_lr = lr_model.predict(X_test_tfidf)
lr_inference_time = time.time() - start

y_test = test_df['label'].values

lr_accuracy  = accuracy_score(y_test, y_pred_lr)
lr_precision = precision_score(y_test, y_pred_lr)
lr_recall    = recall_score(y_test, y_pred_lr)
lr_f1        = f1_score(y_test, y_pred_lr)

print("=" * 55)
print("  TF-IDF + LOGISTIC REGRESSION - TEST SET RESULTS")
print("=" * 55)
print(f"  Accuracy:  {lr_accuracy:.4f}")
print(f"  Precision: {lr_precision:.4f}")
print(f"  Recall:    {lr_recall:.4f}")
print(f"  F1-Score:  {lr_f1:.4f}")
print(f"  Inference time: {lr_inference_time:.4f}s for {len(y_test):,} samples")
print("=" * 55)
print(f"\n{classification_report(y_test, y_pred_lr, target_names=['Negative', 'Positive'])}")

# ============================================================
# Cell 11: Confusion Matrix - Logistic Regression
# ============================================================
fig, ax = plt.subplots(figsize=(7, 6))
cm_lr = confusion_matrix(y_test, y_pred_lr)
disp = ConfusionMatrixDisplay(cm_lr, display_labels=['Negative', 'Positive'])
disp.plot(cmap='Blues', ax=ax, values_format='d')
ax.set_title('TF-IDF + Logistic Regression - Confusion Matrix', fontsize=14)
plt.tight_layout()
plt.show()

# Print the raw numbers
print(f"\nTrue Negatives:  {cm_lr[0][0]:,}")
print(f"False Positives: {cm_lr[0][1]:,}")
print(f"False Negatives: {cm_lr[1][0]:,}")
print(f"True Positives:  {cm_lr[1][1]:,}")

# ============================================================
# Cell 12: Top 20 Positive and Negative Features
# ============================================================
feature_names = np.array(tfidf.get_feature_names_out())
coefs = lr_model.coef_[0]

top_pos_idx = np.argsort(coefs)[-20:][::-1]
top_neg_idx = np.argsort(coefs)[:20]

fig, axes = plt.subplots(1, 2, figsize=(16, 7))

axes[0].barh(range(20), coefs[top_pos_idx], color='#27ae60', edgecolor='black')
axes[0].set_yticks(range(20))
axes[0].set_yticklabels(feature_names[top_pos_idx])
axes[0].set_title('Top 20 Positive Sentiment Features', fontsize=14)
axes[0].set_xlabel('Coefficient Weight')
axes[0].invert_yaxis()

axes[1].barh(range(20), coefs[top_neg_idx], color='#e74c3c', edgecolor='black')
axes[1].set_yticks(range(20))
axes[1].set_yticklabels(feature_names[top_neg_idx])
axes[1].set_title('Top 20 Negative Sentiment Features', fontsize=14)
axes[1].set_xlabel('Coefficient Weight')
axes[1].invert_yaxis()

plt.tight_layout()
plt.show()

# ============================================================
# Cell 13: Prepare Data Subsets and Tokenise
# ============================================================

# Use 5,000 samples for manageable training time on free Colab GPU
TRAIN_SIZE = 5000
TEST_SIZE  = 5000
MAX_LENGTH = 256
MODEL_NAME = 'distilbert-base-uncased'

# Sample subsets
train_bert_df = train_df.sample(n=TRAIN_SIZE, random_state=SEED).reset_index(drop=True)
test_bert_df  = test_df.sample(n=TEST_SIZE, random_state=SEED).reset_index(drop=True)

# Split 10% of training data as validation set
val_size = int(0.1 * len(train_bert_df))
val_bert_df   = train_bert_df.iloc[:val_size].reset_index(drop=True)
train_bert_df = train_bert_df.iloc[val_size:].reset_index(drop=True)

print(f"Training samples:   {len(train_bert_df)}")
print(f"Validation samples: {len(val_bert_df)}")
print(f"Test samples:       {len(test_bert_df)}")

# Load tokeniser
print(f"\nLoading tokeniser: {MODEL_NAME}")
tokenizer = DistilBertTokenizerFast.from_pretrained(MODEL_NAME)

# Tokenise all splits
print("Tokenising training data...")
train_encodings = tokenizer(
    train_bert_df['bert_text'].tolist(),
    truncation=True, padding=True,
    max_length=MAX_LENGTH, return_tensors='pt'
)
print("Tokenising validation data...")
val_encodings = tokenizer(
    val_bert_df['bert_text'].tolist(),
    truncation=True, padding=True,
    max_length=MAX_LENGTH, return_tensors='pt'
)
print("Tokenising test data...")
test_encodings = tokenizer(
    test_bert_df['bert_text'].tolist(),
    truncation=True, padding=True,
    max_length=MAX_LENGTH, return_tensors='pt'
)

print(f"\nTokenised input shape: {train_encodings['input_ids'].shape}")
print("Done!")

# ============================================================
# Cell 14: Custom PyTorch Dataset and Model Initialisation
# ============================================================
class IMDBDataset(Dataset):
    """Custom PyTorch Dataset wrapping tokenised IMDB reviews."""
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

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

    def __getitem__(self, idx):
        item = {key: val[idx] for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx], dtype=torch.long)
        return item

train_dataset = IMDBDataset(train_encodings, train_bert_df['label'].tolist())
val_dataset   = IMDBDataset(val_encodings, val_bert_df['label'].tolist())
test_dataset  = IMDBDataset(test_encodings, test_bert_df['label'].tolist())

# Load the pre-trained model with a classification head
model = DistilBertForSequenceClassification.from_pretrained(
    MODEL_NAME, num_labels=2
)

total_params     = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f"Datasets created:")
print(f"  Train: {len(train_dataset)}")
print(f"  Val:   {len(val_dataset)}")
print(f"  Test:  {len(test_dataset)}")
print(f"\nModel: {MODEL_NAME}")
print(f"Total parameters:     {total_params:,}")
print(f"Trainable parameters: {trainable_params:,}")

# ============================================================
# Cell 15: Fine-tune DistilBERT
# ============================================================

def compute_metrics(eval_pred):
    """Compute metrics for the Trainer."""
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    return {
        'accuracy':  accuracy_score(labels, preds),
        'precision': precision_score(labels, preds),
        'recall':    recall_score(labels, preds),
        'f1':        f1_score(labels, preds)
    }

training_args = TrainingArguments(
    output_dir='./distilbert_imdb',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    learning_rate=2e-5,
    weight_decay=0.01,
    warmup_steps=500,
    eval_strategy='epoch',
    save_strategy='epoch',
    load_best_model_at_end=True,
    metric_for_best_model='f1',
    logging_dir='./logs',
    logging_steps=50,
    seed=SEED,
    fp16=True,            # mixed precision for faster training
    report_to='none'
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=2)]
)

print("Starting DistilBERT fine-tuning...")
print("Config: 3 epochs, batch_size=16, lr=2e-5, fp16=True")
print("-" * 50)
start = time.time()
train_result = trainer.train()
bert_train_time = time.time() - start

print(f"\n{'='*50}")
print(f"Training completed in {bert_train_time:.1f}s ({bert_train_time/60:.1f} min)")
print(f"Final training loss: {train_result.training_loss:.4f}")
print(f"{'='*50}")

# ============================================================
# Cell 16: Plot Training and Validation Curves
# ============================================================
log_history = trainer.state.log_history

train_losses = [x['loss'] for x in log_history
                if 'loss' in x and 'eval_loss' not in x]
eval_entries  = [x for x in log_history if 'eval_loss' in x]
eval_losses   = [x['eval_loss'] for x in eval_entries]
eval_f1s      = [x['eval_f1'] for x in eval_entries]
eval_accs     = [x['eval_accuracy'] for x in eval_entries]

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Training loss over steps
axes[0].plot(train_losses, color='#2980b9', linewidth=1.5)
axes[0].set_title('Training Loss Over Steps', fontsize=14)
axes[0].set_xlabel('Logging Step')
axes[0].set_ylabel('Loss')
axes[0].grid(True, alpha=0.3)

# Validation metrics per epoch
epochs = range(1, len(eval_losses) + 1)
axes[1].plot(epochs, eval_losses, 'o-', label='Eval Loss', color='#e74c3c', linewidth=2)
axes[1].plot(epochs, eval_f1s, 's-', label='Eval F1', color='#27ae60', linewidth=2)
axes[1].plot(epochs, eval_accs, '^-', label='Eval Accuracy', color='#f39c12', linewidth=2)
axes[1].set_title('Validation Metrics Per Epoch', fontsize=14)
axes[1].set_xlabel('Epoch')
axes[1].legend(fontsize=11)
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# Print epoch results
print("\nEpoch-by-epoch validation results:")
for i, entry in enumerate(eval_entries):
    print(f"  Epoch {i+1}: Loss={entry['eval_loss']:.4f}, "
          f"F1={entry['eval_f1']:.4f}, Acc={entry['eval_accuracy']:.4f}")

# ============================================================
# Cell 17: Evaluate DistilBERT on Test Set
# ============================================================
start = time.time()
bert_results = trainer.predict(test_dataset)
bert_inference_time = time.time() - start

y_pred_bert  = np.argmax(bert_results.predictions, axis=-1)
y_test_bert  = test_bert_df['label'].values

bert_accuracy  = accuracy_score(y_test_bert, y_pred_bert)
bert_precision = precision_score(y_test_bert, y_pred_bert)
bert_recall    = recall_score(y_test_bert, y_pred_bert)
bert_f1        = f1_score(y_test_bert, y_pred_bert)

print("=" * 55)
print("     DISTILBERT - TEST SET RESULTS")
print("=" * 55)
print(f"  Accuracy:  {bert_accuracy:.4f}")
print(f"  Precision: {bert_precision:.4f}")
print(f"  Recall:    {bert_recall:.4f}")
print(f"  F1-Score:  {bert_f1:.4f}")
print(f"  Inference time: {bert_inference_time:.2f}s for {len(y_test_bert):,} samples")
print("=" * 55)
print(f"\n{classification_report(y_test_bert, y_pred_bert, target_names=['Negative', 'Positive'])}")

# ============================================================
# Cell 18: Confusion Matrix - DistilBERT
# ============================================================
fig, ax = plt.subplots(figsize=(7, 6))
cm_bert = confusion_matrix(y_test_bert, y_pred_bert)
disp = ConfusionMatrixDisplay(cm_bert, display_labels=['Negative', 'Positive'])
disp.plot(cmap='Oranges', ax=ax, values_format='d')
ax.set_title('DistilBERT - Confusion Matrix', fontsize=14)
plt.tight_layout()
plt.show()

print(f"\nTrue Negatives:  {cm_bert[0][0]:,}")
print(f"False Positives: {cm_bert[0][1]:,}")
print(f"False Negatives: {cm_bert[1][0]:,}")
print(f"True Positives:  {cm_bert[1][1]:,}")

# ============================================================
# Cell 19: Model Comparison Summary
# ============================================================
comparison = pd.DataFrame({
    'Metric': ['Accuracy', 'Precision', 'Recall', 'F1-Score',
               'Train Time (s)', 'Inference Time (s)', 'Parameters'],
    'TF-IDF + LR': [
        f"{lr_accuracy:.4f}", f"{lr_precision:.4f}",
        f"{lr_recall:.4f}", f"{lr_f1:.4f}",
        f"{lr_train_time:.1f}", f"{lr_inference_time:.4f}",
        f"{X_train_tfidf.shape[1]:,} features"
    ],
    'DistilBERT': [
        f"{bert_accuracy:.4f}", f"{bert_precision:.4f}",
        f"{bert_recall:.4f}", f"{bert_f1:.4f}",
        f"{bert_train_time:.1f}", f"{bert_inference_time:.2f}",
        f"{trainable_params:,}"
    ]
})

print("\n" + "=" * 65)
print("          MODEL COMPARISON SUMMARY")
print("=" * 65)
print(comparison.to_string(index=False))
print("=" * 65)

# ============================================================
# Cell 20: Visual Performance Comparison
# ============================================================
metrics_names = ['Accuracy', 'Precision', 'Recall', 'F1-Score']
lr_scores   = [lr_accuracy, lr_precision, lr_recall, lr_f1]
bert_scores = [bert_accuracy, bert_precision, bert_recall, bert_f1]

x = np.arange(len(metrics_names))
width = 0.35

fig, ax = plt.subplots(figsize=(10, 6))
bars1 = ax.bar(x - width/2, lr_scores, width, label='TF-IDF + LR',
               color='#3498db', edgecolor='black')
bars2 = ax.bar(x + width/2, bert_scores, width, label='DistilBERT',
               color='#e67e22', edgecolor='black')

for bars in [bars1, bars2]:
    for bar in bars:
        h = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2., h + 0.003,
                f'{h:.3f}', ha='center', va='bottom',
                fontsize=11, fontweight='bold')

ax.set_ylim(0.78, 1.02)
ax.set_ylabel('Score', fontsize=12)
ax.set_title('Model Performance Comparison', fontsize=15)
ax.set_xticks(x)
ax.set_xticklabels(metrics_names, fontsize=12)
ax.legend(fontsize=12)
ax.grid(axis='y', alpha=0.3)
plt.tight_layout()
plt.show()

# ============================================================
# Cell 21: SHAP Analysis - Global Feature Importance
# ============================================================
print("Computing SHAP values for Logistic Regression...")
print("(This may take 1-2 minutes)\n")

background = X_train_tfidf[:100]
explainer  = shap.LinearExplainer(lr_model, background)
shap_values = explainer.shap_values(X_test_tfidf[:500])

# Global importance: mean |SHAP value|
mean_shap    = np.abs(shap_values).mean(axis=0)
top_shap_idx = np.argsort(mean_shap)[-20:][::-1]

fig, ax = plt.subplots(figsize=(10, 7))
ax.barh(range(20), mean_shap[top_shap_idx], color='#8e44ad', edgecolor='black')
ax.set_yticks(range(20))
ax.set_yticklabels(feature_names[top_shap_idx])
ax.set_title('Top 20 Features by Mean |SHAP Value|', fontsize=14)
ax.set_xlabel('Mean |SHAP Value|')
ax.invert_yaxis()
plt.tight_layout()
plt.show()

# ============================================================
# Cell 22: SHAP Beeswarm Plot
# ============================================================
shap_display = shap.Explanation(
    values=shap_values[:200],
    data=X_test_tfidf[:200].toarray(),
    feature_names=feature_names
)

plt.figure(figsize=(10, 8))
shap.plots.beeswarm(shap_display, max_display=20, show=False)
plt.title('SHAP Beeswarm Plot - Feature Impact on Sentiment', fontsize=13)
plt.tight_layout()
plt.show()

# ============================================================
# Cell 23: SHAP Local Explanation for a Single Review
# ============================================================
sample_idx  = 0
sample_text = test_df['clean_text'].iloc[sample_idx]
sample_true = test_df['label'].iloc[sample_idx]
sample_pred = y_pred_lr[sample_idx]

print(f"Review (first 200 chars): {sample_text[:200]}...")
print(f"True label:  {'Positive' if sample_true else 'Negative'}")
print(f"Predicted:   {'Positive' if sample_pred else 'Negative'}\n")

# Get non-zero SHAP values for this sample
sample_shap = shap_values[sample_idx]
nonzero_idx = np.nonzero(X_test_tfidf[sample_idx].toarray()[0])[0]
sample_features  = feature_names[nonzero_idx]
sample_shap_vals = sample_shap[nonzero_idx]
sort_idx = np.argsort(np.abs(sample_shap_vals))[::-1][:15]

fig, ax = plt.subplots(figsize=(10, 6))
colors = ['#27ae60' if v > 0 else '#e74c3c' for v in sample_shap_vals[sort_idx]]
ax.barh(range(len(sort_idx)), sample_shap_vals[sort_idx], color=colors, edgecolor='black')
ax.set_yticks(range(len(sort_idx)))
ax.set_yticklabels(sample_features[sort_idx])
ax.set_title('SHAP Values - Single Review Explanation', fontsize=14)
ax.set_xlabel('SHAP Value (green = pushes Positive, red = pushes Negative)')
ax.invert_yaxis()
ax.axvline(x=0, color='black', linewidth=0.8)
plt.tight_layout()
plt.show()

# ============================================================
# Cell 24: Attention Visualisation for DistilBERT (FIXED v3)
# ============================================================
import copy

# Step 1: Save trained weights
trained_weights = copy.deepcopy(model.state_dict())

# Step 2: Reload model from scratch with eager attention
model = DistilBertForSequenceClassification.from_pretrained(
    'distilbert-base-uncased',
    num_labels=2,
    attn_implementation="eager",
    output_attentions=True
)

# Step 3: Load our trained weights back in
model.load_state_dict(trained_weights, strict=False)
model.cpu()
model.eval()
print("Model reloaded with eager attention successfully!")

def visualise_attention(text, model, tokenizer, layer=-1, head=0):
    """Extract [CLS] attention weights from a DistilBERT layer."""
    inputs = tokenizer(text, return_tensors='pt', truncation=True,
                       max_length=256, padding=True)

    with torch.no_grad():
        outputs = model(**inputs)

    attention = outputs.attentions[layer][0, head].cpu().numpy()
    cls_attention = attention[0]
    tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])

    valid_mask = [t not in ['[CLS]', '[SEP]', '[PAD]'] for t in tokens]
    filtered_tokens = [t for t, m in zip(tokens, valid_mask) if m]
    filtered_attn   = cls_attention[valid_mask]
    filtered_attn   = filtered_attn / filtered_attn.sum()

    return filtered_tokens, filtered_attn

pos_review = test_bert_df[test_bert_df['label']==1]['bert_text'].iloc[0]
neg_review = test_bert_df[test_bert_df['label']==0]['bert_text'].iloc[0]

fig, axes = plt.subplots(2, 1, figsize=(16, 10))

for idx, (review, label) in enumerate(
        [(pos_review, 'Positive'), (neg_review, 'Negative')]):
    tokens, attn = visualise_attention(review[:500], model, tokenizer)
    n_show  = min(25, len(tokens))
    top_idx = np.argsort(attn)[-n_show:][::-1]
    colours_map = plt.cm.YlOrRd(attn[top_idx] / attn[top_idx].max())

    axes[idx].barh(range(n_show), attn[top_idx],
                   color=colours_map, edgecolor='gray')
    axes[idx].set_yticks(range(n_show))
    axes[idx].set_yticklabels([tokens[i] for i in top_idx], fontsize=9)
    axes[idx].set_title(
        f'DistilBERT [CLS] Attention - {label} Review (Last Layer, Head 0)',
        fontsize=13)
    axes[idx].set_xlabel('Attention Weight')
    axes[idx].invert_yaxis()

plt.tight_layout()
plt.show()

# ============================================================
# Cell 25: Multi-Head Attention Heatmap
# ============================================================

def attention_heatmap(text, model, tokenizer, layer=-1, n_tokens=20):
    """Heatmap of all attention heads for [CLS] token."""
    model.eval()
    inputs = tokenizer(text, return_tensors='pt', truncation=True,
                       max_length=256, padding=True)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model(**inputs, output_attentions=True)

    all_heads = outputs.attentions[layer][0, :, 0, :].cpu().numpy()
    tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])

    valid = [t not in ['[CLS]', '[SEP]', '[PAD]'] for t in tokens]
    valid_idx = [i for i, m in enumerate(valid) if m]
    filt_tokens = [tokens[i] for i in valid_idx[:n_tokens]]
    filt_attn   = all_heads[:, valid_idx[:n_tokens]]
    return filt_tokens, filt_attn

tokens, attn_matrix = attention_heatmap(pos_review[:400], model, tokenizer)

fig, ax = plt.subplots(figsize=(14, 5))
im = ax.imshow(attn_matrix, cmap='YlOrRd', aspect='auto')
ax.set_xticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=45, ha='right', fontsize=9)
ax.set_yticks(range(attn_matrix.shape[0]))
ax.set_yticklabels([f'Head {i}' for i in range(attn_matrix.shape[0])], fontsize=10)
ax.set_title('DistilBERT Multi-Head Attention Heatmap (Last Layer)', fontsize=14)
plt.colorbar(im, ax=ax, label='Attention Weight')
plt.tight_layout()
plt.show()

# ============================================================
# Cell 26: Analyse Misclassified Reviews
# ============================================================
lr_wrong   = y_pred_lr != y_test
bert_wrong = y_pred_bert != y_test_bert

print(f"Logistic Regression errors: {lr_wrong.sum():,} / {len(y_test):,} "
      f"({lr_wrong.mean()*100:.1f}%)")
print(f"DistilBERT errors:          {bert_wrong.sum():,} / {len(y_test_bert):,} "
      f"({bert_wrong.mean()*100:.1f}%)")

# Show examples
wrong_idx = np.where(lr_wrong)[0][:3]
print("\n" + "=" * 60)
print("Sample Misclassified Reviews (Logistic Regression)")
print("=" * 60)
for i in wrong_idx:
    true_lbl = 'Positive' if y_test[i] == 1 else 'Negative'
    pred_lbl = 'Positive' if y_pred_lr[i] == 1 else 'Negative'
    print(f"\n[Review {i}] True: {true_lbl} | Predicted: {pred_lbl}")
    print(f"  {test_df['text'].iloc[i][:200]}...")

# ============================================================
# Cell 27: Accuracy by Review Length
# ============================================================
test_df['word_count'] = test_df['text'].apply(lambda x: len(x.split()))
test_df['lr_correct'] = ~lr_wrong

bins   = [0, 100, 200, 300, 500, 1000, 5000]
labels = ['<100', '100-200', '200-300', '300-500', '500-1k', '1000+']
test_df['length_bin'] = pd.cut(test_df['word_count'], bins=bins, labels=labels)

acc_by_len = test_df.groupby('length_bin')['lr_correct'].mean()

fig, ax = plt.subplots(figsize=(10, 5))
acc_by_len.plot(kind='bar', ax=ax, color='#3498db', edgecolor='black')
for i, v in enumerate(acc_by_len.values):
    ax.text(i, v + 0.003, f'{v:.3f}', ha='center', fontweight='bold')
ax.set_title('Logistic Regression Accuracy by Review Length', fontsize=14)
ax.set_ylabel('Accuracy')
ax.set_xlabel('Review Word Count')
ax.set_ylim(0.80, 1.0)
plt.xticks(rotation=0)
plt.tight_layout()
plt.show()

# ============================================================
# Cell 28: Final Summary
# ============================================================
print("=" * 65)
print("   EXPERIMENT COMPLETE - FINAL SUMMARY")
print("=" * 65)
print(f"\n{'Model':<25} {'Accuracy':>10} {'F1':>10} {'Train(s)':>10} {'Infer(s)':>10}")
print("-" * 65)
print(f"{'TF-IDF + Log. Reg.':<25} {lr_accuracy:>10.4f} {lr_f1:>10.4f} "
      f"{lr_train_time:>10.1f} {lr_inference_time:>10.4f}")
print(f"{'DistilBERT':<25} {bert_accuracy:>10.4f} {bert_f1:>10.4f} "
      f"{bert_train_time:>10.1f} {bert_inference_time:>10.2f}")
print("-" * 65)
improvement = (bert_f1 - lr_f1) / lr_f1 * 100
print(f"\nDistilBERT F1 improvement over LR: +{improvement:.1f}%")
print(f"DistilBERT training time ratio:    {bert_train_time/lr_train_time:.0f}x slower")
print(f"DistilBERT inference time ratio:   {bert_inference_time/lr_inference_time:.0f}x slower")
print("\nNotebook execution complete!")