main.py

import pandas as pd
import numpy as np
import time
import random
import nltk
import sys
import os
from sklearn.datasets import fetch_20newsgroups

try:
    nltk.data.find('tokenizers/punkt')
except LookupError:
    nltk.download('punkt')
    nltk.download('stopwords')
    nltk.download('wordnet')

from data.data_loader import load_data, split_data
from data.data_model import EmailDataModel
from data.preprocessing import EmailPreprocessor
from utils.feature_extraction import TfidfFeatureExtractor
from models.chained_model import ChainedClassifier
from models.hierarchical_model import HierarchicalClassifier
from models.classifiers.random_forest import RandomForestModel
from models.classifiers.svm import SVMModel
from config.config import TYPE2_CATEGORIES, TYPE3_CATEGORIES, TYPE4_CATEGORIES

def create_sample_data(n_samples=500):
    print(f"\nGenerating sample data with {n_samples} samples...")
    start = time.time()
    
    data = fetch_20newsgroups(subset='all', remove=('headers', 'footers', 'quotes'))
    texts = data.data[:n_samples]
    
    print("Preprocessing texts...")
    preprocessor = EmailPreprocessor()
    processed_texts = []
    for i, text in enumerate(texts):
        if i % 100 == 0:
            print(f"Processing text {i}/{len(texts)}")
        try:
            processed = preprocessor.preprocess(text)
            processed_texts.append(processed)
        except Exception as e:
            processed_texts.append(preprocessor.clean_email(text))
    
    type2_labels = [random.choice(TYPE2_CATEGORIES) for _ in range(len(processed_texts))]
    type3_labels = [random.choice(TYPE3_CATEGORIES) for _ in range(len(processed_texts))]
    type4_labels = [random.choice(TYPE4_CATEGORIES) for _ in range(len(processed_texts))]
    
    df = pd.DataFrame({
        'email_text': processed_texts,
        'type2_label': type2_labels,
        'type3_label': type3_labels,
        'type4_label': type4_labels
    })
    
    end = time.time()
    print(f"Sample data created in {end-start:.2f} seconds")
    return df

def main():
    print("\n*** EMAIL CLASSIFICATION SYSTEM ***")
    print("Starting at", time.strftime("%H:%M:%S", time.localtime()))
    
    try:
        print("\nTrying to load data from file...")
        data = load_data()
        if data is None or len(data) < 10:
            print("No good data found. Creating sample data...")
            data = create_sample_data()
    except Exception as e:
        print(f"Error loading data: {e}")
        data = create_sample_data()
    
    print(f"\nData loaded with {len(data)} samples")
    
    train_data, test_data = split_data(data)
    print(f"Training set: {len(train_data['X'])} samples")
    print(f"Testing set: {len(test_data['X'])} samples")
    
    data_model = EmailDataModel.from_dict(train_data, test_data)
    
    print("\nExtracting features from text...")
    feature_extractor = TfidfFeatureExtractor(max_feat=1000)
    X_train = feature_extractor.fit_transform(data_model.X_train)
    X_test = feature_extractor.transform(data_model.X_test)
    
    print("\n" + "=" * 50)
    print("DESIGN 1: CHAINED CLASSIFICATION")
    print("=" * 50)
    
    chained_model = ChainedClassifier(base_classifier=RandomForestModel)
    chained_model.train(
        X_train, 
        data_model.y_train_type2, 
        data_model.y_train_type3, 
        data_model.y_train_type4
    )
    
    type2_pred, type3_pred, type4_pred = chained_model.predict(X_test)
    
    chained_results = chained_model.print_results(
        data_model.y_test_type2, 
        data_model.y_test_type3, 
        data_model.y_test_type4,
        type2_pred, 
        type3_pred, 
        type4_pred
    )
    
    print("\n" + "=" * 50)
    print("DESIGN 2: HIERARCHICAL CLASSIFICATION")
    print("=" * 50)
    
    hierarchical_model = HierarchicalClassifier(base_classifier=SVMModel)
    hierarchical_model.train(
        X_train, 
        data_model.y_train_type2, 
        data_model.y_train_type3, 
        data_model.y_train_type4
    )
    
    type2_pred, type3_pred, type4_pred = hierarchical_model.predict(X_test)
    
    hierarchical_results = hierarchical_model.print_results(
        data_model.y_test_type2, 
        data_model.y_test_type3, 
        data_model.y_test_type4,
        type2_pred, 
        type3_pred, 
        type4_pred
    )
    
    print("\n" + "=" * 50)
    print("FINAL COMPARISON")
    print("=" * 50)
    print(f"Chained Overall Accuracy: {chained_results['overall_acc']:.4f}")
    print(f"Hierarchical Overall Accuracy: {hierarchical_results['overall_acc']:.4f}")
    print(f"Winner: {'Chained' if chained_results['overall_acc'] > hierarchical_results['overall_acc'] else 'Hierarchical'}")
    print("\nClassification completed at", time.strftime("%H:%M:%S", time.localtime()))

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"ERROR: {e}")
        print("Something went wrong! Check the error above.")