pattern_discovery.py

from typing import List, Dict, Any, Tuple, Optional
from collections import Counter
import numpy as np
import re

class PatternDiscoveryAgent:
    """
    Agent for discovering hidden patterns and correlations in historical documents
    """

    def __init__(self, rag_system):
        self.rag_system = rag_system
        self.pattern_cache = {}
        self.baseline_cache = {}

    def analyze_concept_correlation(
            self,
            concept_a: str,
            concept_b: str,
            date_range: Tuple[int, int],
            sample_size: int = 500,
            confidence_level: float = 0.95,
            effect_size_threshold: float = 0.2
    ) -> Dict[str, Any]:
        """
        Enhanced correlation analysis with proper statistical tests
        """
        start_year, end_year = date_range

        # Get documents containing concept_a
        docs_a = self.rag_system.temporal_analyzer.retrieve_with_temporal_filter(
            query=concept_a,
            start_year=start_year,
            end_year=end_year,
            limit=sample_size
        )

        if not docs_a:
            return {"error": f"No documents found containing '{concept_a}'"}

        # Create binary indicators for co-occurrence
        co_occurrences = []
        for doc in docs_a:
            has_b = self._contains_concept(doc['text'], concept_b)
            co_occurrences.append(1 if has_b else 0)

        # Calculate observed statistics
        n_total = len(docs_a)
        n_cooccur = sum(co_occurrences)
        observed_rate = n_cooccur / n_total

        # Get baseline rate from larger random sample
        baseline_rate = self._calculate_baseline_correlation(
            concept_b, date_range, sample_size * 2
        )

        if n_cooccur == 0:
            return {
                "concept_a": concept_a,
                "concept_b": concept_b,
                "date_range": date_range,
                "error": f"No co-occurrences found between '{concept_a}' and '{concept_b}' in the specified time period",
                "sample_statistics": {
                    "documents_analyzed": n_total,
                    "co_occurrences": 0,
                    "observed_rate": 0.0,
                    "baseline_rate": baseline_rate
                },
                "note": "This is expected for anachronistic combinations (e.g., historical concepts with modern terms)"
            }

        # Chi-square test for association
        from scipy.stats import chi2_contingency, fisher_exact

        # Create contingency table
        expected_cooccur = baseline_rate * n_total
        expected_no_cooccur = n_total - expected_cooccur

        # Also add this check before chi-square test
        if expected_cooccur < 1 or n_total < 5:
            return {
                "concept_a": concept_a,
                "concept_b": concept_b,
                "date_range": date_range,
                "error": "Insufficient data for statistical analysis",
                "sample_statistics": {
                    "documents_analyzed": n_total,
                    "co_occurrences": n_cooccur,
                    "observed_rate": observed_rate,
                    "baseline_rate": baseline_rate
                },
                "note": "Need more documents or higher co-occurrence rate for reliable statistical tests"
            }

        contingency_table = np.array([
            [n_cooccur, n_total - n_cooccur],  # With concept_a
            [expected_cooccur, expected_no_cooccur]  # Expected baseline
        ])

        # Perform statistical tests
        if n_total >= 5 and expected_cooccur >= 5:
            chi2, p_value_chi2, dof, expected = chi2_contingency(contingency_table)
            test_used = "chi-square"
        else:
            # Use Fisher's exact test for small samples
            odds_ratio, p_value_chi2 = fisher_exact(contingency_table)
            chi2 = None
            test_used = "fisher_exact"

        # Calculate effect size (Cramér's V for chi-square)
        if chi2 is not None:
            cramers_v = np.sqrt(chi2 / (n_total * (min(contingency_table.shape) - 1)))
        else:
            cramers_v = None

        # Calculate confidence interval for the difference in proportions
        from statsmodels.stats.proportion import proportions_ztest, proportion_confint

        # Z-test for difference in proportions
        count = np.array([n_cooccur, int(baseline_rate * sample_size * 2)])
        nobs = np.array([n_total, sample_size * 2])

        z_stat, p_value_prop = proportions_ztest(count, nobs)

        # Confidence interval for observed rate
        ci_lower, ci_upper = proportion_confint(n_cooccur, n_total,
                                                alpha=1 - confidence_level,
                                                method='wilson')

        # Effect size interpretation
        if cramers_v is not None:
            if cramers_v < 0.1:
                effect_size = "negligible"
            elif cramers_v < 0.3:
                effect_size = "small"
            elif cramers_v < 0.5:
                effect_size = "medium"
            else:
                effect_size = "large"
        else:
            effect_size = "unknown"

        # Calculate lift with confidence interval
        lift = observed_rate / baseline_rate if baseline_rate > 0 else float('inf')

        # Multiple testing correction (Bonferroni)
        alpha = 1 - confidence_level
        corrected_alpha = alpha / 2  # Two tests performed

        # Determine statistical significance
        is_significant = p_value_chi2 < corrected_alpha

        results = {
            "concept_a": concept_a,
            "concept_b": concept_b,
            "date_range": date_range,
            "sample_statistics": {
                "documents_analyzed": n_total,
                "co_occurrences": n_cooccur,
                "observed_rate": observed_rate,
                "baseline_rate": baseline_rate,
                "lift": lift,
                "confidence_interval": [ci_lower, ci_upper]
            },
            "statistical_tests": {
                "test_used": test_used,
                "chi_square": chi2,
                "p_value": p_value_chi2,
                "z_statistic": z_stat,
                "p_value_proportion_test": p_value_prop,
                "corrected_alpha": corrected_alpha
            },
            "effect_size": {
                "cramers_v": cramers_v,
                "interpretation": effect_size,
                "meets_threshold": cramers_v > effect_size_threshold if cramers_v else False
            },
            "significance": {
                "is_statistically_significant": is_significant,
                "is_practically_significant": cramers_v > effect_size_threshold if cramers_v else False,
                "confidence_level": confidence_level
            },
            "supporting_documents": [doc for doc, cooccur in zip(docs_a, co_occurrences) if cooccur][:5]
        }

        return self._ensure_json_serializable(results)

    def detect_temporal_anomalies(
            self,
            concept: str,
            date_range: Tuple[int, int],
            window_size: int = 10,
            sample_size: int = 100,
            confidence_level: float = 0.95,
            min_documents_per_window: int = 5
    ) -> Dict[str, Any]:
        """
        Enhanced anomaly detection with multiple statistical methods
        """
        from scipy import stats
        from scipy.stats import jarque_bera

        start_year, end_year = date_range

        # Create time windows
        windows = []
        current_year = start_year

        while current_year < end_year:
            window_end = min(current_year + window_size - 1, end_year)
            windows.append((current_year, window_end))
            current_year += window_size

        # Analyze each window
        window_data = []
        frequencies = []
        document_counts = []

        for window_start, window_end in windows:
            docs = self.rag_system.temporal_analyzer.retrieve_with_temporal_filter(
                query=concept,
                start_year=window_start,
                end_year=window_end,
                limit=sample_size
            )

            # Estimate total documents in period (more sophisticated)
            total_docs = self._estimate_total_docs_sophisticated(window_start, window_end)
            concept_frequency = len(docs) / total_docs if total_docs > 0 else 0

            window_info = {
                "period": f"{window_start}-{window_end}",
                "start_year": window_start,
                "end_year": window_end,
                "concept_mentions": len(docs),
                "total_docs_estimated": total_docs,
                "frequency": concept_frequency,
                "top_documents": docs[:3]
            }

            window_data.append(window_info)

            # Only include windows with sufficient data
            if len(docs) >= min_documents_per_window:
                frequencies.append(concept_frequency)
                document_counts.append(len(docs))

        if len(frequencies) < 3:
            return {
                "error": "Insufficient data for anomaly detection",
                "concept": concept,
                "windows_with_data": len(frequencies)
            }

        frequencies = np.array(frequencies)
        document_counts = np.array(document_counts)

        # Statistical analysis of the time series
        freq_mean = np.mean(frequencies)
        freq_std = np.std(frequencies, ddof=1)
        freq_median = np.median(frequencies)
        freq_mad = np.median(np.abs(frequencies - freq_median))  # Median Absolute Deviation

        # Test for normality
        shapiro_stat, shapiro_p = stats.shapiro(frequencies)
        jb_stat, jb_p = jarque_bera(frequencies)

        # Calculate control limits (statistical process control)
        alpha = 1 - confidence_level

        # Method 1: Classical z-score
        z_upper = freq_mean + stats.norm.ppf(1 - alpha / 2) * freq_std
        z_lower = freq_mean - stats.norm.ppf(1 - alpha / 2) * freq_std

        # Method 2: Robust MAD-based detection
        mad_multiplier = stats.norm.ppf(1 - alpha / 2) * 1.4826  # MAD to std conversion
        mad_upper = freq_median + mad_multiplier * freq_mad
        mad_lower = freq_median - mad_multiplier * freq_mad

        # Method 3: Interquartile Range (IQR) method
        q1, q3 = np.percentile(frequencies, [25, 75])
        iqr = q3 - q1
        iqr_multiplier = 1.5  # Standard outlier detection
        iqr_upper = q3 + iqr_multiplier * iqr
        iqr_lower = q1 - iqr_multiplier * iqr

        # Method 4: Modified Z-score (using MAD)
        modified_z_scores = 0.6745 * (frequencies - freq_median) / freq_mad if freq_mad > 0 else np.zeros_like(
            frequencies)
        modified_z_threshold = 3.5

        # Detect anomalies using multiple methods
        anomalies = []
        valid_window_idx = 0

        for i, window in enumerate(window_data):
            if window["concept_mentions"] < min_documents_per_window:
                continue

            freq = frequencies[valid_window_idx]
            mod_z = modified_z_scores[valid_window_idx]

            # Check each method
            methods_triggered = []

            if freq > z_upper or freq < z_lower:
                methods_triggered.append("z_score")

            if freq > mad_upper or freq < mad_lower:
                methods_triggered.append("mad_robust")

            if freq > iqr_upper or freq < iqr_lower:
                methods_triggered.append("iqr")

            if abs(mod_z) > modified_z_threshold:
                methods_triggered.append("modified_z")

            # Require at least 2 methods to agree for anomaly
            if len(methods_triggered) >= 2:
                anomaly_type = "spike" if freq > freq_median else "drop"

                # Calculate severity
                z_score = (freq - freq_mean) / freq_std if freq_std > 0 else 0
                severity_score = abs(z_score)

                if severity_score > 3:
                    severity = "extreme"
                elif severity_score > 2:
                    severity = "strong"
                else:
                    severity = "moderate"

                anomalies.append({
                    "period": window["period"],
                    "type": anomaly_type,
                    "frequency": freq,
                    "z_score": z_score,
                    "modified_z_score": mod_z,
                    "severity": severity,
                    "methods_triggered": methods_triggered,
                    "documents": window["top_documents"],
                    "confidence_level": confidence_level
                })

            valid_window_idx += 1

        # Trend analysis
        if len(frequencies) >= 4:
            # Linear trend test
            x = np.arange(len(frequencies))
            slope, intercept, r_value, p_value_trend, std_err = stats.linregress(x, frequencies)

            # Mann-Kendall trend test (non-parametric)
            mk_trend, mk_p_value = self._mann_kendall_test(frequencies)
        else:
            slope, r_value, p_value_trend = None, None, None
            mk_trend, mk_p_value = None, None

        # Generate enhanced insights
        insights = self._generate_anomaly_insights(
            concept, anomalies, window_data, frequencies, shapiro_p < 0.05
        )

        results = {
            "concept": concept,
            "date_range": date_range,
            "window_size": window_size,
            "statistical_summary": {
                "windows_analyzed": len(window_data),
                "windows_with_sufficient_data": len(frequencies),
                "mean_frequency": freq_mean,
                "std_frequency": freq_std,
                "median_frequency": freq_median,
                "mad_frequency": freq_mad,
                "coefficient_of_variation": freq_std / freq_mean if freq_mean > 0 else None
            },
            "normality_tests": {
                "shapiro_wilk_p": shapiro_p,
                "jarque_bera_p": jb_p,
                "is_normal": shapiro_p > 0.05 and jb_p > 0.05
            },
            "control_limits": {
                "z_score": {"lower": z_lower, "upper": z_upper},
                "mad_robust": {"lower": mad_lower, "upper": mad_upper},
                "iqr": {"lower": iqr_lower, "upper": iqr_upper}
            },
            "trend_analysis": {
                "linear_slope": slope,
                "r_squared": r_value ** 2 if r_value else None,
                "trend_p_value": p_value_trend,
                "mann_kendall_trend": mk_trend,
                "mann_kendall_p": mk_p_value
            },
            "anomalies": anomalies,
            "temporal_data": window_data,
            "insights": insights
        }

        return self._ensure_json_serializable(results)

    def _mann_kendall_test(self, data):
        """Mann-Kendall trend test implementation"""
        from scipy.stats import norm

        n = len(data)
        s = 0

        for i in range(n - 1):
            for j in range(i + 1, n):
                if data[j] > data[i]:
                    s += 1
                elif data[j] < data[i]:
                    s -= 1

        # Calculate variance
        var_s = n * (n - 1) * (2 * n + 5) / 18

        if s > 0:
            z = (s - 1) / np.sqrt(var_s)
        elif s < 0:
            z = (s + 1) / np.sqrt(var_s)
        else:
            z = 0

        p_value = 2 * (1 - norm.cdf(abs(z)))

        # Trend interpretation
        if p_value < 0.05:
            if s > 0:
                trend = "increasing"
            else:
                trend = "decreasing"
        else:
            trend = "no_trend"

        return trend, p_value

    def _estimate_total_docs_sophisticated(self, start_year: int, end_year: int) -> int:
        """Robust document count estimation using larger samples"""

        # First try: Use loaded metadata if we have enough coverage
        if hasattr(self.rag_system, 'metadata') and len(self.rag_system.metadata) > 1000:
            count = 0
            total_with_dates = 0

            for metadata in self.rag_system.metadata.values():
                date_str = metadata.get('date', '')
                if date_str:
                    total_with_dates += 1
                    year_match = re.search(r'\b(1[4-7]\d{2})\b', str(date_str))
                    if year_match:
                        year = int(year_match.group(1))
                        if start_year <= year <= end_year:
                            count += 1

            # If we have good metadata coverage, use it
            if total_with_dates > 500:
                print(f"Using metadata: {count} docs found in {start_year}-{end_year}")
                return max(count, 10)

        # Second try: Sample-based estimation with larger sample
        try:
            # Get multiple batches to increase sample size
            all_samples = []
            batch_queries = ["text", "and", "the", "in", "of"]  # Different starting points

            for i, query in enumerate(batch_queries):
                try:
                    batch = self.rag_system.milvus_client.query(
                        collection_name=self.rag_system.collection_name,
                        filter="",
                        output_fields=["filename", "text"],
                        limit=200,  # Larger batches
                        offset=i * 150  # Different offsets
                    )
                    all_samples.extend(batch)
                except:
                    continue

            if len(all_samples) < 100:
                # Try retrieve method as fallback
                for query in ["a", "to", "with", "for"]:
                    try:
                        docs = self.rag_system.retrieve(query=query, limit=100)
                        all_samples.extend(docs)
                    except:
                        continue

            # Remove duplicates and ensure metadata loaded
            seen = set()
            unique_samples = []
            for doc in all_samples:
                filename = doc.get("filename", "")
                base_filename = filename.split('_chunk_')[0] if '_chunk_' in filename else filename
                if base_filename not in seen:
                    seen.add(base_filename)
                    unique_samples.append(doc)

            print(f"Collected {len(unique_samples)} unique documents for sampling")

            if len(unique_samples) >= 200:  # Good sample size
                # Load metadata for samples
                sample_filenames = [doc["filename"] for doc in unique_samples]
                self.rag_system._ensure_metadata_loaded(sample_filenames)

                # Count documents in date range
                in_range = 0
                total_with_dates = 0

                for doc in unique_samples:
                    base_filename = doc["filename"].split('_chunk_')[0] if '_chunk_' in doc["filename"] else doc[
                        "filename"]
                    metadata = self.rag_system.metadata.get(base_filename, {})
                    date_str = metadata.get('date', '')

                    if date_str:
                        total_with_dates += 1
                        year_match = re.search(r'\b(1[4-7]\d{2})\b', str(date_str))
                        if year_match:
                            year = int(year_match.group(1))
                            if start_year <= year <= end_year:
                                in_range += 1

                if total_with_dates > 50:  # Enough dated documents
                    proportion = in_range / total_with_dates

                    # Estimate total collection size (chunks to documents)
                    try:
                        stats = self.rag_system.milvus_client.get_collection_stats(
                            collection_name=self.rag_system.collection_name
                        )
                        total_chunks = stats.get('row_count', len(unique_samples) * 10)
                        estimated_docs = total_chunks // 3  # Assume ~3 chunks per doc
                    except:
                        estimated_docs = len(unique_samples) * 10  # Conservative estimate

                    final_estimate = int(estimated_docs * proportion)

                    print(f"Sample-based estimate: {in_range}/{total_with_dates} = {proportion:.3f}")
                    print(f"Estimated {estimated_docs} total docs → {final_estimate} in range")

                    return max(final_estimate, in_range)  # At least what we found

        except Exception as e:
            print(f"Sampling failed: {e}")

        # Fallback: Historical averages
        years = end_year - start_year + 1

        if start_year >= 1650:
            base_rate = 45  # Later 17th century - more printing
        elif start_year >= 1600:
            base_rate = 35  # Early 17th century
        elif start_year >= 1550:
            base_rate = 30  # Late 16th century
        elif start_year >= 1500:
            base_rate = 25  # Early 16th century
        elif start_year >= 1450:
            base_rate = 20  # Late 15th century
        else:
            base_rate = 15  # Earlier periods

        fallback = years * base_rate
        print(f"Using fallback estimate: {years} years × {base_rate} = {fallback}")

        return fallback

    def _get_embeddings_for_documents(self, documents: List[Dict]) -> Optional[np.ndarray]:
        """Retrieve dense embeddings from Milvus for documents"""
        try:
            embeddings = []
            filenames = [doc['filename'] for doc in documents]

            # Query Milvus in batches
            batch_size = 100
            for i in range(0, len(filenames), batch_size):
                batch_filenames = filenames[i:i + batch_size]

                # Build filter for batch
                filter_parts = [f"filename == '{fn.replace(chr(39), chr(39) + chr(39))}'" for fn in batch_filenames]
                filter_expr = " or ".join(filter_parts)

                results = self.rag_system.milvus_client.query(
                    collection_name=self.rag_system.collection_name,
                    filter=filter_expr,
                    output_fields=["filename", "dense"],
                    limit=len(batch_filenames)
                )

                # Create mapping for this batch
                filename_to_embedding = {r['filename']: r['dense'] for r in results}

                # Add embeddings in correct order
                for fn in batch_filenames:
                    if fn in filename_to_embedding:
                        embeddings.append(filename_to_embedding[fn])
                    else:
                        print(f"Warning: No embedding found for {fn}")

            if len(embeddings) == 0:
                return None

            return np.array(embeddings)

        except Exception as e:
            print(f"Error retrieving embeddings: {e}")
            return None

    def _reduce_dimensions(self, embeddings: np.ndarray, n_components: int = 50) -> np.ndarray:
        """Reduce dimensionality using UMAP"""
        try:
            import umap

            reducer = umap.UMAP(
                n_components=min(n_components, embeddings.shape[0] - 1),
                n_neighbors=min(15, embeddings.shape[0] - 1),
                min_dist=0.1,
                metric='cosine',
                random_state=42
            )

            reduced = reducer.fit_transform(embeddings)
            print(f"Reduced from {embeddings.shape[1]} to {reduced.shape[1]} dimensions")
            return reduced

        except ImportError:
            print("UMAP not available, falling back to PCA")
            from sklearn.decomposition import PCA

            n_components = min(n_components, embeddings.shape[0] - 1, embeddings.shape[1])
            pca = PCA(n_components=n_components, random_state=42)
            reduced = pca.fit_transform(embeddings)
            print(
                f"PCA: Reduced to {reduced.shape[1]} dims, explained variance: {pca.explained_variance_ratio_.sum():.2%}")
            return reduced

    def _cluster_embeddings(self, embeddings: np.ndarray, documents: List[Dict]) -> Dict[int, List[Dict]]:
        """Cluster embeddings using HDBSCAN or fallback"""
        try:
            import hdbscan

            clusterer = hdbscan.HDBSCAN(
                min_cluster_size=max(5, len(embeddings) // 20),
                min_samples=3,
                metric='euclidean',
                cluster_selection_method='eom'
            )

            labels = clusterer.fit_predict(embeddings)

            # Group documents by cluster
            clusters = {}
            for i, label in enumerate(labels):
                if label == -1:  # Noise
                    continue
                if label not in clusters:
                    clusters[label] = []
                clusters[label].append(documents[i])

            print(f"HDBSCAN found {len(clusters)} clusters, {sum(labels == -1)} noise points")
            return clusters

        except ImportError:
            print("HDBSCAN not available, using K-means fallback")
            from sklearn.cluster import KMeans

            n_clusters = max(3, min(8, len(embeddings) // 30))
            kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
            labels = kmeans.fit_predict(embeddings)

            clusters = {}
            for i, label in enumerate(labels):
                if label not in clusters:
                    clusters[label] = []
                clusters[label].append(documents[i])

            print(f"K-means created {n_clusters} clusters")
            return clusters

    def _analyze_embedding_cluster(self, cluster_docs: List[Dict], cluster_id: int) -> Dict[str, Any]:
        """Analyze cluster using TF-IDF and entity extraction"""
        from sklearn.feature_extraction.text import TfidfVectorizer

        # Extract texts
        texts = [doc['text'][:2000] for doc in cluster_docs]  # Limit text length

        # TF-IDF for topic terms
        try:
            vectorizer = TfidfVectorizer(
                max_features=50,
                stop_words='english',
                ngram_range=(1, 2),
                min_df=2
            )
            tfidf_matrix = vectorizer.fit_transform(texts)
            feature_names = vectorizer.get_feature_names_out()

            # Get top terms
            tfidf_scores = np.asarray(tfidf_matrix.sum(axis=0)).flatten()
            top_indices = tfidf_scores.argsort()[-10:][::-1]
            top_terms = [feature_names[i] for i in top_indices]

        except Exception as e:
            print(f"TF-IDF failed: {e}")
            top_terms = []

        # Extract entities from cluster
        all_entities = {'people': [], 'places': [], 'organizations': [], 'dates': []}

        for doc in cluster_docs[:20]:  # Sample for entity extraction
            entities = self._extract_entities(doc['text'][:1000])
            for entity_type, entity_list in entities.items():
                all_entities[entity_type].extend(entity_list)

        # Count most common entities
        common_entities = {}
        for entity_type, entity_list in all_entities.items():
            if entity_list:
                entity_counts = Counter(entity_list)
                common_entities[entity_type] = [entity for entity, count in entity_counts.most_common(5)]
            else:
                common_entities[entity_type] = []

        # Get date range
        all_dates = []
        for doc in cluster_docs:
            metadata = doc.get('metadata', {})
            if metadata.get('date'):
                try:
                    year = int(str(metadata['date'])[:4])
                    all_dates.append(year)
                except:
                    pass

        date_range = f"{min(all_dates)}-{max(all_dates)}" if all_dates else "Unknown"

        # Generate cluster label from top terms and entities
        cluster_label = self._generate_cluster_label(top_terms, common_entities)

        # Representative documents
        representative_docs = []
        for doc in cluster_docs[:5]:
            rep_doc = {
                'filename': doc['filename'],
                'text_preview': doc['text'][:300] + "..." if len(doc['text']) > 300 else doc['text'],
                'metadata': doc.get('metadata', {}),
                'score': doc.get('score', 0)
            }
            representative_docs.append(rep_doc)

        return {
            'cluster_id': cluster_id,
            'size': len(cluster_docs),
            'label': cluster_label,
            'top_terms': top_terms[:5],
            'common_entities': common_entities,
            'date_range': date_range,
            'representative_documents': representative_docs,
            'cluster_summary': self._generate_cluster_summary_embedding(cluster_label, common_entities,
                                                                        len(cluster_docs))
        }

    def _generate_cluster_label(self, top_terms: List[str], entities: Dict) -> str:
        """Generate a descriptive label from terms and entities"""
        # Combine top terms and key entities
        label_parts = []

        if top_terms:
            # Take top 2-3 terms
            label_parts.extend(top_terms[:3])

        if entities.get('people') and len(entities['people']) > 0:
            label_parts.append(entities['people'][0])

        if entities.get('places') and len(entities['places']) > 0:
            label_parts.append(entities['places'][0])

        # Create readable label
        if not label_parts:
            return "Miscellaneous Documents"

        return " | ".join(label_parts[:3])

    def _generate_cluster_summary_embedding(self, label: str, entities: Dict, doc_count: int) -> str:
        """Generate summary for embedding-based cluster"""
        summary_parts = [f"This cluster of {doc_count} documents focuses on: {label}."]

        if entities.get('people'):
            people_str = ', '.join(entities['people'][:3])
            summary_parts.append(f"Key figures: {people_str}.")

        if entities.get('places'):
            places_str = ', '.join(entities['places'][:3])
            summary_parts.append(f"Geographic focus: {places_str}.")

        return ' '.join(summary_parts)

    def _generate_embedding_cluster_insights(self, clusters: List[Dict], total_docs: int) -> List[str]:
        """Generate insights from embedding-based clustering"""
        insights = []

        if not clusters:
            insights.append("No clusters identified from the document sample.")
            return insights

        insights.append(
            f"Identified {len(clusters)} semantic clusters from {total_docs} documents using embedding-based analysis.")

        # Cluster size distribution
        cluster_sizes = [c['size'] for c in clusters]
        avg_size = sum(cluster_sizes) / len(cluster_sizes)
        largest_cluster = max(clusters, key=lambda x: x['size'])

        insights.append(f"Average cluster size: {avg_size:.1f} documents.")
        insights.append(f"Largest cluster: '{largest_cluster['label']}' with {largest_cluster['size']} documents.")

        # Entity analysis
        all_people = set()
        all_places = set()
        for cluster in clusters:
            all_people.update(cluster['common_entities'].get('people', []))
            all_places.update(cluster['common_entities'].get('places', []))

        if all_people:
            insights.append(f"Found references to {len(all_people)} distinct historical figures across clusters.")

        if all_places:
            insights.append(f"Geographic references span {len(all_places)} different locations.")

        # Well-defined clusters
        well_defined = [c for c in clusters if c['size'] >= max(5, total_docs // 50)]
        if well_defined:
            insights.append(
                f"{len(well_defined)} clusters show strong semantic coherence with sufficient document representation.")

        return insights

    def _contains_concept(self, text: str, concept: str) -> bool:
        """Check if text contains concept with word boundaries"""
        concept_lower = concept.lower()
        text_lower = text.lower()

        # Simple word boundary check
        import re
        pattern = r'\b' + re.escape(concept_lower) + r'\b'
        return bool(re.search(pattern, text_lower))

    def _calculate_baseline_correlation(
            self,
            concept: str,
            date_range: Tuple[int, int],
            sample_size: int
    ) -> float:
        """Enhanced baseline calculation with stratified sampling"""
        cache_key = f"{concept}_{date_range}_{sample_size}_enhanced"

        if cache_key in self.baseline_cache:
            return self.baseline_cache[cache_key]

        # Stratified sampling across time periods
        start_year, end_year = date_range
        period_length = max(5, (end_year - start_year) // 5)  # 5 time strata

        all_samples = []
        for period_start in range(start_year, end_year, period_length):
            period_end = min(period_start + period_length - 1, end_year)
            period_samples = self._get_random_sample_docs(
                (period_start, period_end),
                sample_size // 5
            )
            all_samples.extend(period_samples)

        if not all_samples:
            return 0.0

        # Count occurrences
        count = sum(1 for doc in all_samples if self._contains_concept(doc['text'], concept))
        baseline = count / len(all_samples)

        self.baseline_cache[cache_key] = baseline
        return baseline

    def _get_random_sample_docs(self, date_range: Tuple[int, int], sample_size: int) -> List[Dict]:
        """Get truly random sample of documents from time period using pagination"""
        start_year, end_year = date_range

        try:
            # First, get total document count estimate by querying with no filter
            # Use a very broad query to get documents, then filter by date
            all_docs = []

            # Get documents in batches using pagination-like approach
            batch_size = min(100, sample_size * 3)  # Get more than needed
            offset_attempts = 5  # Try different "starting points"

            for attempt in range(offset_attempts):
                try:
                    # Use different broad search terms to get different document sets
                    broad_queries = ["text", "and", "the", "in", "of", "to", "a", "with", "for", "by"]
                    query_term = broad_queries[attempt % len(broad_queries)]

                    # Query with different approaches to get variety
                    if attempt < 3:
                        # Use semantic search with very common terms
                        docs = self.rag_system.retrieve(
                            query=query_term,
                            limit=batch_size
                        )
                    else:
                        # Use direct milvus query for more variety
                        docs = self.rag_system.milvus_client.query(
                            collection_name=self.rag_system.collection_name,
                            filter="",
                            output_fields=["filename", "text"],
                            limit=batch_size,
                            offset=attempt * batch_size  # Different starting point
                        )
                        # Convert to expected format
                        docs = [{"filename": doc["filename"], "text": doc["text"], "score": 1.0}
                                for doc in docs]

                    all_docs.extend(docs)

                except Exception as e:
                    print(f"Batch {attempt} failed: {e}")
                    continue

            if not all_docs:
                print("No documents retrieved, falling back to minimal sample")
                return []

            # Remove duplicates based on filename
            seen_files = set()
            unique_docs = []
            for doc in all_docs:
                base_filename = doc["filename"].split('_chunk_')[0] if '_chunk_' in doc["filename"] else doc["filename"]
                if base_filename not in seen_files:
                    seen_files.add(base_filename)
                    unique_docs.append(doc)

            print(f"Retrieved {len(unique_docs)} unique documents before temporal filtering")

            # Now apply temporal filtering
            filtered_docs = self.rag_system.temporal_analyzer.filter_documents_by_date(
                unique_docs, start_year, end_year
            )

            print(f"After temporal filtering ({start_year}-{end_year}): {len(filtered_docs)} documents")

            if not filtered_docs:
                # If no documents in date range, try without temporal filter
                print("No documents in date range, using all retrieved documents")
                filtered_docs = unique_docs

            # Truly randomize the results
            import random
            random.shuffle(filtered_docs)

            # Return requested sample size
            return filtered_docs[:sample_size]

        except Exception as e:
            print(f"Error in improved random sampling: {e}")
            return []

    def _estimate_total_docs_in_period(self, start_year: int, end_year: int) -> int:
        """Rough estimate of total documents in time period"""
        # This is a placeholder - could be improved with actual corpus statistics
        return max(100, (end_year - start_year) * 50)  # Assume ~50 docs per year

    def _generate_correlation_insights(
            self,
            concept_a: str,
            concept_b: str,
            correlation: float,
            baseline: float,
            supporting_docs: List[Dict],
            date_range: Tuple[int, int]
    ) -> List[str]:
        """Generate human-readable insights about correlations"""
        insights = []

        if correlation > baseline * 3:
            insights.append(
                f"Strong correlation: '{concept_a}' and '{concept_b}' appear together {correlation * 100:.1f}% of the time, which is {correlation / baseline:.1f}x more than expected")
        elif correlation > baseline * 2:
            insights.append(
                f"Moderate correlation: '{concept_a}' and '{concept_b}' co-occur {correlation / baseline:.1f}x more than random chance")

        if supporting_docs:
            # Analyze metadata patterns
            authors = [doc.get('metadata', {}).get('author', '') for doc in supporting_docs]
            common_authors = [a for a in authors if a and authors.count(a) > 1]
            if common_authors:
                insights.append(f"Multiple documents by {common_authors[0]} discuss both concepts")

        return insights

    def _generate_clustering_insights(self, associations: List[Dict],
                                               concept_counts: Dict, total_docs: int) -> List[str]:
        """Generate enhanced insights about concept clustering"""
        insights = []

        if not associations:
            insights.append("No significant concept associations detected in this time period")
            return insights

        # Analyze association strengths
        high_lift = [a for a in associations if a['basic_measures']['lift'] > 3.0]
        strong_correlation = [a for a in associations if
                              a.get('similarity_measures', {}).get('phi_coefficient', 0) > 0.3]

        if high_lift:
            insights.append(f"Found {len(high_lift)} concept pairs with strong lift (>3.0x expected co-occurrence)")

            top_lift = max(high_lift, key=lambda x: x['basic_measures']['lift'])
            insights.append(
                f"Strongest association: '{top_lift['concept_a']}' and '{top_lift['concept_b']}' (lift: {top_lift['basic_measures']['lift']:.1f})")

        if strong_correlation:
            insights.append(f"Detected {len(strong_correlation)} concept pairs with strong correlation (phi > 0.3)")

        # Analyze concept popularity
        popular_concepts = sorted(concept_counts.items(), key=lambda x: x[1], reverse=True)[:5]
        if popular_concepts:
            insights.append(
                f"Most frequent concepts: {', '.join([f'{concept} ({count} docs)' for concept, count in popular_concepts])}")

        # Statistical insights
        significant_count = len([a for a in associations if a['statistical_tests']['is_significant']])
        if significant_count > 0:
            insights.append(
                f"{significant_count} associations are statistically significant after multiple testing correction")

        # Effect size insights
        large_effects = [a for a in associations if a['statistical_tests'].get('effect_size') == 'large']
        if large_effects:
            insights.append(f"Found {len(large_effects)} associations with large effect sizes (Cramér's V > 0.5)")

        return insights

    def _generate_anomaly_insights(self, concept: str, anomalies: List[Dict],
                                   window_data: List[Dict], frequencies: np.ndarray,
                                   is_non_normal: bool) -> List[str]:
        """Generate enhanced insights about temporal anomalies"""
        insights = []

        if not anomalies:
            insights.append(f"No significant temporal anomalies detected for '{concept}' using enhanced methods")
            if is_non_normal:
                insights.append("Data distribution is non-normal - robust detection methods were prioritized")
            return insights

        # Analyze anomaly types and severity
        spikes = [a for a in anomalies if a['type'] == 'spike']
        drops = [a for a in anomalies if a['type'] == 'drop']
        extreme_anomalies = [a for a in anomalies if a.get('severity') == 'extreme']

        if spikes:
            insights.append(f"Detected {len(spikes)} unusual spikes in '{concept}' references")
            if extreme_anomalies:
                extreme_spike = next((a for a in extreme_anomalies if a['type'] == 'spike'), None)
                if extreme_spike:
                    insights.append(
                        f"Extreme spike in {extreme_spike['period']} - {abs(extreme_spike['z_score']):.1f} standard deviations above normal")

        if drops:
            insights.append(f"Found {len(drops)} significant drops in '{concept}' mentions")
            extreme_drop = next((a for a in extreme_anomalies if a['type'] == 'drop'), None)
            if extreme_drop:
                insights.append(
                    f"Extreme drop in {extreme_drop['period']} - {abs(extreme_drop['z_score']):.1f} standard deviations below normal")

        # Method consensus analysis
        method_counts = {}
        for anomaly in anomalies:
            for method in anomaly.get('methods_triggered', []):
                method_counts[method] = method_counts.get(method, 0) + 1

        if method_counts:
            most_common_method = max(method_counts.items(), key=lambda x: x[1])
            insights.append(
                f"Most anomalies detected by {most_common_method[0]} method ({most_common_method[1]} cases)")

        # Temporal clustering of anomalies
        if len(anomalies) >= 2:
            anomaly_years = []
            for anomaly in anomalies:
                period = anomaly['period']
                if '-' in period:
                    start_year = int(period.split('-')[0])
                    anomaly_years.append(start_year)

            if len(anomaly_years) >= 2:
                year_gaps = [anomaly_years[i + 1] - anomaly_years[i] for i in range(len(anomaly_years) - 1)]
                avg_gap = sum(year_gaps) / len(year_gaps)
                if avg_gap <= 20:
                    insights.append(f"Anomalies appear clustered in time (average gap: {avg_gap:.1f} years)")

        # Statistical distribution insights
        if is_non_normal:
            insights.append("Non-normal data distribution detected - used robust statistical methods for detection")
        else:
            insights.append("Data follows normal distribution - parametric methods provided reliable detection")

        # Frequency variation insights
        if len(frequencies) > 3:
            cv = np.std(frequencies) / np.mean(frequencies) if np.mean(frequencies) > 0 else 0
            if cv > 1.0:
                insights.append(f"High variability in '{concept}' frequency over time (CV = {cv:.2f})")
            elif cv < 0.3:
                insights.append(f"Relatively stable '{concept}' frequency with periodic anomalies (CV = {cv:.2f})")

        return insights

    def _ensure_json_serializable(self, obj):
        """Convert numpy types and handle NaN/Inf values"""
        import math

        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            if math.isnan(obj) or math.isinf(obj):
                return None
            return float(obj)
        elif isinstance(obj, np.bool_):
            return bool(obj)
        elif isinstance(obj, dict):
            return {k: self._ensure_json_serializable(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [self._ensure_json_serializable(v) for v in obj]
        elif isinstance(obj, np.ndarray):
            return [self._ensure_json_serializable(x) for x in obj.tolist()]
        elif obj is None:
            return None
        elif isinstance(obj, (int, float, str, bool)):
            if isinstance(obj, float) and (math.isnan(obj) or math.isinf(obj)):
                return None
            return obj
        else:
            try:
                return str(obj)
            except:
                return None