### SECTION 1: INSTALL LIBRARIES ###
# ------------------------------------
print("--- Section 1: Installing Libraries (TM & KG) ---")
!pip install pypdf nltk scikit-learn wordcloud matplotlib seaborn pandas spacy networkx gradio_client requests beautifulsoup4 --quiet
print("Libraries installed successfully.")

# Download NLTK & spaCy models
print("\n--- Downloading Language Models ---")
import nltk
# Define function to download NLTK data robustly
def download_nltk_data(resource):
    try:
        nltk.data.find(resource)
        print(f"NLTK resource '{resource.split('/')[-1]}' already downloaded.")
    except LookupError:
        print(f"Downloading NLTK resource '{resource.split('/')[-1]}'...")
        nltk.download(resource.split('/')[-1], quiet=True)
        print(f"NLTK resource '{resource.split('/')[-1]}' downloaded.")

try:
    download_nltk_data('corpora/wordnet')
    download_nltk_data('corpora/stopwords')
    download_nltk_data('tokenizers/punkt')
except Exception as e:
    print(f"❌ ERROR: NLTK download failed: {e}. Subsequent steps might fail.")
    # Decide if you want to stop: raise SystemExit("NLTK download failed.")

# Download and load spaCy model
spacy_model_name = 'en_core_web_sm'
nlp = None # Initialize nlp
try:
    import spacy
    try:
        nlp = spacy.load(spacy_model_name)
        print(f"spaCy model '{spacy_model_name}' loaded successfully.")
    except OSError:
        print(f"spaCy model '{spacy_model_name}' not found. Downloading...")
        # Use spacy.cli.download for better handling
        spacy.cli.download(spacy_model_name)
        nlp = spacy.load(spacy_model_name) # Try loading again
        print(f"spaCy model '{spacy_model_name}' downloaded and loaded.")
except ImportError:
    print("❌ ERROR: spaCy library not found. Please ensure installation was successful.")
except Exception as e:
    print(f"❌ ERROR: spaCy download/load failed: {e}")

if nlp is None:
     print("⚠️ Warning: spaCy model could not be loaded. Knowledge Graph generation will be skipped.")

print("-" * 30, "\n")





### SECTION 2: IMPORT LIBRARIES ###
# ------------------------------------
print("--- Section 2: Importing Libraries ---")
import os
import re
import warnings
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import json
import pickle
import time
import io
from collections import Counter, defaultdict # <<< FIXED: Added Counter import

from google.colab import files
from pypdf import PdfReader
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
from nltk.tokenize import word_tokenize, sent_tokenize # Ensure sent_tokenize is imported

from sklearn.feature_extraction.text import CountVectorizer
from sklearn.decomposition import LatentDirichletAllocation
from sklearn.metrics.pairwise import cosine_similarity
from wordcloud import WordCloud
import scipy.cluster.hierarchy as sch

import networkx as nx
import requests
from bs4 import BeautifulSoup
from IPython.display import HTML, display, Markdown

# Import and check Gradio Client
try:
    from gradio_client import Client
    GRADIO_AVAILABLE = True
except ImportError:
    GRADIO_AVAILABLE = False
    print("⚠️ Warning: Gradio client library not found. Graph-Mind API functionality disabled.")

# Suppress specific warnings
warnings.filterwarnings("ignore", category=UserWarning, module='pypdf._reader')
warnings.filterwarnings("ignore", category=FutureWarning)

# Set plot style & define globals
sns.set_style('whitegrid')
topic_colors = plt.cm.tab10.colors # Color palette for topics
output_dir = '/content/results' # Central output directory
os.makedirs(output_dir, exist_ok=True)
print(f"Results will be saved in: {output_dir}")

# Ensure spaCy model is accessible (even if None)
if 'nlp' not in locals():
     print("Re-checking spaCy model status...")
     if nlp is None: print("spaCy 'nlp' variable still None.")
     else: print("spaCy 'nlp' variable seems ok.")

