visualized_app_v2_2.py

# /// script
# [tool.marimo.display]
# custom_css = ["./custom_header_font.css"]
# ///

import marimo

__generated_with = "0.13.14"
app = marimo.App(width="full")

with app.setup:
    # Initialization code that runs before all other cells
    import marimo as mo
    from typing import Dict, Optional, List, Union, Any
    from ibm_watsonx_ai import APIClient, Credentials
    from pathlib import Path
    import pandas as pd
    import mimetypes
    import requests
    import zipfile
    import tempfile
    import certifi
    import base64
    import polars
    import nltk
    import time
    import json
    import ast
    import os
    import io
    import re

    from dotenv import load_dotenv
    load_dotenv()

    def get_iam_token(api_key):
        return requests.post(
            'https://iam.cloud.ibm.com/identity/token',
            headers={'Content-Type': 'application/x-www-form-urlencoded'},
            data={'grant_type': 'urn:ibm:params:oauth:grant-type:apikey', 'apikey': api_key},
            verify=certifi.where()
        ).json()['access_token']

    def setup_task_credentials(client):
        # Get existing task credentials
        existing_credentials = client.task_credentials.get_details()

        # Delete existing credentials if any
        if "resources" in existing_credentials and existing_credentials["resources"]:
            for cred in existing_credentials["resources"]:
                cred_id = client.task_credentials.get_id(cred)
                client.task_credentials.delete(cred_id)

        # Store new credentials
        return client.task_credentials.store()



@app.cell
def client_variables(client_instantiation_form):
    client_setup = client_instantiation_form.value or None

    ### Extract Credential Variables:
    if client_setup:
        wx_url = client_setup["wx_region"] if client_setup["wx_region"] else "EU"
        wx_api_key = client_setup["wx_api_key"].strip() if client_setup["wx_api_key"] else None
        os.environ["WATSONX_APIKEY"] = wx_api_key or ""

        project_id = client_setup["project_id"].strip() if client_setup["project_id"] else None
        space_id = client_setup["space_id"].strip() if client_setup["space_id"] else None
    else:
        os.environ["WATSONX_APIKEY"] = ""
        project_id = space_id = wx_api_key = wx_url = None
    return client_setup, project_id, space_id, wx_api_key, wx_url


@app.cell
def _():
    from baked_in_credentials.creds import credentials
    from base_variables import wx_regions, wx_platform_url
    from helper_functions.helper_functions import wrap_with_spaces, get_key_by_value, markdown_spacing, get_cell_values, create_parameter_table, get_cred_value
    return (
        create_parameter_table,
        credentials,
        get_cell_values,
        get_cred_value,
        get_key_by_value,
        wrap_with_spaces,
        wx_regions,
    )


@app.cell
def client_instantiation(
    client_setup,
    project_id,
    space_id,
    wx_api_key,
    wx_url,
):
    ### Instantiate the watsonx.ai client
    if client_setup:
        try:
            wx_credentials = Credentials(url=wx_url, api_key=wx_api_key)
            project_client = (
                APIClient(credentials=wx_credentials, project_id=project_id)
                if project_id
                else None
            )
            deployment_client = (
                APIClient(credentials=wx_credentials, space_id=space_id)
                if space_id
                else None
            )
            instantiation_success = True
            instantiation_error = None
        except Exception as e:
            instantiation_success = False
            instantiation_error = str(e)
            wx_credentials = project_client = deployment_client = None
    else:
        wx_credentials = project_client = deployment_client = None
        instantiation_success = None
        instantiation_error = None

    return (
        deployment_client,
        instantiation_error,
        instantiation_success,
        project_client,
    )


@app.cell
def _():
    mo.md(
        r"""
    #watsonx.ai Embedding Visualizer - Marimo Notebook

    #### This marimo notebook can be used to develop a more intuitive understanding of how vector embeddings work by creating a 3D visualization of vector embeddings based on chunked PDF document pages. 

    #### It can also serve as a useful tool for identifying gaps in model choice, chunking strategy or contents used in building collections by showing how far you are from what you want.
    <br>

    /// admonition
    Created by ***Milan Mrdenovic*** [milan.mrdenovic@ibm.com] for IBM Ecosystem Client Engineering, NCEE - ***version 5.3** - 20.04.2025*
    ///


    >Licensed under apache 2.0, users hold full accountability for any use or modification of the code.
    ><br>This asset is part of a set meant to support IBMers, IBM Partners, Clients in developing understanding of how to better utilize various watsonx features and generative AI as a subject matter.

    <br>
    """
    )
    return


@app.cell
def _():
    mo.md("""###Part 1 - Client Setup, File Preparation and Chunking""")
    return


@app.cell
def accordion_client_setup(
    client_selector,
    client_stack,
    current_mode,
    switch_file_loader_type,
):
    ui_accordion_part_1_1 = mo.accordion(
        {
            "Instantiate Client": mo.vstack([
                client_stack, 
                mo.hstack([client_selector, switch_file_loader_type], justify="space-around", gap=2),
                current_mode
                ], align="center"),
        }
    )

    ui_accordion_part_1_1
    return


@app.cell
def _(switch_file_loader_type):
    if switch_file_loader_type.value:
        current_mode = mo.md("**Current Mode:** Using pre-made embedding/text files.")
    else:
        current_mode = mo.md("**Current Mode:** Using loaded pdf files and chunking.")
    return (current_mode,)


@app.cell
def accordion_file_upload(select_stack):
    ui_accordion_part_1_2 = mo.accordion(
        {
            "Select Model & Upload Files": select_stack
        }
    )

    ui_accordion_part_1_2
    return


@app.cell
def loaded_texts(
    create_temp_files_from_uploads,
    file_loader,
    pdf_reader,
    run_upload_button,
    set_text_state,
    switch_file_loader_type,
):
    if file_loader.value is not None and run_upload_button.value:
        filepaths = create_temp_files_from_uploads(file_loader.value)
        if switch_file_loader_type.value:
            loaded_texts = load_json_csv_data_with_progress(filepaths, file_loader.value, show_progress=True)
        else:
            loaded_texts = load_pdf_data_with_progress(pdf_reader, filepaths, file_loader.value, show_progress=True)

        set_text_state(loaded_texts)
    else:
        filepaths = None
        loaded_texts = None
    return (loaded_texts,)


@app.cell
def _(chunker_setup, file_column_setup, switch_file_loader_type):
    if switch_file_loader_type.value:
        ui_accordion_part_1_3 = mo.accordion(
            {
                "Column Selector": file_column_setup
            }
        )
    else:
        ui_accordion_part_1_3 = mo.accordion(
            {
                "Chunker Setup": chunker_setup
            }
        )    

    ui_accordion_part_1_3
    return


@app.cell
def accordion_chunker_setup():
    # ui_accordion_part_1_3 = mo.accordion(
    #     {
    #         "Chunker Setup": chunker_setup
    #     }
    # )

    # ui_accordion_part_1_3
    return


@app.cell
def chunk_documents_to_nodes(
    get_text_state,
    sentence_splitter,
    sentence_splitter_config,
    set_chunk_state,
):
    if sentence_splitter_config.value and sentence_splitter and get_text_state() is not None:
        chunked_texts = chunk_documents(get_text_state(), sentence_splitter, show_progress=True)
        set_chunk_state(chunked_texts)
    else:
        chunked_texts = None
    return (chunked_texts,)


@app.cell
def _():
    mo.md(r"""###Part 2 - Query Setup and Visualization""")
    return


@app.cell
def accordion_chunk_range():
    # ui_accordion_part_2_1 = mo.accordion(
    #     {
    #         "Chunk Range Selection": chart_range_selection
    #     }
    # )
    # ui_accordion_part_2_1
    return


@app.cell
def _(chart_range_selection, switch_file_loader_type):
    ui_accordion_part_2_1 = mo.accordion(
        {
            "Chunk Range Selection": chart_range_selection
        }
    )
    ui_accordion_part_2_1 if switch_file_loader_type.value == False else None
    return


@app.cell
def chunk_embedding(
    chunks_to_process,
    embedding,
    sentence_splitter_config,
    set_embedding_state,
):
    if sentence_splitter_config.value is not None and chunks_to_process is not None:
        with mo.status.spinner(title="Embedding Documents...", remove_on_exit=True) as _spinner:
            output_embeddings = embedding.embed_documents(chunks_to_process)
            _spinner.update("Almost Done")
            time.sleep(1.5)
            set_embedding_state(output_embeddings)
            _spinner.update("Documents Embedded")
    else:
        output_embeddings = None
    return


@app.cell
def preview_chunks(chunks_dict):
    if chunks_dict is not None:
        stats = create_stats(chunks_dict, 
                             bordered=True, 
                             object_names=['text','text'], 
                             group_by_row=True, 
                             items_per_row=5, 
                             gap=1,
                             label="Chunk")
        ui_chunk_viewer = mo.accordion(
            {
                "View Chunks": stats,
            }
        )
    else:
        ui_chunk_viewer = None

    ui_chunk_viewer
    return


