import numpy as np
from openai import OpenAI
import tiktoken
import os

# Connect to the local LM Studio server using OpenAI client.
client = OpenAI(base_url="http://localhost:1234/v1", api_key="lm-studio")  # Ensure LM Studio is running on this URL.

# Define your model identifiers based on LM Studio setup
EMBEDDING_MODEL = "embedding-model-name"  # Replace with actual embedding model name
CHAT_COMPLETION_MODEL = "chat-completion-model-name"  # Replace with actual chat completion model name


# Step 1: Load data from a text file
def load_corpus(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            return file.read()
    except FileNotFoundError:
        print(f"Error: The file '{file_path}' was not found.")
        return None
    except Exception as e:
        print(f"An error occurred while reading the file: {e}")
        return None


# Step 2: Split the text into manageable chunks
def split_into_chunks(text, chunk_size=50, overlap=10):
    encoding = tiktoken.get_encoding("cl100k_base")
    tokens = encoding.encode(text)
    chunks = []
    for i in range(0, len(tokens), chunk_size - overlap):
        chunk = tokens[i:i + chunk_size]
        decoded_chunk = encoding.decode(chunk)
        chunks.append(decoded_chunk)
    return chunks


# Step 3: Get embeddings for the text chunks
def get_embeddings(texts):
    try:
        print("Sending chunks to the LLM for embeddings...")
        response = client.embeddings.create(
            model=EMBEDDING_MODEL,  # Use your actual embeddings model
            input=texts
        )
        print("Received embeddings from the LLM.")
        return [embedding.embedding for embedding in response.data]
    except Exception as e:
        print(f"An error occurred while getting embeddings: {e}")
        return None


# Step 4: Calculate cosine similarity
def cosine_similarity(vec1, vec2):
    vec1 = np.array(vec1)
    vec2 = np.array(vec2)
    norm1 = np.linalg.norm(vec1)
    norm2 = np.linalg.norm(vec2)
    if norm1 == 0 or norm2 == 0:
        return 0.0
    return np.dot(vec1, vec2) / (norm1 * norm2)


# Step 5: Retrieve relevant text chunks
def retrieve_relevant_chunks(query, chunks, chunk_embeddings, top_n=10):
    print(f"\nRetrieving relevant chunks for the query: '{query}'")
    query_embedding = get_embeddings([query])
    if query_embedding is None:
        return []
    query_embedding = query_embedding[0]
    scores = [(cosine_similarity(query_embedding, chunk_embedding), chunk)
              for chunk, chunk_embedding in zip(chunks, chunk_embeddings)]
    scores.sort(reverse=True, key=lambda x: x[0])
    top_chunks = [(score, chunk) for score, chunk in scores[:top_n] if score > 0]
    print(f"Found {len(top_chunks)} relevant chunks.")
    for score, chunk in top_chunks:
        print(f"Score: {score:.4f}, Chunk: \"{chunk}\"")
    return top_chunks


# Step 6: Check for contradictions
def check_for_contradictions(user_statement, contextual_chunks):
    print(f"\nChecking for contradictions against the user statement: '{user_statement}'")
    context = "\n".join(chunk for _, chunk in contextual_chunks)

    messages = [
        {"role": "system", "content": "Determine if the user statement contradicts the provided context."},
        {"role": "user", "content": f"Context: {context}\nUser Statement: {user_statement}"}
    ]

    print(f"Complete Prompt:\n{messages}")

    try:
        completion = client.chat.completions.create(
            model=CHAT_COMPLETION_MODEL,  # Use your actual chat completion model
            messages=messages,
            temperature=0.5,
        )
        response_message = completion.choices[0].message.content
        print("Contradiction check complete.")
        return response_message
    except Exception as e:
        print(f"An error occurred during contradiction check: {e}")
        return "Error occurred during contradiction check."


# Step 7: Calculate word importance
def word_importance_calculation(query, original_top_chunks, chunks, chunk_embeddings):
    query_words = query.split()
    word_scores = {}
    print("\nCalculating word importance...")
    for word in query_words:
        modified_query = " ".join(w for w in query_words if w.lower() != word.lower())
        modified_top_chunks = retrieve_relevant_chunks(modified_query, chunks, chunk_embeddings, top_n=5)
        common_chunks = set(chunk for _, chunk in modified_top_chunks).intersection(
            set(chunk for _, chunk in original_top_chunks))
        word_scores[word] = len(common_chunks)
    importance_words = sorted(word_scores, key=word_scores.get, reverse=True)[:5]
    print("Word importance calculated.")
    return importance_words


# Main function for user interaction
def main():
    print("Welcome to the RAG interactive console!")
    print("Type your query to get relevant information or 'exit' to quit.\n")

    # Load and preprocess the corpus
    corpus_file = 'corpus_text.txt'
    text = load_corpus(corpus_file)
    if text is None:
        return

    chunks = split_into_chunks(text)
    print(f"\nTotal chunks created: {len(chunks)}")
    chunk_embeddings = get_embeddings(chunks)
    if chunk_embeddings is None:
        return

    while True:
        query = input("\nPlease enter your query: ").strip()
        if query.lower() == 'exit':
            print("Exiting the program. Goodbye!")
            break
        if not query:
            print("Empty query received. Please enter a valid query.")
            continue

        # Retrieve the top relevant chunks for the user's query
        top_chunks = retrieve_relevant_chunks(query, chunks, chunk_embeddings, top_n=10)

        if not top_chunks:
            print("No relevant chunks found for your query.")
            continue

        # Print the top 3 chunks with their similarity scores
        print("\nTop 3 Relevant Chunks:")
        for score, chunk in top_chunks[:3]:
            print(f"Score: {score:.4f}, Chunk: \"{chunk}\"")

        # Ask for a user statement to check for contradictions
        user_statement = input("\nEnter a statement to check for contradictions against the context: ").strip()
        if not user_statement:
            print("Empty statement received. Skipping contradiction check.")
        else:
            contradiction_response = check_for_contradictions(user_statement, top_chunks)
            print(f"\nContradiction Check Response: {contradiction_response}")

        # Calculate and print the most important words associated with the query
        most_important_words = word_importance_calculation(query, top_chunks, chunks, chunk_embeddings)
        print("\nMost Important Words in Your Query:")
        for word in most_important_words:
            print(word)


if __name__ == "__main__":
    main()

simplefy this code as much as piossible so a beginner can understand it well also add some more print outputs for the user to see whats going aon and add enough ocomments for a gbeginner tto understand the code easily