print("Libraries imported.")
print("-" * 30, "\n")


### SECTION 3: UPLOAD PDF FILES ###
# ------------------------------------
print("--- Section 3: Upload PDF Files ---")
upload_dir = 'uploaded_pdfs_temp' # Use a distinct temp dir name
if not os.path.exists(upload_dir): os.makedirs(upload_dir)

# Clear previous uploads
if os.path.exists(upload_dir):
    for f in os.listdir(upload_dir):
        file_path = os.path.join(upload_dir, f)
        if os.path.isfile(file_path): os.remove(file_path)
    print("Cleared previous uploads from temporary directory.")

print(f"\nPlease upload your PDF journal articles.")
# Use Colab's upload widget
uploaded = files.upload()

pdf_files_paths = [] # Store paths to uploaded files in the temp dir
uploaded_filenames_map = {} # Map temp path to original filename

if uploaded:
    for filename, content in uploaded.items():
        # Sanitize filename slightly for path safety
        safe_filename = re.sub(r'[\\/*?:"<>|]', '_', filename)
        filepath = os.path.join(upload_dir, safe_filename)
        uploaded_filenames_map[filepath] = filename # Store original name mapping
        with open(filepath, 'wb') as f: f.write(content)
        # Check extension again after potential sanitization
        if filepath.lower().endswith('.pdf'):
             pdf_files_paths.append(filepath)
             print(f'Saved temporary file: {safe_filename} (Original: {filename})')
        else:
             print(f'Skipped non-PDF file: {filename}')
    print(f"\nSuccessfully processed {len(pdf_files_paths)} PDF files.")
else:
    print("\nNo files were uploaded.")

print("-" * 30, "\n")



### SECTION 4: EXTRACT TEXT, PREPROCESS (LDA), STORE RAW & SENTENCES (KG) ###
# ------------------------------------
print("--- Section 4: Extracting & Processing Text ---")
import nltk
nltk.download('punkt_tab')

# Initialize data storage dictionaries
pdf_contents_raw = {}    # {original_filename: raw_extracted_text}
document_sentences = {} # {original_filename: [sentence1, sentence2,...]}
extracted_texts_lda = [] # [processed_text_doc1_for_lda, ...]
processed_file_labels_lda = [] # [original_filename_doc1, ...] - maps LDA results back

# --- LDA Preprocessing Setup ---
lemmatizer = WordNetLemmatizer()
stop_words_set = set(stopwords.words('english'))
custom_stopwords = {'ieee', 'url', 'content', 'malla', 'code', 'example', 'et', 'al', 'fig', 'figure', 'table', 'abstract', 'introduction', 'conclusion', 'references', 'doi', 'journal', 'vol', 'published', 'author', 'copyright', 'reserved', 'rights', 'university', 'department', 'paper', 'article', 'research', 'study', 'method', 'result', 'analysis', 'model', 'datum', 'based', 'using', 'show', 'however', 'provide', 'propose', 'discuss', 'present'}
stop_words_set.update(custom_stopwords)
# --- End LDA Setup ---

if not pdf_files_paths:
    print("❌ No PDF files found to process. Upload files in Section 3.")