@app.cell
def accordion_query_view(chart_visualization, query_stack):
    ui_accordion_part_2_2 = mo.accordion(
        {
            "Query": mo.vstack([query_stack, mo.hstack([chart_visualization])], align="center", gap=3)
        }
    )
    ui_accordion_part_2_2
    return


@app.cell
def chunker_setup(sentence_splitter_config):
    chunker_setup = mo.hstack([sentence_splitter_config], justify="space-around", align="center", widths=[0.55])
    return (chunker_setup,)


@app.cell
def file_and_model_select(
    file_loader,
    get_embedding_model_list,
    run_upload_button,
):
    select_stack = mo.hstack([get_embedding_model_list(), mo.vstack([mo.md("Drag & Drop or Double Click to select PDFs, then press **Load Files**"),file_loader, run_upload_button], align="center")], justify="space-around", align="center", widths=[0.3,0.3])
    return (select_stack,)


@app.cell
def client_instantiation_form(credentials, get_cred_value, wx_regions):
    baked_in_creds = credentials
    # Create a form with multiple elements
    client_instantiation_form = (
        mo.md('''
        ###**watsonx.ai credentials:**

        {wx_region}

        {wx_api_key}

        {project_id}

        {space_id}

        > You can add either a project_id, space_id or both, **only one is required**.  
        > If you provide both you can switch the active one in the dropdown.
    ''')
        .batch(
            wx_region = mo.ui.dropdown(
                wx_regions, 
                label="Select your watsonx.ai region:", 
                value=get_cred_value('region', creds_var_name='baked_in_creds') or "EU",
                searchable=True
            ),
            wx_api_key = mo.ui.text(
                placeholder="Add your IBM Cloud api-key...", 
                label="IBM Cloud Api-key:", 
                kind="password", 
                value=get_cred_value('api_key', creds_var_name='baked_in_creds')
            ),
            project_id = mo.ui.text(
                placeholder="Add your watsonx.ai project_id...", 
                label="Project_ID:", 
                kind="text", 
                value=get_cred_value('project_id', creds_var_name='baked_in_creds')
            ),
            space_id = mo.ui.text(
                placeholder="Add your watsonx.ai space_id...", 
                label="Space_ID:", 
                kind="text",
                value=get_cred_value('space_id', creds_var_name='baked_in_creds')
            )
        ,)
        .form(show_clear_button=True, bordered=False)
    )
    return (client_instantiation_form,)


@app.cell
def instantiation_status(
    client_callout_kind,
    client_instantiation_form,
    client_status,
):
    client_callout = mo.callout(client_status, kind=client_callout_kind)
    client_stack = mo.hstack([client_instantiation_form, client_callout], align="center", justify="space-around", gap=10)
    return (client_stack,)


@app.cell
def _(
    client_key,
    client_options,
    client_selector,
    client_setup,
    get_key_by_value,
    instantiation_error,
    instantiation_success,
    wrap_with_spaces,
):
    active_client_name = get_key_by_value(client_options, client_key)

    if client_setup:
        if instantiation_success:
            client_status = mo.md(
                f"### ✅ Client Instantiation Successful ✅\n\n"
                f"{client_selector}\n\n"
                f"**Active Client:**{wrap_with_spaces(active_client_name, prefix_spaces=5)}"
            )
            client_callout_kind = "success"
        else:
            client_status = mo.md(
                f"### ❌ Client Instantiation Failed\n**Error:** {instantiation_error}\n\nCheck your region selection and credentials"
            )
            client_callout_kind = "danger"
    else:
        client_status = mo.md(
            f"### Client Instantiation Status will turn Green When Ready\n\n"
            f"{client_selector}\n\n"
            f"**Active Client:**{wrap_with_spaces(active_client_name, prefix_spaces=5)}"
        )
        client_callout_kind = "neutral"

    return client_callout_kind, client_status


@app.cell
def client_selector(deployment_client, project_client):
    if project_client is not None and deployment_client is not None:
        client_options = {"Project Client":project_client, "Deployment Client":deployment_client}

    elif project_client is not None:
        client_options = {"Project Client":project_client}

    elif deployment_client is not None:
        client_options = {"Deployment Client":deployment_client}

    else:
        client_options = {"No Client": "Instantiate a Client"}

    default_client = next(iter(client_options))
    client_selector = mo.ui.dropdown(client_options, value=default_client, label="**Switch your active client:**")
    return client_options, client_selector


@app.cell
def active_client(client_selector):
    client_key = client_selector.value
    if client_key == "Instantiate a Client":
        client = None
    else:
        client = client_key
    return client, client_key


@app.cell
def emb_model_selection(client, set_embedding_model_list):
    if client is not None:
        model_specs = client.foundation_models.get_embeddings_model_specs()
        # model_specs = client.foundation_models.get_model_specs()
        resources = model_specs["resources"]
        # Define embedding models reference data
        embedding_models = {
            "ibm/granite-embedding-107m-multilingual": {"max_tokens": 512, "embedding_dimensions": 384},
            "ibm/granite-embedding-278m-multilingual": {"max_tokens": 512, "embedding_dimensions": 768},
            "ibm/slate-125m-english-rtrvr-v2": {"max_tokens": 512, "embedding_dimensions": 768},
            "ibm/slate-125m-english-rtrvr": {"max_tokens": 512, "embedding_dimensions": 768},
            "ibm/slate-30m-english-rtrvr-v2": {"max_tokens": 512, "embedding_dimensions": 384},
            "ibm/slate-30m-english-rtrvr": {"max_tokens": 512, "embedding_dimensions": 384},
            "sentence-transformers/all-minilm-l6-v2": {"max_tokens": 128, "embedding_dimensions": 384},
            "sentence-transformers/all-minilm-l12-v2": {"max_tokens": 128, "embedding_dimensions": 384},
            "intfloat/multilingual-e5-large": {"max_tokens": 512, "embedding_dimensions": 1024}
        }

        # Get model IDs from resources
        model_id_list = []
        for resource in resources:
            model_id_list.append(resource["model_id"])

        # Create enhanced model data for the table
        embedding_model_data = []
        for model_id in model_id_list:
            model_entry = {"model_id": model_id}

            # Add properties if model exists in our reference, otherwise use 0
            if model_id in embedding_models:
                model_entry["max_tokens"] = embedding_models[model_id]["max_tokens"]
                model_entry["embedding_dimensions"] = embedding_models[model_id]["embedding_dimensions"]
            else:
                model_entry["max_tokens"] = 0
                model_entry["embedding_dimensions"] = 0

            embedding_model_data.append(model_entry)

        embedding_model_selection = mo.ui.table(
            embedding_model_data,
            selection="single",  # Only allow selecting one row
            label="Select an embedding model to use.",
            page_size=30,
            initial_selection=[1]
        )
        set_embedding_model_list(embedding_model_selection)
    else:
        default_model_data = [{
            "model_id": "ibm/granite-embedding-107m-multilingual",
            "max_tokens": 512,
            "embedding_dimensions": 384
        }]

        set_embedding_model_list(create_emb_model_selection_table(default_model_data, initial_selection=0, selection_type="single", label="Select a model to use."))
    return


@app.function
def create_emb_model_selection_table(model_data, initial_selection=0, selection_type="single", label="Select a model to use."):
    embedding_model_selection = mo.ui.table(
        model_data,
        selection=selection_type,  # Only allow selecting one row
        label=label,
        page_size=30,
        initial_selection=[initial_selection]
    )
    return embedding_model_selection


@app.cell
def embedding_model():
    get_embedding_model_list, set_embedding_model_list = mo.state(None)
    return get_embedding_model_list, set_embedding_model_list


@app.cell
def emb_model_parameters(emb_model_max_tk, embedding_model):
    from ibm_watsonx_ai.metanames import EmbedTextParamsMetaNames as EmbedParams
    if embedding_model is not None:
        embed_params = {
             EmbedParams.TRUNCATE_INPUT_TOKENS: emb_model_max_tk,
             EmbedParams.RETURN_OPTIONS: {
             'input_text': True
             }
         }
    else:
        embed_params = {
             EmbedParams.TRUNCATE_INPUT_TOKENS: 128,
             EmbedParams.RETURN_OPTIONS: {
             'input_text': True
             }
         }      
    return (embed_params,)


@app.cell
def emb_model_state(get_embedding_model_list):
    embedding_model = get_embedding_model_list()
    return (embedding_model,)


@app.cell
def emb_model_setup(embedding_model):
    if embedding_model is not None:
        emb_model = embedding_model.value[0]['model_id']
        emb_model_max_tk = embedding_model.value[0]['max_tokens']
        emb_model_emb_dim = embedding_model.value[0]['embedding_dimensions']
    else:
        emb_model = None
        emb_model_max_tk = None
        emb_model_emb_dim = None
    return emb_model, emb_model_emb_dim, emb_model_max_tk


@app.cell
def emb_model_instantiation(client, emb_model, embed_params):
    from ibm_watsonx_ai.foundation_models import Embeddings
    if client is not None:
        embedding = Embeddings(
             model_id=emb_model,
             api_client=client,
             params=embed_params,
             batch_size=1000,
             concurrency_limit=10
             )
    else:
        embedding = None
    return (embedding,)


