# Изисквания: pip install numpy pandas scikit-learn matplotlib
import numpy as np
import pandas as pd
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
import matplotlib.pyplot as plt

# 1) Генериране на синтетични данни (симулиращи мрежови трафик)
# Тук използваме make_classification: всеки ред = сесия/пакет; target = клас (напр. "злонамерен"/"нормален")
X, y = make_classification(
    n_samples=2000,       # брой примери
    n_features=10,        # общ брой признаци (можеш да увеличиш по-нататък)
    n_informative=6,      # колко признака са информативни за класа
    n_redundant=2,        # линейно зависими признаци
    n_repeated=0,
    n_classes=2,          # бинарна класификация (например 0 = нормален, 1 = опасен)
    weights=[0.8, 0.2],   # лек дисбаланс (80% нормални, 20% атаки)
    class_sep=1.2,
    random_state=42
)

# 2) Зареждане в DataFrame чрез pandas
feature_names = [f'feat_{i}' for i in range(X.shape[1])]
df = pd.DataFrame(X, columns=feature_names)
df['label'] = y
print(df.head())

# 3) Разделяне на тренировъчен и тестов набор
X = df[feature_names].values
y = df['label'].values
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, stratify=y, random_state=42
)

# 4) Използване на данните за обучение на Decision Tree
clf = DecisionTreeClassifier(
    criterion='gini',
    max_depth=5,        # контролира простотата на дървото; пробвай различни стойности
    min_samples_split=10,
    random_state=42
)
clf.fit(X_train, y_train)

# 5) Начертаване на дървото
plt.figure(figsize=(16,8))
plot_tree(clf, feature_names=feature_names, class_names=['normal','attack'], filled=True, fontsize=8)
plt.title("Decision Tree")
plt.show()

# 6) Оценка на представянето върху тестовия масив
y_pred = clf.predict(X_test)
print("Accuracy:", accuracy_score(y_test, y_pred))
print("Confusion matrix:\n", confusion_matrix(y_test, y_pred))
print("Classification report:\n", classification_report(y_test, y_pred))

# 7) Проба със същото, но увеличаване на броя признаци
# Генерираме нов набор с повече признаци (напр. 30)
X2, y2 = make_classification(
    n_samples=2000,
    n_features=30,
    n_informative=15,
    n_redundant=10,
    n_classes=2,
    weights=[0.8, 0.2],
    class_sep=1.0,
    random_state=1
)
df2 = pd.DataFrame(X2, columns=[f'feat_{i}' for i in range(X2.shape[1])])
df2['label'] = y2

X2 = df2.drop(columns=['label']).values
y2 = df2['label'].values
X2_train, X2_test, y2_train, y2_test = train_test_split(X2, y2, test_size=0.3, stratify=y2, random_state=1)

clf2 = DecisionTreeClassifier(max_depth=6, min_samples_split=10, random_state=1)
clf2.fit(X2_train, y2_train)

y2_pred = clf2.predict(X2_test)
print("With 30 features - Accuracy:", accuracy_score(y2_test, y2_pred))
print("With 30 features - Classification report:\n", classification_report(y2_test, y2_pred))