# -- Code Cell --
import nltk
import numpy as np
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer

nltk.download('punkt')
nltk.download('punkt_tab')


def extract_text_features(texts):
    # 1. Tokenization Example
    tokens = nltk.word_tokenize(texts[0])
    print("Sample Tokens:", tokens)

    # 2. Bag of Words
    count_vectorizer = CountVectorizer()
    bow_matrix = count_vectorizer.fit_transform(texts)

    # 3. TF-IDF
    tfidf_vectorizer = TfidfVectorizer()
    tfidf_matrix = tfidf_vectorizer.fit_transform(texts)

    # 4. Basic Text Statistics
    word_counts = [len(nltk.word_tokenize(text)) for text in texts]
    avg_word_length = [
        np.mean([len(word) for word in nltk.word_tokenize(text)])
        for text in texts
    ]

    feature_summary = {
        "vocabulary_size": len(count_vectorizer.vocabulary_),
        "bow_shape": bow_matrix.shape,
        "tfidf_shape": tfidf_matrix.shape,
        "word_counts": word_counts,
        "avg_word_length": avg_word_length
    }

    return feature_summary


if __name__ == "__main__":
    texts = [
        "Machine learning is amazing.",
        "Feature extraction converts data into numerical format.",
        "Python is widely used in AI."
    ]

    features = extract_text_features(texts)

    print("\nExtracted Text Features:")
    for key, value in features.items():
        print(f"{key}: {value}")


# -- Code Cell --
import torch
from transformers import AutoTokenizer, AutoModel


def extract_bert_features(texts):
    # Load pretrained tokenizer & model
    model_name = "bert-base-uncased"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name)

    model.eval()  # evaluation mode

    # Tokenize input
    inputs = tokenizer(
        texts,
        padding=True,
        truncation=True,
        return_tensors="pt"
    )

    # Forward pass (no gradient needed)
    with torch.no_grad():
        outputs = model(**inputs)

    # CLS token embedding (sentence representation)
    cls_embeddings = outputs.last_hidden_state[:, 0, :]

    print("Embedding shape:", cls_embeddings.shape)
    return cls_embeddings


if __name__ == "__main__":
    texts = [
        "Deep learning is powerful.",
        "Natural language processing uses neural networks.",
        "Transformers changed NLP."
    ]

    features = extract_bert_features(texts)

    print("\nSample Feature Vector (first sentence):")
    print(features[0])


# -- Code Cell --
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np


def get_bert_embeddings(texts, model_name="bert-base-uncased"):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name)
    model.eval()

    inputs = tokenizer(
        texts,
        padding=True,
        truncation=True,
        return_tensors="pt"
    )

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

    # CLS token embedding
    embeddings = outputs.last_hidden_state[:, 0, :]
    return embeddings


def compute_distances(embeddings):
    embeddings_np = embeddings.numpy()

    # Cosine Similarity
    cos_sim_matrix = cosine_similarity(embeddings_np)

    # Euclidean Distance
    euclidean_matrix = np.linalg.norm(
        embeddings_np[:, None] - embeddings_np, axis=2
    )

    return cos_sim_matrix, euclidean_matrix


if __name__ == "__main__":

    texts = [
        "I love machine learning.",
        "Artificial intelligence is fascinating.",
        "The weather is sunny today.",
        "Deep learning models are powerful."
    ]

    embeddings = get_bert_embeddings(texts)

    cos_sim, euclidean_dist = compute_distances(embeddings)

    print("\nCosine Similarity Matrix:\n")
    print(np.round(cos_sim, 3))

    print("\nEuclidean Distance Matrix:\n")
    print(np.round(euclidean_dist, 3))

    # Find most similar pair (excluding self similarity)
    np.fill_diagonal(cos_sim, -1)
    most_similar = np.unravel_index(np.argmax(cos_sim), cos_sim.shape)

    print("\nMost similar texts:")
    print(f"Text 1: {texts[most_similar[0]]}")
    print(f"Text 2: {texts[most_similar[1]]}")


# -- Code Cell --
import torch
import torch.nn as nn
import torch.optim as optim
from transformers import AutoTokenizer, AutoModel
from sklearn.model_selection import train_test_split


# -----------------------
# 1. Sample Dataset
# -----------------------
texts = [
    "Machine learning is amazing",
    "Deep learning models are powerful",
    "Artificial intelligence is the future",
    "The weather is sunny today",
    "I love going to the beach",
    "It is raining outside"
]

# Labels: 1 = AI-related, 0 = Not AI-related
labels = torch.tensor([1, 1, 1, 0, 0, 0])


# -----------------------
# 2. Model Definition
# -----------------------
class BERTClassifier(nn.Module):
    def __init__(self, model_name="bert-base-uncased", num_classes=2):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name)
        self.fc = nn.Linear(768, num_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
        cls_output = outputs.last_hidden_state[:, 0, :]
        logits = self.fc(cls_output)
        return logits


# -----------------------
# 3. Prepare Data
# -----------------------
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

encodings = tokenizer(
    texts,
    padding=True,
    truncation=True,
    return_tensors="pt"
)

input_ids = encodings["input_ids"]
attention_mask = encodings["attention_mask"]

# Train-test split
X_train_ids, X_test_ids, y_train, y_test, train_mask, test_mask = train_test_split(
    input_ids, labels, attention_mask, test_size=0.3, random_state=42
)


# -----------------------
# 4. Training Setup
# -----------------------
model = BERTClassifier()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=2e-5)

model.train()

epochs = 15

for epoch in range(epochs):
    optimizer.zero_grad()

    outputs = model(X_train_ids, train_mask)
    loss = criterion(outputs, y_train)

    loss.backward()
    optimizer.step()

    print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f}")


# -----------------------
# 5. Evaluation
# -----------------------
model.eval()
with torch.no_grad():
    outputs = model(X_test_ids, test_mask)
    predictions = torch.argmax(outputs, dim=1)

print("\nTest Predictions:", predictions)
print("True Labels:     ", y_test)


# -----------------------
# 6. Test New Sentence
# -----------------------
def predict(text):
    model.eval()
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model(
            inputs["input_ids"],
            inputs["attention_mask"]
        )
    prediction = torch.argmax(outputs, dim=1).item()
    return prediction


new_text = "Neural networks are part of AI"
prediction = predict(new_text)

print("\nNew Text:", new_text)
print("Predicted Class:", prediction)


# -- Code Cell --