else:
    print(f"Processing {len(pdf_files_paths)} PDF files...")
    for pdf_path in pdf_files_paths:
        original_filename = uploaded_filenames_map.get(pdf_path, os.path.basename(pdf_path))
        raw_text = "" # Initialize for this file
        try:
            print(f"  Reading: {original_filename}...")
            reader = PdfReader(pdf_path)
            if reader.is_encrypted:
                 try: reader.decrypt('')
                 except Exception: print(f"    ⚠️ Warning: Encrypted PDF {original_filename} couldn't be decrypted. Skipping."); continue

            for page in reader.pages:
                try: page_text = page.extract_text(); raw_text += (page_text + " ") if page_text else ""
                except Exception as page_error: print(f"    ⚠️ Warning: Can't extract text from a page in {original_filename}. Error: {page_error}")

            raw_text = raw_text.strip()
            if not raw_text: print(f"    ⚠️ Warning: No text extracted from {original_filename}. Skipping."); continue

            # 1. Store Raw Text (for KG)
            pdf_contents_raw[original_filename] = raw_text
            print(f"    Stored raw text ({len(raw_text)} chars).")

            # 2. Store Sentences (for KG)
            text_for_sentences = re.sub(r'\s+', ' ', raw_text) # Normalize whitespace first
            try:
                sentences = sent_tokenize(text_for_sentences)
                valid_sentences = [s.strip() for s in sentences if len(s.strip().split()) > 3 and len(s.strip()) > 15]
                if valid_sentences:
                     document_sentences[original_filename] = valid_sentences
                     print(f"    Stored {len(valid_sentences)} sentences.")
                else:
                     print(f"    ⚠️ Warning: No valid sentences found after filtering for {original_filename}. Storing raw text chunk as single 'sentence'.")
                     document_sentences[original_filename] = [raw_text] if raw_text else [] # Fallback, ensure list
            except Exception as e_sent:
                 print(f"    ❌ Error tokenizing sentences for {original_filename}: {e_sent}. Storing raw text chunk as single 'sentence'.")
                 document_sentences[original_filename] = [raw_text] if raw_text else [] # Fallback, ensure list


            # 3. Preprocess for LDA
            text_lda = raw_text.lower()
            text_lda = re.sub(r'[^a-z\s]', '', text_lda) # Keep only letters and spaces
            text_lda = re.sub(r'\s+', ' ', text_lda).strip()
            words = text_lda.split()
            meaningful_words = [lemmatizer.lemmatize(word) for word in words if word not in stop_words_set and len(word) > 2]

            if not meaningful_words: print(f"    ⚠️ Warning: No meaningful words left in {original_filename} after LDA preprocessing. Skipping this doc for LDA."); continue

            # 4. Store Processed Text & Label (for LDA)
            extracted_texts_lda.append(" ".join(meaningful_words))
            processed_file_labels_lda.append(original_filename)
            print(f"    Stored preprocessed text for LDA.")

        except Exception as e: print(f"    ❌ UNEXPECTED ERROR processing {original_filename}: {e}. Skipping this file.")

    # --- Save intermediate results ---
    try:
         intermediate_data = {
             'pdf_contents_raw': pdf_contents_raw,
             'document_sentences': document_sentences,
             'extracted_texts_lda': extracted_texts_lda,
             'processed_file_labels_lda': processed_file_labels_lda
         }
         save_path = os.path.join(output_dir, 'intermediate_data.pkl')
         with open(save_path, 'wb') as f: pickle.dump(intermediate_data, f)
         print(f"\n✅ Intermediate data (raw text, sentences, LDA text) saved to {save_path}")
    except Exception as e_save: print(f"\n⚠️ Warning: Could not save intermediate data: {e_save}")

    if not extracted_texts_lda: print("\n❌ No text could be successfully processed for LDA Topic Modeling. Cannot proceed with LDA.")
    elif len(extracted_texts_lda) != len(processed_file_labels_lda): print("\n❌ Error: Mismatch between processed texts and labels for LDA. Check processing logs.")
    else: print(f"\n✅ Successfully processed text from {len(extracted_texts_lda)} files for LDA.")

print("-" * 30, "\n")



### SECTION 5: TOPIC MODELING (LDA) ###
# ------------------------------------
print("--- Section 5: Performing Topic Modeling (LDA) ---")

# !!! --- IMPORTANT: SET THIS VARIABLE --- !!!
# Adjust this based on the number of distinct themes you expect in your documents.
NUMBER_OF_TOPICS = 6 # <<< Example: Set to 5, 8, 10, etc.
# !!! ------------------------------------ !!!