@app.cell
def _():
    get_embedding_state, set_embedding_state = mo.state(None)
    return get_embedding_state, set_embedding_state


@app.cell
def _():
    get_query_state, set_query_state = mo.state(None)
    return get_query_state, set_query_state


@app.cell
def file_loader_input(switch_file_loader_type):
    if switch_file_loader_type.value:
        file_loader = mo.ui.file(
                    kind="area",
                    filetypes=[".json",".csv"],
                    label="   Load pre-made embedding/text pair files (.json,.csv)  ",
                    multiple=False
                )
    else:
        file_loader = mo.ui.file(
                kind="area",
                filetypes=[".pdf"],
                label="   Load .pdf files  ",
                multiple=True
            )
    return (file_loader,)


@app.cell
def file_loader_run(file_loader):
    if file_loader.value:
        run_upload_button = mo.ui.run_button(label="Load Files")
    else:
        run_upload_button = mo.ui.run_button(disabled=True, label="Load Files")
    return (run_upload_button,)


@app.cell
def helper_function_tempfiles():
    def create_temp_files_from_uploads(upload_results) -> List[str]:
        """
        Creates temporary files from a tuple of FileUploadResults objects and returns their paths.
        Args:
            upload_results: Object containing a value attribute that is a tuple of FileUploadResults
        Returns:
            List of temporary file paths
        """
        temp_file_paths = []

        # Get the number of items in the tuple
        num_items = len(upload_results)

        # Process each item by index
        for i in range(num_items):
            result = upload_results[i]  # Get item by index

            # Create a temporary file with the original filename
            temp_dir = tempfile.gettempdir()
            file_name = result.name
            temp_path = os.path.join(temp_dir, file_name)
            # Write the contents to the temp file
            with open(temp_path, 'wb') as temp_file:
                temp_file.write(result.contents)
            # Add the path to our list
            temp_file_paths.append(temp_path)

        return temp_file_paths

    def cleanup_temp_files(temp_file_paths: List[str]) -> None:
        """Delete temporary files after use."""
        for path in temp_file_paths:
            if os.path.exists(path):
                os.unlink(path)
    return (create_temp_files_from_uploads,)


@app.function
def load_pdf_data_with_progress(pdf_reader, filepaths, file_loader_value, show_progress=True):
    """
    Loads PDF data for each file path and organizes results by original filename.
    Args:
        pdf_reader: The PyMuPDFReader instance
        filepaths: List of temporary file paths
        file_loader_value: The original upload results value containing file information
        show_progress: Whether to show a progress bar during loading (default: False)
    Returns:
        Dictionary mapping original filenames to their loaded text content
    """
    results = {}

    # Process files with or without progress bar
    if show_progress:
        import marimo as mo
        # Use progress bar with the length of filepaths as total
        with mo.status.progress_bar(
            total=len(filepaths),
            title="Loading PDFs",
            subtitle="Processing documents...",
            completion_title="PDF Loading Complete",
            completion_subtitle=f"{len(filepaths)} documents processed",
            remove_on_exit=True
        ) as bar:
            # Process each file path
            for i, file_path in enumerate(filepaths):

                original_file_name = file_loader_value[i].name
                bar.update(subtitle=f"Processing {original_file_name}...")
                loaded_text = pdf_reader.load_data(file_path=file_path, metadata=True)

                # Store the result with the original filename as the key
                results[original_file_name] = loaded_text
                # Update progress bar
                bar.update(increment=1)
    else:
        # Original logic without progress bar
        for i, file_path in enumerate(filepaths):
            original_file_name = file_loader_value[i].name
            loaded_text = pdf_reader.load_data(file_path=file_path, metadata=True)
            results[original_file_name] = loaded_text

    return results


@app.cell
def file_readers():
    from llama_index.readers.file import PyMuPDFReader
    from llama_index.readers.file import FlatReader
    from llama_index.core.node_parser import SentenceSplitter

    ### File Readers
    pdf_reader = PyMuPDFReader()
    # flat_file_reader = FlatReader()
    return SentenceSplitter, pdf_reader


@app.cell
def sentence_splitter_setup():
    ### Chunker Setup
    sentence_splitter_config = (
        mo.md('''
        ###**Chunking Setup:**

        > Unless you want to do some advanced sentence splitting, it's best to stick to adjusting only the chunk size and overlap. Changing the other settings might result in unexpected results.

        Separator value is set to **" "** by default, while the paragraph separator is **"\\n\\n\\n"**.

        {chunk_size} 

        {chunk_overlap}

        {separator} {paragraph_separator}

        {secondary_chunking_regex} {include_metadata}

    ''')
        .batch(
            chunk_size = mo.ui.slider(start=100, stop=5000, step=1, label="**Chunk Size:**", value=275, show_value=True, full_width=True),
            chunk_overlap = mo.ui.slider(start=0, stop=1000, step=1, label="**Chunk Overlap** *(Must always be smaller than Chunk Size)* **:**", value=0, show_value=True, full_width=True),
            separator = mo.ui.text(placeholder="Define a separator", label="**Separator:**", kind="text", value=" "),
            paragraph_separator = mo.ui.text(placeholder="Define a paragraph separator", 
                                             label="**Paragraph Separator:**", kind="text", 
                                             value="\n\n\n"),
            secondary_chunking_regex = mo.ui.text(placeholder="Define a secondary chunking regex", 
                                                  label="**Chunking Regex:**", kind="text", 
                                                  value="[^,.;?!]+[,.;?!]?"),
            include_metadata= mo.ui.checkbox(value=True, label="**Include Metadata**")
        )
        .form(show_clear_button=True, bordered=False, submit_button_label="Chunk Documents")
    )
    return (sentence_splitter_config,)


@app.cell
def sentence_splitter_instantiation(
    SentenceSplitter,
    sentence_splitter_config,
):
    ### Chunker/Sentence Splitter 
    def simple_whitespace_tokenizer(text):
        return text.split()

    if sentence_splitter_config.value is not None:
        sentence_splitter_config_values = sentence_splitter_config.value
        validated_chunk_overlap = min(sentence_splitter_config_values.get("chunk_overlap"), 
                     int(sentence_splitter_config_values.get("chunk_size") * 0.3))

        sentence_splitter = SentenceSplitter(
            chunk_size=sentence_splitter_config_values.get("chunk_size"),
            chunk_overlap=validated_chunk_overlap,
            separator=sentence_splitter_config_values.get("separator"),
            paragraph_separator=sentence_splitter_config_values.get("paragraph_separator"),
            secondary_chunking_regex=sentence_splitter_config_values.get("secondary_chunking_regex"),
            include_metadata=sentence_splitter_config_values.get("include_metadata"),
            tokenizer=simple_whitespace_tokenizer
        )

    else:
        sentence_splitter = SentenceSplitter(
            chunk_size=2048,
            chunk_overlap=204,
            separator=" ",
            paragraph_separator="\n\n\n",
            secondary_chunking_regex="[^,.;?!]+[,.;?!]?",
            include_metadata=True,
            tokenizer=simple_whitespace_tokenizer
        )
    return (sentence_splitter,)


@app.cell
def text_state():
    get_text_state, set_text_state = mo.state(None)
    return get_text_state, set_text_state


@app.cell
def chunk_state():
    get_chunk_state, set_chunk_state = mo.state(None)
    return get_chunk_state, set_chunk_state


@app.function
def chunk_documents(loaded_texts, sentence_splitter, show_progress=True):
    """
    Process each document in the loaded_texts dictionary using the sentence_splitter,
    with an optional marimo progress bar tracking progress at document level.

    Args:
        loaded_texts (dict): Dictionary containing lists of Document objects
        sentence_splitter: The sentence splitter object with get_nodes_from_documents method
        show_progress (bool): Whether to show a progress bar during processing

    Returns:
        dict: Dictionary with the same structure but containing chunked texts
    """
    chunked_texts_dict = {}

    # Get the total number of documents across all keys
    total_docs = sum(len(docs) for docs in loaded_texts.values())
    processed_docs = 0

    # Process with or without progress bar
    if show_progress:
        import marimo as mo
        # Use progress bar with the total number of documents as total
        with mo.status.progress_bar(
            total=total_docs,
            title="Processing Documents",
            subtitle="Chunking documents...",
            completion_title="Processing Complete",
            completion_subtitle=f"{total_docs} documents processed",
            remove_on_exit=True
        ) as bar:
            # Process each key-value pair in the loaded_texts dictionary
            for key, documents in loaded_texts.items():
                # Update progress bar subtitle to show current key
                doc_count = len(documents)
                bar.update(subtitle=f"Chunking {key}... ({doc_count} documents)")

                # Apply the sentence splitter to each list of documents
                chunked_texts = sentence_splitter.get_nodes_from_documents(
                    documents, 
                    show_progress=False  # Disable internal progress to avoid nested bars
                )

                # Store the result with the same key
                chunked_texts_dict[key] = chunked_texts
                time.sleep(0.15)

                # Update progress bar with the number of documents in this batch
                bar.update(increment=doc_count)
                processed_docs += doc_count
    else:
        # Process without progress bar
        for key, documents in loaded_texts.items():
            chunked_texts = sentence_splitter.get_nodes_from_documents(
                documents, 
                show_progress=True  # Use the internal progress bar if no marimo bar
            )
            chunked_texts_dict[key] = chunked_texts

    return chunked_texts_dict


