uap_analyzer.py

# Standard library
import concurrent.futures
import json
import logging
import os
import pickle
import textwrap
import time

# Third-party core
import numpy as np
import pandas as pd

# ML / NLP
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from sklearn.metrics import accuracy_score, confusion_matrix
from sklearn.model_selection import train_test_split
from scipy.stats import chi2_contingency
from statsmodels.api import stats
from statsmodels.graphics.mosaicplot import mosaic
import xgboost as xgb
from xgboost import plot_importance
from Levenshtein import distance

# Visualization
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import seaborn as sns
import plotly.graph_objects as go
import plotly.express as px
import squarify

# Deep learning / embeddings
import torch
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import pytorch_cos_sim, pairwise_cos_sim

# UMAP + HDBSCAN
import umap
import fast_hdbscan as hdbscan

# Web / API
from requests.exceptions import HTTPError
from openai import OpenAI
from stqdm import stqdm
stqdm.pandas()
import streamlit as st

torch.set_float32_matmul_precision("medium")

_device = 'cuda' if torch.cuda.is_available() else 'cpu'
_embed_model = None

def get_embed_model():
    global _embed_model
    if _embed_model is None:
        _model_kwargs = {"torch_dtype": torch.bfloat16}
        if _device == "cuda":
            _model_kwargs["device_map"] = "cuda"
        _embed_model = SentenceTransformer(
            "microsoft/harrier-oss-v1-270m", model_kwargs=_model_kwargs
        )
        if _device != "cuda":
            _embed_model.to(_device)
        # Harrier's prompts load automatically from its config; verify with:
        #   print(_embed_model.prompts)
        # Set max_seq_length explicitly — many models silently truncate otherwise.
        _embed_model.max_seq_length = 512
    return _embed_model


# GPU acceleration — cuML (UMAP/HDBSCAN/KMeans/PCA) and cuVS (ANN search)
# Falls back to CPU libraries if RAPIDS is not installed.
def _check_gpu() -> bool:
    try:
        import cupy as cp
        cp.zeros(1)
        return True
    except Exception:
        return False

_GPU_AVAILABLE = _check_gpu()

_CUML_AVAILABLE = False
if _GPU_AVAILABLE:
    try:
        from cuml.manifold import UMAP as _cuUMAP
        from cuml.cluster.hdbscan import HDBSCAN as _cuHDBSCAN
        from cuml.cluster import KMeans as _cuKMeans
        from cuml.decomposition import PCA as _cuPCA
        import rmm
        rmm.reinitialize(pool_allocator=True, initial_pool_size=2**31)  # 2 GB pool
        _CUML_AVAILABLE = True
        logging.getLogger(__name__).info("cuML detected — GPU acceleration enabled for UMAP/HDBSCAN/KMeans/PCA")
    except ImportError:
        logging.getLogger(__name__).warning("cuML not found — install with: uv add --extra-index-url=https://pypi.nvidia.com cuml-cu12")