# Initialize LDA variables
vectorizer = None; dtm = None; lda = None; feature_names = None; doc_topic_dist = None

# Check if prerequisite data exists and is valid
if 'extracted_texts_lda' not in locals() or not extracted_texts_lda or \
   'processed_file_labels_lda' not in locals() or not processed_file_labels_lda or \
   len(extracted_texts_lda) != len(processed_file_labels_lda):
    print("❌ Cannot perform topic modeling: Missing or inconsistent preprocessed text data from Section 4.")
else:
    num_docs_lda = len(extracted_texts_lda)
    min_document_frequency = min(5, max(1, 2 if num_docs_lda >= 10 else 1))
    print(f"Number of documents for LDA: {num_docs_lda}")

    # Validate NUMBER_OF_TOPICS
    if not isinstance(NUMBER_OF_TOPICS, int) or NUMBER_OF_TOPICS <= 0:
        print(f"❌ Error: NUMBER_OF_TOPICS ('{NUMBER_OF_TOPICS}') must be a positive integer."); raise SystemExit("Stopping: Invalid number of topics.")
    if NUMBER_OF_TOPICS > num_docs_lda:
        print(f"⚠️ Warning: Topics ({NUMBER_OF_TOPICS}) > Docs ({num_docs_lda}). Reducing topics to {num_docs_lda}."); NUMBER_OF_TOPICS = num_docs_lda

    print(f"Setting number of topics to: {NUMBER_OF_TOPICS}")
    print(f"Setting min_df to: {min_document_frequency}")
    print(f"Vectorizing text for LDA (max_df=0.90, min_df={min_document_frequency}, max_features=1500)...")

    try:
        vectorizer = CountVectorizer(max_df=0.90, min_df=min_document_frequency, max_features=1500, stop_words='english')
        dtm = vectorizer.fit_transform(extracted_texts_lda)
        print("Vectorization complete."); print(f"DTM shape: {dtm.shape}")

        if dtm.shape[1] == 0: print("\n❌ Error: DTM has 0 features. Check preprocessing/vectorizer settings."); lda = None; feature_names = None
        elif dtm.shape[0] < NUMBER_OF_TOPICS: print(f"\n❌ Error: Docs in DTM ({dtm.shape[0]}) < Topics ({NUMBER_OF_TOPICS}). Reduce topics."); lda = None; feature_names = None
        else:
            print(f"Building LDA model with {NUMBER_OF_TOPICS} topics...")
            lda = LatentDirichletAllocation(n_components=NUMBER_OF_TOPICS, random_state=42, learning_method='online', n_jobs=-1, max_iter=25, evaluate_every=5)
            lda.fit(dtm)
            doc_topic_dist = lda.transform(dtm)
            print("✅ LDA model trained & doc-topic distributions calculated.")
            feature_names = vectorizer.get_feature_names_out()

            # --- Save LDA Model and Vectorizer ---
            try:
                with open(os.path.join(output_dir, 'lda_model.pkl'), 'wb') as f: pickle.dump(lda, f)
                with open(os.path.join(output_dir, 'vectorizer.pkl'), 'wb') as f: pickle.dump(vectorizer, f)
                print(f"✅ LDA model and vectorizer saved to {output_dir}")
            except Exception as e_save_lda: print(f"⚠️ Warning: Could not save LDA model/vectorizer: {e_save_lda}")

    except ValueError as e: print(f"\n❌ Error during vectorization/LDA: {e}"); lda = None; feature_names = None; doc_topic_dist = None

print("-" * 30, "\n")

# ================================================================
# 5-1. GENERATE TOPIC–KEYWORD TABLE (Topic | Keywords)
# ================================================================
#  Added  code on 12Dec2025 to generate Topics and their respective keywords.
print("--- Section 5B: Generating Topic–Keyword Table ---")