@app.cell
def chunked_nodes(chunked_texts, get_chunk_state, sentence_splitter):
    if chunked_texts is not None and sentence_splitter:
        chunked_documents = get_chunk_state()
    else:
        chunked_documents = None
    return (chunked_documents,)


@app.cell
def prep_cumulative_df(chunked_documents, llamaindex_convert_docs_multi):
    if chunked_documents is not None:
        dict_from_nodes = llamaindex_convert_docs_multi(chunked_documents)
        nodes_from_dict = llamaindex_convert_docs_multi(dict_from_nodes)
    else:
        dict_from_nodes = None
        nodes_from_dict = None
    return (dict_from_nodes,)


@app.cell
def chunks_to_process(
    dict_from_nodes,
    document_range_stack,
    get_data_in_range_triplequote,
):
    if dict_from_nodes is not None and document_range_stack is not None:

        chunk_dict_df = create_cumulative_dataframe(dict_from_nodes)

        if document_range_stack.value is not None:
            chunk_start_idx = document_range_stack.value[0]
            chunk_end_idx = document_range_stack.value[1]
        else:
            chunk_start_idx = 0
            chunk_end_idx = len(chunk_dict_df)

        chunk_range_index = [chunk_start_idx, chunk_end_idx]
        chunks_dict = get_data_in_range_triplequote(chunk_dict_df,
                                                    index_range=chunk_range_index, 
                                                    columns_to_include=["text"])

        chunks_to_process = chunks_dict['text'] if 'text' in chunks_dict else []
    else:
        chunk_objects = None
        chunks_dict = None
        chunks_to_process = None
    return chunks_dict, chunks_to_process


@app.cell
def helper_function_doc_formatting():
    def llamaindex_convert_docs_multi(items):
        """
        Automatically convert between document objects and dictionaries.

        This function handles:
        - Converting dictionaries to document objects
        - Converting document objects to dictionaries
        - Processing lists or individual items
        - Supporting dictionary structures where values are lists of documents

        Args:
            items: A document object, dictionary, or list of either.
                  Can also be a dictionary mapping filenames to lists of documents.

        Returns:
            Converted item(s) maintaining the original structure
        """
        # Handle empty or None input
        if not items:
            return []

        # Handle dictionary mapping filenames to document lists (from load_pdf_data)
        if isinstance(items, dict) and all(isinstance(v, list) for v in items.values()):
            result = {}
            for filename, doc_list in items.items():
                result[filename] = llamaindex_convert_docs(doc_list)
            return result

        # Handle single items (not in a list)
        if not isinstance(items, list):
            # Single dictionary to document
            if isinstance(items, dict):
                # Determine document class
                doc_class = None
                if 'doc_type' in items:
                    import importlib
                    module_path, class_name = items['doc_type'].rsplit('.', 1)
                    module = importlib.import_module(module_path)
                    doc_class = getattr(module, class_name)
                if not doc_class:
                    from llama_index.core.schema import Document
                    doc_class = Document
                return doc_class.from_dict(items)
            # Single document to dictionary
            elif hasattr(items, 'to_dict'):
                return items.to_dict()
            # Return as is if can't convert
            return items

        # Handle list input
        result = []

        # Handle empty list
        if len(items) == 0:
            return result

        # Determine the type of conversion based on the first non-None item
        first_item = next((item for item in items if item is not None), None)

        # If we found no non-None items, return empty list
        if first_item is None:
            return result

        # Convert dictionaries to documents
        if isinstance(first_item, dict):
            # Get the right document class from the items themselves
            doc_class = None
            # Try to get doc class from metadata if available
            if 'doc_type' in first_item:
                import importlib
                module_path, class_name = first_item['doc_type'].rsplit('.', 1)
                module = importlib.import_module(module_path)
                doc_class = getattr(module, class_name)
            if not doc_class:
                # Fallback to default Document class from llama_index
                from llama_index.core.schema import Document
                doc_class = Document

            # Convert each dictionary to document
            for item in items:
                if isinstance(item, dict):
                    result.append(doc_class.from_dict(item))
                elif item is None:
                    result.append(None)
                elif isinstance(item, list):
                    result.append(llamaindex_convert_docs(item))
                else:
                    result.append(item)

        # Convert documents to dictionaries
        else:
            for item in items:
                if hasattr(item, 'to_dict'):
                    result.append(item.to_dict())
                elif item is None:
                    result.append(None)
                elif isinstance(item, list):
                    result.append(llamaindex_convert_docs(item))
                else:
                    result.append(item)

        return result

    def llamaindex_convert_docs(items):
        """
        Automatically convert between document objects and dictionaries.

        Args:
            items: A list of document objects or dictionaries

        Returns:
            List of converted items (dictionaries or document objects)
        """
        result = []

        # Handle empty or None input
        if not items:
            return result

        # Determine the type of conversion based on the first item
        if isinstance(items[0], dict):
            # Get the right document class from the items themselves
            # Look for a 'doc_type' or '__class__' field in the dictionary
            doc_class = None

            # Try to get doc class from metadata if available
            if 'doc_type' in items[0]:
                import importlib
                module_path, class_name = items[0]['doc_type'].rsplit('.', 1)
                module = importlib.import_module(module_path)
                doc_class = getattr(module, class_name)

            if not doc_class:
                # Fallback to default Document class from llama_index
                from llama_index.core.schema import Document
                doc_class = Document

            # Convert dictionaries to documents
            for item in items:
                if isinstance(item, dict):
                    result.append(doc_class.from_dict(item))
        else:
            # Convert documents to dictionaries
            for item in items:
                if hasattr(item, 'to_dict'):
                    result.append(item.to_dict())

        return result
    return (llamaindex_convert_docs_multi,)


@app.cell
def helper_function_create_df():
    def create_document_dataframes(dict_from_docs):
        """
        Creates a pandas DataFrame for each file in the dictionary.

        Args:
            dict_from_docs: Dictionary mapping filenames to lists of documents

        Returns:
            List of pandas DataFrames, each representing all documents from a single file
        """
        dataframes = []

        for filename, docs in dict_from_docs.items():
            # Create a list to hold all document records for this file
            file_records = []

            for i, doc in enumerate(docs):
                # Convert the document to a format compatible with DataFrame
                if hasattr(doc, 'to_dict'):
                    doc_data = doc.to_dict()
                elif isinstance(doc, dict):
                    doc_data = doc
                else:
                    doc_data = {'content': str(doc)}

                # Add document index information
                doc_data['doc_index'] = i

                # Add to the list of records for this file
                file_records.append(doc_data)

            # Create a single DataFrame for all documents in this file
            if file_records:
                df = pd.DataFrame(file_records)
                df['filename'] = filename  # Add filename as a column
                dataframes.append(df)

        return dataframes

    def create_dataframe_previews(dataframe_list, page_size=5):
        """
        Creates a list of mo.ui.dataframe components, one for each DataFrame in the input list.

        Args:
            dataframe_list: List of pandas DataFrames (output from create_document_dataframes)
            page_size: Number of rows to show per page for each component

        Returns:
            List of mo.ui.dataframe components
        """
        # Create a list of mo.ui.dataframe components
        preview_components = []

        for df in dataframe_list:
            # Create a mo.ui.dataframe component for this DataFrame
            preview = mo.ui.dataframe(df, page_size=page_size)
            preview_components.append(preview)

        return preview_components
    return


@app.cell
def _():
    switch_file_loader_type = mo.ui.switch(label="**Switch** to pre-made Embedding/Text pairs")
    return (switch_file_loader_type,)


@app.cell
def _():
    import csv

    def csv_to_json(csv_file_path, json_file_path):
        """
        Convert CSV file to JSON format.
    
        Args:
            csv_file_path (str): Path to input CSV file
            json_file_path (str): Path to output JSON file
        """
        with open(csv_file_path, 'r', encoding='utf-8') as csv_file:
            csv_reader = csv.DictReader(csv_file)
            data = list(csv_reader)
    
        with open(json_file_path, 'w', encoding='utf-8') as json_file:
            json.dump(data, json_file, indent=2)

    return