_CUVS_AVAILABLE = False
if _GPU_AVAILABLE:
    try:
        from cuvs.neighbors import cagra as _cagra
        _CUVS_AVAILABLE = True
        logging.getLogger(__name__).info("cuVS detected — GPU ANN search enabled")
    except ImportError:
        logging.getLogger(__name__).warning("cuVS not found — install with: uv add --extra-index-url=https://pypi.nvidia.com cuvs-cu12")

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class UAPAnalyzer:
    """
    A class for analyzing and clustering textual data within a pandas DataFrame using 
    Natural Language Processing (NLP) techniques and machine learning models.
    
    Attributes:
        data (pd.DataFrame): The dataset containing textual data for analysis.
        column (str): The name of the column in the DataFrame to be analyzed.
        embeddings (np.ndarray): The vector representations of textual data.
        reduced_embeddings (np.ndarray): The dimensionality-reduced embeddings.
        cluster_labels (np.ndarray): The labels assigned to each data point after clustering.
        cluster_terms (list): The list of terms associated with each cluster.
        tfidf_matrix (sparse matrix): The Term Frequency-Inverse Document Frequency (TF-IDF) matrix.
        models (dict): A dictionary to store trained machine learning models.
        evaluations (dict): A dictionary to store evaluation results of models.
        data_nums (pd.DataFrame): The DataFrame with numerical encoding of categorical data.
    """

    def __init__(self, data, column, has_embeddings=False):
        """
        Initializes the UAPAnalyzer with a dataset and a specified column for analysis.
        
        Args:
            data (pd.DataFrame): The dataset for analysis.
            column (str): The column within the dataset to analyze.
        """
        assert isinstance(data, pd.DataFrame), "Data must be a pandas DataFrame"
        assert column in data.columns, f"Column '{column}' not found in DataFrame"
        self.has_embeddings = has_embeddings        
        self.data = data
        self.column = column
        self.embeddings = None
        self.reduced_embeddings = None
        self.cluster_labels = None
        self.cluster_names = None
        self.cluster_terms = None 
        self.cluster_terms_embeddings = None
        self.tfidf_matrix = None
        self.models = {}  # To store trained models
        self.evaluations = {}  # To store evaluation results
        self.data_nums = None  # Encoded numerical data
        self.x_train = None
        self.y_train = None
        self.x_test = None
        self.y_test = None
        self.preds = None
        self.new_dataset = None
        #self.cluster_names_ = pd.DataFrame()

        logging.info("UAPAnalyzer initialized")

    def preprocess_data(self, trim=False, has_embeddings=False, top_n=32,):
        """
        Preprocesses the data by optionally trimming the dataset to include only the top N labels and extracting embeddings.
        
        Args:
            trim (bool): Whether to trim the dataset to include only the top N labels.
            top_n (int): The number of top labels to retain if trimming is enabled.
        """
        logging.info("Preprocessing data")

        # if trim is True
        if trim:
            # Identify the top labels based on value counts
            top_labels = self.data[self.column].value_counts().nlargest(top_n).index.tolist()
            # Revise the column data, setting values to 'Other' if they are not in the top labels
            self.data[f'{self.column}_revised'] = np.where(self.data[self.column].isin(top_labels), self.data[self.column], 'Other')
        # Convert the column data to string type before passing to _extract_embeddings
        # This is useful especially if the data type of the column is not originally string
        string_data = self.data[f'{self.column}'].astype(str)
        # Extract embeddings from the revised and string-converted column data
        if has_embeddings:
            self.embeddings = self.data['embeddings'].to_list()
        else:
            model = get_embed_model()
            pool = model.start_multi_process_pool(target_devices=[_device])
            try:
                self.embeddings = self._extract_embeddings(string_data, pool=pool)
            finally:
                model.stop_multi_process_pool(pool)
        logging.info("Data preprocessing complete")


    def _extract_embeddings(self, data_column, pool=None):
        """
        Extracts embeddings from the given data column.

        Args:
            data_column (pd.Series): The column from which to extract embeddings.
            pool: Optional multi-process pool from model.start_multi_process_pool().

        Returns:
            np.ndarray: float32 embeddings of shape (n_samples, dim).
        """
        logging.info("Extracting embeddings")
        embs = get_embed_model().encode_document(
            data_column.tolist(),
            show_progress_bar=True,
            batch_size=256,
            pool=pool,
            chunk_size=5000,
        )
        return np.array(embs, dtype=np.float32)

    def reduce_dimensionality(self, method='UMAP', n_components=2,
                               pca_preprocess=True, pca_components=50, **kwargs):
        """
        Reduces the dimensionality of embeddings using the specified method.

        When a GPU is available and cuML is installed, UMAP and PCA run on the GPU
        (60-300x faster than CPU).  An optional PCA pre-reduction step compresses
        high-dimensional embeddings (e.g. 768-d) to `pca_components` before UMAP,
        which dramatically speeds up the UMAP graph construction.

        Args:
            method (str): 'UMAP' or 'PCA'.
            n_components (int): Target dimensionality.
            pca_preprocess (bool): Apply PCA before UMAP when embeddings are wide.
            pca_components (int): PCA target dims used only as a pre-reduction step.
            **kwargs: Forwarded to the UMAP/PCA constructor.
        """
        logging.info(f"Reducing dimensionality using {method}")
        embeddings = np.array(self.embeddings, dtype=np.float32)

        # PCA pre-reduction: collapse 768→50 before UMAP to cut cost ~50x
        if pca_preprocess and method == 'UMAP' and embeddings.shape[1] > pca_components:
            n_pca = min(pca_components, embeddings.shape[0] - 1, embeddings.shape[1])
            if _CUML_AVAILABLE:
                pre_pca = _cuPCA(n_components=n_pca)
            else:
                pre_pca = PCA(n_components=n_pca)
            embeddings = pre_pca.fit_transform(embeddings)
            if hasattr(embeddings, 'get'):
                embeddings = embeddings.get()
            embeddings = np.array(embeddings, dtype=np.float32)
            logging.info(f"PCA pre-reduction to {n_pca} dims complete")

        if method == 'UMAP':
            if _CUML_AVAILABLE:
                reducer = _cuUMAP(n_components=n_components, **kwargs)
            else:
                reducer = umap.UMAP(n_components=n_components, **kwargs)
        elif method == 'PCA':
            if _CUML_AVAILABLE:
                reducer = _cuPCA(n_components=n_components)
            else:
                reducer = PCA(n_components=n_components)
        else:
            raise ValueError("Unsupported dimensionality reduction method")

        result = reducer.fit_transform(embeddings)
        if hasattr(result, 'get'):      # CuPy → NumPy
            result = result.get()
        self.reduced_embeddings = np.array(result, dtype=np.float32)
        logging.info(f"Dimensionality reduced using {method} ({'GPU' if _CUML_AVAILABLE else 'CPU'})")

    def cluster_data(self, method='HDBSCAN', **kwargs):
        """
        Clusters the reduced dimensionality data using the specified clustering method.

        When cuML is available and the dataset is large enough (n*d > 50k), HDBSCAN
        and KMeans run on GPU.  For small/low-dim inputs the CPU implementation is used
        because GPU kernel launch overhead dominates.

        Args:
            method (str): 'HDBSCAN' or 'KMeans'.
            **kwargs: Forwarded to the clusterer constructor.
        """
        logging.info(f"Clustering data using {method}")
        embeddings = np.array(self.reduced_embeddings, dtype=np.float32)

        # cuML HDBSCAN/KMeans only beats CPU when both N and dim are large.
        # Benchmarks: 50k×2D → GPU 40× slower; 300k×50D → GPU 13× faster.
        # Crossover is ~100k rows AND dim > 5.
        n, d = embeddings.shape
        _use_gpu_cluster = _CUML_AVAILABLE and (n > 100_000 and d > 5)

        if method == 'HDBSCAN':
            clusterer = _cuHDBSCAN(**kwargs) if _use_gpu_cluster else hdbscan.HDBSCAN(**kwargs)
        elif method == 'KMeans':
            clusterer = _cuKMeans(**kwargs) if _use_gpu_cluster else KMeans(**kwargs)
        else:
            raise ValueError("Unsupported clustering method")

        clusterer.fit(embeddings)
        labels = clusterer.labels_
        if hasattr(labels, 'get'):      # CuPy → NumPy
            labels = labels.get()
        self.cluster_labels = np.array(labels)
        logging.info(f"Data clustering complete using {method} ({'GPU' if _use_gpu_cluster else 'CPU'})")

    # ------------------------------------------------------------------
    # GPU-accelerated ANN search (cuVS CAGRA)
    # ------------------------------------------------------------------

    def build_search_index(self):
        """
        Build a GPU CAGRA ANN index over the raw embeddings for fast similarity search.

        Requires cuVS (uv add --extra-index-url=https://pypi.nvidia.com cuvs-cu12).
        Embeddings are L2-normalised before indexing so inner-product == cosine sim.
        """
        if not _CUVS_AVAILABLE:
            logging.warning("cuVS not available — search index not built")
            return
        import cupy as cp
        embs = cp.array(self.embeddings, dtype=cp.float32)
        norms = cp.linalg.norm(embs, axis=1, keepdims=True)
        embs_norm = embs / (norms + 1e-8)
        index_params = _cagra.IndexParams(metric="inner_product")
        self._search_index = _cagra.build(index_params, embs_norm)
        self._search_embs_norm = embs_norm
        logging.info(f"CAGRA search index built on GPU ({embs.shape[0]} vectors, dim={embs.shape[1]})")

    def find_similar_reports(self, query_texts: list, k: int = 10):
        """
        Return the indices and cosine-similarity scores of the k most similar
        stored reports for each query string.

        Args:
            query_texts: List of query strings to encode and search.
            k: Number of nearest neighbours to return.

        Returns:
            distances: np.ndarray shape (len(query_texts), k) — cosine similarities.
            neighbors: np.ndarray shape (len(query_texts), k) — row indices into self.data.
        """
        if not _CUVS_AVAILABLE or not hasattr(self, '_search_index'):
            raise RuntimeError("Call build_search_index() first (requires cuVS).")
        import cupy as cp
        query_embs = get_embed_model().encode_query(query_texts, batch_size=256)
        q = cp.array(query_embs, dtype=cp.float32)
        norms = cp.linalg.norm(q, axis=1, keepdims=True)
        q_norm = q / (norms + 1e-8)
        distances, neighbors = _cagra.search(_cagra.SearchParams(), self._search_index, q_norm, k=k)
        return distances.get(), neighbors.get()

    def get_tf_idf_clusters(self, top_n=2):
        """
        Names clusters using the most frequent terms based on TF-IDF analysis.

        Args:
            top_n (int): The number of top terms to consider for naming each cluster.
        """
        logging.info("Naming clusters based on top TF-IDF terms.")

        # Ensure data has been clustered
        assert self.cluster_labels is not None, "Data has not been clustered yet."
        vectorizer = TfidfVectorizer(max_features=1000, stop_words='english')

        # Fit the vectorizer to the text data and transform it into a TF-IDF matrix
        tfidf_matrix = vectorizer.fit_transform(self.data[f'{self.column}'].astype(str))

        # Initialize an empty list to store the cluster terms
        self.cluster_terms = []

        for cluster_id in np.unique(self.cluster_labels):
            # Skip noise if present (-1 in HDBSCAN)
            if cluster_id == -1:
                continue

            # Find indices of documents in the current cluster
            indices = np.where(self.cluster_labels == cluster_id)[0]

            # Compute the mean TF-IDF score for each term in the cluster
            cluster_tfidf_mean = np.mean(tfidf_matrix[indices], axis=0)

            # Use the matrix directly for indexing if it does not support .toarray()
            # Ensure it's in a format that supports indexing, convert if necessary
            if hasattr(cluster_tfidf_mean, "toarray"):
                dense_mean = cluster_tfidf_mean.toarray().flatten()
            else:
                dense_mean = np.asarray(cluster_tfidf_mean).flatten()

            # Get the indices of the top_n terms
            top_n_indices = np.argsort(dense_mean)[-top_n:]

            # Get the corresponding terms for these top indices
            terms = vectorizer.get_feature_names_out()
            top_terms = [terms[i] for i in top_n_indices]

            # Join the top_n terms with a hyphen
            cluster_name = '-'.join(top_terms)

            # Append the cluster name to the list
            self.cluster_terms.append(cluster_name)

        # Convert the list of cluster terms to a categorical data type
        self.cluster_terms = pd.Categorical(self.cluster_terms)
        logging.info("Cluster naming completed.")



    def cluster_levenshtein(self, cluster_terms, cluster_labels, char_diff_threshold=3):
        from Levenshtein import distance  # Make sure to import the correct distance function

        merge_map = {}
        # Iterate over term pairs and decide on merging based on the distance
        for idx, term1 in enumerate(cluster_terms):
            for jdx, term2 in enumerate(cluster_terms):
                if idx < jdx and distance(term1, term2) <= char_diff_threshold:  
                    labels_to_merge = [label for label, term_index in enumerate(cluster_labels) if term_index == jdx]
                    for label in labels_to_merge:
                        merge_map[label] = idx  # Map the label to use the term index of term1
                    logging.info(f"Merging '{term2}' into '{term1}'")
                    st.write(f"Merging '{term2}' into '{term1}'")
        # Update the cluster labels
        updated_cluster_labels = [merge_map.get(label, label) for label in cluster_labels]
        # Update string labels to reflect merged labels
        updated_string_labels = [cluster_terms[label] for label in updated_cluster_labels]
        return updated_string_labels

    def cluster_cosine(self, cluster_terms, cluster_labels, similarity_threshold):
        from sklearn.metrics.pairwise import cosine_similarity

        cluster_terms_embeddings = get_embed_model().encode_document(cluster_terms)
        # Compute cosine similarity matrix in a vectorized form
        cos_sim_matrix = cosine_similarity(cluster_terms_embeddings, cluster_terms_embeddings)

        merge_map = {}
        n_terms = len(cluster_terms)
        # Iterate only over upper triangular matrix excluding diagonal to avoid redundant computations and self-comparison
        for idx in range(n_terms):
            for jdx in range(idx + 1, n_terms):
                if cos_sim_matrix[idx, jdx] >= similarity_threshold:
                    labels_to_merge = [label for label, term_index in enumerate(cluster_labels) if term_index == jdx]
                    for label in labels_to_merge:
                        merge_map[label] = idx
                    st.write(f"Merging '{cluster_terms[jdx]}' into '{cluster_terms[idx]}'")
                    logging.info(f"Merging '{cluster_terms[jdx]}' into '{cluster_terms[idx]}'")
        # Update the cluster labels
        updated_cluster_labels = [merge_map.get(label, label) for label in cluster_labels]
        # Update string labels to reflect merged labels
        updated_string_labels = [cluster_terms[label] for label in updated_cluster_labels]
        # make a dataframe with index, cluster label and cluster term
        return updated_string_labels

    def merge_similar_clusters(self, cluster_terms, cluster_labels, distance_type='cosine', char_diff_threshold=3, similarity_threshold=0.92):
        if distance_type == 'levenshtein':
            return self.cluster_levenshtein(cluster_terms, cluster_labels, char_diff_threshold)
        elif distance_type == 'cosine':
            return self.cluster_cosine(cluster_terms, cluster_labels, similarity_threshold)


    def plot_embeddings2(self, title=None):
        assert self.reduced_embeddings is not None, "Dimensionality reduction has not been performed yet."
        assert self.cluster_terms is not None, "Cluster TF-IDF analysis has not been performed yet."

        logging.info("Plotting embeddings with TF-IDF colors")

        fig = go.Figure()

        unique_cluster_terms = np.unique(self.cluster_terms)

        for cluster_term in unique_cluster_terms:
            if cluster_term != 'Noise':
                indices = np.where(np.array(self.cluster_terms) == cluster_term)[0]

                # Plot points in the current cluster
                fig.add_trace(
                    go.Scatter(
                        x=self.reduced_embeddings[indices, 0],
                        y=self.reduced_embeddings[indices, 1],
                        mode='markers',
                        marker=dict(
                            size=5,
                            opacity=0.8,
                        ),
                        name=cluster_term,
                        text=self.data[f'{self.column}'].iloc[indices], 
                        hoverinfo='text',
                    )
                )
            else:
                # Plot noise points differently if needed
                fig.add_trace(
                    go.Scatter(
                        x=self.reduced_embeddings[indices, 0],
                        y=self.reduced_embeddings[indices, 1],
                        mode='markers',
                        marker=dict(
                            size=5,
                            opacity=0.5,
                            color='grey'
                        ),
                        name='Noise',
                        text=[self.data[f'{self.column}'][i] for i in indices],  # Adjusted for potential pandas use
                        hoverinfo='text',
                    )
                )
            # else:
            #     indices = np.where(np.array(self.cluster_terms) == 'Noise')[0]

            #     # Plot noise points
            #     fig.add_trace(
            #         go.Scatter(
            #             x=self.reduced_embeddings[indices, 0],
            #             y=self.reduced_embeddings[indices, 1],
            #             mode='markers',
            #             marker=dict(
            #                 size=5,
            #                 opacity=0.8,
            #             ),
            #             name='Noise',
            #             text=self.data[f'{self.column}'].iloc[indices],
            #             hoverinfo='text',
            #         )
            #     )

        fig.update_layout(title=title, showlegend=True, legend_title_text='Top TF-IDF Terms')
        #return fig
        st.plotly_chart(fig, use_container_width=True)  
       #fig.show()
        #logging.info("Embeddings plotted with TF-IDF colors")

    def plot_embeddings3(self, title=None):
        assert self.reduced_embeddings is not None, "Dimensionality reduction has not been performed yet."
        assert self.cluster_terms is not None, "Cluster TF-IDF analysis has not been performed yet."

        logging.info("Plotting embeddings with TF-IDF colors")

        fig = go.Figure()

        unique_cluster_terms = np.unique(self.cluster_terms)

        terms_order = {term: i for i, term in enumerate(np.unique(self.cluster_terms, return_index=True)[0])}
        #indices = np.argsort([terms_order[term] for term in self.cluster_terms])

        # Handling color assignment, especially for noise
        colors = {term: ('grey' if term == 'Noise' else None) for term in unique_cluster_terms}
        color_map = px.colors.qualitative.Plotly  # Default color map from Plotly Express for consistency

        # Apply a custom color map, handling 'Noise' specifically
        color_idx = 0
        for cluster_term in unique_cluster_terms:
            indices = np.where(np.array(self.cluster_terms) == cluster_term)[0]
            if cluster_term != 'Noise':
                marker_color = color_map[color_idx % len(color_map)]
                color_idx += 1
            else:
                marker_color = 'grey'

            fig.add_trace(
                go.Scatter(
                    x=self.reduced_embeddings[indices, 0],
                    y=self.reduced_embeddings[indices, 1],
                    mode='markers',
                    marker=dict(
                        size=5,
                        opacity=(0.5 if cluster_term == 'Noise' else 0.8),
                        color=marker_color
                    ),
                    name=cluster_term,
                    text=self.data[f'{self.column}'].iloc[indices],
                    hoverinfo='text'
                )
            )
        fig.data = sorted(fig.data, key=lambda trace: terms_order[trace.name])
        fig.update_layout(title=title if title else "Embeddings Visualized", showlegend=True, legend_title_text='Top TF-IDF Terms')
        st.plotly_chart(fig, use_container_width=True)


    def plot_embeddings(self, title=None):
        """
        Plots the reduced dimensionality embeddings with clusters indicated.
        
        Args:
            title (str): The title of the plot.
        """
        # Ensure dimensionality reduction and TF-IDF based cluster naming have been performed
        assert self.reduced_embeddings is not None, "Dimensionality reduction has not been performed yet."
        assert self.cluster_terms is not None, "Cluster TF-IDF analysis has not been performed yet."

        logging.info("Plotting embeddings with TF-IDF colors")
        
        fig = go.Figure()
        
        #for i, term in enumerate(self.cluster_terms):
            # Indices of points in the current cluster
        #unique_cluster_ids = np.unique(self.cluster_labels[self.cluster_labels != -1])  # Exclude noise
        unique_cluster_terms = np.unique(self.cluster_terms)
        unique_cluster_labels = np.unique(self.cluster_labels)
            
        for i, (cluster_id, cluster_terms) in enumerate(zip(unique_cluster_labels, unique_cluster_terms)):
            indices = np.where(self.cluster_labels == cluster_id)[0]
            #indices = np.where(self.cluster_labels == i)[0]
            
            # Plot points in the current cluster
            fig.add_trace(
                go.Scatter(
                    x=self.reduced_embeddings[indices, 0],
                    y=self.reduced_embeddings[indices, 1],
                    mode='markers',
                    marker=dict(
                        #color=i,
                        #colorscale='rainbow',
                        size=5,
                        opacity=0.8,
                    ),
                    name=cluster_terms,
                    text=self.data[f'{self.column}'].iloc[indices],
                    hoverinfo='text',
                )
            )
            
        
        fig.update_layout(title=title, showlegend=True, legend_title_text='Top TF-IDF Terms')
        st.plotly_chart(fig, use_container_width=True)
        logging.info("Embeddings plotted with TF-IDF colors")

    def plot_embeddings4(self, title=None, cluster_terms=None, cluster_labels=None, reduced_embeddings=None, column=None, data=None):
        """
        Plots the reduced dimensionality embeddings with clusters indicated.
        
        Args:
            title (str): The title of the plot.
        """
        # Ensure dimensionality reduction and TF-IDF based cluster naming have been performed
        assert reduced_embeddings is not None, "Dimensionality reduction has not been performed yet."
        assert cluster_terms is not None, "Cluster TF-IDF analysis has not been performed yet."

        logging.info("Plotting embeddings with TF-IDF colors")
        
        fig = go.Figure()
        
        # Determine unique cluster IDs and terms, and ensure consistent color mapping
        unique_cluster_ids = np.unique(cluster_labels)
        unique_cluster_terms = [cluster_terms[i] for i in unique_cluster_ids]#if i != -1]  # Exclude noise by ID

        color_map = px.colors.qualitative.Plotly  # Using Plotly Express's qualitative colors for consistency
        color_idx = 0
        
        # Map each cluster ID to a color
        cluster_colors = {}
        for cid in unique_cluster_ids:
            #if cid != -1:  # Exclude noise
                cluster_colors[cid] = color_map[color_idx % len(color_map)]
                color_idx += 1
            #else:
            #    cluster_colors[cid] = 'grey'  # Noise or outliers in grey

        for cluster_id, cluster_term in zip(unique_cluster_ids, unique_cluster_terms):
            indices = np.where(cluster_labels == cluster_id)[0]
            fig.add_trace(
                go.Scatter(
                    x=reduced_embeddings[indices, 0],
                    y=reduced_embeddings[indices, 1],
                    mode='markers',
                    marker=dict(
                        color=cluster_colors[cluster_id],
                        size=5,
                        opacity=0.8#if cluster_id != -1 else 0.5,
                    ),
                    name=cluster_term,
                    text=data[f'{column}'].iloc[indices],
                    hoverinfo='text',
                )
            )
            
        fig.update_layout(
            title=title if title else "Embeddings Visualized",
            showlegend=True,
            legend_title_text='Top TF-IDF Terms',
            legend=dict(
                traceorder='normal',  # 'normal' or 'reversed'; ensures that traces appear in the order they are added
                itemsizing='constant'
            )
        )
        st.plotly_chart(fig, use_container_width=True)
        logging.info("Embeddings plotted with TF-IDF colors")