# Ensure LDA model is ready
if lda is None or feature_names is None:
    print("❌ Cannot generate topic table: LDA model or feature names missing.")
else:
    TOP_N_KEYWORDS = 15   # Number of keywords per topic (adjust as needed)

    topic_keyword_rows = []

    for topic_idx, topic in enumerate(lda.components_):
        # Get indices of top weighted words
        top_indices = topic.argsort()[-TOP_N_KEYWORDS:][::-1]

        # Convert to keyword list
        keywords = [feature_names[i] for i in top_indices]

        # Join words into a comma-separated string
        keyword_string = ", ".join(keywords)

        # Append row
        topic_keyword_rows.append([f"Topic {topic_idx + 1}", keyword_string])

    # Create final DataFrame
    topic_keyword_df = pd.DataFrame(topic_keyword_rows, columns=["Topic", "Keywords"])

    # PRINT WITHOUT INDEX
    print("\nTopic–Keyword Table:")
    print(topic_keyword_df.to_string(index=False))

    # Optional: display in notebooks without index
    try:
        from IPython.display import display
        display(topic_keyword_df.style.hide(axis="index"))
    except:
        pass

### SECTION 5A: GOODNESS-OF-FIT EVALUATION (TOPIC OPTIMIZATION) ###
# ----------------------------------------------------------------
print("--- Section 5A: Evaluating Model Fit with Different Number of Topics ---")

from sklearn.decomposition import LatentDirichletAllocation

min_topics = 2
max_topics = min(15, len(extracted_texts_lda))  # Avoid going beyond document count
step = 1
log_likelihoods = []
perplexities = []
topic_range = list(range(min_topics, max_topics + 1))

print(f"Evaluating topics from {min_topics} to {max_topics}...")

for num_topics in topic_range:
    print(f"Training LDA for {num_topics} topics...")
    lda_model = LatentDirichletAllocation(n_components=num_topics,
                                          max_iter=10,
                                          learning_method='online',
                                          random_state=42,
                                          evaluate_every=1,
                                          n_jobs=-1)
    lda_model.fit(dtm)
    log_likelihood = lda_model.score(dtm)
    perplexity = lda_model.perplexity(dtm)
    log_likelihoods.append(log_likelihood)
    perplexities.append(perplexity)

# Plot Log Likelihood
plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
plt.plot(topic_range, log_likelihoods, marker='o')
plt.title("Log Likelihood vs Number of Topics")
plt.xlabel("Number of Topics")
plt.ylabel("Log Likelihood")

# Plot Perplexity
plt.subplot(1, 2, 2)
plt.plot(topic_range, perplexities, marker='o', color='red')
plt.title("Perplexity vs Number of Topics")
plt.xlabel("Number of Topics")
plt.ylabel("Perplexity (lower is better)")

plt.tight_layout()
goodness_path = os.path.join(output_dir, 'lda_goodness_of_fit.png')
plt.savefig(goodness_path, dpi=300)
plt.show()

print(f"✅ Goodness-of-Fit plots saved to {goodness_path}")
print("🔍 Interpretation: Look for elbow point (log-likelihood plateaus or perplexity bottoms out).")



### SECTION 6: DISPLAY TOPICS (Keywords) ###
# ------------------------------------
print("--- Section 6: Displaying Topics (Keywords) ---")

NUM_TOP_WORDS = 15 # Show more words
topic_keywords = {} # Store keywords

# Check prerequisites for this section
if 'lda' not in locals() or lda is None or \
   'feature_names' not in locals() or feature_names is None or \
   len(feature_names) == 0:
    print("❌ Cannot display topics: LDA model or features not available from Section 5.")