@app.function
def load_json_csv_data_with_progress(filepaths, file_loader_value, show_progress=True):
    """
    Loads CSV (converted to JSON) or JSON data for a single file.
    Returns the raw JSON content without filename mapping.
    """
    import csv
    import json
    
    filepath = filepaths[0]  # Only process first file
    original_file_name = file_loader_value[0].name
    
    if show_progress:
        import marimo as mo
        with mo.status.progress_bar(
            total=1,
            title="Loading File",
            subtitle=f"Processing {original_file_name}...",
            completion_title="File Loading Complete",
            completion_subtitle="1 file processed",
            remove_on_exit=True
        ) as bar:
            if filepath.lower().endswith('.csv'):
                with open(filepath, 'r', encoding='utf-8') as csv_file:
                    csv_reader = csv.DictReader(csv_file)
                    result = list(csv_reader)
            elif filepath.lower().endswith('.json'):
                with open(filepath, 'r', encoding='utf-8') as json_file:
                    result = json.load(json_file)
            bar.update(increment=1)
    else:
        if filepath.lower().endswith('.csv'):
            with open(filepath, 'r', encoding='utf-8') as csv_file:
                csv_reader = csv.DictReader(csv_file)
                result = list(csv_reader)
        elif filepath.lower().endswith('.json'):
            with open(filepath, 'r', encoding='utf-8') as json_file:
                result = json.load(json_file)
    
    return result


@app.function
def organize_data_by_columns(loaded_texts, columns_to_use):
    """
    Organizes loaded text data into specified column groups based on configuration.
    
    Args:
        loaded_texts: List of dictionaries containing the data
        columns_to_use: Dictionary mapping column group names to their field configurations
    
    Returns:
        Dictionary with column group names as keys and lists of field values as values
    """
    result = {}
    
    for group_name, field_config in columns_to_use.items():
        result[group_name] = []
        
        # Get fields that are marked as True
        selected_fields = [field for field, include in field_config.items() if include]
        
        for record in loaded_texts:
            # Add each selected field's value directly to the list
            for field in selected_fields:
                if field in record:
                    result[group_name].append(record[field])
    
    return result


@app.cell
def _(
    create_parameter_table,
    get_text_state,
    loaded_texts,
    switch_file_loader_type,
):
    if switch_file_loader_type.value and loaded_texts:
        column_list = list(get_text_state()[0])
        text_column = create_parameter_table(
                                           label="Select the Embedded Text Column",
                                           input_list=column_list,
                                           column_name="Text Column", 
                                           selection_type="single-cell", 
                                           text_justify="center"
                                          )
        embedding_column = create_parameter_table(
                                           label="Select the Corresponding Embeddings Column",
                                           input_list=column_list,
                                           column_name="Embedding Column", 
                                           selection_type="single-cell", 
                                           text_justify="center"
                                          )
        column_selection_stack = mo.hstack([text_column, embedding_column], justify="space-around", widths=[0.4,0.4])
        file_column_setup = mo.hstack([column_selection_stack], justify="space-around", align="center", widths=[0.75])
    else:
        text_column = embedding_column = column_selection_stack = file_column_setup = None
    return embedding_column, file_column_setup, text_column


@app.cell
def _(embedding_column, get_cell_values, switch_file_loader_type, text_column):
    if switch_file_loader_type.value:
        text_col_value = get_cell_values(text_column)
        emb_col_value = get_cell_values(embedding_column)
        columns_to_use = {
            "texts": text_col_value,
            "embeddings": emb_col_value
        }
    else:
        text_col_value = emb_col_value = columns_to_use = None
    return columns_to_use, text_col_value


@app.cell
def _(columns_to_use, loaded_texts, text_col_value):
    if text_col_value and columns_to_use and loaded_texts:
        premade_documents = organize_data_by_columns(columns_to_use=columns_to_use, loaded_texts=loaded_texts)
        text_col_state = validate_value(premade_documents["texts"])
        emb_col_state = validate_value(premade_documents["embeddings"])
        columns_selected = all_true(text_col_state, emb_col_state)
    else:
        premade_documents = text_col_state = emb_col_state = emb_col_state = columns_selected = None
    return columns_selected, premade_documents


@app.function
def validate_value(value):
    """
    Check if a value is not None or an empty data object.
    """
    return value is not None and bool(value)


@app.cell
def helper_function_chart_preparation():
    import altair as alt
    import numpy as np
    import plotly.express as px
    from sklearn.manifold import TSNE

    def prepare_embedding_data(embeddings, texts, model_id=None, embedding_dimensions=None):
        """
        Prepare embedding data for visualization

        Args:
            embeddings: List of embeddings arrays
            texts: List of text strings
            model_id: Embedding model ID (optional)
            embedding_dimensions: Embedding dimensions (optional)

        Returns:
            DataFrame with processed data and metadata
        """
        # Flatten embeddings (in case they're nested)
        flattened_embeddings = []
        for emb in embeddings:
            if isinstance(emb, list) and len(emb) > 0 and isinstance(emb[0], list):
                flattened_embeddings.append(emb[0])  # Take first element if nested
            else:
                flattened_embeddings.append(emb)

        # Convert to numpy array
        embedding_array = np.array(flattened_embeddings)

        # Apply dimensionality reduction (t-SNE)
        tsne = TSNE(n_components=2, random_state=42, perplexity=min(30, len(embedding_array)-1))
        reduced_embeddings = tsne.fit_transform(embedding_array)

        # Create truncated texts for display
        truncated_texts = [text[:50] + "..." if len(text) > 50 else text for text in texts]

        # Create dataframe for visualization
        df = pd.DataFrame({
            "x": reduced_embeddings[:, 0],
            "y": reduced_embeddings[:, 1],
            "text": truncated_texts,
            "full_text": texts,
            "index": range(len(texts))
        })

        # Add metadata
        metadata = {
            "model_id": model_id,
            "embedding_dimensions": embedding_dimensions
        }

        return df, metadata

    def create_embedding_chart(df, metadata=None):
        """
        Create an Altair chart for embedding visualization

        Args:
            df: DataFrame with x, y coordinates and text
            metadata: Dictionary with model_id and embedding_dimensions

        Returns:
            Altair chart
        """
        model_id = metadata.get("model_id") if metadata else None
        embedding_dimensions = metadata.get("embedding_dimensions") if metadata else None

        selection = alt.selection_multi(fields=['index'])

        base = alt.Chart(df).encode(
            x=alt.X("x:Q", title="Dimension 1"),
            y=alt.Y("y:Q", title="Dimension 2"),
            tooltip=["text", "index"]
        )

        points = base.mark_circle(size=100).encode(
            color=alt.Color("index:N", legend=None),
            opacity=alt.condition(selection, alt.value(1), alt.value(0.2))
        ).add_selection(selection)  # Add this line to apply the selection

        text = base.mark_text(align="left", dx=7).encode(
            text="index:N"
        )

        return (points + text).properties(
            width=700, 
            height=500,
            title=f"Embedding Visualization{f' - Model: {model_id}' if model_id else ''}{f' ({embedding_dimensions} dimensions)' if embedding_dimensions else ''}"
        ).interactive()

    def show_selected_text(indices, texts):
        """
        Create markdown display for selected texts

        Args:
            indices: List of selected indices
            texts: List of all texts

        Returns:
            Markdown string
        """
        if not indices:
            return "No text selected"

        selected_texts = [texts[i] for i in indices if i < len(texts)]
        return "\n\n".join([f"**Document {i}**:\n{text}" for i, text in zip(indices, selected_texts)])

    def prepare_embedding_data_3d(embeddings, texts, model_id=None, embedding_dimensions=None):
            """
            Prepare embedding data for 3D visualization

            Args:
                embeddings: List of embeddings arrays
                texts: List of text strings
                model_id: Embedding model ID (optional)
                embedding_dimensions: Embedding dimensions (optional)

            Returns:
                DataFrame with processed data and metadata
            """
            # Flatten embeddings (in case they're nested)
            flattened_embeddings = []
            for emb in embeddings:
                if isinstance(emb, list) and len(emb) > 0 and isinstance(emb[0], list):
                    flattened_embeddings.append(emb[0])
                else:
                    flattened_embeddings.append(emb)

            # Convert to numpy array
            embedding_array = np.array(flattened_embeddings)

            # Handle the case of a single embedding differently
            if len(embedding_array) == 1:
                # For a single point, we don't need t-SNE, just use a fixed position
                reduced_embeddings = np.array([[0.0, 0.0, 0.0]])
            else:
                # Apply dimensionality reduction to 3D
                # Fix: Ensure perplexity is at least 1.0
                perplexity_value = max(1.0, min(30, len(embedding_array)-1))
                tsne = TSNE(n_components=3, random_state=42, perplexity=perplexity_value)
                reduced_embeddings = tsne.fit_transform(embedding_array)

            # Format texts for display
            formatted_texts = []
            for text in texts:
                # Truncate if needed
                if len(text) > 500:
                    text = text[:500] + "..."

                # Insert line breaks for wrapping
                wrapped_text = ""
                for i in range(0, len(text), 50):
                    wrapped_text += text[i:i+50] + "<br>"

                formatted_texts.append("<b>"+wrapped_text+"</b>")

            # Create dataframe for visualization
            df = pd.DataFrame({
                "x": reduced_embeddings[:, 0],
                "y": reduced_embeddings[:, 1],
                "z": reduced_embeddings[:, 2],
                "text": formatted_texts,
                "full_text": texts,
                "index": range(len(texts)),
                "embedding": flattened_embeddings  # Store the original embeddings for later use
            })

            # Add metadata
            metadata = {
                "model_id": model_id,
                "embedding_dimensions": embedding_dimensions
            }

            return df, metadata

    def create_3d_embedding_chart(df, metadata=None, chart_width=1200, chart_height=800, marker_size_var: int=3):
        """
        Create a 3D Plotly chart for embedding visualization with proximity-based coloring
        """
        model_id = metadata.get("model_id") if metadata else None
        embedding_dimensions = metadata.get("embedding_dimensions") if metadata else None

        # Calculate the proximity between points
        from scipy.spatial.distance import pdist, squareform
        # Get the coordinates as a numpy array
        coords = df[['x', 'y', 'z']].values

        # Calculate pairwise distances
        dist_matrix = squareform(pdist(coords))

        # For each point, find its average distance to all other points
        avg_distances = np.mean(dist_matrix, axis=1)

        # Add this to the dataframe - smaller values = closer to other points
        df['proximity'] = avg_distances

        # Create 3D scatter plot with proximity-based coloring
        fig = px.scatter_3d(
            df, 
            x='x', 
            y='y', 
            z='z',
            # x='petal_length',  # Changed from 'x' to 'petal_length'
            # y='petal_width',   # Changed from 'y' to 'petal_width'
            # z='petal_height',
            color='proximity',  # Color based on proximity
            color_continuous_scale='Viridis_r',  # Reversed so closer points are warmer colors
            hover_data=['text', 'index', 'proximity'],
            labels={'x': 'Dimension 1', 'y': 'Dimension 2', 'z': 'Dimension 3', 'proximity': 'Avg Distance'},
            # labels={'x': 'Dimension 1', 'y': 'Dimension 2', 'z': 'Dimension 3', 'proximity': 'Avg Distance'},
            title=f"<b>3D Embedding Visualization</b>{f' - Model: <b>{model_id}</b>' if model_id else ''}{f'  <i>({embedding_dimensions} dimensions)</i>' if embedding_dimensions else ''}",
            text='index',
            # size_max=marker_size_var
        )

        # Update marker size and layout
        # fig.update_traces(marker=dict(size=3), selector=dict(mode='markers'))
        fig.update_traces(
            marker=dict(
                size=marker_size_var,              # Very small marker size
                opacity=0.7,         # Slightly transparent 
                symbol="diamond",      # Use circle markers (other options: "square", "diamond", "cross", "x")
                line=dict(
                    width=0.5,       # Very thin border
                    color="white"    # White outline makes small dots more visible
                )
            ),
            textfont=dict(
                color="rgba(255, 255, 255, 0.3)",
                size=8
            ),
            # hovertemplate="<b>index=%{text}</b><br>%{customdata[0]}<br><br>Avg Distance=%{customdata[2]:.4f}<extra></extra>", ### Hover Changes
            hovertemplate="text:<br><b>%{customdata[0]}</b><br>index: <b>%{text}</b><br><br>Avg Distance: <b>%{customdata[2]:.4f}</b><extra></extra>",
            hoverinfo="text+name",
            hoverlabel=dict(
                bgcolor="white",     # White background for hover labels
                font_size=12         # Font size for hover text
            ),
            selector=dict(type='scatter3d')
        )

        # Keep your existing layout settings
        fig.update_layout(
            scene=dict(
                xaxis=dict(
                    title='Dimension 1',
                    nticks=40,
                    backgroundcolor="rgb(10, 10, 20, 0.1)",
                    gridcolor="white",
                    showbackground=True,
                    gridwidth=0.35,
                    zerolinecolor="white",
                ),
                yaxis=dict(
                    title='Dimension 2',
                    nticks=40,
                    backgroundcolor="rgb(10, 10, 20, 0.1)",
                    gridcolor="white",
                    showbackground=True,
                    gridwidth=0.35,
                    zerolinecolor="white",
                ),
                zaxis=dict(
                    title='Dimension 3',
                    nticks=40,
                    backgroundcolor="rgb(10, 10, 20, 0.1)",
                    gridcolor="white",
                    showbackground=True,
                    gridwidth=0.35,
                    zerolinecolor="white",
                ),
                # Control camera view angle
                camera=dict(
                    up=dict(x=0, y=0, z=1),
                    center=dict(x=0, y=0, z=0),
                    eye=dict(x=1.25, y=1.25, z=1.25),
                ),
                aspectratio=dict(x=1, y=1, z=1),
                aspectmode='data'
            ),
            width=int(chart_width),
            height=int(chart_height),
            margin=dict(r=20, l=10, b=10, t=50),
            paper_bgcolor="rgb(0, 0, 0)",
            plot_bgcolor="rgb(0, 0, 0)",
            coloraxis_colorbar=dict(
                title="Average Distance",
                thicknessmode="pixels", thickness=20,
                lenmode="pixels", len=400,
                yanchor="top", y=1,
                ticks="outside",
                dtick=0.1
            )
        )

        return fig
    return create_3d_embedding_chart, prepare_embedding_data_3d