def analyze_and_predict(data, analyzers, col_names):
    """
    Performs analysis on the data using provided analyzers and makes predictions on specified columns.

    Args:
        data (pd.DataFrame): The dataset for analysis.
        analyzers (list): A list of UAPAnalyzer instances.
        col_names (list): Column names to be analyzed and predicted.
    """
    new_data = pd.DataFrame()
    for i, (column, analyzer) in enumerate(zip(col_names, analyzers)):
        new_data[f'Analyzer_{column}'] = analyzer.__dict__['cluster_terms']
        logging.info(f"Cluster terms extracted for {column}")

    new_data = new_data.fillna('null').astype('category')
    data_nums = new_data.apply(lambda x: x.cat.codes)

    for col in data_nums.columns:
        try:
            categories = new_data[col].cat.categories
            x_train, x_test, y_train, y_test = train_test_split(data_nums.drop(columns=[col]), data_nums[col], test_size=0.2, random_state=42)
            bst, accuracy, preds = train_xgboost(x_train, y_train, x_test, y_test, len(categories))
            plot_results(new_data, bst, x_test, y_test, preds, categories, accuracy, col)
        except Exception as e:
            logging.error(f"Error processing {col}: {e}")
    return new_data

def train_xgboost(x_train, y_train, x_test, y_test, num_classes):
    """
    Trains an XGBoost model and evaluates its performance.

    Uses GPU training (tree_method='hist', device='cuda') when a CUDA GPU is
    available.  Early stopping avoids over-training and reduces wall time.

    Args:
        x_train (pd.DataFrame): Training features.
        y_train (pd.Series): Training labels.
        x_test (pd.DataFrame): Test features.
        y_test (pd.Series): Test labels.
        num_classes (int): The number of unique classes in the target variable.

    Returns:
        bst (Booster): The trained XGBoost model.
        accuracy (float): The accuracy of the model on the test set.
        preds (np.ndarray): Predictions on the test set.
    """
    dtrain = xgb.DMatrix(x_train, label=y_train, enable_categorical=True)
    dtest = xgb.DMatrix(x_test, label=y_test)

    params = {
        'objective': 'multi:softmax',
        'num_class': num_classes,
        'max_depth': 6,
        'eta': 0.3,
        'tree_method': 'hist',
        'device': 'cuda' if _GPU_AVAILABLE else 'cpu',
        'nthread': -1,
    }
    bst = xgb.train(
        dtrain=dtrain,
        params=params,
        num_boost_round=100,
        evals=[(dtest, 'eval')],
        early_stopping_rounds=10,
        verbose_eval=False,
    )
    preds = bst.predict(dtest)
    accuracy = accuracy_score(y_test, preds)

    logging.info(f"XGBoost trained with accuracy: {accuracy:.2f} ({'GPU' if _GPU_AVAILABLE else 'CPU'})")
    return bst, accuracy, preds