else:
    # Ensure NUMBER_OF_TOPICS reflects the actual components in the trained model
    ACTUAL_NUMBER_OF_TOPICS = lda.n_components
    print(f"\nTop {NUM_TOP_WORDS} words for each of the {ACTUAL_NUMBER_OF_TOPICS} topics found:\n")
    for topic_idx, topic_weights in enumerate(lda.components_):
        num_words_to_show = min(NUM_TOP_WORDS, len(feature_names))
        top_word_indices = topic_weights.argsort()[:-num_words_to_show - 1:-1]
        top_words = [feature_names[i] for i in top_word_indices]
        topic_keywords[topic_idx] = top_words
        print(f"Topic #{topic_idx + 1}: {', '.join(top_words)}")

    print("\n\n--- Interpretation ---")
    print("Review keywords. Adjust NUMBER_OF_TOPICS in Section 5 if needed and re-run.")
    print("Consider refining custom_stopwords in Section 4 if needed.")

    # Save topic keywords to a file
    try:
        df_topics = pd.DataFrame([{'Topic': f"Topic {i+1}", 'Keywords': ", ".join(words)} for i, words in topic_keywords.items()])
        topics_path = os.path.join(output_dir, 'topic_keywords.csv')
        df_topics.to_csv(topics_path, index=False)
        print(f"\n✅ Topic keywords saved to {topics_path}")
    except Exception as e_save_topics:
         print(f"\n⚠️ Warning: Could not save topic keywords: {e_save_topics}")

print("\n" + "-" * 30)
print("Run Section 7 next for Topic Modeling visualizations.")
print("-" * 30, "\n")



### SECTION 7: TOPIC MODELING VISUALIZATIONS ###
# ------------------------------------
print("--- Section 7: Generating Topic Modeling Visualizations ---")

# Check if necessary components exist for visualization
if 'lda' not in locals() or lda is None or \
   'dtm' not in locals() or dtm is None or \
   'feature_names' not in locals() or feature_names is None or len(feature_names) == 0 or \
   'doc_topic_dist' not in locals() or doc_topic_dist is None or \
   'processed_file_labels_lda' not in locals() or not processed_file_labels_lda or \
   len(processed_file_labels_lda) != doc_topic_dist.shape[0]: # Check alignment
    print("❌ Cannot generate TM visualizations: missing or inconsistent data from previous steps.")
