import pandas as pd
import numpy as np
import os
import joblib
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.model_selection import train_test_split

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATASET_PATH = os.path.join(BASE_DIR, "datasets", "dataset_2.csv")

MODEL_DIR = os.path.join(BASE_DIR, "ml_models")
os.makedirs(MODEL_DIR, exist_ok=True)

MODEL_PATH = os.path.join(MODEL_DIR, "device_classifier.pkl")
SCALER_PATH = os.path.join(BASE_DIR, "datasets", "scaler.pkl")
LE_PATH = os.path.join(BASE_DIR, "datasets", "label_encoder.pkl")

def train():
    print("Loading dataset...")
    df = pd.read_csv(DATASET_PATH)
    
    print("Columns available:", list(df.columns))

    print("Preprocessing data...")
    # Drop NaNs
    df.dropna(subset=['device_type', 'packet_size', 'protocol_type', 'bandwidth_usage', 'transmission_time'], inplace=True)
    
    # Encode protocol (TCP/UDP mapping)
    print("Encoding protocol...")
    df['protocol'] = df['protocol_type'].apply(lambda x: 0 if str(x).upper() == 'TCP' else (1 if str(x).upper() == 'UDP' else 2))
    
    # Map features
    df['flow_duration'] = df['transmission_time']
    df['packet_rate'] = df['bandwidth_usage'] # Simple mapping
    
    features = ['packet_size', 'protocol', 'packet_rate', 'flow_duration']
    
    X = df[features]
    y = df['device_type']
    
    print("Encoding Labels...")
    le = LabelEncoder()
    y_encoded = le.fit_transform(y)
    
    print("Scaling Features...")
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    print("Splitting...")
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y_encoded, test_size=0.2, random_state=42)
    
    print("Training Model...")
    clf = RandomForestClassifier(n_estimators=100, random_state=42, n_jobs=-1)
    clf.fit(X_train, y_train)
    
    print(f"Accuracy: {clf.score(X_test, y_test):.4f}")
    
    print("Saving artifacts...")
    pd.DataFrame(X_test, columns=features).to_csv(os.path.join(BASE_DIR, "datasets", "X_test.csv"), index=False)
    pd.DataFrame(y_test, columns=['device_type']).to_csv(os.path.join(BASE_DIR, "datasets", "y_test.csv"), index=False)
    joblib.dump(clf, MODEL_PATH)
    joblib.dump(scaler, SCALER_PATH)
    joblib.dump(le, LE_PATH)
    print(f"Saved model to {MODEL_PATH}")
    print(f"Saved scaler to {SCALER_PATH}")
    print(f"Saved label encoder to {LE_PATH}")
    print("Done!")

if __name__ == "__main__":
    train()