def plot_results(new_data, bst, x_test, y_test, preds, categories, accuracy, col):
    """
    Plots the feature importance, confusion matrix, and contingency table.

    Args:
        bst (Booster): The trained XGBoost model.
        x_test (pd.DataFrame): Test features.
        y_test (pd.Series): Test labels.
        preds (np.array): Predictions made by the model.
        categories (Index): Category names for the target variable.
        accuracy (float): The accuracy of the model on the test set.
        col (str): The target column name being analyzed and predicted.
    """
    fig, axs = plt.subplots(1, 3, figsize=(25, 5), dpi=300)
    fig.suptitle(f'{col.split(sep=".")[-1]} prediction', fontsize=35)

    plot_importance(bst, ax=axs[0], importance_type='gain', show_values=False)
    conf_matrix = confusion_matrix(y_test, preds)
    sns.heatmap(conf_matrix, annot=True, fmt='g', cmap='Blues', xticklabels=categories, yticklabels=categories, ax=axs[1])
    axs[1].set_title(f'Confusion Matrix\nAccuracy: {accuracy * 100:.2f}%')
    # make axes rotated
    axs[1].set_yticklabels(axs[1].get_yticklabels(), rotation=30, ha='right')
    sorted_features = sorted(bst.get_score(importance_type="gain").items(), key=lambda x: x[1], reverse=True)
    # The most important feature is the first element in the sorted list
    most_important_feature = sorted_features[0][0]
    # Create a contingency table
    contingency_table = pd.crosstab(new_data[col], new_data[most_important_feature])

    # resid pearson is used to calculate the residuals, which 
    table = stats.Table(contingency_table).resid_pearson
    #print(table)
    # Perform the chi-squared test
    chi2, p, dof, expected = chi2_contingency(contingency_table)
    # Print the results
    logging.info(f"Chi-squared test for {col} and {most_important_feature}: p-value = {p}")
    
    sns.heatmap(table, annot=True, cmap='Greens', ax=axs[2])
    # make axis rotated
    axs[2].set_yticklabels(axs[2].get_yticklabels(), rotation=30, ha='right')
    axs[2].set_title(f'Contingency Table between {col.split(sep=".")[-1]} and {most_important_feature.split(sep=".")[-1]}\np-value = {p}')    

    plt.tight_layout()
    #plt.savefig(f"{col}_{accuracy:.2f}_prediction_XGB.jpeg", dpi=300)
    return plt

def cramers_v(confusion_matrix):
    """Calculate Cramer's V statistic for categorical-categorical association."""
    try:
        chi2 = chi2_contingency(confusion_matrix)[0]
    except ValueError:
        return 0.0
    n = confusion_matrix.sum().sum()
    if n <= 1:
        return 0.0
    phi2 = chi2 / n
    r, k = confusion_matrix.shape
    phi2corr = max(0, phi2 - ((k-1)*(r-1))/(n-1))
    r_corr = r - ((r-1)**2)/(n-1)
    k_corr = k - ((k-1)**2)/(n-1)
    denom = min((k_corr-1), (r_corr-1))
    if denom <= 0:
        return 0.0
    return np.sqrt(phi2corr / denom)

def plot_cramers_v_heatmap(data, significance_level=0.05):
    """Plot heatmap of Cramer's V statistic for each pair of categorical variables in a DataFrame."""
    # Initialize a DataFrame to store Cramer's V values
    cramers_v_df = pd.DataFrame(index=data.columns, columns=data.columns, data=np.nan)
    
    # Compute Cramer's V for each pair of columns
    for col1 in data.columns:
        for col2 in data.columns:
            if col1 != col2:  # Avoid self-comparison
                confusion_matrix = pd.crosstab(data[col1], data[col2])
                chi2, p, dof, expected = chi2_contingency(confusion_matrix)
                # Check if the p-value is less than the significance level
                #if p < significance_level:
                #    cramers_v_df.at[col1, col2] = cramers_v(confusion_matrix)
                # alternatively, you can use the following line to include all pairs
                cramers_v_df.at[col1, col2] = cramers_v(confusion_matrix)
    
    # Plot the heatmap
    plt.figure(figsize=(12, 10), dpi=200)
    mask = np.triu(np.ones_like(cramers_v_df, dtype=bool))  # Mask for the upper triangle
    # make a max and min of the cmap
    sns.heatmap(cramers_v_df, annot=True, fmt=".2f", cmap='coolwarm', cbar=True, mask=mask, square=True)
    plt.title(f"Heatmap of Cramér's V (p < {significance_level})")
    return plt


