import os
import pandas as pd
import requests
import random
import numpy as np
from tqdm.auto import tqdm
from sentence_transformers import SentenceTransformer, util
import torch
import textwrap

# Read the CSV file
csv_path = "IPL2020 Commentary Data.csv"
df = pd.read_csv(csv_path)

# Inspect the data to check the columns
print("DataFrame columns:", df.columns)

# Preprocess the text data
def text_formatter(text: str) -> str:
    """Performs minor formatting on text."""
    if isinstance(text, str):
        return text.replace("\n", " ").strip()
    return ""

df['formatted_text'] = df['long_text'].apply(text_formatter)

# Verify that the formatted_text column has been created
if 'formatted_text' not in df.columns:
    raise KeyError("Column 'formatted_text' was not created in the DataFrame.")

# Check the first few rows to confirm the formatting
print("First few rows after formatting:\n", df[['formatted_text']].head())

# Define a function to split text into chunks
def split_list(input_list: list, slice_size: int) -> list[list[str]]:
    """Splits the input_list into sublists of size slice_size (or as close as possible)."""
    return [input_list[i:i + slice_size] for i in range(0, len(input_list), slice_size)]

# Split the commentary into chunks of sentences
num_sentence_chunk_size = 10

# Add sentence chunks to the dataframe
df['sentence_chunks'] = df['formatted_text'].apply(lambda x: split_list(x.split('. '), num_sentence_chunk_size))
df['num_chunks'] = df['sentence_chunks'].apply(len)

# Check available columns before expanding rows
print("Available columns before expanding rows:", df.columns)

# Ensure that the necessary columns 'over' and 'ball' exist, if not, handle the missing columns appropriately
if 'over' not in df.columns:
    df['over'] = df.get('over', np.nan)

# Expand the dataframe to have each chunk in a separate row
expanded_rows = []

for _, row in tqdm(df.iterrows(), total=len(df)):
    for sentence_chunk in row['sentence_chunks']:
        chunk_dict = {
            "over": row["over"],
            "sentence_chunk": " ".join(sentence_chunk),
            "chunk_char_count": len(" ".join(sentence_chunk)),
            "chunk_word_count": len(" ".join(sentence_chunk).split(" ")),
            "chunk_token_count": len(" ".join(sentence_chunk)) / 4
        }
        expanded_rows.append(chunk_dict)

df_chunks = pd.DataFrame(expanded_rows)
print(df_chunks.head())

# Load the embedding model
embedding_model = SentenceTransformer(model_name_or_path="all-mpnet-base-v2", device="cuda" if torch.cuda.is_available() else "cpu")

# Create embeddings for the sentence chunks
df_chunks['embedding'] = df_chunks['sentence_chunk'].apply(lambda x: embedding_model.encode(x, convert_to_tensor=True))

# Save embeddings to file
embeddings_save_path = "ipl_commentary_embeddings.csv"
df_chunks.to_csv(embeddings_save_path, index=False)

# Load the embeddings from the file
df_chunks_loaded = pd.read_csv(embeddings_save_path)
df_chunks_loaded['embedding'] = df_chunks_loaded['embedding'].apply(lambda x: np.fromstring(x.strip("[]"), sep=" "))

# Convert embeddings to tensor
embeddings = torch.tensor(np.vstack(df_chunks_loaded['embedding'].values), dtype=torch.float32).to(embedding_model.device)
print(embeddings.shape)

# Define helper function to print wrapped text
def print_wrapped(text, wrap_length=80):
    wrapped_text = textwrap.fill(text, wrap_length)
    print(wrapped_text)

# Perform similarity search with a query
query = "excellent batting performance"
print(f"Query: {query}")

# Embed the query
query_embedding = embedding_model.encode(query, convert_to_tensor=True)

# Calculate similarity scores
dot_scores = util.dot_score(query_embedding, embeddings)[0]

# Get the top-k results
top_k = 5
top_results = torch.topk(dot_scores, k=top_k)

# Display the results
print(f"Query: '{query}'\n")
print("Results:")

for score, idx in zip(top_results[0], top_results[1]):
    print(f"Score: {score:.4f}")
    print("Text:")
    print_wrapped(df_chunks_loaded.iloc[idx]['sentence_chunk'])
    print(f"Over: {df_chunks_loaded.iloc[idx]['over']}\n")
RuntimeError                              Traceback (most recent call last)
<ipython-input-37-09c4cbafdc6b> in <cell line: 2>()
      1 # Calculate similarity scores
----> 2 dot_scores = util.dot_score(query_embedding, embeddings)[0]
      3 
      4 # Get the top-k results
      5 top_k = 5

/usr/local/lib/python3.10/dist-packages/sentence_transformers/util.py in dot_score(a, b)
    129     b = _convert_to_batch_tensor(b)
    130 
--> 131     return torch.mm(a, b.transpose(0, 1))
    132 
    133 

RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x768 and 0x14447)