else:
    # Use the actual number of topics from the trained model
    ACTUAL_NUMBER_OF_TOPICS_VIZ = lda.n_components
    num_docs_viz = len(processed_file_labels_lda)
    print(f"Preparing TM visualizations for {num_docs_viz} documents and {ACTUAL_NUMBER_OF_TOPICS_VIZ} topics...")

    # --- Visualization 1: Word Clouds per Topic ---
    print("\nGenerating Word Clouds...")
    topic_word_distributions = lda.components_ / lda.components_.sum(axis=1)[:, np.newaxis]
    cols = 3; rows = int(np.ceil(ACTUAL_NUMBER_OF_TOPICS_VIZ / cols)); plt.figure(figsize=(6 * cols, 5 * rows))
    for topic_idx, topic_weights in enumerate(topic_word_distributions):
        top_indices = topic_weights.argsort()[:-51:-1]
        top_word_freq = {feature_names[i]: topic_weights[i] for i in top_indices if i < len(feature_names)}
        ax = plt.subplot(rows, cols, topic_idx + 1)
        try:
            if not top_word_freq: raise ValueError("No words found.")
            wordcloud = WordCloud(width=500, height=400, background_color='white', max_words=50, colormap='viridis', prefer_horizontal=0.95, random_state=42).generate_from_frequencies(top_word_freq)
            ax.imshow(wordcloud, interpolation='bilinear'); ax.set_title(f'Topic #{topic_idx + 1}', fontsize=14)
        except ValueError as e: print(f"  ⚠️ Warning: WC Error T{topic_idx + 1}: {e}"); ax.text(0.5, 0.5, f'Topic {topic_idx + 1}\n(Error)', ha='center', va='center')
        ax.axis('off')
    plt.tight_layout(pad=3.0); plt.suptitle("Word Clouds for Each Topic", fontsize=18, y=1.03);
    wc_path = os.path.join(output_dir, 'topic_wordclouds.png'); plt.savefig(wc_path, dpi=300, bbox_inches='tight'); plt.show()
    print(f"✅ Word Clouds saved to {wc_path}")

    # --- Visualization 2: Topic Prevalence Bar Chart ---
    print("\nGenerating Topic Prevalence Chart...")
    topic_prevalence = doc_topic_dist.mean(axis=0)
    plt.figure(figsize=(12, 6)); topic_nums = [f"Topic {i+1}" for i in range(ACTUAL_NUMBER_OF_TOPICS_VIZ)]
    sns.barplot(x=topic_nums, y=topic_prevalence, palette="viridis")
    plt.title('Average Topic Prevalence Across All Documents', fontsize=16); plt.xlabel('Topic', fontsize=12); plt.ylabel('Average Prevalence Score', fontsize=12)
    plt.xticks(rotation=45, ha='right'); plt.tight_layout();
    tp_path = os.path.join(output_dir, 'topic_prevalence.png'); plt.savefig(tp_path, dpi=300, bbox_inches='tight'); plt.show()
    print(f"✅ Topic Prevalence Chart saved to {tp_path}")

    # --- Visualization 3: Thematic Map (Topic Relationship Clustermap) ---
    print("\nGenerating Thematic Map (Clustermap)...")
    if ACTUAL_NUMBER_OF_TOPICS_VIZ > 1:
        try:
            topic_similarity_matrix = cosine_similarity(topic_word_distributions)
            topic_labels = [f"Topic {i+1}" for i in range(ACTUAL_NUMBER_OF_TOPICS_VIZ)]
            df_topic_similarity = pd.DataFrame(topic_similarity_matrix, index=topic_labels, columns=topic_labels)
            print("  Calculating clusters and plotting heatmap...")
            cluster_map = sns.clustermap(df_topic_similarity, method='ward', cmap="viridis", linewidths=0.5, linecolor='lightgray', annot=True, fmt=".2f", figsize=(max(8, ACTUAL_NUMBER_OF_TOPICS_VIZ*0.9), max(8, ACTUAL_NUMBER_OF_TOPICS_VIZ*0.9)))
            cluster_map.fig.suptitle('Thematic Map: Topic Similarity Clustermap', y=1.02, fontsize=16)
            plt.setp(cluster_map.ax_heatmap.get_xticklabels(), rotation=45, ha='right'); plt.setp(cluster_map.ax_heatmap.get_yticklabels(), rotation=0)
            tm_path = os.path.join(output_dir, 'thematic_map_clustermap.png'); plt.savefig(tm_path, dpi=300, bbox_inches='tight');
            plt.show()
            print(f"✅ Thematic Map (Clustermap) saved to {tm_path}")
            print("  Interpretation: Topics clustered together are similar. Brighter cells indicate higher similarity.")
        except Exception as e: print(f"  ❌ Error generating Clustermap: {e}")
    else: print("  Skipping Thematic Map: requires > 1 topic.")

    # --- Visualization 4: Document-Topic Distribution (Stacked Bar Chart) ---
    print("\nGenerating Document-Topic Distribution Chart...")
    try:
        topic_names = [f"Topic {i+1}" for i in range(ACTUAL_NUMBER_OF_TOPICS_VIZ)]
        df_doc_topic = pd.DataFrame(doc_topic_dist, columns=topic_names, index=processed_file_labels_lda) # Use LDA labels
        plot_kind = 'bar'; fig_height = max(7, num_docs_viz * 0.5); fig_width = 14
        if num_docs_viz > 25: print("  Info: Using horizontal bars for >25 documents."); plot_kind = 'barh'; fig_height, fig_width = fig_width, fig_height
        ax = df_doc_topic.plot(kind=plot_kind, stacked=True, figsize=(fig_width, fig_height), colormap='viridis', width=0.8 if plot_kind=='bar' else None, fontsize=10)
        ax.set_title('Topic Distribution Across Documents', fontsize=16)
        if plot_kind == 'bar':
            ax.set_xlabel('Documents', fontsize=12); ax.set_ylabel('Topic Proportion', fontsize=12); plt.xticks(rotation=90)
            plt.tight_layout(rect=[0, 0.03, 1, 0.95]); ax.legend(title='Topics', bbox_to_anchor=(1.02, 1), loc='upper left', fontsize='medium')
        else:
            ax.set_xlabel('Topic Proportion', fontsize=12); ax.set_ylabel('Documents', fontsize=12); plt.gca().invert_yaxis()
            plt.tight_layout(rect=[0.1, 0, 0.85, 0.95]); ax.legend(title='Topics', bbox_to_anchor=(1.02, 1), loc='upper left', fontsize='medium')
        dtd_path = os.path.join(output_dir, 'document_topic_distribution.png'); plt.savefig(dtd_path, dpi=300, bbox_inches='tight'); plt.show()
        print(f"✅ Document-Topic Distribution chart saved to {dtd_path}")
    except Exception as e_dtd:
         print(f"  ❌ Error generating Document-Topic chart: {e_dtd}")