class UAPVisualizer:
    def __init__(self, data=None):
        pass  # Initialization can be added if needed

    def analyze_and_predict(self, data, analyzers, col_names):
        new_data = pd.DataFrame()
        for i, (column, analyzer) in enumerate(zip(col_names, analyzers)):
            new_data[f'Analyzer_{column}'] = analyzer.__dict__['cluster_terms']
            logging.info(f"Cluster terms extracted for {column}")

        new_data = new_data.fillna('null').astype('category')
        data_nums = new_data.apply(lambda x: x.cat.codes)

        for col in data_nums.columns:
            try:
                categories = new_data[col].cat.categories
                x_train, x_test, y_train, y_test = train_test_split(data_nums.drop(columns=[col]), data_nums[col], test_size=0.2, random_state=42)
                bst, accuracy, preds = self.train_xgboost(x_train, y_train, x_test, y_test, len(categories))
                self.plot_results(new_data, bst, x_test, y_test, preds, categories, accuracy, col)
            except Exception as e:
                logging.error(f"Error processing {col}: {e}")

    def train_xgboost(self, x_train, y_train, x_test, y_test, num_classes):
        dtrain = xgb.DMatrix(x_train, label=y_train, enable_categorical=True)
        dtest = xgb.DMatrix(x_test, label=y_test)

        params = {
            'objective': 'multi:softmax',
            'num_class': num_classes,
            'max_depth': 6,
            'eta': 0.3,
            'tree_method': 'hist',
            'device': 'cuda' if _GPU_AVAILABLE else 'cpu',
            'nthread': -1,
        }
        bst = xgb.train(
            dtrain=dtrain,
            params=params,
            num_boost_round=100,
            evals=[(dtest, 'eval')],
            early_stopping_rounds=10,
            verbose_eval=False,
        )
        preds = bst.predict(dtest)
        accuracy = accuracy_score(y_test, preds)

        logging.info(f"XGBoost trained with accuracy: {accuracy:.2f} ({'GPU' if _GPU_AVAILABLE else 'CPU'})")
        return bst, accuracy, preds

    def plot_results(self, new_data, bst, x_test, y_test, preds, categories, accuracy, col):
        fig, axs = plt.subplots(1, 3, figsize=(25, 5))
        fig.suptitle(f'{col.split(sep=".")[-1]} prediction', fontsize=35)

        plot_importance(bst, ax=axs[0], importance_type='gain', show_values=False)
        conf_matrix = confusion_matrix(y_test, preds)
        sns.heatmap(conf_matrix, annot=True, fmt='g', cmap='Blues', xticklabels=categories, yticklabels=categories, ax=axs[1])
        axs[1].set_title(f'Confusion Matrix\nAccuracy: {accuracy * 100:.2f}%')

        sorted_features = sorted(bst.get_score(importance_type="gain").items(), key=lambda x: x[1], reverse=True)
        most_important_feature = sorted_features[0][0]
        contingency_table = pd.crosstab(new_data[col], new_data[most_important_feature])
        chi2, p, dof, expected = chi2_contingency(contingency_table)
        logging.info(f"Chi-squared test for {col} and {most_important_feature}: p-value = {p}")

        sns.heatmap(contingency_table, annot=True, cmap='Greens', ax=axs[2])
        axs[2].set_title(f'Contingency Table between {col.split(sep=".")[-1]} and {most_important_feature.split(sep=".")[-1]}\np-value = {p}')    

        plt.tight_layout()
        plt.savefig(f"{col}_{accuracy:.2f}_prediction_XGB.jpeg", dpi=300)
        plt.show()

    @staticmethod
    def cramers_v(confusion_matrix):
        try:
            chi2 = chi2_contingency(confusion_matrix)[0]
        except ValueError:
            return 0.0
        n = confusion_matrix.sum().sum()
        if n <= 1:
            return 0.0
        phi2 = chi2 / n
        r, k = confusion_matrix.shape
        phi2corr = max(0, phi2 - ((k-1)*(r-1))/(n-1))
        r_corr = r - ((r-1)**2)/(n-1)
        k_corr = k - ((k-1)**2)/(n-1)
        denom = min((k_corr-1), (r_corr-1))
        if denom <= 0:
            return 0.0
        return np.sqrt(phi2corr / denom)

    def plot_cramers_v_heatmap(self, data, significance_level=0.05):
        cramers_v_df = pd.DataFrame(index=data.columns, columns=data.columns, data=np.nan)
        
        for col1 in data.columns:
            for col2 in data.columns:
                if col1 != col2:
                    confusion_matrix = pd.crosstab(data[col1], data[col2])
                    chi2, p, dof, expected = chi2_contingency(confusion_matrix)
                    if p < significance_level:
                        cramers_v_df.at[col1, col2] = UAPVisualizer.cramers_v(confusion_matrix)
        
        plt.figure(figsize=(10, 8)),# facecolor="black")
        mask = np.triu(np.ones_like(cramers_v_df, dtype=bool))
        #sns.set_theme(style="dark", rc={"axes.facecolor": "black", "grid.color": "white", "xtick.color": "white", "ytick.color": "white", "axes.labelcolor": "white", "axes.titlecolor": "white"})
        # ax = sns.heatmap(cramers_v_df, annot=True, fmt=".1f", linewidths=.5, linecolor='white', cmap='coolwarm', annot_kws={"color":"white"}, cbar=True, mask=mask, square=True)
        # Customizing the color of the ticks and labels to white
        # plt.xticks(color='white')
        # plt.yticks(color='white')
        sns.heatmap(cramers_v_df, annot=True, fmt=".2f", cmap='coolwarm', cbar=True, mask=mask, square=True)
        plt.title(f"Heatmap of Cramér's V (p < {significance_level})")
        plt.show()

    
    def plot_treemap(self, df, column, top_n=32):
        # Get the value counts and the top N labels
        value_counts = df[column].value_counts()
        top_labels = value_counts.iloc[:top_n].index
        

        # Use np.where to replace all values not in the top N with 'Other'
        revised_column = f'{column}_revised'
        df[revised_column] = np.where(df[column].isin(top_labels), df[column], 'Other')

        # Get the value counts including the 'Other' category
        sizes = df[revised_column].value_counts().values
        labels = df[revised_column].value_counts().index

        # Get a gradient of colors
        colors = list(mcolors.TABLEAU_COLORS.values())

        # Get % of each category
        percents = sizes / sizes.sum()

        # Prepare labels with percentages
        labels = [f'{label}\n {percent:.1%}' for label, percent in zip(labels, percents)]

        # Plot the treemap
        squarify.plot(sizes=sizes, label=labels, alpha=0.7, pad=True, color=colors, text_kwargs={'fontsize': 10})

        ax = plt.gca()

        # Iterate over text elements and rectangles (patches) in the axes for color adjustment
        for text, rect in zip(ax.texts, ax.patches):
            background_color = rect.get_facecolor()
            r, g, b, _ = mcolors.to_rgba(background_color)
            brightness = np.average([r, g, b])
            text.set_color('white' if brightness < 0.5 else 'black')

            # Adjust font size based on rectangle's area and wrap long text
            coef = 0.8
            font_size = np.sqrt(rect.get_width() * rect.get_height()) * coef
            text.set_fontsize(font_size)
            wrapped_text = textwrap.fill(text.get_text(), width=20)
            text.set_text(wrapped_text)

        plt.axis('off')
        plt.gca().invert_yaxis()
        plt.gcf().set_size_inches(20, 12)
        plt.show()



