import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import os
import gc

def visualize_raw_data(input_path):
    """
    Gera visualizações para o dataset completo (distribuição de uma feature e das classes),
    processando todas as linhas do arquivo.
    """

    print(f"Carregando o dataset completo de: {input_path}")
    
    # Carregar o dataset inteiro
    df_raw = pd.read_csv(input_path, low_memory=False)  # Carrega o dataset inteiro
    print(f"Dataset carregado com {df_raw.shape[0]} linhas e {df_raw.shape[1]} colunas.")

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

    # Criar diretório de saída
    os.makedirs("visualizations", exist_ok=True)

    # Detectar coluna de rótulo
    label_col = None
    for col in df_raw.columns:
        if col.strip().lower() in ["label", "attack", "class"]:
            label_col = col
            break

    if label_col is None:
        print("❌ Nenhuma coluna de rótulo encontrada.")
        return

    print(f"✅ Coluna de rótulo detectada: \'{label_col}\'")

    # Tentativa de conversão forçada para float (ignora erros e deixa NaN)
    df_converted = df_raw.copy()
    for col in df_converted.columns:
        if col != label_col:  # Não converter a coluna de label
            df_converted[col] = pd.to_numeric(df_converted[col], errors='coerce')

    # Seleciona a primeira coluna numérica válida (com poucos NaNs)
    numeric_cols = df_converted.select_dtypes(include=["float", "int"]).columns
    if len(numeric_cols) == 0:
        print("❌ Ainda nenhuma coluna numérica detectada após conversão.")
        return

    selected_feature = None
    for feature in numeric_cols:
        if df_converted[feature].notna().sum() > 1000:  # Pelo menos 1000 valores não-NaN
            selected_feature = feature
            break
    else:
        print("❌ Nenhuma feature numérica com dados suficientes encontrada.")
        return

    print(f"📊 Coluna numérica selecionada para distribuição: \'{selected_feature}\'")

    # Histograma da feature numérica
    plt.figure(figsize=(10, 6))
    sns.histplot(df_converted[selected_feature].dropna(), kde=True, bins=50)
    plt.title(f"Distribuição da Feature \'{selected_feature}\' (Dados Brutos Completo)")
    plt.xlabel(selected_feature)
    plt.ylabel("Frequência")
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("visualizations/raw_data_feature_distribution.png", dpi=300)
    plt.close()
    print("✅ Gráfico de distribuição salvo: visualizations/raw_data_feature_distribution.png")

    # Gráfico de distribuição da classe
    plt.figure(figsize=(10, 5))
    sns.countplot(x=df_raw[label_col].astype(str).str.strip())
    plt.title("Distribuição das Classes (Dados Brutos Completo)")
    plt.xlabel("Classe")
    plt.ylabel("Contagem")
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    plt.savefig("visualizations/raw_data_label_distribution.png", dpi=300)
    plt.close()
    print("✅ Gráfico de distribuição das classes salvo: visualizations/raw_data_label_distribution.png")

if __name__ == '__main__':
    input_file = 'cicids2017.csv'
    if not os.path.exists(input_file):
        print(f"❌ Arquivo '{input_file}' não encontrado.")
    else:
        visualize_raw_data(input_file)