print("\n" + "-" * 30)
print("✅ Topic Modeling & Visualization Phase Complete.")
print("-" * 30, "\n")



########################
###KG 
##########################3

### SECTION 8: KNOWLEDGE GRAPH GENERATION ###
# -------------------------------------------
print("--- Section 8: Generating Knowledge Graph ---")

import networkx as nx
import matplotlib.pyplot as plt

# --- CONFIGURABLE ---
TOP_N_WORDS = 10  # limit top keywords per topic
LDA_MODEL_PATH = os.path.join(output_dir, 'lda_model.pkl')
VECTORIZER_PATH = os.path.join(output_dir, 'vectorizer.pkl')

# --- LOAD MODEL & FEATURES ---
with open(LDA_MODEL_PATH, 'rb') as f:
    lda = pickle.load(f)
with open(VECTORIZER_PATH, 'rb') as f:
    vectorizer = pickle.load(f)

feature_names = vectorizer.get_feature_names_out()
topic_keywords = {}

for topic_idx, topic in enumerate(lda.components_):
    top_indices = topic.argsort()[:-TOP_N_WORDS - 1:-1]
    keywords = [feature_names[i] for i in top_indices]
    topic_keywords[f"Topic {topic_idx + 1}"] = keywords

# --- BUILD GRAPH ---
G = nx.Graph()

# Add nodes and edges
for topic, keywords in topic_keywords.items():
    G.add_node(topic, type='topic')
    for kw in keywords:
        G.add_node(kw, type='keyword')
        G.add_edge(topic, kw)


# --- VISUALIZE ---
plt.figure(figsize=(12, 8))
pos = nx.spring_layout(G, k=0.5, iterations=50, seed=42)

# Node coloring
node_colors = ['skyblue' if G.nodes[n]['type'] == 'topic' else 'lightgreen' for n in G.nodes]
node_sizes = [1000 if G.nodes[n]['type'] == 'topic' else 500 for n in G.nodes]

nx.draw_networkx_nodes(G, pos, node_size=node_sizes, node_color=node_colors, alpha=0.9)
nx.draw_networkx_edges(G, pos, alpha=0.5)
nx.draw_networkx_labels(G, pos, font_size=10)

plt.title(f"Topic-Keyword Knowledge Graph (Top {TOP_N_WORDS} keywords)", fontsize=14)
plt.axis('off')
plt.tight_layout()

kg_path = os.path.join(output_dir, 'topic_keyword_graph.png')
plt.savefig(kg_path, dpi=300, bbox_inches='tight')
plt.show()
print(f"✅ Knowledge Graph saved to {kg_path}")