import pandas as pd
import nltk
import re
import pickle
import joblib
import json
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.naive_bayes import MultinomialNB
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
from collections import Counter
import os

# Download NLTK resources
try:
    nltk.download('punkt')
    nltk.download('stopwords')
    nltk.download('wordnet')
    nltk.download('punkt_tab')
except Exception as e:
    print(f"Warning: Could not download NLTK data: {e}")

# Text Preprocessing Function
def preprocess_text(text):
    # Lowercase
    text = text.lower()
    # Remove special characters and digits
    text = re.sub(r'[^a-zA-Z\s]', '', text)
    # Tokenize
    tokens = word_tokenize(text)
    # Remove stopwords
    stop_words = set(stopwords.words('english'))
    tokens = [word for word in tokens if word not in stop_words]
    # Lemmatize
    lemmatizer = WordNetLemmatizer()
    tokens = [lemmatizer.lemmatize(word) for word in tokens]
    # Join tokens back to string
    return ' '.join(tokens)

# Load and Preprocess Dataset
def load_and_preprocess_data(file_path="IMDB Dataset.csv"):
    try:
        df = pd.read_csv(file_path)
        print(f"Loaded {len(df)} reviews")
        # Apply preprocessing to reviews
        print("Preprocessing reviews...")
        df['cleaned_review'] = df['review'].apply(preprocess_text)
        # Convert sentiment to binary (1 for positive, 0 for negative)
        df['sentiment'] = df['sentiment'].replace({'positive': 1, 'negative': 0})
        return df
    except FileNotFoundError:
        print(f"Error: Could not find file '{file_path}'")
        return None

# Train Models and Save
def train_and_save_models(df, model_dir="saved_models"):
    # Create model directory
    os.makedirs(model_dir, exist_ok=True)
    
    # Convert text to TF-IDF features
    print("Vectorizing text data...")
    vectorizer = TfidfVectorizer(max_features=5000)
    X = vectorizer.fit_transform(df['cleaned_review'])
    y = df['sentiment']
    
    # Split data
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    # Train Logistic Regression
    print("Training Logistic Regression...")
    lr_model = LogisticRegression(max_iter=1000, random_state=42)
    lr_model.fit(X_train, y_train)
    lr_predictions = lr_model.predict(X_test)
    lr_accuracy = accuracy_score(y_test, lr_predictions)
    
    # Train Naive Bayes
    print("Training Naive Bayes...")
    nb_model = MultinomialNB()
    nb_model.fit(X_train, y_train)
    nb_predictions = nb_model.predict(X_test)
    nb_accuracy = accuracy_score(y_test, nb_predictions)
    
    # Save models
    print("Saving models...")
    
    # Save vectorizer
    joblib.dump(vectorizer, os.path.join(model_dir, 'tfidf_vectorizer.pkl'))
    
    # Save Logistic Regression model
    joblib.dump(lr_model, os.path.join(model_dir, 'logistic_regression_model.pkl'))
    
    # Save Naive Bayes model
    joblib.dump(nb_model, os.path.join(model_dir, 'naive_bayes_model.pkl'))
    
    # Save model metadata
    metadata = {
        'vectorizer_features': vectorizer.get_feature_names_out().tolist(),
        'max_features': 5000,
        'lr_accuracy': float(lr_accuracy),
        'nb_accuracy': float(nb_accuracy),
        'training_samples': X_train.shape[0],
        'test_samples': X_test.shape[0],
        'preprocessing_steps': [
            'lowercase',
            'remove_special_chars',
            'tokenization',
            'stopword_removal',
            'lemmatization'
        ]
    }
    
    with open(os.path.join(model_dir, 'model_metadata.json'), 'w') as f:
        json.dump(metadata, f, indent=2)
    
    # Print results
    print("\n" + "="*50)
    print("MODEL TRAINING RESULTS")
    print("="*50)
    print(f"Logistic Regression Accuracy: {lr_accuracy:.4f}")
    print(f"Naive Bayes Accuracy: {nb_accuracy:.4f}")
    print(f"Models saved to: {model_dir}/")
    print("="*50)
    
    return vectorizer, lr_model, nb_model, lr_accuracy, nb_accuracy