@app.cell
def helper_function_text_preparation():
    def convert_table_to_json_docs(df, selected_columns=None):
        """
        Convert a pandas DataFrame or dictionary to a list of JSON documents.
        Dynamically includes columns based on user selection.
        Column names are standardized to lowercase with underscores instead of spaces
        and special characters removed.

        Args:
            df: The DataFrame or dictionary to process
            selected_columns: List of column names to include in the output documents

        Returns:
            list: A list of dictionaries, each representing a row as a JSON document
        """
        import pandas as pd
        import re

        def standardize_key(key):
            """Convert a column name to lowercase with underscores instead of spaces and no special characters"""
            if not isinstance(key, str):
                return str(key).lower()
            # Replace spaces with underscores and convert to lowercase
            key = key.lower().replace(' ', '_')
            # Remove special characters (keeping alphanumeric and underscores)
            return re.sub(r'[^\w]', '', key)

        # Handle case when input is a dictionary
        if isinstance(df, dict):
            # Filter the dictionary to include only selected columns
            if selected_columns:
                return [{standardize_key(k): df.get(k, None) for k in selected_columns}]
            else:
                # If no columns selected, return all key-value pairs with standardized keys
                return [{standardize_key(k): v for k, v in df.items()}]

        # Handle case when df is None
        if df is None:
            return []

        # Ensure df is a DataFrame
        if not isinstance(df, pd.DataFrame):
            try:
                df = pd.DataFrame(df)
            except:
                return []  # Return empty list if conversion fails

        # Now check if DataFrame is empty
        if df.empty:
            return []

        # If no columns are specifically selected, use all available columns
        if not selected_columns or not isinstance(selected_columns, list) or len(selected_columns) == 0:
            selected_columns = list(df.columns)

        # Determine which columns exist in the DataFrame
        available_columns = []
        columns_lower = {col.lower(): col for col in df.columns if isinstance(col, str)}

        for col in selected_columns:
            if col in df.columns:
                available_columns.append(col)
            elif isinstance(col, str) and col.lower() in columns_lower:
                available_columns.append(columns_lower[col.lower()])

        # If no valid columns found, return empty list
        if not available_columns:
            return []

        # Process rows
        json_docs = []
        for _, row in df.iterrows():
            doc = {}
            for col in available_columns:
                value = row[col]
                # Standardize the column name when adding to document
                std_col = standardize_key(col)
                doc[std_col] = None if pd.isna(value) else value
            json_docs.append(doc)

        return json_docs

    def get_column_values(df, columns_to_include):
        """
        Extract values from specified columns of a dataframe as lists.

        Args:
            df: A pandas DataFrame
            columns_to_include: A list of column names to extract

        Returns:
            Dictionary with column names as keys and their values as lists
        """
        result = {}

        # Validate that columns exist in the dataframe
        valid_columns = [col for col in columns_to_include if col in df.columns]
        invalid_columns = set(columns_to_include) - set(valid_columns)

        if invalid_columns:
            print(f"Warning: These columns don't exist in the dataframe: {list(invalid_columns)}")

        # Extract values for each valid column
        for col in valid_columns:
            result[col] = df[col].tolist()

        return result

    def get_data_in_range(doc_dict_df, index_range, columns_to_include):
        """
        Extract values from specified columns of a dataframe within a given index range.

        Args:
            doc_dict_df: The pandas DataFrame to extract data from
            index_range: An integer specifying the number of rows to include (from 0 to index_range-1)
            columns_to_include: A list of column names to extract

        Returns:
            Dictionary with column names as keys and their values (within the index range) as lists
        """
        # Validate the index range
        max_index = len(doc_dict_df)
        if index_range <= 0:
            print(f"Warning: Invalid index range {index_range}. Must be positive.")
            return {}

        # Adjust index_range if it exceeds the dataframe length
        if index_range > max_index:
            print(f"Warning: Index range {index_range} exceeds dataframe length {max_index}. Using maximum length.")
            index_range = max_index

        # Slice the dataframe to get rows from 0 to index_range-1
        df_subset = doc_dict_df.iloc[:index_range]

        # Use the provided get_column_values function to extract column data
        return get_column_values(df_subset, columns_to_include)

    def get_data_in_range_triplequote(doc_dict_df, index_range, columns_to_include):
        """
        Extract values from specified columns of a dataframe within a given index range.
        Wraps string values with triple quotes and escapes URLs.

        Args:
            doc_dict_df: The pandas DataFrame to extract data from
            index_range: A list of two integers specifying the start and end indices of rows to include
                        (e.g., [0, 10] includes rows from index 0 to 9 inclusive)
            columns_to_include: A list of column names to extract
        """
        # Validate the index range
        start_idx, end_idx = index_range
        max_index = len(doc_dict_df)

        # Validate start index
        if start_idx < 0:
            print(f"Warning: Invalid start index {start_idx}. Using 0 instead.")
            start_idx = 0

        # Validate end index
        if end_idx <= start_idx:
            print(f"Warning: End index {end_idx} must be greater than start index {start_idx}. Using {start_idx + 1} instead.")
            end_idx = start_idx + 1

        # Adjust end index if it exceeds the dataframe length
        if end_idx > max_index:
            print(f"Warning: End index {end_idx} exceeds dataframe length {max_index}. Using maximum length.")
            end_idx = max_index

        # Slice the dataframe to get rows from start_idx to end_idx-1
        # Using .loc with slice to preserve original indices
        df_subset = doc_dict_df.iloc[start_idx:end_idx]

        # Use the provided get_column_values function to extract column data
        result = get_column_values(df_subset, columns_to_include)

        # Process each string result to wrap in triple quotes
        for col in result:
            if isinstance(result[col], list):
                # Create a new list with items wrapped in triple quotes
                processed_items = []
                for item in result[col]:
                    if isinstance(item, str):
                        # Replace http:// and https:// with escaped versions
                        item = item.replace("http://", "http\\://").replace("https://", "https\\://")
                        # processed_items.append('"""' + item + '"""')
                        processed_items.append(item)
                    else:
                        processed_items.append(item)
                result[col] = processed_items
        return result
    return (get_data_in_range_triplequote,)


