import os
import pandas as pd
import torch

from settings import parquet_file

import tiktoken  # tokenizer library for use with OpenAI LLMs 
from llama_index.legacy.text_splitter import SentenceSplitter
from sentence_transformers import SentenceTransformer

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# create tensors on GPU if available
if torch.cuda.is_available():
    torch.set_default_tensor_type('torch.cuda.FloatTensor')


def chunk_vectorize(doc_content: dict = None, 
                    chunk_size: int = 256,    # limit for 'all-mpnet-base-v2'
                    chunk_overlap: int = 20,  # some overlap to link the chunks
                    encoder: str = 'gpt-3.5-turbo-0613',
                    model_name: str = 'sentence-transformers/all-mpnet-base-v2'):  # can try all-MiniLM-L6-v2
    # see tests in chunking_indexing.ipynb for more details

    encoding = tiktoken.encoding_for_model(encoder)

    splitter = SentenceSplitter(chunk_size=chunk_size, 
                                tokenizer=encoding.encode, 
                                chunk_overlap=chunk_overlap)

    # let's create the splits for every document
    contents_splits = {}
    for fname, content in doc_content.items():
        splits = [splitter.split_text(page) for page in content]
        contents_splits[fname] = [split for sublist in splits for split in sublist]
        
    model = SentenceTransformer(model_name)

    content_emb = {}
    for fname, splits in contents_splits.items():
        content_emb[fname] = [(split, model.encode(split)) for split in splits]

    # save fname since it carries information, and could be used as a property in Weaviate
    text_vector_tuples = [(fname, split, emb.tolist()) for fname, splits_emb in content_emb.items() for split, emb in splits_emb]

    new_df = pd.DataFrame(
        text_vector_tuples, 
        columns=['file', 'content', 'content_embedding']
    )
    
    # load the existing parquet file if it exists and update it 
    if os.path.exists(parquet_file):
        new_df = pd.concat([pd.read_parquet(parquet_file), new_df])

    # no optimization here (zipping etc) since the data is small
    new_df.to_parquet(parquet_file, index=False)        
    
    return

# TODO 
# import unittest
# from unitesting_utils import load_impact_theory_data

# class TestSplitContents(unittest.TestCase):
#     '''
#     Unit test to ensure proper functionality of split_contents function
#     '''
    
#     def test_split_contents(self):
#         import tiktoken
#         from llama_index.text_splitter import SentenceSplitter
        
#         data = load_impact_theory_data()
                
#         subset = data[:3]
#         chunk_size = 256
#         chunk_overlap = 0
#         encoding = tiktoken.encoding_for_model('gpt-3.5-turbo-0613')
#         gpt35_txt_splitter = SentenceSplitter(chunk_size=chunk_size, tokenizer=encoding.encode, chunk_overlap=chunk_overlap)
#         results = split_contents(subset, gpt35_txt_splitter)
#         self.assertEqual(len(results), 3)
#         self.assertEqual(len(results[0]), 83)
#         self.assertEqual(len(results[1]), 178)
#         self.assertEqual(len(results[2]), 144)
#         self.assertTrue(isinstance(results, list))
#         self.assertTrue(isinstance(results[0], list))
#         self.assertTrue(isinstance(results[0][0], str))
# unittest.TextTestRunner().run(unittest.TestLoader().loadTestsFromTestCase(TestSplitContents))