# GENAI Experiment 8: Chatbot: tokenizer, TF-IDF + cosine similarity, context memory, exit loop
# ===========================================================================

# # Part 1 — Knowledge base and tokenization

import re, math, random
import numpy as np
random.seed(0)

KB = {
    "what is artificial intelligence": "Artificial intelligence is the science of making machines perform tasks that need human intelligence.",
    "what is machine learning": "Machine learning lets computers learn patterns from data instead of being explicitly programmed.",
    "what is deep learning": "Deep learning uses neural networks with many layers to learn features directly from data.",
    "what is a neural network": "A neural network is layers of connected neurons that compute weighted sums followed by activations.",
    "what is a gan generative adversarial network": "A GAN has a generator that makes fake data and a discriminator that tries to tell real from fake.",
    "what is a vae variational autoencoder": "A VAE encodes data to a distribution in latent space and decodes samples from it to generate new data.",
    "what is a transformer attention": "A transformer uses self-attention so every word can look at every other word to understand context.",
    "what is a diffusion model": "A diffusion model adds noise to data step by step and learns to remove it, turning noise into new samples.",
    "how do i study for exams tips": "Make a timetable, revise in short sessions, practise previous papers and sleep well.",
    "hello hi hey good morning": "Hello! I am your GenAI lab assistant. Ask me about AI topics.",
    "how are you": "I am just a program, but I am working fine. How can I help you?",
    "thank you thanks": "You are welcome!",
}
MORE = {   # second-level answers used for "tell me more" (context)
    "what is machine learning": "There are three main types: supervised, unsupervised and reinforcement learning.",
    "what is deep learning": "Popular deep networks are CNNs for images, RNNs/LSTMs for sequences and transformers for language.",
    "what is a gan generative adversarial network": "Training alternates: update the discriminator on real and fake, then update the generator to fool it.",
    "what is a vae variational autoencoder": "Its loss is reconstruction loss plus KL divergence, and it uses the reparameterization trick z = mu + sigma * eps.",
    "what is a diffusion model": "The network is trained to predict the added noise; sampling runs the denoising steps from t = T down to 1.",
}
STOP = {"what", "is", "a", "an", "the", "of", "me", "about", "please", "can", "you", "tell", "do", "i", "to", "how", "it"}

def tokenize(text):
    return [w for w in re.findall(r"[a-z]+", text.lower()) if w not in STOP]

print(tokenize("What is Machine Learning, please?"))

# # Part 2 — TF-IDF vectors and cosine similarity (the "model")

questions = list(KB.keys())
vocab = sorted({w for q in questions for w in tokenize(q)})
widx = {w: i for i, w in enumerate(vocab)}
df = np.zeros(len(vocab))
for q in questions:
    for w in set(tokenize(q)): df[widx[w]] += 1
idf = np.log((1 + len(questions)) / (1 + df)) + 1                 # inverse document frequency

def tfidf(text):
    v = np.zeros(len(vocab))
    for w in tokenize(text):
        if w in widx: v[widx[w]] += 1                              # term frequency
    v *= idf
    n = np.linalg.norm(v); return v / n if n > 0 else v

Q = np.stack([tfidf(q) for q in questions])                      # matrix: one row per known question
def best_match(text):
    sims = Q @ tfidf(text)                                       # cosine similarity (vectors are normalised)
    i = int(np.argmax(sims)); return questions[i], float(sims[i])

print(best_match("explain neural networks"))

# # Part 3 — Response generation with context and a confidence threshold

state = {"last_topic": None}
FALLBACK = ["Sorry, I don't know that yet. Try asking about AI, ML, GANs, VAEs or diffusion.",
            "I am not sure. Could you rephrase your question?"]

def chatbot_reply(user_text, threshold=0.3):
    t = user_text.lower()
    if re.search(r"\b(more|elaborate|explain it|what about it|continue)\b", t) and state["last_topic"]:
        return MORE.get(state["last_topic"], "That's all I know about this topic.")      # context
    q, score = best_match(user_text)
    if score < threshold:
        return random.choice(FALLBACK)                                                  # low confidence
    state["last_topic"] = q
    return KB[q]

for q in ["Hi!", "What is deep learning?", "Tell me more", "What's a GAN?", "can you elaborate",
          "what is the capital of France", "how do I prepare for exams", "thanks"]:
    print(f"You: {q}\nBot: {chatbot_reply(q)}\n")

# # Part 4 — Interactive loop (type exit to stop)

def run_chatbot():
    print("GenAI lab chatbot - type 'exit' to quit")
    while True:
        try:
            user_text = input("You: ")
        except Exception:                      # no keyboard (e.g. auto-run) -> short demo
            for user_text in ["what is a vae", "tell me more", "exit"]:
                print("You:", user_text)
                if user_text == "exit": print("Bot: Goodbye!"); return
                print("Bot:", chatbot_reply(user_text))
            return
        if user_text.strip().lower() == "exit":
            print("Bot: Goodbye!"); break
        print("Bot:", chatbot_reply(user_text))

run_chatbot()

# Exam explanation: the input is tokenized (lowercase words without stop-words), turned into a TF-IDF
# vector and compared with every known question by cosine similarity. The most similar answer is
# returned. A threshold rejects unrelated questions (e.g. capitals of countries) instead of
# hallucinating, and a last-topic memory gives context-aware follow-ups ("tell me more"). Compared
# with DistilGPT2: retrieval answers are always correct for known topics but cannot invent new
# sentences, while the pretrained model is fluent but often wrong.