@app.cell
def prepare_doc_select(sentence_splitter_config):
    def prepare_document_selection(node_dict):
        """
        Creates document selection UI component.
        Args:
            node_dict: Dictionary mapping filenames to lists of documents
        Returns:
            mo.ui component for document selection
        """
        # Calculate total number of documents across all files
        total_docs = sum(len(docs) for docs in node_dict.values())

        # Create a combined DataFrame of all documents for table selection
        all_docs_records = []
        doc_index_global = 0
        for filename, docs in node_dict.items():
            for i, doc in enumerate(docs):
                # Convert the document to a format compatible with DataFrame
                if hasattr(doc, 'to_dict'):
                    doc_data = doc.to_dict()
                elif isinstance(doc, dict):
                    doc_data = doc
                else:
                    doc_data = {'content': str(doc)}

                # Add metadata
                doc_data['filename'] = filename
                doc_data['doc_index'] = i
                doc_data['global_index'] = doc_index_global
                all_docs_records.append(doc_data)
                doc_index_global += 1

        # Create UI component
        stop_value = max(total_docs, 1)
        llama_docs = mo.ui.range_slider(
            start=1, 
            stop=stop_value, 
            step=1, 
            full_width=True, 
            show_value=True, 
            label="**Select a Range of Chunks to Visualize:**"
        ).form(submit_button_disabled=check_state(sentence_splitter_config.value), submit_button_label="Change Document View Range")

        return llama_docs
    return (prepare_document_selection,)


@app.cell
def document_range_selection(
    dict_from_nodes,
    prepare_document_selection,
    set_range_slider_state,
):
    if dict_from_nodes is not None:
        llama_docs = prepare_document_selection(dict_from_nodes)
        set_range_slider_state(llama_docs)
    else:
        bare_dict = {}
        llama_docs = prepare_document_selection(bare_dict)
    return


@app.function
def create_cumulative_dataframe(dict_from_docs):
    """
    Creates a cumulative DataFrame from a nested dictionary of documents.

    Args:
        dict_from_docs: Dictionary mapping filenames to lists of documents

    Returns:
        DataFrame with all documents flattened with global indices
    """
    # Create a list to hold all document records
    all_records = []
    global_idx = 1  # Start from 1 to match range slider expectations

    for filename, docs in dict_from_docs.items():
        for i, doc in enumerate(docs):
            # Convert the document to a dict format
            if hasattr(doc, 'to_dict'):
                doc_data = doc.to_dict()
            elif isinstance(doc, dict):
                doc_data = doc.copy()
            else:
                doc_data = {'content': str(doc)}

            # Add additional metadata
            doc_data['filename'] = filename
            doc_data['doc_index'] = i
            doc_data['global_index'] = global_idx

            # If there's 'content' but no 'text', create a 'text' field
            if 'content' in doc_data and 'text' not in doc_data:
                doc_data['text'] = doc_data['content']

            all_records.append(doc_data)
            global_idx += 1

    # Create DataFrame from all records
    return pd.DataFrame(all_records)


@app.function
def create_stats(texts_dict, bordered=False, object_names=None, group_by_row=False, items_per_row=6, gap=2, label="Chunk"):
    """
    Create a list of stat objects for each item in the specified dictionary.

    Parameters:
    - texts_dict (dict): Dictionary containing the text data
    - bordered (bool): Whether the stats should be bordered
    - object_names (list or tuple): Two object names to use for label and value
                                   [label_object, value_object]
    - group_by_row (bool): Whether to group stats in rows (horizontal stacks)
    - items_per_row (int): Number of stat objects per row when group_by_row is True

    Returns:
    - object: A vertical stack of stat objects or rows of stat objects
    """
    if not object_names or len(object_names) < 2:
        raise ValueError("You must provide two object names as a list or tuple")

    label_object = object_names[0]
    value_object = object_names[1]

    # Validate that both objects exist in the dictionary
    if label_object not in texts_dict:
        raise ValueError(f"Label object '{label_object}' not found in texts_dict")
    if value_object not in texts_dict:
        raise ValueError(f"Value object '{value_object}' not found in texts_dict")

    # Determine how many items to process (based on the label object length)
    num_items = len(texts_dict[label_object])

    # Create individual stat objects
    individual_stats = []
    for i in range(num_items):
        stat = mo.stat(
            label=texts_dict[label_object][i],
            value=f"{label} Number: {len(texts_dict[value_object][i])}",
            bordered=bordered
        )
        individual_stats.append(stat)

    # If grouping is not enabled, just return a vertical stack of all stats
    if not group_by_row:
        return mo.vstack(individual_stats, wrap=False)

    # Group stats into rows based on items_per_row
    rows = []
    for i in range(0, num_items, items_per_row):
        # Get a slice of stats for this row (up to items_per_row items)
        row_stats = individual_stats[i:i+items_per_row]
        # Create a horizontal stack for this row
        widths = [0.35] * len(row_stats)
        row = mo.hstack(row_stats, gap=gap, align="start", justify="center", widths=widths)
        rows.append(row)

    # Return a vertical stack of all rows
    return mo.vstack(rows)


@app.cell
def prepare_chart_embeddings(
    chunks_to_process,
    emb_model,
    emb_model_emb_dim,
    get_embedding_state,
    prepare_embedding_data_3d,
):
    # chart_dataframe, chart_metadata = None, None
    if chunks_to_process is not None and get_embedding_state() is not None:
        chart_dataframe, chart_metadata = prepare_embedding_data_3d(
            get_embedding_state(), 
            chunks_to_process,
            model_id=emb_model,
            embedding_dimensions=emb_model_emb_dim
        )
    else:
        chart_dataframe, chart_metadata = None, None
    return chart_dataframe, chart_metadata


@app.function
def all_true(*args):
    """
    Check if all provided boolean arguments are True.
    """
    return all(args)


@app.cell
def _(chart_dataframe_prem, columns_selected):
    print(columns_selected,chart_dataframe_prem)
    return


@app.cell
def _(
    columns_selected,
    emb_model,
    emb_model_emb_dim,
    premade_documents,
    prepare_embedding_data_3d,
):
    if premade_documents and columns_selected:
        chart_dataframe_prem, chart_metadata_prem = prepare_embedding_data_3d(
            premade_documents["embeddings"], 
            premade_documents["texts"],
            model_id=emb_model,
            embedding_dimensions=emb_model_emb_dim
        )
    else:
        chart_dataframe_prem = chart_metadata_prem = None
    return chart_dataframe_prem, chart_metadata_prem


@app.cell
def chart_dims():
    chart_dimensions = (
        mo.md('''
        > **Adjust Chart Window**

        {chart_height}

        {chat_width}

        ''').batch(
                chart_height = mo.ui.slider(start=500, step=30, stop=1000, label="**Height:**", value=800, show_value=True),
                chat_width = mo.ui.slider(start=900, step=50, stop=1400, label="**Width:**", value=1200, show_value=True)
            )
    )
    return (chart_dimensions,)


@app.cell
def chart_dim_values(chart_dimensions):
    chart_height = chart_dimensions.value['chart_height']
    chart_width = chart_dimensions.value['chat_width']
    return chart_height, chart_width