# Create inference script
def create_inference_script():
    inference_code = '''import joblib
import json
import re
import nltk
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer

# Download NLTK resources
try:
    nltk.download('punkt')
    nltk.download('stopwords')
    nltk.download('wordnet')
except:
    pass

class SentimentAnalyzer:
    def __init__(self, model_dir="saved_models"):
        # Load models
        self.vectorizer = joblib.load(f"{model_dir}/tfidf_vectorizer.pkl")
        self.lr_model = joblib.load(f"{model_dir}/logistic_regression_model.pkl")
        self.nb_model = joblib.load(f"{model_dir}/naive_bayes_model.pkl")
        
        # Load metadata
        with open(f"{model_dir}/model_metadata.json", 'r') as f:
            self.metadata = json.load(f)
    
    def preprocess_text(self, text):
        # Lowercase
        text = text.lower()
        # Remove special characters and digits
        text = re.sub(r'[^a-zA-Z\\s]', '', text)
        # Tokenize
        tokens = word_tokenize(text)
        # Remove stopwords
        stop_words = set(stopwords.words('english'))
        tokens = [word for word in tokens if word not in stop_words]
        # Lemmatize
        lemmatizer = WordNetLemmatizer()
        tokens = [lemmatizer.lemmatize(word) for word in tokens]
        # Join tokens back to string
        return ' '.join(tokens)
    
    def predict(self, text, model_type='both'):
        # Preprocess text
        cleaned_text = self.preprocess_text(text)
        
        # Vectorize
        text_vector = self.vectorizer.transform([cleaned_text])
        
        results = {}
        
        if model_type in ['lr', 'both']:
            lr_pred = self.lr_model.predict(text_vector)[0]
            lr_prob = self.lr_model.predict_proba(text_vector)[0]
            results['logistic_regression'] = {
                'prediction': 'positive' if lr_pred == 1 else 'negative',
                'confidence': float(max(lr_prob)),
                'probabilities': {
                    'negative': float(lr_prob[0]),
                    'positive': float(lr_prob[1])
                }
            }
        
        if model_type in ['nb', 'both']:
            nb_pred = self.nb_model.predict(text_vector)[0]
            nb_prob = self.nb_model.predict_proba(text_vector)[0]
            results['naive_bayes'] = {
                'prediction': 'positive' if nb_pred == 1 else 'negative',
                'confidence': float(max(nb_prob)),
                'probabilities': {
                    'negative': float(nb_prob[0]),
                    'positive': float(nb_prob[1])
                }
            }
        
        return results

# Example usage
if __name__ == "__main__":
    analyzer = SentimentAnalyzer()
    
    # Test with sample reviews
    test_reviews = [
        "This movie was absolutely fantastic! I loved every minute of it.",
        "Terrible film, waste of time. Don't watch it.",
        "It was okay, nothing special but not bad either."
    ]
    
    for review in test_reviews:
        print(f"\\nReview: {review}")
        results = analyzer.predict(review)
        for model, result in results.items():
            print(f"{model}: {result['prediction']} (confidence: {result['confidence']:.2f})")
'''
    
    with open('inference.py', 'w') as f:
        f.write(inference_code)
    
    print("Created inference.py for model deployment")

# Create requirements.txt
def create_requirements():
    requirements = '''pandas>=1.3.0
nltk>=3.6
scikit-learn>=1.0.0
joblib>=1.1.0
numpy>=1.21.0
streamlit>=1.0.0
matplotlib>=3.5.0
seaborn>=0.11.0'''
    
    with open('requirements.txt', 'w') as f:
        f.write(requirements)
    
    print("Created requirements.txt")

# Create README
def create_readme():
    readme = '''# IMDb Sentiment Analysis Model

This repository contains a trained sentiment analysis model for IMDb movie reviews.

## Model Performance
- Logistic Regression: ~88.47% accuracy
- Naive Bayes: ~85.2% accuracy

## Files
- `saved_models/`: Directory containing trained models
- `inference.py`: Script for making predictions
- `train_and_save_model.py`: Script to train and save models
- `requirements.txt`: Python dependencies

## Usage

### Load and Use the Model
```python
from inference import SentimentAnalyzer

# Initialize analyzer
analyzer = SentimentAnalyzer()

# Make prediction
result = analyzer.predict("This movie was amazing!")
print(result)
```

### Deploy on Streamlit
```bash
streamlit run streamlit_deployment.py
```

## Model Details
- **Vectorizer**: TF-IDF with 5000 features
- **Preprocessing**: Lowercase, special char removal, tokenization, stopword removal, lemmatization
- **Models**: Logistic Regression and Naive Bayes
'''
    
    with open('README.md', 'w') as f:
        f.write(readme)
    
    print("Created README.md")

if __name__ == "__main__":
    print("Training and saving sentiment analysis models...")
    
    # Load data
    df = load_and_preprocess_data()
    
    if df is not None:
        # Train and save models
        train_and_save_models(df)
        
        # Create deployment files
        create_inference_script()
        create_requirements()
        create_readme()
        
        print("\n✅ Model training and saving completed!")
        print("📁 Models saved in 'saved_models/' directory")
        print("🚀 Ready for deployment on Hugging Face, Kaggle, or other platforms")
    else:
        print("❌ Failed to load data. Please check if 'IMDB Dataset.csv' exists.")