"""
NLP Comprehensive Analyzer - Main Module

This module contains the ComprehensiveNLP class that provides all NLP functionality
including text preprocessing, semantic analysis, and information extraction.
"""

import nltk
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from collections import Counter
from nltk.tokenize import word_tokenize, sent_tokenize
from nltk.corpus import stopwords, wordnet
from nltk.stem import WordNetLemmatizer, PorterStemmer
from nltk.tag import pos_tag
from nltk.chunk import ne_chunk
from nltk.util import ngrams
from nltk.probability import LaplaceProbDist, FreqDist
import math
import warnings
warnings.filterwarnings('ignore')


class ComprehensiveNLP:
    """
    A comprehensive NLP analyzer that covers three main areas:
    
    1. Text Preprocessing & Representation (Q1)
    2. Semantic Understanding & Language Modeling (Q2)  
    3. Information Extraction & Sentiment Analysis (Q3)
    
    Perfect for educational purposes and NLP learning.
    """
    
    def __init__(self, download_data=True):
        """
        Initialize NLP tools and resources
        
        Args:
            download_data (bool): Whether to download NLTK data automatically
        """
        if download_data:
            self._download_nltk_data()
        
        self.lemmatizer = WordNetLemmatizer()
        self.stemmer = PorterStemmer()
        
        try:
            self.stop_words = set(stopwords.words('english'))
        except:
            # Fallback stopwords if NLTK data not available
            self.stop_words = {
                'i', 'me', 'my', 'myself', 'we', 'our', 'ours', 'ourselves', 'you', 'your', 
                'yours', 'yourself', 'yourselves', 'he', 'him', 'his', 'himself', 'she', 
                'her', 'hers', 'herself', 'it', 'its', 'itself', 'they', 'them', 'their', 
                'theirs', 'themselves', 'what', 'which', 'who', 'whom', 'this', 'that', 
                'these', 'those', 'am', 'is', 'are', 'was', 'were', 'be', 'been', 'being', 
                'have', 'has', 'had', 'having', 'do', 'does', 'did', 'doing', 'a', 'an', 
                'the', 'and', 'but', 'if', 'or', 'because', 'as', 'until', 'while', 'of', 
                'at', 'by', 'for', 'with', 'through', 'during', 'before', 'after', 'above', 
                'below', 'up', 'down', 'in', 'out', 'on', 'off', 'over', 'under', 'again', 
                'further', 'then', 'once'
            }
        
        # Default sample texts for demonstration
        self.sample_texts = [
            "I feel overwhelmed with anxiety and stress lately. Need professional help.",
            "The therapy sessions have been very helpful for my mental health recovery.",
            "Depression affects many students during exam periods. Support is crucial.",
            "Meditation and exercise help reduce stress and improve wellbeing significantly.",
            "Mental health awareness campaigns are important in universities and colleges.",
            "Counseling services provide excellent support for students facing difficulties."
        ]
        
        print("✓ NLP Comprehensive Analyzer initialized successfully!")
    
    def _download_nltk_data(self):
        """Download required NLTK data"""
        required_data = [
            'punkt', 'stopwords', 'wordnet', 'averaged_perceptron_tagger',
            'maxent_ne_chunker', 'words', 'omw-1.4'
        ]
        
        print("Downloading NLTK data...")
        for data in required_data:
            try:
                nltk.download(data, quiet=True)
            except Exception as e:
                print(f"Warning: Could not download {data}: {e}")
        print("✓ NLTK data download complete!")
    
    # ========== Q1: TEXT PREPROCESSING & REPRESENTATION ==========
    
    def preprocess_text(self, texts=None):
        """
        Complete text preprocessing pipeline
        
        Args:
            texts (list): List of texts to process. If None, uses sample texts.
            
        Returns:
            tuple: (processed_data, all_tokens)
        """
        if texts is None:
            texts = self.sample_texts
        
        print("\n" + "="*60)
        print("Q1: TEXT PREPROCESSING & REPRESENTATION")
        print("="*60)
        
        processed_data = []
        all_tokens = []
        
        for i, text in enumerate(texts, 1):
            print(f"\nProcessing Text {i}: {text[:50]}...")
            
            # Step 1: Tokenization
            tokens = word_tokenize(text.lower())
            
            # Step 2: Remove punctuation and stopwords
            filtered_tokens = [word for word in tokens 
                             if word.isalpha() and word not in self.stop_words]
            
            # Step 3: Lemmatization
            lemmatized = [self.lemmatizer.lemmatize(word) for word in filtered_tokens]
            
            # Step 4: Stemming (for comparison)
            stemmed = [self.stemmer.stem(word) for word in filtered_tokens]
            
            processed_data.append({
                'original': text,
                'tokens': tokens,
                'filtered': filtered_tokens,
                'lemmatized': lemmatized,
                'stemmed': stemmed
            })
            
            all_tokens.extend(lemmatized)
            
            print(f"  Original length: {len(tokens)}")
            print(f"  After filtering: {len(filtered_tokens)}")
            print(f"  Lemmatized sample: {lemmatized[:5]}")
        
        # Text Representation
        self.create_text_representations(processed_data, all_tokens)
        return processed_data, all_tokens
    
    def create_text_representations(self, processed_data, all_tokens):
        """Create BoW and TF-IDF representations"""
        print(f"\n--- TEXT REPRESENTATION METHODS ---")
        
        # Create vocabulary
        vocabulary = sorted(set(all_tokens))
        print(f"Vocabulary size: {len(vocabulary)}")
        
        # Bag of Words
        bow_matrix = []
        for data in processed_data:
            bow_vector = [data['lemmatized'].count(word) for word in vocabulary]
            bow_matrix.append(bow_vector)
        
        print(f"Bag-of-Words matrix shape: {len(bow_matrix)} x {len(bow_matrix[0])}")
        
        # TF-IDF
        tfidf_matrix = self.calculate_tfidf(processed_data, vocabulary)
        print(f"TF-IDF matrix shape: {len(tfidf_matrix)} x {len(tfidf_matrix[0])}")
        
        # Show top words by TF-IDF
        word_scores = {}
        for i, word in enumerate(vocabulary):
            total_score = sum(doc[i] for doc in tfidf_matrix)
            word_scores[word] = total_score
        
        top_words = sorted(word_scores.items(), key=lambda x: x[1], reverse=True)[:5]
        print(f"Top 5 words by TF-IDF: {[word for word, _ in top_words]}")
        
        return bow_matrix, tfidf_matrix, vocabulary
    
    def calculate_tfidf(self, processed_data, vocabulary):
        """Calculate TF-IDF scores"""
        num_docs = len(processed_data)
        tfidf_matrix = []
        
        # Document frequency
        doc_freq = {}
        for word in vocabulary:
            doc_freq[word] = sum(1 for data in processed_data 
                               if word in data['lemmatized'])
        
        for data in processed_data:
            doc_tokens = data['lemmatized']
            doc_length = len(doc_tokens)
            
            tfidf_vector = []
            for word in vocabulary:
                # TF
                tf = doc_tokens.count(word) / doc_length if doc_length > 0 else 0
                # IDF
                idf = math.log(num_docs / (doc_freq[word] + 1))
                # TF-IDF
                tfidf = tf * idf
                tfidf_vector.append(tfidf)
            
            tfidf_matrix.append(tfidf_vector)
        
        return tfidf_matrix
    
    # ========== Q2: SEMANTIC UNDERSTANDING & LANGUAGE MODELING ==========
    
    def semantic_analysis(self, tokens):
        """
        WordNet analysis for semantic understanding
        
        Args:
            tokens (list): List of tokens to analyze
            
        Returns:
            dict: Semantic analysis results
        """
        print("\n" + "="*60)
        print("Q2: SEMANTIC UNDERSTANDING & LANGUAGE MODELING")
        print("="*60)
        
        # Analyze key words
        key_words = ['stress', 'help', 'support', 'health', 'therapy']
        
        print(f"\n--- WORDNET SEMANTIC ANALYSIS ---")
        semantic_data = {}
        
        for word in key_words:
            if word in tokens:
                try:
                    synsets = wordnet.synsets(word)
                    if synsets:
                        main_synset = synsets[0]
                        
                        # Get synonyms
                        synonyms = set()
                        antonyms = set()
                        for syn in synsets[:2]:  # First 2 synsets
                            for lemma in syn.lemmas():
                                synonyms.add(lemma.name().replace('_', ' '))
                                for ant in lemma.antonyms():
                                    antonyms.add(ant.name().replace('_', ' '))
                        
                        # Get hypernyms
                        hypernyms = [hyp.name().split('.')[0] for hyp in main_synset.hypernyms()[:3]]
                        
                        semantic_data[word] = {
                            'definition': main_synset.definition(),
                            'synonyms': list(synonyms)[:5],
                            'antonyms': list(antonyms)[:3],
                            'hypernyms': hypernyms
                        }
                        
                        print(f"\n{word.upper()}:")
                        print(f"  Definition: {semantic_data[word]['definition']}")
                        print(f"  Synonyms: {semantic_data[word]['synonyms']}")
                        if semantic_data[word]['antonyms']:
                            print(f"  Antonyms: {semantic_data[word]['antonyms']}")
                        print(f"  Hypernyms: {semantic_data[word]['hypernyms']}")
                except Exception as e:
                    print(f"Error analyzing {word}: {e}")
        
        # Language Modeling
        self.build_language_model(tokens)
        return semantic_data
    
    def build_language_model(self, tokens):
        """Build N-gram model with Laplace smoothing"""
        print(f"\n--- N-GRAM LANGUAGE MODEL WITH LAPLACE SMOOTHING ---")
        
        # Add sentence boundaries
        tokens_with_boundaries = ['<START>'] + tokens + ['<END>']
        
        # Create bigrams
        bigrams = list(ngrams(tokens_with_boundaries, 2))
        unigrams = list(ngrams(tokens_with_boundaries, 1))
        
        # Count frequencies
        bigram_counts = Counter(bigrams)
        unigram_counts = Counter(unigrams)
        
        print(f"Total bigrams: {len(bigrams)}")
        print(f"Unique bigrams: {len(bigram_counts)}")
        print(f"Vocabulary size: {len(unigram_counts)}")
        
        # Show sample bigrams
        print(f"\nTop 5 bigrams:")
        for bigram, count in bigram_counts.most_common(5):
            print(f"  {bigram}: {count}")
        
        # Laplace smoothing function
        def laplace_probability(w1, w2, alpha=1):
            bigram_count = bigram_counts.get((w1, w2), 0)
            unigram_w1_count = unigram_counts.get((w1,), 0)
            vocab_size = len(unigram_counts)
            
            return (bigram_count + alpha) / (unigram_w1_count + alpha * vocab_size)
        
        # Test predictions
        test_words = ['stress', 'mental', 'support']
        predict_words = ['help', 'health', 'therapy']
        
        print(f"\nLaplace Smoothing Predictions:")
        for w1 in test_words:
            if (w1,) in unigram_counts:
                best_prob = 0
                best_word = None
                
                for w2 in predict_words:
                    prob = laplace_probability(w1, w2)
                    print(f"  P({w2}|{w1}) = {prob:.6f}")
                    
                    if prob > best_prob:
                        best_prob = prob
                        best_word = w2
                
                if best_word:
                    print(f"  → Best prediction: '{w1} {best_word}' ({best_prob:.6f})")
                print()
        
        return bigram_counts, unigram_counts
    
    # ========== Q3: INFORMATION EXTRACTION & SENTIMENT ANALYSIS ==========
    
    def information_extraction_and_sentiment(self, texts=None):
        """
        NER, Sentiment Analysis, and Word Embeddings
        
        Args:
            texts (list): List of texts to analyze. If None, uses sample texts.
            
        Returns:
            tuple: (entities, sentiments)
        """
        if texts is None:
            texts = self.sample_texts
        
        print("\n" + "="*60)
        print("Q3: INFORMATION EXTRACTION & SENTIMENT ANALYSIS")
        print("="*60)
        
        # Named Entity Recognition
        print(f"\n--- NAMED ENTITY RECOGNITION ---")
        all_entities = {'PERSON': [], 'ORGANIZATION': [], 'GPE': []}
        
        for i, text in enumerate(texts, 1):
            tokens = word_tokenize(text)
            pos_tagged = pos_tag(tokens)
            
            try:
                entities = ne_chunk(pos_tagged)
                
                for subtree in entities:
                    if hasattr(subtree, 'label'):
                        entity_name = ' '.join([token for token, pos in subtree.leaves()])
                        entity_label = subtree.label()
                        
                        if entity_label in all_entities:
                            all_entities[entity_label].append(entity_name)
                
                print(f"Text {i}: Found {len([e for e in entities if hasattr(e, 'label')])} entities")
            
            except Exception as e:
                print(f"Text {i}: Using simple entity detection - {e}")
                # Simple fallback entity detection
                entities = self.simple_entity_detection(text)
                all_entities['ORGANIZATION'].extend(entities.get('organizations', []))
        
        print(f"All entities found:")
        for entity_type, entities in all_entities.items():
            if entities:
                print(f"  {entity_type}: {list(set(entities))}")
        
        # Sentiment Analysis
        sentiments = self.analyze_sentiment(texts)
        
        # Word Similarity (Simple implementation)
        self.word_similarity_analysis(texts)
        
        return all_entities, sentiments
    
    def simple_entity_detection(self, text):
        """Simple entity detection fallback"""
        organizations = ['therapy', 'university', 'college', 'hospital', 'clinic']
        found_orgs = [org for org in organizations if org.lower() in text.lower()]
        return {'organizations': found_orgs}
    
    def analyze_sentiment(self, texts):
        """Rule-based sentiment analysis"""
        print(f"\n--- SENTIMENT ANALYSIS ---")
        
        positive_words = {
            'helpful', 'excellent', 'good', 'better', 'improve', 'support', 
            'successful', 'effective', 'positive', 'recovery', 'wellbeing'
        }
        
        negative_words = {
            'overwhelmed', 'anxiety', 'stress', 'depression', 'difficult', 
            'problem', 'trouble', 'worried', 'sad', 'negative'
        }
        
        sentiments = []
        
        for i, text in enumerate(texts, 1):
            tokens = word_tokenize(text.lower())
            
            positive_score = sum(1 for word in tokens if word in positive_words)
            negative_score = sum(1 for word in tokens if word in negative_words)
            
            if positive_score > negative_score:
                sentiment = "Positive"
            elif negative_score > positive_score:
                sentiment = "Negative"
            else:
                sentiment = "Neutral"
            
            sentiments.append(sentiment)
            print(f"Text {i}: {sentiment} (pos:{positive_score}, neg:{negative_score})")
        
        # Sentiment distribution
        sentiment_counts = Counter(sentiments)
        print(f"\nSentiment Distribution:")
        for sentiment, count in sentiment_counts.items():
            percentage = (count / len(sentiments)) * 100
            print(f"  {sentiment}: {count} texts ({percentage:.1f}%)")
        
        return sentiments
    
    def word_similarity_analysis(self, texts):
        """Simple word similarity using co-occurrence"""
        print(f"\n--- WORD SIMILARITY ANALYSIS ---")
        
        # Target words for similarity
        target_words = ['stress', 'help', 'therapy', 'support']
        
        # Create co-occurrence matrix
        all_tokens = []
        for text in texts:
            tokens = word_tokenize(text.lower())
            filtered = [word for word in tokens if word.isalpha() and word not in self.stop_words]
            all_tokens.extend(filtered)
        
        # Calculate co-occurrence
        window_size = 3
        cooccurrence = {word: Counter() for word in target_words}
        
        for i, word in enumerate(all_tokens):
            if word in target_words:
                start = max(0, i - window_size)
                end = min(len(all_tokens), i + window_size + 1)
                
                for j in range(start, end):
                    if j != i and all_tokens[j] in target_words:
                        cooccurrence[word][all_tokens[j]] += 1
        
        print(f"Word Co-occurrence Analysis:")
        for word in target_words:
            if cooccurrence[word]:
                top_cooccur = cooccurrence[word].most_common(3)
                print(f"  {word}: co-occurs with {dict(top_cooccur)}")
        
        # Simple similarity using Jaccard
        print(f"\nPairwise Similarity (Jaccard):")
        for i in range(len(target_words)):
            for j in range(i + 1, len(target_words)):
                word1, word2 = target_words[i], target_words[j]
                
                set1 = set(cooccurrence[word1].keys())
                set2 = set(cooccurrence[word2].keys())
                
                if set1 and set2:
                    intersection = len(set1.intersection(set2))
                    union = len(set1.union(set2))
                    similarity = intersection / union if union > 0 else 0
                    print(f"  {word1} - {word2}: {similarity:.3f}")
        
        return cooccurrence
    
    def run_complete_analysis(self, texts=None):
        """
        Run all analyses - the main method to use
        
        Args:
            texts (list): List of texts to analyze. If None, uses sample texts.
            
        Returns:
            dict: Complete analysis results
        """
        if texts is None:
            texts = self.sample_texts
        
        print("COMPREHENSIVE NLP ANALYSIS - ALL TASKS")
        print("Mental Health Text Analysis Domain")
        print("="*80)
        
        # Q1: Text Preprocessing & Representation
        processed_data, all_tokens = self.preprocess_text(texts)
        
        # Q2: Semantic Understanding & Language Modeling
        semantic_data = self.semantic_analysis(all_tokens)
        
        # Q3: Information Extraction & Sentiment Analysis
        entities, sentiments = self.information_extraction_and_sentiment(texts)
        
        # Summary
        print(f"\n" + "="*60)
        print("ANALYSIS SUMMARY")
        print("="*60)
        print(f"✓ Processed {len(texts)} texts")
        print(f"✓ Total vocabulary: {len(set(all_tokens))} unique words")
        print(f"✓ Semantic analysis: {len(semantic_data)} key words analyzed")
        print(f"✓ Entities extracted: {sum(len(ents) for ents in entities.values())} total")
        print(f"✓ Sentiment distribution: {Counter(sentiments)}")
        print(f"✓ All NLP tasks completed successfully!")
        
        return {
            'processed_data': processed_data,
            'all_tokens': all_tokens,
            'semantic_data': semantic_data,
            'entities': entities,
            'sentiments': sentiments
        }


# Utility functions for standalone usage
def quick_analysis(texts):
    """Quick analysis function for simple usage"""
    nlp = ComprehensiveNLP()
    return nlp.run_complete_analysis(texts)


def preprocess_only(texts):
    """Only run preprocessing"""
    nlp = ComprehensiveNLP()
    return nlp.preprocess_text(texts)


def sentiment_only(texts):
    """Only run sentiment analysis"""
    nlp = ComprehensiveNLP()
    return nlp.analyze_sentiment(texts)