@app.cell
def create_baseline_chart(
    chart_dataframe,
    chart_dataframe_prem,
    chart_height,
    chart_metadata,
    chart_metadata_prem,
    chart_width,
    create_3d_embedding_chart,
):
    if chart_dataframe is not None and chart_metadata is not None:
        emb_plot = create_3d_embedding_chart(chart_dataframe, chart_metadata, chart_width, chart_height, marker_size_var=9)
        chart = mo.ui.plotly(emb_plot)
        chart_reference = chart_dataframe
    
    elif chart_dataframe_prem is not None and chart_metadata_prem is not None:
        emb_plot = create_3d_embedding_chart(chart_dataframe_prem, chart_metadata_prem, chart_width, chart_height, marker_size_var=9)
        chart = mo.ui.plotly(emb_plot)
        chart_reference = chart_dataframe_prem
    
    else:
        emb_plot = chart = chart_reference = None
    return chart, chart_reference, emb_plot


@app.cell
def test_query(get_chunk_state, premade_documents, switch_file_loader_type):
    placeholder = """How can i use watsonx.data to perform vector search?"""
    if switch_file_loader_type.value:
        query = mo.ui.text_area(label="**Write text to check:**", full_width=True, rows=8, value=placeholder).form(show_clear_button=True, 
                                                                                                                   submit_button_disabled=check_state(premade_documents),  
                                                                                                                   submit_button_label="Query and View Visualization")
    else:
        query = mo.ui.text_area(label="**Write text to check:**", full_width=True, rows=8, value=placeholder).form(show_clear_button=True, 
                                                                                                                   submit_button_disabled=check_state(get_chunk_state()),  
                                                                                                                   submit_button_label="Query and View Visualization")
    return (query,)


@app.cell
def query_stack(chart_dimensions, query):
    # query_stack = mo.hstack([query], justify="space-around", align="center", widths=[0.65])
    query_stack = mo.hstack([query, chart_dimensions], justify="space-around", align="center", gap=15)
    return (query_stack,)


@app.function
def check_state(variable):
    return variable is None


@app.cell
def helper_function_add_query_to_chart():
    def add_query_to_embedding_chart(existing_chart, query_coords, query_text, marker_size=12):
        """
        Add a query point to an existing 3D embedding chart as a large red dot.

        Args:
            existing_chart: The existing plotly figure or chart data
            query_coords: Dictionary with 'x', 'y', 'z' coordinates for the query point
            query_text: Text of the query to display on hover
            marker_size: Size of the query marker (default: 18, typically 2x other markers)

        Returns:
            A modified plotly figure with the query point added as a red dot
        """
        import plotly.graph_objects as go

        # Create a deep copy of the existing chart to avoid modifying the original
        import copy
        chart_copy = copy.deepcopy(existing_chart)

        # Handle case where chart_copy is a dictionary or list (from mo.ui.plotly)
        if isinstance(chart_copy, (dict, list)):
            # Create a new plotly figure from the data
            import plotly.graph_objects as go

            if isinstance(chart_copy, list):
                # If it's a list, assume it's a list of traces
                fig = go.Figure(data=chart_copy)
            else:
                # If it's a dict with 'data' and 'layout'
                fig = go.Figure(data=chart_copy.get('data', []), layout=chart_copy.get('layout', {}))

            chart_copy = fig

        # Create the query trace
        query_trace = go.Scatter3d(
            x=[query_coords['x']],
            y=[query_coords['y']],
            z=[query_coords['z']],
            mode='markers',
            name='Query',
            marker=dict(
                size=marker_size,  # Typically 2x the size of other markers
                color='red',       # Bright red color
                symbol='circle',   # Circle shape
                opacity=0.70,       # Fully opaque
                line=dict(
                    width=1,       # Thin white border
                    color='white'
                )
            ),
            # text=['Query: ' + query_text],
            text=['<b>Query:</b><br>' + '<br>'.join([query_text[i:i+50] for i in range(0, len(query_text), 50)])], ### Text Wrapping
            hoverinfo="text+name"
        )

        # Add the query trace to the chart copy
        chart_copy.add_trace(query_trace)

        return chart_copy


    def get_query_coordinates(reference_embeddings=None, query_embedding=None):
        """
        Calculate appropriate coordinates for a query point based on reference embeddings.

        This function handles several scenarios:
        1. If both reference embeddings and query embedding are provided, it places the
           query near similar documents.
        2. If only reference embeddings are provided, it places the query at a visible 
           location near the center of the chart.
        3. If neither are provided, it returns default origin coordinates.

        Args:
            reference_embeddings: DataFrame with x, y, z coordinates from the main chart
            query_embedding: The embedding vector of the query

        Returns:
            Dictionary with x, y, z coordinates for the query point
        """
        import numpy as np

        # Default coordinates (origin with slight offset)
        default_coords = {'x': 0.0, 'y': 0.0, 'z': 0.0}

        # If we don't have reference embeddings, return default
        if reference_embeddings is None or len(reference_embeddings) == 0:
            return default_coords

        # If we have reference embeddings but no query embedding, 
        # position at a visible location near the center
        if query_embedding is None:
            center_coords = {
                'x': reference_embeddings['x'].mean(),
                'y': reference_embeddings['y'].mean(),
                'z': reference_embeddings['z'].mean()
            }
            return center_coords

        # If we have both reference embeddings and query embedding,
        # try to position near similar documents
        try:
            from sklearn.metrics.pairwise import cosine_similarity

            # Check if original embeddings are in the dataframe
            if 'embedding' in reference_embeddings.columns:
                # Get all document embeddings as a 2D array
                if isinstance(reference_embeddings['embedding'].iloc[0], list):
                    doc_embeddings = np.array(reference_embeddings['embedding'].tolist())
                else:
                    doc_embeddings = np.array([emb for emb in reference_embeddings['embedding'].values])

                # Reshape query embedding for comparison
                query_emb_array = np.array(query_embedding)
                if query_emb_array.ndim == 1:
                    query_emb_array = query_emb_array.reshape(1, -1)

                # Calculate cosine similarities
                similarities = cosine_similarity(query_emb_array, doc_embeddings)[0]

                # Find the closest document
                closest_idx = np.argmax(similarities)

                # Use the position of the closest document, with slight offset for visibility
                query_coords = {
                    'x': reference_embeddings['x'].iloc[closest_idx] + 0.2,
                    'y': reference_embeddings['y'].iloc[closest_idx] + 0.2,
                    'z': reference_embeddings['z'].iloc[closest_idx] + 0.2
                }
                return query_coords
        except Exception as e:
            print(f"Error positioning query near similar documents: {e}")

        # Fallback to center position if similarity calculation fails
        center_coords = {
            'x': reference_embeddings['x'].mean(),
            'y': reference_embeddings['y'].mean(),
            'z': reference_embeddings['z'].mean()
        }
        return center_coords
    return add_query_to_embedding_chart, get_query_coordinates


@app.cell
def combined_chart_visualization(
    add_query_to_embedding_chart,
    chart,
    chart_reference,
    emb_plot,
    embedding,
    get_query_coordinates,
    get_query_state,
    query,
    set_chart_state,
    set_query_state,
):
    # Usage with highlight_closest=True
    if chart is not None and query.value:
        with mo.status.spinner(title="Embedding Query...", remove_on_exit=True) as _spinner:
                query_emb = embedding.embed_documents([query.value])
                set_query_state(query_emb)

                _spinner.update("Preparing Query Coordinates") # --- --- ---
                time.sleep(1.0)

                # Get appropriate coordinates for the query
                query_coords = get_query_coordinates(
                    reference_embeddings=chart_reference,
                    query_embedding=get_query_state()
                )

                _spinner.update("Adding Query to Chart") # --- --- ---
                time.sleep(1.0)

                # Add the query to the chart with closest points highlighted
                result = add_query_to_embedding_chart(
                    existing_chart=emb_plot,
                    query_coords=query_coords,
                    query_text=query.value,
                )

                chart_with_query = result

                _spinner.update("Preparing Visualization") # --- --- ---
                time.sleep(1.0)

                # Create the visualization
                combined_viz = mo.ui.plotly(chart_with_query)
                set_chart_state(combined_viz)

                _spinner.update("Done") # --- --- ---
    else:
        combined_viz = None
    return


@app.cell
def _():
    get_range_slider_state, set_range_slider_state = mo.state(None)
    return get_range_slider_state, set_range_slider_state


@app.cell
def _(get_range_slider_state):
    if get_range_slider_state() is not None:
        document_range_stack = get_range_slider_state()
    else:
        document_range_stack = None
    return (document_range_stack,)


@app.cell
def _():
    get_chart_state, set_chart_state = mo.state(None)
    return get_chart_state, set_chart_state


@app.cell
def _(get_chart_state, query):
    if query.value is not None:
        chart_visualization = get_chart_state()
    else:
        chart_visualization = None
    return (chart_visualization,)


@app.cell
def c(document_range_stack):
    chart_range_selection = mo.hstack([document_range_stack], justify="space-around", align="center", widths=[0.65])
    return (chart_range_selection,)


if __name__ == "__main__":
    app.run()