import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
import gc

def preprocess_cicids2017_full(input_path, output_path):
    """
    Realiza o pré-processamento e limpeza do dataset CICIDS2017 completo.

    Args:
        input_path (str): Caminho para o arquivo CSV de entrada do dataset.
        output_path (str): Caminho para salvar o arquivo CSV pré-processado.
    """
    print(f"Carregando o dataset completo de: {input_path}")
    
    try:
        df = pd.read_csv(input_path, low_memory=False)
        print(f"Dataset completo carregado com {df.shape[0]} linhas e {df.shape[1]} colunas")
    except Exception as e:
        print(f"Erro ao carregar o dataset completo: {e}")
        print("Por favor, verifique se há memória suficiente ou se o arquivo está corrompido.")
        return

    # Renomear colunas para remover espaços e caracteres problemáticos
    df.columns = df.columns.str.strip().str.replace(' ', '_').str.replace('/', '_').str.replace('(', '', regex=False).str.replace(')', '', regex=False)

    # Verifica se a coluna 'Label' existe
    if 'Label' not in df.columns:
        raise ValueError("A coluna 'Label' não foi encontrada no dataset. Verifique se o arquivo contém os rótulos.")

    print(f"Distribuição das classes antes da limpeza:")
    print(df['Label'].value_counts())

    # Converte todas as colunas (exceto 'Label') para numérico
    feature_cols = df.columns[df.columns != 'Label']
    print(f"Convertendo {len(feature_cols)} colunas para numérico...")
    
    for col in feature_cols:
        df[col] = pd.to_numeric(df[col], errors='coerce')

    # Substitui valores infinitos por NaN e remove linhas inválidas
    df.replace([np.inf, -np.inf], np.nan, inplace=True)
    initial_rows = df.shape[0]
    df.dropna(inplace=True)
    print(f"Removidas {initial_rows - df.shape[0]} linhas com valores NaN ou infinitos.")

    # Remover colunas com baixa variância (constantes ou quase constantes)
    numeric_cols = df.select_dtypes(include=[np.number]).columns
    low_variance_cols = [col for col in numeric_cols if df[col].nunique() <= 1]
    if low_variance_cols:
        df.drop(columns=low_variance_cols, inplace=True)
        print(f"Removidas colunas com baixa variância: {low_variance_cols}")
    else:
        print("Nenhuma coluna com baixa variância encontrada.")

    # Separar X e y
    X = df.drop('Label', axis=1)
    y = df['Label']

    print(f"Características finais: {X.shape[1]} colunas")
    print(f"Amostras finais: {X.shape[0]} linhas")
    print(f"Distribuição das classes após limpeza:")
    print(y.value_counts())

    # Normalização Min-Max
    print("Aplicando normalização Min-Max...")
    scaler = MinMaxScaler()
    X_scaled = scaler.fit_transform(X)
    X_processed = pd.DataFrame(X_scaled, columns=X.columns)

    # Combina X normalizado e y
    df_processed = pd.concat([X_processed, y.reset_index(drop=True)], axis=1)

    # Salva o resultado
    df_processed.to_csv(output_path, index=False)
    print(f"Dataset pré-processado salvo em: {output_path}")
    print(f"Formato final: {df_processed.shape}")

    # Limpeza de memória
    del df, X, y, X_scaled, X_processed
    gc.collect()

if __name__ == '__main__':
    input_file = 'cicids2017.csv'  # Caminho real do dataset bruto
    output_file = 'cicids2017_preprocessed.csv'

    preprocess_cicids2017_full(input_file, output_file)