script.py

import json
import pdfplumber
import spacy
import numpy as np
from datetime import datetime
from typing import List, Dict, Tuple, Any, Set
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from keybert import KeyBERT
from sentence_transformers import SentenceTransformer, util
import re
import os


class JobAwarePDFAnalyzer:
    def __init__(self):
        """Initialize the analyzer with better embedding models and hybrid similarity."""
        # Load spaCy model
        try:
            self.nlp = spacy.load("en_core_web_sm")
        except OSError:
            print("Downloading spaCy model 'en_core_web_sm'...")
            from spacy.cli import download
            download("en_core_web_sm")
            self.nlp = spacy.load("en_core_web_sm")

        # Initialize with better embedding models
        print("Loading improved embedding models...")
        
        try:
            model_name = 'sentence-transformers/all-mpnet-base-v2'
            self.sentence_model = SentenceTransformer(model_name)
            print(f"Successfully loaded model: {model_name}")
        except Exception as e:
            print(f"Error loading model: {e}")
            # Fallback to a simpler model
            self.sentence_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
            print("Loaded fallback model: paraphrase-MiniLM-L6-v2")
        
        self.kw_model = KeyBERT(self.sentence_model)
        
        # TF-IDF vectorizer (now actively used for hybrid similarity)
        self.tfidf_vectorizer = TfidfVectorizer(
            stop_words="english",
            max_features=5000,  # Increased for better coverage
            ngram_range=(1, 3),
            min_df=1,
            max_df=0.95,  # Remove terms that appear in >95% of documents
            sublinear_tf=True  # Apply sublinear tf scaling
        )

        # Extended stop words and non-keyword words to filter out
        self.extended_stop_words = {
            'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by',
            'is', 'are', 'was', 'were', 'be', 'been', 'being', 'have', 'has', 'had', 'do', 'does', 'did',
            'will', 'would', 'could', 'should', 'may', 'might', 'can', 'must', 'shall',
            'this', 'that', 'these', 'those', 'i', 'you', 'he', 'she', 'it', 'we', 'they',
            'me', 'him', 'her', 'us', 'them', 'my', 'your', 'his', 'our', 'their',
            'myself', 'yourself', 'himself', 'herself', 'itself', 'ourselves', 'themselves',
            'what', 'which', 'who', 'whom', 'whose', 'where', 'when', 'why', 'how',
            'all', 'any', 'both', 'each', 'few', 'more', 'most', 'other', 'some', 'such',
            'no', 'nor', 'not', 'only', 'own', 'same', 'so', 'than', 'very', 'too',
            'just', 'now', 'here', 'there', 'then', 'once', 'again', 'also', 'still',
            'back', 'down', 'off', 'over', 'under', 'up', 'out', 'through', 'during',
            'before', 'after', 'above', 'below', 'between', 'into', 'onto', 'upon',
            'get', 'got', 'make', 'made', 'take', 'took', 'come', 'came', 'go', 'went',
            'see', 'saw', 'know', 'knew', 'think', 'thought', 'say', 'said', 'tell', 'told',
            'work', 'worked', 'use', 'used', 'find', 'found', 'give', 'gave', 'put', 'keep',
            'let', 'help', 'show', 'try', 'call', 'ask', 'need', 'want', 'like', 'look',
            'feel', 'seem', 'become', 'leave', 'move', 'play', 'run', 'walk', 'talk',
            'bring', 'write', 'read', 'hear', 'watch', 'follow', 'stop', 'create',
            'include', 'provide', 'allow', 'add', 'start', 'turn', 'change', 'end',
            'build', 'buy', 'pay', 'sell', 'send', 'receive', 'meet', 'join', 'win',
            'lose', 'die', 'kill', 'break', 'fix', 'open', 'close', 'cut', 'clear',
            'develop', 'grow', 'draw', 'fall', 'push', 'pull', 'carry', 'pick'
        }

        # Document statistics for normalization
        self.document_stats = {}

    # ==============================================================================
    # SECTION 1: ENHANCED HEADING EXTRACTION WITH DOCUMENT STATISTICS TRACKING
    # ==============================================================================

    def extract_text_from_pdf(self, pdf_path: str) -> Tuple[str, List[Dict]]:
        """
        Extract text from PDF with structural information and document statistics tracking.
        """
        full_text = ""
        sections = []
        
        # Document statistics for normalization
        page_count = 0
        total_words = 0
        total_chars = 0

        try:
            with pdfplumber.open(pdf_path) as pdf:
                page_count = len(pdf.pages)
                
                for page_num, page in enumerate(pdf.pages, 1):
                    page_text = page.extract_text()
                    if page_text:
                        full_text += page_text + "\n"
                        page_sections = self._extract_headings_from_page(page, page_num)
                        sections.extend(page_sections)
                        
                        # Track page statistics
                        page_words = len(page_text.split())
                        total_words += page_words
                        total_chars += len(page_text)
                
        except Exception as e:
            print(f"Error extracting from {pdf_path}: {e}")
            return "", []
        
        # Store document statistics for normalization
        filename = os.path.basename(pdf_path)
        self.document_stats[filename] = {
            'page_count': page_count,
            'total_words': total_words,
            'total_chars': total_chars,
            'section_count': len(sections),
            'avg_words_per_page': total_words / page_count if page_count > 0 else 0,
            'density_factor': total_words / (page_count * 500) if page_count > 0 else 1.0  # Assuming 500 words per average page
        }
        
        print(f"DEBUG: Document stats for {filename}: {self.document_stats[filename]}")
        
        sections = self._filter_and_rank_headings(sections)
        return full_text, sections

    def _extract_headings_from_page(self, page, page_num: int) -> List[Dict]:
        """
        Extract headings from a single page using multiple detection methods.
        """
        headings = []
        headings.extend(self._extract_font_based_headings(page, page_num))
        headings.extend(self._extract_position_based_headings(page, page_num))
        headings.extend(self._extract_content_pattern_headings(page, page_num))
        return headings

    def _is_valid_heading(self, text: str) -> bool:
        """
        Enhanced validation to filter out headings with unwanted punctuation.
        """
        if not text or len(text.strip()) < 3:
            return False
            
        text = text.strip()
        
        # Filter out headings with problematic punctuation
        problematic_chars = {',', '?', ';', ':', '!'}
        if any(char in text for char in problematic_chars):
            return False
            
        # Filter out very long headings (likely not headings)
        if len(text) > 120:
            return False
            
        # Filter out headings that are mostly numbers or special characters
        if len(re.sub(r'[^a-zA-Z\s]', '', text)) < len(text) * 0.6:
            return False
            
        # Filter out single words that are too common
        common_words = {'the', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by'}
        words = text.lower().split()
        if len(words) == 1 and words[0] in common_words:
            return False
            
        return True

    def _extract_font_based_headings(self, page, page_num: int) -> List[Dict]:
        headings = []
        try:
            chars = page.chars
            if not chars: return headings
            
            lines = {}
            for char in chars:
                y_pos = round(char['y0'], 1)
                if y_pos not in lines: lines[y_pos] = []
                lines[y_pos].append(char)

            font_sizes = [char['size'] for char in chars]
            avg_font_size = np.mean(font_sizes) if font_sizes else 12
            large_font_threshold = avg_font_size * 1.2

            for y_pos, line_chars in lines.items():
                line_text = ''.join([char['text'] for char in line_chars]).strip()
                
                if not self._is_valid_heading(line_text):
                    continue

                line_font_sizes = [char['size'] for char in line_chars]
                avg_line_font_size = np.mean(line_font_sizes)
                font_names = [char.get('fontname', '') for char in line_chars]
                is_bold = any('bold' in name.lower() for name in font_names)
                is_large_font = avg_line_font_size > large_font_threshold
                is_isolated = self._is_isolated_line(y_pos, lines)

                if (is_large_font or is_bold) and is_isolated:
                    headings.append({
                        'text': line_text, 'page_number': page_num, 'type': 'font_heading',
                        'confidence': 0.8 if is_large_font and is_bold else 0.6
                    })
        except Exception as e:
            print(f"Font-based heading extraction error: {e}")
        return headings

    def _extract_position_based_headings(self, page, page_num: int) -> List[Dict]:
        headings = []
        try:
            page_text = page.extract_text()
            if not page_text: return headings
            lines = page_text.split('\n')

            for i, line in enumerate(lines):
                line = line.strip()
                
                if not self._is_valid_heading(line):
                    continue

                is_short_line = len(line) < 60
                is_title_case = self._is_title_case(line)
                is_standalone = self._is_standalone_line(lines, i)
                has_no_punctuation_end = not line.endswith(('.', ',', ';', '!', '?'))
                
                score = sum([is_short_line * 0.3, is_title_case * 0.3, is_standalone * 0.2, has_no_punctuation_end * 0.2])

                if score >= 0.6:
                    headings.append({'text': line, 'page_number': page_num, 'type': 'position_heading', 'confidence': score})
        except Exception as e:
            print(f"Position-based heading extraction error: {e}")
        return headings

    def _extract_content_pattern_headings(self, page, page_num: int) -> List[Dict]:
        headings = []
        try:
            page_text = page.extract_text()
            if not page_text: return headings
            lines = page_text.split('\n')
            
            heading_patterns = [
                r'^\d+\.?\s+[A-Z][A-Za-z\s]+$',  # Numbered headings
                r'^[A-Z][A-Z\s]+$',              # ALL CAPS headings
                r'^[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*$',  # Title Case (no colon)
                r'^Chapter\s+\d+',               # Chapter headings
                r'^Section\s+\d+',               # Section headings
                r'^Part\s+[A-Z]',                # Part headings
            ]
            
            for line in lines:
                line = line.strip()
                
                if not self._is_valid_heading(line):
                    continue

                pattern_score = 0.7 if any(re.match(p, line) for p in heading_patterns) else 0
                semantic_score = self._calculate_semantic_heading_score(line)
                combined_score = max(pattern_score, semantic_score)

                if combined_score >= 0.5:
                    headings.append({'text': line, 'page_number': page_num, 'type': 'content_heading', 'confidence': combined_score})
        except Exception as e:
            print(f"Content pattern heading extraction error: {e}")
        return headings

    def _calculate_semantic_heading_score(self, text: str) -> float:
        try:
            doc = self.nlp(text)
            score = 0.0
            if not doc or len(doc) == 0: return 0.0
            
            important_pos = ['NOUN', 'PROPN', 'ADJ']
            complex_pos = ['CCONJ', 'SCONJ', 'ADP']
            topic_indicators = ['overview', 'introduction', 'methodology', 'results', 'conclusion', 
                              'summary', 'analysis', 'guide', 'tips', 'planning', 'recommendations']
            
            if len([tok for tok in doc if tok.pos_ in important_pos]) / len(doc) > 0.5: score += 0.2
            if len([tok for tok in doc if tok.pos_ in complex_pos]) / len(doc) < 0.2: score += 0.15
            if not text.endswith(('.', '!', '?')): score += 0.1
            if any(indicator in text.lower() for indicator in topic_indicators): score += 0.25
            
            words = text.split()
            if len(words) > 1 and sum(1 for w in words if w and w[0].isupper()) / len(words) > 0.7: score += 0.15
            
            return min(score, 1.0)
        except Exception as e:
            print(f"Semantic score calculation error: {e}")
            return 0.0

    def _is_isolated_line(self, y_pos: float, lines: Dict) -> bool:
        nearby_threshold = 15
        return sum(1 for other_y in lines.keys() if abs(other_y - y_pos) < nearby_threshold and other_y != y_pos) < 2

    def _is_title_case(self, text: str) -> bool:
        words = text.split()
        if not words: return False
        if len(words) < 2: return text[0].isupper()
        return sum(1 for word in words if word and word[0].isupper()) / len(words) >= 0.7

    def _is_standalone_line(self, lines: List[str], i: int) -> bool:
        curr = lines[i].strip()
        prev = lines[i - 1].strip() if i > 0 else ""
        next = lines[i + 1].strip() if i < len(lines) - 1 else ""
        return (not prev or abs(len(prev) - len(curr)) > len(curr) * 0.5) and \
               (not next or abs(len(next) - len(curr)) > len(curr) * 0.5)

    def _filter_and_rank_headings(self, sections: List[Dict]) -> List[Dict]:
        if not sections: return []
        unique_sections = []
        seen_texts = set()
        
        for section in sorted(sections, key=lambda x: x.get('confidence', 0), reverse=True):
            text_clean = re.sub(r'[^\w\s]', '', section['text'].lower()).strip()
            if text_clean and text_clean not in seen_texts:
                is_duplicate = False
                for existing in unique_sections:
                    if self._are_similar_headings(section['text'], existing['text']):
                        is_duplicate = True
                        break
                if not is_duplicate:
                    unique_sections.append(section)
                    seen_texts.add(text_clean)
        return unique_sections[:30]

    def _are_similar_headings(self, text1: str, text2: str) -> bool:
        text1_clean = re.sub(r'[^\w\s]', '', text1.lower())
        text2_clean = re.sub(r'[^\w\s]', '', text2.lower())
        if text1_clean in text2_clean or text2_clean in text1_clean: return True
        words1, words2 = set(text1_clean.split()), set(text2_clean.split())
        if not words1 or not words2: return False
        return len(words1.intersection(words2)) / max(len(words1), len(words2)) > 0.7

    # ==============================================================================
    # SECTION 2: ENHANCED KEYWORD GENERATION WITH STOP WORD FILTERING
    # ==============================================================================

    def _filter_keywords(self, keywords: List[str]) -> List[str]:
        """
        Filter out stop words and non-keyword words from the keyword list.
        """
        filtered_keywords = []
        
        for keyword in keywords:
            if not keyword or len(keyword) < 2:
                continue
                
            keyword_lower = keyword.lower().strip()
            
            # Skip if it's a stop word
            if keyword_lower in self.extended_stop_words:
                continue
                
            # Skip if it's mostly numbers or special characters
            if len(re.sub(r'[^a-zA-Z]', '', keyword)) < len(keyword) * 0.7:
                continue
                
            # Skip very short words that are likely not meaningful
            if len(keyword_lower) < 3 and keyword_lower not in ['ai', 'ml', 'it', 'ui', 'ux']:
                continue
                
            # Use spaCy to check if it's a meaningful word
            try:
                doc = self.nlp(keyword)
                if doc and len(doc) > 0:
                    # Keep words that are nouns, proper nouns, adjectives, or verbs
                    meaningful_pos = ['NOUN', 'PROPN', 'ADJ', 'VERB']
                    if any(token.pos_ in meaningful_pos for token in doc):
                        filtered_keywords.append(keyword_lower)
            except Exception:
                # If spaCy fails, keep the keyword as a fallback
                filtered_keywords.append(keyword_lower)
        
        return list(set(filtered_keywords))  # Remove duplicates

    def generate_job_context_keywords(self, persona: str, job_task: str) -> List[str]:
        """
        Generate contextual keywords based on persona and job-to-be-done using NLP with enhanced filtering.
        """
        job_keywords = set()
        combined_context = f"{persona} {job_task}"

        print(f"DEBUG: Generating keywords for context: '{combined_context}'")

        # KeyBERT extraction
        try:
            kws = self.kw_model.extract_keywords(combined_context, keyphrase_ngram_range=(1, 3), top_n=25, stop_words='english')
            raw_keywords = [kw[0].lower() for kw in kws]
            filtered_keybert = self._filter_keywords(raw_keywords)
            job_keywords.update(filtered_keybert)
            print(f"KeyBERT generated {len(raw_keywords)} raw keywords, filtered to {len(filtered_keybert)}")
        except Exception as e: 
            print(f"KeyBERT error: {e}")

        # spaCy extraction
        try:
            doc = self.nlp(combined_context)
            # Extract noun chunks
            noun_chunks = [chunk.text.lower().strip() for chunk in doc.noun_chunks if len(chunk.text.strip()) > 2]
            filtered_chunks = self._filter_keywords(noun_chunks)
            job_keywords.update(filtered_chunks)
            
            # Extract individual meaningful tokens
            tokens = [tok.lemma_.lower() for tok in doc if not tok.is_stop and not tok.is_punct and tok.pos_ in ['NOUN', 'VERB', 'ADJ']]
            filtered_tokens = self._filter_keywords(tokens)
            job_keywords.update(filtered_tokens)
            print(f"spaCy extracted {len(noun_chunks + tokens)} raw terms, filtered to {len(filtered_chunks + filtered_tokens)}")
        except Exception as e: 
            print(f"spaCy error: {e}")

        # Semantic expansion
        try:
            seed_kws = list(job_keywords)[:15]
            expanded_kws = self._expand_keywords_semantically(seed_kws, combined_context)
            filtered_expanded = self._filter_keywords(expanded_kws)
            job_keywords.update(filtered_expanded)
            print(f"Semantic expansion generated {len(expanded_kws)} keywords, filtered to {len(filtered_expanded)}")
        except Exception as e: 
            print(f"Semantic expansion error: {e}")

        # Contextual terms
        try:
            contextual_terms = self._generate_contextual_terms(persona, job_task)
            filtered_contextual = self._filter_keywords(contextual_terms)
            job_keywords.update(filtered_contextual)
            print(f"Contextual terms generated {len(contextual_terms)} keywords, filtered to {len(filtered_contextual)}")
        except Exception as e: 
            print(f"Contextual term generation error: {e}")

        final_keywords = list(job_keywords)
        print(f"Final filtered keyword count: {len(final_keywords)}")
        print(f"DEBUG: Generated keywords: {final_keywords}")
        return final_keywords[:60]  # Return top 60 filtered keywords

    def _generate_contextual_terms(self, persona: str, job_task: str) -> List[str]:
        """
        Generate contextual terms by creating semantic queries from the persona and task.
        """
        contextual_terms = set()
        queries = [
            f"{persona} specializes in {job_task}", 
            f"what a {persona} needs for {job_task}",
            f"{job_task} from the perspective of a {persona}",
            f"{persona} working on {job_task} requires",
            f"essential tools for {persona} doing {job_task}"
        ]
        
        for query in queries:
            try:
                kws = self.kw_model.extract_keywords(query, keyphrase_ngram_range=(1, 3), top_n=12, stop_words='english')
                contextual_terms.update(kw[0].lower() for kw in kws)
            except Exception as e: 
                print(f"Contextual query error: {e}")
        return list(contextual_terms)

    def _expand_keywords_semantically(self, seed_keywords: List[str], context: str) -> List[str]:
        """
        Expand keywords by finding semantically similar terms within the context.
        """
        expanded = []
        if not seed_keywords or not context: return expanded
        try:
            context_doc = self.nlp(context)
            context_tokens = list(set(tok.lemma_.lower() for tok in context_doc 
                                    if tok.pos_ in ['NOUN', 'VERB', 'ADJ'] and not tok.is_stop 
                                    and not tok.is_punct and len(tok.text) > 2))
            if not context_tokens: return expanded

            seed_embeddings = self.sentence_model.encode(seed_keywords)
            token_embeddings = self.sentence_model.encode(context_tokens)
            similarities = util.cos_sim(seed_embeddings, token_embeddings)

            for i in range(len(seed_keywords)):
                similar_indices = (similarities[i] > 0.35).nonzero(as_tuple=True)[0]
                for idx in similar_indices[:5]:
                    token = context_tokens[idx]
                    if token not in seed_keywords and token not in expanded:
                        expanded.append(token)
        except Exception as e:
            print(f"Semantic expansion error: {e}")
        return expanded

    # ==============================================================================
    # SECTION 3: ENHANCED SIMILARITY WITH NORMALIZATION FACTORS
    # ==============================================================================

    def filter_relevant_sections(self, sections: List[Dict], job_keywords: List[str], similarity_threshold: float = 0.05) -> List[Dict]:
        """
        Enhanced hybrid filtering with document length normalization.
        """
        if not sections or not job_keywords: 
            print("DEBUG: No sections or keywords provided")
            return []
            
        print(f"DEBUG: Processing {len(sections)} sections with {len(job_keywords)} keywords")
        
        relevant_sections = []
        try:
            section_texts = [s.get('text', '') for s in sections]
            if not any(section_texts): 
                print("DEBUG: No section texts found")
                return []
            
            combined_keywords = ' '.join(job_keywords)
            
            # ===== TF-IDF SIMILARITY =====
            corpus = section_texts + [combined_keywords]
            tfidf_matrix = self.tfidf_vectorizer.fit_transform(corpus)
            keyword_tfidf_vector = tfidf_matrix[-1]
            section_tfidf_vectors = tfidf_matrix[:-1]
            tfidf_similarities = cosine_similarity(section_tfidf_vectors, keyword_tfidf_vector).flatten()
            
            # ===== SENTENCE TRANSFORMER SIMILARITY =====
            section_embeddings = self.sentence_model.encode(section_texts)
            keyword_embedding = self.sentence_model.encode([combined_keywords])
            semantic_similarities = util.cos_sim(section_embeddings, keyword_embedding).numpy().flatten()
            
            # ===== INDIVIDUAL KEYWORD MATCHING =====
            job_embeddings = self.sentence_model.encode(job_keywords)
            individual_similarities = util.cos_sim(section_embeddings, job_embeddings)
            max_individual_sims = individual_similarities.max(dim=1)[0].numpy()
            mean_individual_sims = individual_similarities.mean(dim=1).numpy()
            high_sim_counts = (individual_similarities > 0.20).sum(dim=1).numpy()
            
            for i, section in enumerate(sections):
                tfidf_score = float(tfidf_similarities[i])
                semantic_score = float(semantic_similarities[i])
                max_individual_sim = float(max_individual_sims[i])
                mean_individual_sim = float(mean_individual_sims[i])
                high_sim_bonus = min(float(high_sim_counts[i]) * 0.05, 0.25)
                
                individual_score = (0.5 * max_individual_sim) + (0.3 * mean_individual_sim) + (0.2 * high_sim_bonus)
                base_score = (0.3 * tfidf_score) + (0.4 * semantic_score) + (0.3 * individual_score)
                
                # ===== ENHANCED SECTION TITLE RELEVANCE CHECK =====
                section_title_relevance = self._calculate_title_keyword_relevance(
                    section.get('text', ''),
                    job_keywords,
                    title_type="section"
                )
                
                # ===== DOCUMENT TITLE RELEVANCE CHECK =====
                document_title_relevance = self._calculate_title_keyword_relevance(
                    section.get('document', ''),
                    job_keywords,
                    title_type="document"
                )
                
                # ===== COMBINED RELEVANCE SCORE =====
                combined_relevance_score = (0.7 * section_title_relevance) + (0.3 * document_title_relevance)
                
                # ===== PERFECT MATCH BONUS SYSTEM =====
                perfect_match_bonus = self._calculate_perfect_match_bonus(
                    section.get('text', ''),
                    section.get('document', ''),
                    job_keywords
                )
                
                # Apply penalty
                penalty = self._calculate_keyword_similarity_penalty(
                    section.get('text', ''),
                    section.get('document', ''),
                    job_keywords
                )
                
                # ===== NEW: DOCUMENT LENGTH NORMALIZATION =====
                normalization_factor = self._calculate_normalization_factor(section.get('document', ''))
                
                # Final score with normalization
                relevance_multiplier = max(0.1, combined_relevance_score)
                pre_normalized_score = (base_score * relevance_multiplier) + perfect_match_bonus - (0.2 * penalty)
                final_score = pre_normalized_score * normalization_factor
                
                print(f"DEBUG: '{section.get('text', '')[:30]}' - Base:{base_score:.3f}, Relevance:{combined_relevance_score:.3f}, PreNorm:{pre_normalized_score:.3f}, NormFactor:{normalization_factor:.3f}, Final:{final_score:.3f}")
                
                if final_score > similarity_threshold:
                    section['similarity_score'] = final_score
                    section['pre_normalized_score'] = pre_normalized_score
                    section['normalization_factor'] = normalization_factor
                    section['tfidf_score'] = tfidf_score
                    section['semantic_score'] = semantic_score
                    section['individual_score'] = individual_score
                    section['section_title_relevance'] = section_title_relevance
                    section['document_title_relevance'] = document_title_relevance
                    section['combined_relevance_score'] = combined_relevance_score
                    section['perfect_match_bonus'] = perfect_match_bonus
                    section['base_similarity'] = base_score
                    section['keyword_penalty'] = penalty
                    section['max_individual_sim'] = max_individual_sim
                    relevant_sections.append(section)
                    
        except Exception as e:
            print(f"Hybrid section filtering error: {e}")
            import traceback
            traceback.print_exc()
        
        print(f"DEBUG: Found {len(relevant_sections)} relevant sections after filtering")
        return sorted(relevant_sections, key=lambda x: x['similarity_score'], reverse=True)

    def _calculate_normalization_factor(self, document_name: str) -> float:
        """
        Calculate normalization factor based on document length, page count, and section count.
        This ensures shorter, more focused documents aren't penalized against longer ones.
        """
        if not document_name or document_name not in self.document_stats:
            return 1.0  # Default normalization if no stats available
        
        stats = self.document_stats[document_name]
        
        # Get all document stats for relative comparison
        all_stats = list(self.document_stats.values())
        if len(all_stats) <= 1:
            return 1.0  # No comparison possible with single document
        
        # Calculate normalization components
        page_counts = [s['page_count'] for s in all_stats]
        word_counts = [s['total_words'] for s in all_stats]
        section_counts = [s['section_count'] for s in all_stats]
        
        # Calculate percentiles for robust normalization
        page_median = np.median(page_counts)
        word_median = np.median(word_counts)
        section_median = np.median(section_counts)
        
        # Current document metrics
        current_pages = stats['page_count']
        current_words = stats['total_words']
        current_sections = stats['section_count']
        
        # Calculate normalization factors (inverse relationship - longer docs get lower factors)
        # Using log scaling to prevent extreme penalties
        
        # 1. Page count normalization (30% weight)
        if current_pages > page_median:
            page_factor = page_median / current_pages
            page_factor = max(0.7, page_factor)  # Don't penalize too much
        else:
            page_factor = min(1.2, current_pages / page_median) if page_median > 0 else 1.0
        
        # 2. Word count normalization (40% weight)
        if current_words > word_median:
            word_factor = word_median / current_words
            word_factor = max(0.6, word_factor)  # Don't penalize too much
        else:
            word_factor = min(1.3, current_words / word_median) if word_median > 0 else 1.0
        
        # 3. Section count normalization (30% weight)
        if current_sections > section_median:
            section_factor = section_median / current_sections
            section_factor = max(0.7, section_factor)  # Don't penalize too much
        else:
            section_factor = min(1.2, current_sections / section_median) if section_median > 0 else 1.0
        
        # Weighted combination
        normalization_factor = (0.3 * page_factor) + (0.4 * word_factor) + (0.3 * section_factor)
        
        # Apply smoothing to prevent extreme values
        normalization_factor = max(0.5, min(1.5, normalization_factor))
        
        print(f"DEBUG: Normalization for {document_name}: pages={current_pages}(med={page_median:.0f}), words={current_words}(med={word_median:.0f}), sections={current_sections}(med={section_median:.0f}) -> factor={normalization_factor:.3f}")
        
        return normalization_factor

    def _calculate_title_keyword_relevance(self, title_text: str, job_keywords: List[str], title_type: str = "section") -> float:
        """
        Calculate strict keyword relevance score for either section title or document title.
        Penalizes heavily for any irrelevant content in the title.
        """
        if not job_keywords or not title_text:
            return 0.1
        
        # Clean the title text (remove file extensions for documents)
        if title_type == "document":
            title_text = re.sub(r'\.(pdf|doc|docx|txt)$', '', title_text, flags=re.IGNORECASE)
        
        # Clean and tokenize
        title_tokens = set(re.findall(r'\b[a-zA-Z]{3,}\b', title_text.lower()))
        
        # Remove stop words for cleaner analysis
        clean_title_tokens = title_tokens - self.extended_stop_words
        clean_keyword_set = set(job_keywords) - self.extended_stop_words
        
        if not clean_title_tokens:
            return 0.1  # Very low relevance for empty content
        
        # Calculate different types of relevance
        relevance_score = 1.0  # Start with perfect relevance
        
        # 1. DIRECT KEYWORD MATCHES (Positive factor)
        direct_matches = clean_title_tokens & clean_keyword_set
        match_ratio = len(direct_matches) / len(clean_title_tokens) if clean_title_tokens else 0
        
        # Boost score based on direct matches (more generous for document titles)
        if title_type == "document":
            # More lenient for document titles since they're often less specific
            if match_ratio > 0.3:
                relevance_score *= 1.15  # 15% boost for good match ratio in doc titles
            elif match_ratio > 0.2:
                relevance_score *= 1.05  # 5% boost for moderate match ratio
            elif match_ratio > 0.1:
                relevance_score *= 1.0   # No change for low matches
            else:
                relevance_score *= 0.8   # 20% reduction for no matches
        else:
            # Stricter for section titles
            if match_ratio > 0.5:
                relevance_score *= 1.2   # 20% boost for high match ratio
            elif match_ratio > 0.3:
                relevance_score *= 1.1   # 10% boost for good match ratio
            elif match_ratio > 0.1:
                relevance_score *= 1.0   # No change for moderate matches
            else:
                relevance_score *= 0.7   # 30% reduction for low matches
        
        # 2. CONTRADICTORY/IRRELEVANT TERMS (Heavy penalty)
        contradictory_terms = self._identify_contradictory_terms(clean_title_tokens, job_keywords)
        if contradictory_terms:
            # More severe penalty for section titles
            penalty_multiplier = 0.2 if title_type == "section" else 0.15  # 20% vs 15% penalty per term
            contradiction_penalty = len(contradictory_terms) * penalty_multiplier
            relevance_score *= max(0.1, 1.0 - contradiction_penalty)
            print(f"DEBUG: Found contradictory terms in {title_type} '{title_text[:30]}': {contradictory_terms}, penalty: {contradiction_penalty:.3f}")
        
        # 3. SEMANTIC RELEVANCE CHECK
        try:
            title_embedding = self.sentence_model.encode([title_text])
            keyword_embeddings = self.sentence_model.encode(job_keywords[:10])  # Limit for performance
            
            if len(keyword_embeddings) > 0:
                semantic_similarities = util.cos_sim(title_embedding, keyword_embeddings)
                avg_semantic_similarity = semantic_similarities.mean().item()
                
                # Apply semantic multiplier (more lenient for document titles)
                if title_type == "document":
                    if avg_semantic_similarity > 0.5:
                        relevance_score *= 1.05  # 5% boost for high semantic similarity
                    elif avg_semantic_similarity > 0.3:
                        relevance_score *= 1.0   # No change for moderate similarity
                    elif avg_semantic_similarity > 0.15:
                        relevance_score *= 0.9   # 10% reduction for low similarity
                    else:
                        relevance_score *= 0.7   # 30% reduction for very low similarity
                else:
                    # Stricter for section titles
                    if avg_semantic_similarity > 0.6:
                        relevance_score *= 1.1   # 10% boost for high semantic similarity
                    elif avg_semantic_similarity > 0.4:
                        relevance_score *= 1.0   # No change for moderate similarity
                    elif avg_semantic_similarity > 0.2:
                        relevance_score *= 0.8   # 20% reduction for low similarity
                    else:
                        relevance_score *= 0.5   # 50% reduction for very low similarity
                        
        except Exception as e:
            print(f"Semantic relevance calculation error for {title_type}: {e}")
        
        # 4. TITLE LENGTH PENALTY 
        title_length = len(title_text.split())
        if title_length < 2:
            relevance_score *= 0.8   # Penalty for very short titles
        elif title_length > 15:
            relevance_score *= 0.9   # Small penalty for very long titles
        
        # 5. HIGH-IMPORTANCE KEYWORD PRESENCE
        high_importance_keywords = self._get_high_importance_keywords(job_keywords)
        important_matches = clean_title_tokens & high_importance_keywords
        if important_matches:
            importance_boost = len(important_matches) * 0.1  # 10% boost per important keyword
            relevance_score *= (1.0 + importance_boost)
        
        # 6. EXACT PHRASE MATCHING BONUS
        title_text_lower = title_text.lower()
        for keyword in job_keywords:
            if len(keyword.split()) > 1 and keyword in title_text_lower:
                relevance_score *= 1.15  # 15% boost for exact phrase matches
                break
        
        return max(0.05, min(1.5, relevance_score))  # Clamp between 0.05 and 1.5

    def _identify_contradictory_terms(self, content_tokens: Set[str], job_keywords: List[str]) -> Set[str]:
        """
        Identify terms in content that contradict the job keywords.
        """
        contradictory_terms = set()
        
        # Enhanced contradictory mappings based on common job contexts
        contradictions = {
            'vegetarian': ['chicken', 'beef', 'pork', 'meat', 'fish', 'seafood', 'turkey', 'lamb', 'bacon', 'ham', 'sausage', 'steak'],
            'vegan': ['chicken', 'beef', 'pork', 'meat', 'fish', 'seafood', 'turkey', 'lamb', 'bacon', 'ham', 'cheese', 'milk', 'butter', 'egg', 'dairy', 'cream', 'yogurt'],
            'gluten': ['wheat', 'barley', 'rye', 'bread', 'pasta'] if 'free' in ' '.join(job_keywords) else [],
            'corporate': ['casual', 'informal', 'family', 'home', 'picnic', 'barbecue'],
            'buffet': ['plated', 'served', 'individual', 'personal', 'single'],
            'breakfast': ['dinner', 'lunch', 'evening'] if any('breakfast' in kw for kw in job_keywords) else [],
            'lunch': ['breakfast', 'dinner', 'morning', 'evening'] if any('lunch' in kw for kw in job_keywords) else [],
            'dinner': ['breakfast', 'lunch', 'morning'] if any('dinner' in kw for kw in job_keywords) else [],
            'healthy': ['fried', 'deep', 'greasy', 'fatty'],
            'light': ['heavy', 'rich', 'cream', 'butter'],
            'fresh': ['frozen', 'canned', 'preserved'],
        }
        
        # Check for contradictions
        for keyword in job_keywords:
            keyword_lower = keyword.lower()
            if keyword_lower in contradictions:
                contradictory_in_content = content_tokens & set(contradictions[keyword_lower])
                contradictory_terms.update(contradictory_in_content)
        
        return contradictory_terms

    def _calculate_perfect_match_bonus(self, section_text: str, document_name: str, job_keywords: List[str]) -> float:
        """
        Calculate bonus points for perfect keyword matches in section titles and document names.
        """
        if not job_keywords:
            return 0.0
        
        bonus = 0.0
        
        # Clean document name (remove extensions)
        clean_document_name = re.sub(r'\.(pdf|doc|docx|txt)$', '', document_name, flags=re.IGNORECASE)
        
        # Clean and tokenize
        section_tokens = set(re.findall(r'\b[a-zA-Z]{3,}\b', section_text.lower()))
        doc_tokens = set(re.findall(r'\b[a-zA-Z]{3,}\b', clean_document_name.lower()))
        keyword_set = set(job_keywords)
        
        # Remove stop words for cleaner matching
        clean_section_tokens = section_tokens - self.extended_stop_words
        clean_doc_tokens = doc_tokens - self.extended_stop_words
        clean_keyword_set = keyword_set - self.extended_stop_words
        
        # ===== PERFECT MATCH BONUSES =====
        
        # 1. Exact keyword matches in section title
        section_matches = clean_section_tokens & clean_keyword_set
        bonus += len(section_matches) * 0.15  # 0.15 bonus per perfect match in title
        
        # 2. Exact keyword matches in document name
        doc_matches = clean_doc_tokens & clean_keyword_set
        bonus += len(doc_matches) * 0.10  # 0.10 bonus per perfect match in filename
        
        # 3. High-importance keyword perfect matches (extra bonus)
        high_importance_keywords = self._get_high_importance_keywords(job_keywords)
        section_important_matches = clean_section_tokens & high_importance_keywords
        doc_important_matches = clean_doc_tokens & high_importance_keywords
        
        bonus += len(section_important_matches) * 0.20  # Extra bonus for important keywords in title
        bonus += len(doc_important_matches) * 0.15     # Extra bonus for important keywords in filename
        
        # 4. Multiple keyword phrase matches
        section_text_lower = section_text.lower()
        document_name_lower = clean_document_name.lower()
        
        for keyword in job_keywords:
            if len(keyword.split()) > 1:  # Multi-word keywords
                if keyword in section_text_lower:
                    bonus += 0.25  # Big bonus for multi-word phrase match in title
                if keyword in document_name_lower:
                    bonus += 0.20  # Big bonus for multi-word phrase match in filename
        
        # 5. Keyword density bonus for section titles
        if clean_section_tokens:
            keyword_density = len(section_matches) / len(clean_section_tokens)
            if keyword_density > 0.5:  # More than 50% of title words are keywords
                bonus += 0.10
            elif keyword_density > 0.3:  # More than 30% of title words are keywords
                bonus += 0.05
        
        # 6. Document title keyword density bonus
        if clean_doc_tokens:
            doc_keyword_density = len(doc_matches) / len(clean_doc_tokens)
            if doc_keyword_density > 0.4:  # More than 40% of document name words are keywords
                bonus += 0.08
            elif doc_keyword_density > 0.2:  # More than 20% of document name words are keywords
                bonus += 0.04
        
        # 7. Complete keyword coverage bonus
        important_coverage = len((clean_section_tokens | clean_doc_tokens) & high_importance_keywords)
        total_important = len(high_importance_keywords)
        if total_important > 0:
            coverage_ratio = important_coverage / total_important
            if coverage_ratio > 0.7:  # Contains >70% of important keywords
                bonus += 0.15
            elif coverage_ratio > 0.5:  # Contains >50% of important keywords
                bonus += 0.10
        
        return min(bonus, 0.8)  # Cap maximum bonus at 0.8

    def _calculate_keyword_similarity_penalty(self, section_text: str, document_name: str, job_keywords: List[str]) -> float:
        """
        Calculate penalty for sections that don't match keywords (reduced since we have relevance scoring).
        """
        if not job_keywords:
            return 0.0
        
        # Clean document name
        clean_document_name = re.sub(r'\.(pdf|doc|docx|txt)$', '', document_name, flags=re.IGNORECASE)
            
        # Tokenize and clean section title and document name separately
        section_tokens = set(re.findall(r'\b[a-zA-Z]{3,}\b', section_text.lower()))
        doc_tokens = set(re.findall(r'\b[a-zA-Z]{3,}\b', clean_document_name.lower()))
        
        # Remove stop words from tokens
        clean_section_tokens = section_tokens - self.extended_stop_words
        clean_doc_tokens = doc_tokens - self.extended_stop_words
        keyword_set = set(job_keywords) - self.extended_stop_words
        
        if not clean_section_tokens and not clean_doc_tokens:
            return 0.15  # Reduced penalty since we have relevance scoring
        
        # Calculate separate overlaps for section title and document name
        section_overlap = clean_section_tokens & keyword_set
        doc_overlap = clean_doc_tokens & keyword_set
        
        # Calculate overlap ratios
        section_overlap_ratio = len(section_overlap) / len(clean_section_tokens) if clean_section_tokens else 0
        doc_overlap_ratio = len(doc_overlap) / len(clean_doc_tokens) if clean_doc_tokens else 0
        
        # REDUCED PENALTY SYSTEM (since relevance scoring handles most of this)
        penalty = 0.0
        
        # 1. Section Title Penalty (Reduced)
        if section_overlap_ratio == 0:
            penalty += 0.05  # Reduced from higher values
        elif section_overlap_ratio < 0.15:
            penalty += 0.03
        
        # 2. Document Name Penalty (Reduced)
        if doc_overlap_ratio == 0:
            penalty += 0.04  # Reduced from higher values
        elif doc_overlap_ratio < 0.15:
            penalty += 0.02
        
        # 3. Additional penalty if BOTH title and filename have no matches
        if section_overlap_ratio == 0 and doc_overlap_ratio == 0:
            penalty += 0.06  # Reduced from higher values
        
        return min(penalty, 0.2)  # Cap maximum penalty at 0.2

    def _get_high_importance_keywords(self, job_keywords: List[str]) -> Set[str]:
        """
        Identify high-importance keywords that should strongly influence bonuses/penalties.
        """
        high_importance_patterns = [
            'vegetarian', 'vegan', 'plant', 'buffet', 'gluten', 'free', 
            'corporate', 'catering', 'dinner', 'menu', 'food', 'meal',
            'healthy', 'fresh', 'organic', 'breakfast', 'lunch'
        ]
        
        high_importance = set()
        for keyword in job_keywords:
            # Check if keyword contains any high-importance pattern
            if any(pattern in keyword.lower() for pattern in high_importance_patterns):
                high_importance.add(keyword)
            # Also include longer keywords (likely more specific/important)
            elif len(keyword) > 6:
                high_importance.add(keyword)
        
        return high_importance

    def _extract_comprehensive_content(self, full_text: str, section_title: str, max_length: int = 2500) -> str:
        """
        Extract comprehensive content following a section title with better context understanding.
        """
        if not section_title or not full_text: return ""
        
        content_lines = []
        try:
            lines = full_text.split('\n')
            in_section = False
            section_found = False
            
            # Find the section with fuzzy matching
            for i, line in enumerate(lines):
                stripped_line = line.strip()
                
                # Check for section match (exact or fuzzy)
                if not section_found:
                    if (section_title.lower() in stripped_line.lower() or 
                        stripped_line.lower() in section_title.lower() or
                        self._fuzzy_match(section_title.lower(), stripped_line.lower())):
                        in_section = True
                        section_found = True
                        continue
                
                if in_section:
                    # Stop if we hit another heading
                    if stripped_line and len(content_lines) > 5 and self._is_likely_heading(line):
                        break
                    
                    # Add meaningful content
                    if stripped_line and not self._is_noise_line(stripped_line):
                        content_lines.append(stripped_line)
                    
                    # Stop if we have enough content
                    if len(' '.join(content_lines)) >= max_length:
                        break
                        
        except Exception as e:
            print(f"Error extracting comprehensive content: {e}")
        
        # Join and clean up the content
        content = ' '.join(content_lines)
        return self._clean_extracted_content(content)[:max_length]

    def _fuzzy_match(self, text1: str, text2: str, threshold: float = 0.7) -> bool:
        """Simple fuzzy matching for section titles."""
        words1 = set(text1.split())
        words2 = set(text2.split())
        if not words1 or not words2: return False
        
        intersection = words1.intersection(words2)
        union = words1.union(words2)
        return len(intersection) / len(union) >= threshold

    def _is_noise_line(self, line: str) -> bool:
        """Check if a line is noise (page numbers, headers, etc.)."""
        line = line.strip()
        if len(line) < 3: return True
        
        # Common noise patterns
        noise_patterns = [
            r'^\d+$',  # Just numbers
            r'^page\s+\d+',  # Page numbers
            r'^\d+\s*$',  # Isolated numbers
            r'^copyright',  # Copyright notices
            r'^\w{1,3}$',  # Very short words
        ]
        
        return any(re.match(pattern, line.lower()) for pattern in noise_patterns)

    def _clean_extracted_content(self, content: str) -> str:
        """Clean and format extracted content."""
        if not content: return ""
        
        # Remove excessive whitespace
        content = re.sub(r'\s+', ' ', content)
        
        # Remove all types of bullet points and Unicode characters
        bullet_patterns = [
            r'^\s*[•·▪▫‣⁃]\s*',  # Standard bullet points
            r'^\s*\u2022\s*',     # Unicode bullet point
            r'^\s*\u2023\s*',     # Triangular bullet
            r'^\s*\u25E6\s*',     # White bullet
            r'^\s*\u2043\s*',     # Hyphen bullet
            r'^\s*[-*]\s*',       # Dash and asterisk bullets
        ]
        
        for pattern in bullet_patterns:
            content = re.sub(pattern, '', content, flags=re.MULTILINE)
        
        # Remove other common Unicode artifacts
        content = content.replace('\u2022', '').replace('\u2023', '').replace('\u25E6', '')
        
        # Clean up multiple spaces created by bullet removal
        content = re.sub(r'\s+', ' ', content)
        
        # Ensure proper sentence structure
        sentences = content.split('.')
        cleaned_sentences = []
        
        for sentence in sentences:
            sentence = sentence.strip()
            if len(sentence) > 10:  # Meaningful sentences only
                if sentence and not sentence[0].isupper():
                    sentence = sentence.capitalize()
                cleaned_sentences.append(sentence)
        
        return '. '.join(cleaned_sentences) + '.' if cleaned_sentences else content

    def _is_likely_heading(self, line: str) -> bool:
        """Enhanced heuristics to identify if a line is a potential heading."""
        line = line.strip()
        if not line or len(line) > 100 or len(line.split()) > 12: return False
        
        # Skip lines with problematic punctuation
        if any(char in line for char in [',', '?', ';', ':', '!']):
            return False
            
        if re.match(r'^\d+(\.\d+)*\s+', line): return True
        if line.isupper() and len(line.split()) < 6: return True
        if self._is_title_case(line) and not line.endswith('.'): return True
        return False

    def process_documents(self, input_data: Dict, collection_path: str = None) -> Dict:
        """
        Enhanced main processing function with document length normalization.
        """
        documents = input_data.get('documents', [])
        persona = input_data.get('persona', {}).get('role', 'Analyst')
        job_task = input_data.get('job_to_be_done', {}).get('task', 'Analyze documents')

        result = {
            "metadata": {
                "input_documents": [doc.get('filename') for doc in documents], 
                "persona": persona,
                "job_to_be_done": job_task, 
                "processing_timestamp": datetime.now().isoformat()
            },
            "extracted_sections": [],
            "comprehensive_analysis": []
        }

        job_keywords = self.generate_job_context_keywords(persona, job_task)
        if not job_keywords:
            job_keywords = [p.lower() for p in persona.split()] + [t.lower() for t in job_task.split()]
        print(f"Generated {len(job_keywords)} job keywords. Top 15: {job_keywords[:15]}")

        all_scored_sections = []
        document_texts = {}
        
        # First pass: extract all documents and collect statistics
        for doc in documents:
            filename = doc.get('filename', '')
            if not filename:
                print(f"Warning: No filename provided, skipping document")
                continue
            
            # Look for PDF in the current collection's PDFs folder  
            if collection_path:
                pdf_path = os.path.join(collection_path, 'PDFs', filename)
            else:
                # Fallback: search in common locations
                pdf_path = None
                possible_paths = [
                    os.path.join('Challenge_1b', 'Collection 1', 'PDFs', filename),
                    os.path.join('Challenge_1b', 'Collection 2', 'PDFs', filename),
                    os.path.join('Challenge_1b', 'Collection 3', 'PDFs', filename),
                ]
                
                for path in possible_paths:
                    if os.path.exists(path):
                        pdf_path = path
                        break
            
            if not pdf_path or not os.path.exists(pdf_path):
                print(f"Warning: File not found in collection folder, skipping: {filename}")
                continue
            
            try:
                full_text, sections = self.extract_text_from_pdf(pdf_path)
                if not full_text or not sections: 
                    continue
                document_texts[filename] = full_text
                print(f"DEBUG: Extracted {len(sections)} sections from {filename}")
                
                # Store sections for second pass
                for section in sections:
                    section['document'] = filename
                all_scored_sections.extend(sections)

            except Exception as e:
                print(f"Error processing {filename}: {e}")

        # Second pass: filter with normalization (now that we have all document stats)
        print(f"DEBUG: Processing {len(all_scored_sections)} total sections across {len(self.document_stats)} documents")
        
        # Group sections by document for filtering
        sections_by_doc = {}
        for section in all_scored_sections:
            doc_name = section.get('document', 'Unknown')
            if doc_name not in sections_by_doc:
                sections_by_doc[doc_name] = []
            sections_by_doc[doc_name].append(section)
        
        # Filter sections for each document
        final_scored_sections = []
        for doc_name, doc_sections in sections_by_doc.items():
            relevant_sections = self.filter_relevant_sections(doc_sections, job_keywords)
            print(f"Found {len(relevant_sections)} relevant sections in {doc_name}")
            final_scored_sections.extend(relevant_sections)
        
        # Sort all sections by final normalized score
        final_scored_sections.sort(key=lambda x: x.get('similarity_score', 0), reverse=True)

        # Generate comprehensive responses for top sections
        for rank, section in enumerate(final_scored_sections[:5], 1):
            doc_filename = section.get('document', 'Unknown')
            page_num = section.get('page_number', 1)
            
            result["extracted_sections"].append({
                "document": doc_filename,
                "section_title": section.get('text', 'Untitled Section'),
                "importance_rank": rank,
                "page_number": page_num,
                "similarity_score": round(section.get('similarity_score', 0), 4),
                "pre_normalized_score": round(section.get('pre_normalized_score', 0), 4),
                "normalization_factor": round(section.get('normalization_factor', 1), 4),
                "tfidf_score": round(section.get('tfidf_score', 0), 4),
                "semantic_score": round(section.get('semantic_score', 0), 4),
                "individual_score": round(section.get('individual_score', 0), 4),
                "section_title_relevance": round(section.get('section_title_relevance', 0), 4),
                "document_title_relevance": round(section.get('document_title_relevance', 0), 4),
                "combined_relevance_score": round(section.get('combined_relevance_score', 0), 4),
                "perfect_match_bonus": round(section.get('perfect_match_bonus', 0), 4),
                "keyword_penalty": round(section.get('keyword_penalty', 0), 4),
                "base_similarity": round(section.get('base_similarity', 0), 4),
                "max_individual_sim": round(section.get('max_individual_sim', 0), 4),
                "section_type": section.get('type', 'unknown')
            })

            # Extract comprehensive content
            full_doc_text = document_texts.get(doc_filename, '')
            comprehensive_content = self._extract_comprehensive_content(
                full_doc_text, 
                section.get('text', ''),
                max_length=3000
            )
            
            if comprehensive_content and len(comprehensive_content) > 50:
                comprehensive_content = self._final_content_cleanup(comprehensive_content)
                
                result["comprehensive_analysis"].append({
                    "document": doc_filename,
                    "section_title": section.get('text', 'Untitled Section'),
                    "comprehensive_content": comprehensive_content,
                    "page_number": page_num,
                    "relevance_score": round(section.get('similarity_score', 0), 4),
                    "content_length": len(comprehensive_content)
                })
        
        # Add document statistics to result for transparency
        result["document_statistics"] = self.document_stats
        
        return result

    def _final_content_cleanup(self, content: str) -> str:
        """Final cleanup to ensure no artifacts remain in the content."""
        if not content:
            return ""
        
        # Remove any remaining bullet points and Unicode artifacts
        content = re.sub(r'[•·▪▫‣⁃\u2022\u2023\u25E6\u2043]', '', content)
        
        # Remove multiple spaces
        content = re.sub(r'\s+', ' ', content)
        
        # Remove leading/trailing whitespace
        content = content.strip()
        
        # Ensure proper sentence flow
        sentences = content.split('. ')
        cleaned_sentences = []
        
        for sentence in sentences:
            sentence = sentence.strip()
            if len(sentence) > 5:  # Keep meaningful sentences
                # Ensure proper capitalization
                if sentence and not sentence[0].isupper():
                    sentence = sentence.capitalize()
                # Remove any remaining bullet artifacts at sentence start
                sentence = re.sub(r'^[•·▪▫‣⁃\-\*]\s*', '', sentence)
                if sentence:  # Only add if still has content
                    cleaned_sentences.append(sentence)
        
        result = '. '.join(cleaned_sentences)
        if result and not result.endswith('.'):
            result += '.'
            
        return result


def main():
    """Main function that processes all collections from input directory and outputs to output directory"""
    # Define input and output directories for Docker environment
    input_dir = "/app/input"
    output_dir = "/app/output"
    
    try:
        # Check if input directory exists
        if not os.path.exists(input_dir):
            print(f"Error: Input directory {input_dir} not found!")
            print("Please ensure the input directory exists with collections.")
            return
        
        # Create output directory with proper permissions
        os.makedirs(output_dir, exist_ok=True)
        # Ensure we have write permissions
        try:
            test_file = os.path.join(output_dir, '.test')
            with open(test_file, 'w') as f:
                f.write('test')
            os.remove(test_file)
        except PermissionError:
            print(f"Error: No write permission for output directory {output_dir}")
            print("Please ensure the output directory has proper write permissions")
            return
        
        # Find all collection folders
        collection_folders = []
        for item in os.listdir(input_dir):
            item_path = os.path.join(input_dir, item)
            if os.path.isdir(item_path) and item.startswith("Collection"):
                collection_folders.append(item)
        
        if not collection_folders:
            print(f"No collection folders found in {input_dir}!")
            return
        
        collection_folders.sort()  # Sort to process in order
        print(f"Found {len(collection_folders)} collections: {collection_folders}")
        
        # Initialize analyzer once (reused for all collections)
        print("Initializing PDF analyzer with document length normalization...")
        analyzer = JobAwarePDFAnalyzer()
        
        # Process each collection
        for collection_name in collection_folders:
            print(f"\n{'='*60}")
            print(f"Processing {collection_name}...")
            print(f"{'='*60}")
            
            collection_path = os.path.join(input_dir, collection_name)
            input_file = os.path.join(collection_path, "challenge1b_input.json")
            
            # Create output paths for both collection directory and output directory
            collection_output_dir = os.path.join(collection_path, "challenge1b_output.json")
            docker_output_file = os.path.join(output_dir, f"{collection_name}_output.json")
            
            # Check if input file exists
            if not os.path.exists(input_file):
                print(f"Warning: {input_file} not found, skipping {collection_name}")
                continue
            
            # Check if PDFs folder exists
            pdfs_folder = os.path.join(collection_path, "PDFs")
            if not os.path.exists(pdfs_folder):
                print(f"Warning: PDFs folder not found in {collection_name}, skipping")
                continue
            
            try:
                # Load input data
                with open(input_file, 'r', encoding='utf-8') as f:
                    input_data = json.load(f)
                
                print(f"Successfully loaded input from {input_file}")
                
                # Validate that documents exist in PDFs folder
                missing_files = []
                for doc in input_data.get('documents', []):
                    filename = doc.get('filename', '')
                    if filename:
                        pdf_path = os.path.join(pdfs_folder, filename)
                        if not os.path.exists(pdf_path):
                            missing_files.append(filename)
                
                if missing_files:
                    print(f"Warning: The following files were not found in {collection_name}/PDFs/ folder:")
                    for file in missing_files:
                        print(f"  - {file}")
                    print("Processing will continue with available files...")
                
                # Process documents
                print(f"Processing documents in {collection_name}...")
                result = analyzer.process_documents(input_data, collection_path)
                
                # Add collection metadata
                result["collection_info"] = {
                    "collection_name": collection_name,
                    "collection_path": collection_path,
                    "processing_timestamp": datetime.now().isoformat()
                }
                
                # Write output to both collection directory and output directory
                for output_file in [collection_output_dir, docker_output_file]:
                    try:
                        # Create parent directory if it doesn't exist
                        os.makedirs(os.path.dirname(output_file), exist_ok=True)
                        
                        # Try to write the file
                        with open(output_file, 'w', encoding='utf-8') as f:
                            json.dump(result, f, indent=2, ensure_ascii=False)
                        print(f"Successfully wrote output to {output_file}")
                    except PermissionError as e:
                        print(f"Permission error writing to {output_file}: {e}")
                        print("Skipping this output file and continuing...")
                        continue
                    except Exception as e:
                        print(f"Error writing to {output_file}: {e}")
                        print("Skipping this output file and continuing...")
                        continue
                
                print(f"✓ Processing complete for {collection_name}!")
                print(f"  Found {len(result['extracted_sections'])} relevant sections")
                print(f"  Generated {len(result['comprehensive_analysis'])} comprehensive analyses")
                
            except FileNotFoundError as e:
                print(f"File error in {collection_name}: {e}")
            except json.JSONDecodeError as e:
                print(f"JSON parsing error in {input_file}: {e}")
            except Exception as e:
                print(f"Error during processing {collection_name}: {e}")
                import traceback
                traceback.print_exc()
        
        print(f"\n{'='*60}")
        print("All collections processed!")
        print(f"{'='*60}")
        
    except Exception as e:
        print(f"Error during processing: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()