# ── API pricing (May 2026) ────────────────────────────────────────────────────
# key → (api_model_id, provider, input_$/1M, cached_$/1M, output_$/1M, batch_discount_fraction)
_PRICING: dict[str, tuple] = {
    # (name, provider, price_in/M, price_cached_in/M, price_out/M, batch_discount)
    # OpenAI
    "gpt-5.4-mini":       ("gpt-5.4-mini",       "openai",    0.400, 0.1000, 1.600, 0.50),  # pricing placeholder — verify at platform.openai.com/docs/pricing
    "gpt-4o-mini":        ("gpt-4o-mini",        "openai",    0.150, 0.0750, 0.600, 0.50),
    "gpt-4o":             ("gpt-4o",             "openai",    2.500, 1.2500, 10.00, 0.50),
    "gpt-4.1-mini":       ("gpt-4.1-mini",       "openai",    0.400, 0.1000, 1.600, 0.50),
    "gpt-4.1":            ("gpt-4.1",            "openai",    2.000, 0.5000, 8.000, 0.50),
    "gpt-3.5-turbo-0125": ("gpt-3.5-turbo-0125", "openai",    0.500, 0.2500, 1.500, 0.50),
    # DeepSeek — v4-flash and v4-pro are the current generation
    # deepseek-chat / deepseek-reasoner are deprecated 2026-07-24
    "deepseek-v4-flash":  ("deepseek-v4-flash",  "deepseek",  0.140, 0.0028, 0.280, 0.00),
    "deepseek-v4-pro":    ("deepseek-v4-pro",    "deepseek",  0.435, 0.0036, 0.870, 0.00),
    "deepseek-chat":      ("deepseek-chat",      "deepseek",  0.270, 0.0700, 1.100, 0.00),
    "deepseek-reasoner":  ("deepseek-reasoner",  "deepseek",  0.550, 0.1400, 2.190, 0.00),
}

OPENAI_MODELS    = [k for k, v in _PRICING.items() if v[1] == "openai"]
DEEPSEEK_MODELS  = [k for k, v in _PRICING.items() if v[1] == "deepseek"]
_DEEPSEEK_BASE   = "https://api.deepseek.com"
_BATCH_SENTINELS = {"__batch_ids__", "__batch_descs__"}


def _extract_json(text: str) -> dict | None:
    """
    Robustly extract a JSON object from a model response.

    Handles the common patterns that DeepSeek (and other models) produce even
    when ``response_format={"type": "json_object"}`` is requested:
      - Pure JSON (ideal)
      - Markdown fences: ```json ... ``` or ``` ... ```
      - JSON embedded in prose (find first { … last })
      - Single-quoted keys/values
    Returns None if all strategies fail.
    """
    if not isinstance(text, str):
        return None

    import re as _re

    def _try(s: str) -> dict | None:
        s = s.strip()
        try:
            return json.loads(s)
        except json.JSONDecodeError:
            pass
        # single-quote fallback
        try:
            return json.loads(s.replace("'", '"'))
        except (json.JSONDecodeError, Exception):
            pass
        return None

    # 1. Direct parse
    result = _try(text)
    if result is not None:
        return result

    # 2. Strip markdown fences  ```json\n...\n```  or  ```\n...\n```
    fence = _re.sub(r"^```(?:json)?\s*", "", text.strip(), flags=_re.IGNORECASE)
    fence = _re.sub(r"\s*```$", "", fence.strip())
    result = _try(fence)
    if result is not None:
        return result

    # 3. Extract first { ... last } block (handles prose wrapping)
    start = text.find("{")
    end   = text.rfind("}")
    if start != -1 and end > start:
        result = _try(text[start : end + 1])
        if result is not None:
            return result

    return None


