import pandas as pd
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score

if __name__ == '__main__':
    data = pd.read_csv('final_dataset.csv')
    X = data.drop('Protocol', axis=1)
    y = data['Protocol']

    c1 = LabelEncoder()
    X["Source"] = c1.fit_transform(X["Source"])
    c2 = LabelEncoder()
    X["Destination"] = c2.fit_transform(X["Destination"])

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1)

    clf = DecisionTreeClassifier()
    clf.fit(X_train, y_train)
    y_pred = clf.predict(X_test)

    accuracy = accuracy_score(y_test, y_pred)
    print("Presnost natrenovaneho modelu je: {}".format(accuracy))

    src = "10.22.22.22"
    dst = "224.0.0.10"
    paylen = 114

    src = c1.transform([src])
    dst = c2.transform([dst])
    new_data = [[src[0], dst[0], paylen]]

    new_predict = clf.predict(new_data)
    print("Nove data patria protokolu: {}".format(new_predict))