data_loader.py

import os
import zipfile
import urllib.request
import pandas as pd

class DataLoader:
    def __init__(self, data_dir='dataset'):
        self.data_dir = data_dir
        self.csv_path = os.path.join(data_dir, 'enron_spam_data.csv')
        self.zip_path = os.path.join(data_dir, 'enron_spam_data.zip')
        
    def create_dataset_dir(self):
        """Tạo thư mục dataset nếu chưa có"""
        if not os.path.exists(self.data_dir):
            os.makedirs(self.data_dir)
            print(f"Đã tạo thư mục: {self.data_dir}")
    
    def download_dataset(self, url='https://github.com/MWiechmann/enron_spam_data/raw/refs/heads/master/enron_spam_data.zip'):
        """Tải dataset từ URL"""
        self.create_dataset_dir()
        
        if os.path.exists(self.csv_path):
            print(f"Dataset đã tồn tại: {self.csv_path}")
            return True
            
        try:
            print(f"Đang tải dataset từ: {url}")
            urllib.request.urlretrieve(url, self.zip_path)
            print(f"Đã tải: {self.zip_path}")
            return True
        except Exception as e:
            print(f"Lỗi khi tải dataset: {e}")
            return False
    
    def extract_dataset(self):
        """Giải nén file zip"""
        if not os.path.exists(self.zip_path):
            print(f"Không tìm thấy file zip: {self.zip_path}")
            return False
            
        try:
            print(f"Đang giải nén: {self.zip_path}")
            with zipfile.ZipFile(self.zip_path, 'r') as zip_ref:
                zip_ref.extractall(self.data_dir)
            print(f"Đã giải nén vào: {self.data_dir}")
            
            # Xóa file zip sau khi giải nén
            os.remove(self.zip_path)
            print(f"Đã xóa file zip: {self.zip_path}")
            return True
        except Exception as e:
            print(f"Lỗi khi giải nén: {e}")
            return False
    
    def load_data(self):
        """Tải dữ liệu từ CSV"""
        if not os.path.exists(self.csv_path):
            print(f"Không tìm thấy file CSV: {self.csv_path}")
            return None
            
        try:
            print(f"Đang đọc: {self.csv_path}")
            df = pd.read_csv(self.csv_path)
            print(f"Đã load {len(df)} dòng dữ liệu")
            return df
        except Exception as e:
            print(f"Lỗi khi đọc CSV: {e}")
            return None
    
    def get_data_info(self):
        """Hiển thị thông tin dataset"""
        df = self.load_data()
        if df is None:
            return
            
        print("\nThông tin Dataset:")
        print(f"   - Số lượng email: {len(df)}")
        print(f"   - Cột: {list(df.columns)}")
        
        if 'Spam/Ham' in df.columns:
            spam_count = len(df[df['Spam/Ham'] == 'spam'])
            ham_count = len(df[df['Spam/Ham'] == 'ham'])
            print(f"   - Spam: {spam_count} ({spam_count/len(df)*100:.1f}%)")
            print(f"   - Ham: {ham_count} ({ham_count/len(df)*100:.1f}%)")
        
        print(f"\nMẫu dữ liệu:")
        print(df.head())
    
    def setup_complete_dataset(self):
        """Thiết lập hoàn chỉnh dataset (tải + giải nén + kiểm tra)"""
        print("Bắt đầu thiết lập dataset...")
        
        # Bước 1: Tải dataset
        if not self.download_dataset():
            return False
            
        # Bước 2: Giải nén (nếu cần)
        if os.path.exists(self.zip_path):
            if not self.extract_dataset():
                return False
        
        # Bước 3: Kiểm tra và hiển thị thông tin
        if os.path.exists(self.csv_path):
            self.get_data_info()
            print("\nDataset đã sẵn sàng sử dụng!")
            return True
        else:
            print(f"Không tìm thấy file CSV sau khi giải nén")
            return False

def main():
    """Test DataLoader"""
    loader = DataLoader()
    loader.setup_complete_dataset()

if __name__ == "__main__":
    main()