def _count_tokens(text: str, model: str = "gpt-4o-mini") -> int:
    """Token count via tiktoken when available, char/4 otherwise."""
    try:
        import tiktoken
        try:
            enc = tiktoken.encoding_for_model(model)
        except KeyError:
            enc = tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(text))
    except Exception:
        return max(1, len(text) // 4)


def estimate_cost(
    descriptions: list[str],
    format_long,
    model: str = "gpt-4o-mini",
    use_cache: bool = True,
    use_batch: bool = False,
) -> dict:
    """
    Estimate API cost before running process_descriptions().

    The system prompt + format_long template is identical for every request
    (the cacheable prefix). With prompt caching, only the first request pays
    the full input price; all subsequent ones pay the cached rate for that
    prefix — a major saving at scale.

    Parameters
    ----------
    descriptions : list[str]
        The raw texts you plan to process.
    format_long :
        The JSON template (dict or string) passed to process_descriptions().
    model : str
        Key from _PRICING (e.g. "gpt-4o-mini", "deepseek-v4-flash").
    use_cache : bool
        Assume the prefix is cached after the first request.
    use_batch : bool
        Apply OpenAI Batch API 50% discount (OpenAI models only).

    Returns
    -------
    dict with full cost breakdown and token counts.
    """
    if model not in _PRICING:
        raise ValueError(f"Unknown model {model!r}. Choose from: {sorted(_PRICING)}")

    _, provider, p_in, p_cached, p_out, batch_disc = _PRICING[model]
    fmt_str = format_long if isinstance(format_long, str) else json.dumps(format_long)

    system = (
        "You are a structured-data extraction assistant. "
        "Extract UAP sighting information from the provided report "
        "and return it as JSON matching the given template exactly. "
        "Leave fields empty string or null when the information is absent. "
        "IMPORTANT — engagement_type constraint: exactly ONE field in "
        "engagement_type must be set to \"P\" (the primary engagement type). "
        "Any number of additional fields may be \"S\" (secondary). "
        "All remaining fields must be blank. "
        "Never assign \"P\" to more than one field. "
        "Never assign \"S\" to any field without also assigning \"P\" to exactly one field. "
        "If no engagement type applies, set no_engagement to \"P\" and leave every other field blank."
    )
    prefix_text = (
        f"{system}\n\n"
        f"Parse into this JSON structure (leave missing fields empty): {fmt_str}\n\nOutput:"
    )
    prefix_tokens = _count_tokens(prefix_text, model)
    desc_counts   = [_count_tokens(f"Input report: {d}\n\n", model) for d in descriptions]
    output_per_q  = _count_tokens(fmt_str, model)

    n = len(descriptions)
    total_input  = sum(prefix_tokens + t for t in desc_counts)
    total_output = n * output_per_q

    if use_cache and n > 1:
        cached_tokens   = prefix_tokens * (n - 1)
        uncached_tokens = total_input - cached_tokens
    else:
        cached_tokens   = 0
        uncached_tokens = total_input

    disc = batch_disc if use_batch and provider == "openai" else 0.0

    cost_in     = uncached_tokens / 1_000_000 * p_in     * (1 - disc)
    cost_cached = cached_tokens   / 1_000_000 * p_cached * (1 - disc)
    cost_out    = total_output    / 1_000_000 * p_out    * (1 - disc)
    total_usd   = cost_in + cost_cached + cost_out

    return {
        "total_usd":              round(total_usd, 4),
        "cost_input_usd":         round(cost_in, 4),
        "cost_cached_usd":        round(cost_cached, 4),
        "cost_output_usd":        round(cost_out, 4),
        "total_input_tokens":     total_input,
        "cached_tokens":          cached_tokens,
        "uncached_tokens":        uncached_tokens,
        "output_tokens":          total_output,
        "prefix_tokens":          prefix_tokens,
        "avg_desc_tokens":        int(sum(desc_counts) / max(1, n)),
        "output_tokens_per_query": output_per_q,
        "batch_discount_pct":     int(disc * 100),
        "model":                  model,
        "provider":               provider,
        "n_queries":              n,
    }


class UAPParser:
    """
    Parse raw UAP report texts into structured JSON via OpenAI or DeepSeek.

    Providers
    ---------
    ``provider="openai"``   — concurrent requests via OpenAI Chat API.
                              Set ``use_batch=True`` for the Batch API (50% off,
                              async ≤24 h; batch_ids saved to checkpoint).
    ``provider="deepseek"`` — concurrent requests via DeepSeek (OpenAI-compatible).
                              No batch API; uses dynamic rate-limit back-off.

    Checkpointing
    -------------
    Pass ``checkpoint_path`` to persist every response immediately. On restart
    with the same path, completed requests are skipped. OpenAI Batch IDs are
    also stored so jobs can be resumed after a Streamlit rerun.

    Cost estimation
    ---------------
    Call the module-level ``estimate_cost(descriptions, format_long, model=...)``
    before running to get a full cost breakdown.

    Usage
    -----
        # Concurrent (OpenAI or DeepSeek)
        parser = UAPParser(api_key=KEY, model="gpt-4o-mini",
                           checkpoint_path="ckpt.json")
        parser.process_descriptions(texts, FORMAT_LONG)

        # OpenAI Batch API (async, 50% cheaper)
        parser = UAPParser(api_key=KEY, model="gpt-4o-mini",
                           use_batch=True, checkpoint_path="ckpt.json")
        parser.process_descriptions(texts, FORMAT_LONG)   # submits & polls

        # DeepSeek
        parser = UAPParser(api_key=DS_KEY, model="deepseek-v4-flash",
                           checkpoint_path="ckpt.json")
        parser.process_descriptions(texts, FORMAT_LONG)
    """

    _BATCH_CHUNK = 50_000   # OpenAI Batch API hard limit per job

    def __init__(
        self,
        api_key: str,
        model: str = "gpt-4o-mini",
        provider: str | None = None,
        use_batch: bool = False,
        col=None,
        format_long=None,
        checkpoint_path: str | None = None,
    ):
        import threading as _threading
        # Auto-detect provider from model name
        if provider is None:
            provider = "deepseek" if model in DEEPSEEK_MODELS else "openai"

        self.model    = model
        self.provider = provider
        self.use_batch = use_batch and provider == "openai"
        self.col      = col
        self.checkpoint_path = checkpoint_path
        self.responses: dict[str, str] = {}
        self._responses_lock = _threading.Lock()
        self.last_errors: list[str] = []   # populated by process_descriptions

        if provider == "deepseek":
            self.client = OpenAI(api_key=api_key, base_url=_DEEPSEEK_BASE)
        else:
            os.environ["OPENAI_API_KEY"] = api_key
            self.client = OpenAI()

        if checkpoint_path:
            self._load_checkpoint()

    # ── Checkpoint ────────────────────────────────────────────────────────

    def _load_checkpoint(self) -> None:
        if not self.checkpoint_path:
            return
        try:
            with open(self.checkpoint_path, "r", encoding="utf-8") as f:
                self.responses.update(json.load(f))
            logging.info(f"Checkpoint: {len(self.responses)} entries from {self.checkpoint_path}")
        except FileNotFoundError:
            logging.info(f"No checkpoint at {self.checkpoint_path} — starting fresh")
        except (json.JSONDecodeError, OSError) as e:
            logging.warning(f"Could not read checkpoint: {e}")

    def _save_checkpoint(self) -> None:
        if not self.checkpoint_path:
            return
        try:
            with open(self.checkpoint_path, "w", encoding="utf-8") as f:
                json.dump(self.responses, f, ensure_ascii=False, indent=2)
        except OSError as e:
            logging.warning(f"Could not write checkpoint: {e}")

    def save_checkpoint(self) -> None:
        with self._responses_lock:
            self._save_checkpoint()

    # ── Shared prompt builder ─────────────────────────────────────────────

    def _build_messages(self, description: str, format_long) -> list:
        fmt = format_long if isinstance(format_long, str) else json.dumps(format_long)
        return [
            {"role": "system", "content": (
                "You are a structured-data extraction assistant. "
                "Extract UAP sighting information from the provided report "
                "and return it as JSON matching the given template exactly. "
                "Leave fields empty string or null when the information is absent. "
                "IMPORTANT — engagement_type constraint: exactly ONE field in "
                "engagement_type must be set to \"P\" (the primary engagement type). "
                "Any number of additional fields may be \"S\" (secondary). "
                "All remaining fields must be blank. "
                "Never assign \"P\" to more than one field. "
                "Never assign \"S\" to any field without also assigning \"P\" to exactly one field. "
                "no_engagement may only ever be \"P\" or blank — never \"S\". "
                "Set no_engagement to \"P\" only when no other engagement type applies, and leave every other field blank."
            )},
            {"role": "user", "content": (
                f"Input report: {description}\n\n"
                f"Parse into this JSON structure (leave missing fields empty): "
                f"{fmt}\n\nOutput:"
            )},
        ]

    # ── Single-request fetch (both providers) ─────────────────────────────

    def fetch_response(self, description: str, format_long) -> object | None:
        INITIAL_WAIT, MAX_WAIT, MAX_RETRIES = 5, 600, 10
        wait = INITIAL_WAIT
        last_exc = None
        for _ in range(MAX_RETRIES):
            try:
                return self.client.chat.completions.create(
                    model=self.model,
                    response_format={"type": "json_object"},
                    messages=self._build_messages(description, format_long),
                )
            except HTTPError as e:
                if "TooManyRequests" in str(e):
                    logging.warning(f"Rate limited — waiting {wait}s")
                    time.sleep(wait)
                    wait = min(wait * 2, MAX_WAIT)
                    last_exc = e
                else:
                    raise
            except Exception as e:
                err = str(e)
                if "429" in err or "rate" in err.lower() or "concurren" in err.lower():
                    logging.warning(f"Rate limited — waiting {wait}s")
                    time.sleep(wait)
                    wait = min(wait * 2, MAX_WAIT)
                    last_exc = e
                elif any(code in err for code in ("500", "502", "503", "504")):
                    logging.warning(f"Server error ({err[:80]}) — retrying in {wait}s")
                    time.sleep(wait)
                    wait = min(wait * 2, MAX_WAIT)
                    last_exc = e
                else:
                    raise  # surfaces the real error into process_descriptions.last_errors
        if last_exc:
            raise last_exc
        return None

    # ── OpenAI Batch API ──────────────────────────────────────────────────

    def _submit_batch_openai(self, descriptions: list[str], format_long) -> list[str]:
        """Upload requests in 50k chunks and return list of batch IDs."""
        import io as _io
        batch_ids = []
        for start in range(0, len(descriptions), self._BATCH_CHUNK):
            chunk = descriptions[start : start + self._BATCH_CHUNK]
            lines = [
                json.dumps({
                    "custom_id": str(start + i),
                    "method":    "POST",
                    "url":       "/v1/chat/completions",
                    "body": {
                        "model":           self.model,
                        "response_format": {"type": "json_object"},
                        "messages":        self._build_messages(desc, format_long),
                    },
                })
                for i, desc in enumerate(chunk)
            ]
            content  = "\n".join(lines).encode("utf-8")
            uploaded = self.client.files.create(
                file=(f"batch_{start}.jsonl", _io.BytesIO(content), "application/jsonl"),
                purpose="batch",
            )
            batch = self.client.batches.create(
                input_file_id=uploaded.id,
                endpoint="/v1/chat/completions",
                completion_window="24h",
            )
            batch_ids.append(batch.id)
            logging.info(
                f"Submitted batch {batch.id} for chunk "
                f"[{start}:{start + len(chunk)}] ({len(chunk)} requests)"
            )
        return batch_ids

    def _poll_and_collect(self, batch_ids: list[str], descriptions: list[str]) -> None:
        """Poll all batch jobs until done, then store results in self.responses."""
        remaining = list(batch_ids)
        while remaining:
            still_pending = []
            for bid in remaining:
                batch  = self.client.batches.retrieve(bid)
                status = batch.status
                c      = batch.request_counts
                logging.info(
                    f"Batch {bid}: {status} "
                    f"({c.completed}/{c.total} done, {c.failed} failed)"
                )
                if status == "completed":
                    text = self.client.files.content(batch.output_file_id).text
                    n_saved = 0
                    for line in text.splitlines():
                        if not line.strip():
                            continue
                        result  = json.loads(line)
                        idx     = int(result["custom_id"])
                        desc    = descriptions[idx]
                        choices = (result.get("response") or {}).get("body", {}).get("choices", [])
                        if choices:
                            with self._responses_lock:
                                self.responses[desc] = choices[0]["message"]["content"]
                            n_saved += 1
                    self._save_checkpoint()
                    logging.info(f"Batch {bid} collected — {n_saved} results saved.")
                elif status in ("failed", "expired", "cancelled"):
                    logging.error(f"Batch {bid} ended with status: {status}")
                else:
                    still_pending.append(bid)
            remaining = still_pending
            if remaining:
                logging.info(f"{len(remaining)} batch(es) pending — waiting 60s")
                time.sleep(60)

    def _process_batch_openai(self, pending: list[str], format_long) -> None:
        """Full Batch API flow with checkpoint persistence for resume after restart."""
        raw_ids   = self.responses.pop("__batch_ids__",   None)
        raw_descs = self.responses.pop("__batch_descs__", None)

        if raw_ids and raw_descs:
            logging.info("Resuming existing batch job(s) from checkpoint.")
            batch_ids  = json.loads(raw_ids)
            all_descs  = json.loads(raw_descs)
        else:
            batch_ids = self._submit_batch_openai(pending, format_long)
            all_descs = pending
            with self._responses_lock:
                self.responses["__batch_ids__"]   = json.dumps(batch_ids)
                self.responses["__batch_descs__"] = json.dumps(all_descs)
                self._save_checkpoint()

        self._poll_and_collect(batch_ids, all_descs)

        with self._responses_lock:
            self.responses.pop("__batch_ids__",   None)
            self.responses.pop("__batch_descs__", None)
            self._save_checkpoint()

    def batch_status(self) -> list[dict] | None:
        """Non-blocking check of in-progress batch jobs. Returns None if none active."""
        raw = self.responses.get("__batch_ids__")
        if not raw:
            return None
        result = []
        for bid in json.loads(raw):
            try:
                b = self.client.batches.retrieve(bid)
                result.append({
                    "batch_id":  bid,
                    "status":    b.status,
                    "completed": b.request_counts.completed,
                    "total":     b.request_counts.total,
                    "failed":    b.request_counts.failed,
                })
            except Exception as e:
                result.append({"batch_id": bid, "status": "error", "error": str(e)})
        return result

    # ── Non-blocking Batch API (Streamlit-safe) ───────────────────────────

    def submit_batch_only(self, descriptions: list[str], format_long) -> str:
        """
        Submit to OpenAI Batch API and return immediately without polling.

        Batch IDs and the original description list are written to the checkpoint
        file so any future session can resume via ``fetch_batch_results()``.

        Safe to call from Streamlit — does not block or sleep.
        Returns the checkpoint path (create one if none was given).
        """
        if not self.checkpoint_path:
            import uuid as _uuid
            self.checkpoint_path = f"checkpoints/batch_{_uuid.uuid4().hex[:12]}.json"
            os.makedirs("checkpoints", exist_ok=True)

        pending = [d for d in descriptions
                   if d not in self.responses and d not in _BATCH_SENTINELS]
        if not pending:
            logging.info("All descriptions already processed — nothing to submit")
            return self.checkpoint_path

        if "__batch_ids__" in self.responses:
            logging.info("Batch already submitted — checkpoint has IDs, skipping resubmission")
            return self.checkpoint_path

        batch_ids = self._submit_batch_openai(pending, format_long)
        with self._responses_lock:
            self.responses["__batch_ids__"]   = json.dumps(batch_ids)
            self.responses["__batch_descs__"] = json.dumps(pending)
            self._save_checkpoint()

        logging.info(
            f"Batch submitted: {len(batch_ids)} job(s) covering {len(pending)} requests. "
            f"Checkpoint: {self.checkpoint_path}"
        )
        return self.checkpoint_path

    def fetch_batch_results(self) -> dict:
        """
        Non-blocking single poll of all pending batch jobs.

        Call this on reconnect (or via a Streamlit button) to check progress and
        collect finished results. Does NOT sleep — returns immediately with the
        current state of every batch job.

        Returns a dict:
            status   : "complete" | "in_progress" | "no_batch" | "error"
            batches  : list of per-job status dicts
            n_saved  : total new responses written to self.responses this call
            checkpoint: path to the checkpoint file (for display in the UI)

        When ``status == "complete"`` the sentinel keys are cleared from the
        checkpoint so a subsequent call correctly reports ``no_batch``.
        """
        raw_ids   = self.responses.get("__batch_ids__")
        raw_descs = self.responses.get("__batch_descs__")

        if not raw_ids:
            return {"status": "no_batch",
                    "message": "No active batch found in checkpoint.",
                    "checkpoint": self.checkpoint_path}

        batch_ids = json.loads(raw_ids)
        all_descs = json.loads(raw_descs) if raw_descs else []

        statuses: list[dict] = []
        n_saved_total = 0
        all_terminal  = True

        for bid in batch_ids:
            try:
                b  = self.client.batches.retrieve(bid)
                st = b.status
                c  = b.request_counts
                info: dict = {
                    "batch_id":  bid,
                    "status":    st,
                    "completed": c.completed,
                    "total":     c.total,
                    "failed":    c.failed,
                }
                if st == "completed":
                    text    = self.client.files.content(b.output_file_id).text
                    n_saved = 0
                    for line in text.splitlines():
                        if not line.strip():
                            continue
                        result  = json.loads(line)
                        idx     = int(result["custom_id"])
                        if idx >= len(all_descs):
                            continue
                        desc    = all_descs[idx]
                        choices = (result.get("response") or {}).get("body", {}).get("choices", [])
                        if choices:
                            with self._responses_lock:
                                self.responses[desc] = choices[0]["message"]["content"]
                            n_saved += 1
                    self._save_checkpoint()
                    info["results_saved"] = n_saved
                    n_saved_total += n_saved
                elif st in ("failed", "expired", "cancelled"):
                    info["error"] = f"Batch ended with non-success status: {st}"
                else:
                    all_terminal = False  # validating / in_progress / finalizing
                statuses.append(info)
            except Exception as e:
                statuses.append({"batch_id": bid, "status": "error", "error": str(e)})

        if all_terminal:
            with self._responses_lock:
                self.responses.pop("__batch_ids__",   None)
                self.responses.pop("__batch_descs__", None)
                self._save_checkpoint()

        return {
            "status":     "complete" if all_terminal else "in_progress",
            "batches":    statuses,
            "n_saved":    n_saved_total,
            "checkpoint": self.checkpoint_path,
        }

    # ── Main entry point ──────────────────────────────────────────────────

    def process_descriptions(
        self, descriptions, format_long, max_workers: int = 10,
        progress_callback=None,
    ) -> None:
        """
        Parse all descriptions not already in self.responses.

        Routes to:
        - OpenAI Batch API  if self.use_batch=True  (async, ≤24 h, blocking poll)
        - Concurrent single requests otherwise       (both providers)

        For non-blocking Batch API use (Streamlit), call ``submit_batch_only``
        then ``fetch_batch_results`` from a separate button/page rerun.
        """
        pending = [
            d for d in descriptions
            if d not in self.responses and d not in _BATCH_SENTINELS
        ]
        already_done = len(descriptions) - len(pending)
        if already_done:
            logging.info(f"Skipping {already_done} already-processed descriptions")
        if not pending:
            logging.info("All descriptions already processed — nothing to do")
            return

        if self.use_batch:
            self._process_batch_openai(pending, format_long)
            return

        self.last_errors.clear()
        n_ok = n_fail = 0

        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_desc = {
                executor.submit(self.fetch_response, desc, format_long): desc
                for desc in pending
            }
            for future in stqdm(
                concurrent.futures.as_completed(future_to_desc),
                total=len(pending),
                desc=f"Parsing [{self.provider}/{self.model}]",
            ):
                desc = future_to_desc[future]
                try:
                    response = future.result()
                    text = response.choices[0].message.content if response else None
                    if text:
                        with self._responses_lock:
                            self.responses[desc] = text
                            self._save_checkpoint()
                        n_ok += 1
                    else:
                        n_fail += 1
                        msg = f"No response returned for: {str(desc)[:80]!r}"
                        self.last_errors.append(msg)
                        logging.warning(msg)
                except Exception as exc:
                    n_fail += 1
                    self.last_errors.append(str(exc))
                    logging.error(f"Error processing description: {exc}")
                if progress_callback is not None:
                    progress_callback(n_ok, n_fail, len(pending))

        logging.info(
            f"process_descriptions done — ok={n_ok} fail={n_fail} "
            f"model={self.model} provider={self.provider}"
        )

    # ── Post-processing ───────────────────────────────────────────────────

    def parse_responses(self) -> dict:
        parsed, failed = {}, 0
        for k, v in self.responses.items():
            if k in _BATCH_SENTINELS:
                continue
            result = _extract_json(v)
            if result is not None:
                parsed[k] = result
            else:
                failed += 1
                logging.warning(f"Could not parse response for key (first 120 chars): {str(v)[:120]!r}")
        logging.info(f"Parsed: {len(parsed)}  |  Failed: {failed}")
        return parsed

    def responses_to_df(self, col: str, parsed_responses: dict) -> pd.DataFrame:
        """Expand one top-level JSON group into a flat DataFrame."""
        rows = {k: v[col] for k, v in parsed_responses.items() if col in v}
        df = pd.json_normalize(list(rows.values()))
        df.index = list(rows.keys())
        return df