from chonkie import SemanticChunker
from sentence_transformers import SentenceTransformer
from pathlib import Path
import pandas as pd

data_dir = Path("../datasets")
train_faq_df = pd.read_csv(data_dir / "processed" / "bank_faq" /"train_faq.csv")

# Initialize with an embedding model
chunker = SemanticChunker(
    embedding_model="multi-qa-mpnet-base-dot-v1",
    chunk_size=200
)

chunked_docs = []
for idx, row in train_faq_df.iterrows():
    question = row["Question"]
    ans = row["Answer"]
    domain   = row["Class"]

    if len(ans.split()) < 200:
        # short — keep as single doc, embed q+a together
        chunked_docs.append({
            "text":      row["QA_pair"],  # for embedding
            "question":  question,
            "answer":    ans,
            "domain":    domain,
            "source_id": str(idx),
            "chunk_num": "0"
        })
    else:
        # long — chunk only the answer
        chunks = chunker.chunk(ans)
        for chunk_num, chunk in enumerate(chunks):
            chunked_docs.append({
                "text":      f"Q: {question} A: {chunk.text}",  # question anchors every chunk
                "question":  question,
                "answer":    chunk.text,
                "domain":    domain,
                "source_id": str(idx),
                "chunk_num": str(chunk_num)
            })

print(f"Total chunks: {len(chunked_docs)}") 

print(chunked_docs[0])



model = SentenceTransformer("multi-qa-mpnet-base-dot-v1")

texts = [doc["text"] for doc in chunked_docs]
embeddings = model.encode(texts, batch_size=64, show_progress_bar=True)

embed_faq_df = pd.DataFrame(chunked_docs)
embed_faq_df["embedding"] = embeddings.tolist()

embed_faq_df.to_csv(data_dir / "processed" / "bank_faq" / "embed_faq.csv", index=False)
print(embed_faq_df.shape)
print(embed_faq_df.head())

# for doc in chunked_docs:
#     print(len(doc["text"]), "|", doc["text"][:80])

# cleaned_faq_df["token_count"] = cleaned_faq_df["QA_pair"].apply(lambda x: len(x.split()))
# print(cleaned_faq_df["token_count"].describe())