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

def preprocess_data(input_data, pmap, fmap, features, sc):
    """Preprocesses the input data."""
    input_df = pd.DataFrame([input_data], columns=features)
    input_df['protocol_type'] = input_df['protocol_type'].map(pmap)
    input_df['flag'] = input_df['flag'].map(fmap)
    input_df = input_df[features]
    input_data_scaled = sc.transform(input_df)
    return input_data_scaled.reshape((-1,30,1)) # Reshape for CNN
def predict_attack(preprocessed_data, cnn_model):
    """Predicts the attack type using the CNN model."""
    prediction = cnn_model.predict(preprocessed_data)
    amap = {0: 'dos', 1: 'normal', 2: 'probe', 3: 'r2l', 4: 'u2r'}
    predicted_attack_type = amap[np.argmax(prediction)]
    return predicted_attack_type

def network_attack_pipeline(input_data, cnn_model, pmap, fmap, features, sc):
    """
    A pipeline for predicting network attacks.

    Args:
        input_data (dict): Input data dictionary.
        cnn_model (keras.Model): Trained CNN model.
        pmap (dict): Mapping for protocol_type.
        fmap (dict): Mapping for flag.
        features (list): List of features used in training.
        sc (MinMaxScaler): Scaler object.

    Returns:
        str: Predicted attack type.
    """
    preprocessed_data = preprocess_data(input_data, pmap, fmap, features, sc)
    predicted_attack = predict_attack(preprocessed_data, cnn_model)
    return predicted_attack

pmap = {'icmp': 0, 'tcp': 1, 'udp': 2}
fmap = {'SF': 0, 'S0': 1, 'REJ': 2, 'RSTR': 3, 'RSTO': 4, 'SH': 5, 'S1': 6, 'S2': 7, 'RSTOS0': 8, 'S3': 9, 'OTH': 10}
features = ['duration', 'protocol_type', 'flag', 'src_bytes', 'dst_bytes', 'land',
       'wrong_fragment', 'urgent', 'hot', 'num_failed_logins', 'logged_in',
       'num_compromised', 'root_shell', 'su_attempted', 'num_file_creations',
       'num_shells', 'num_access_files', 'is_guest_login', 'count',
       'srv_count', 'serror_rate', 'rerror_rate', 'same_srv_rate',
       'diff_srv_rate', 'srv_diff_host_rate', 'dst_host_count',
       'dst_host_srv_count', 'dst_host_diff_srv_rate',
       'dst_host_same_src_port_rate', 'dst_host_srv_diff_host_rate']

from tensorflow.keras.models import load_model
import joblib
import pickle
cnn_model = load_model("cnn_model.h5")

# Load the MinMaxScaler
scaler = joblib.load("scaler.pkl")


# Now this section in for frontend....
import streamlit as st
st.caption('Input must be in given sequence: ')
st.caption(features)
st.caption('''Eg: "0, 'tcp','SF',181,5450,0,0,0,0,0,1,0,0,0,0,0,0,0,8,8,0,0,1,0,0,9,9,0,0.11,0"''')
input=st.text_input("Enter the input with comma Seprated value: ")
listVal=input.split(',')
def predict():
    inputList=[]
    inputList.append((int)(listVal[0]))
    inputList.append(listVal[1])
    inputList.append(listVal[2])
    inputList.append((int)(listVal[3]))
    inputList.append((int)(listVal[4]))
    inputList.append((int)(listVal[5]))
    inputList.append((int)(listVal[6]))
    inputList.append((int)(listVal[7]))
    inputList.append((int)(listVal[8]))
    inputList.append((int)(listVal[9]))
    inputList.append((int)(listVal[10]))
    inputList.append((int)(listVal[11]))
    inputList.append((int)(listVal[12]))
    inputList.append((int)(listVal[13]))
    inputList.append((int)(listVal[14]))
    inputList.append((int)(listVal[15]))
    inputList.append((int)(listVal[16]))
    inputList.append((int)(listVal[17]))
    inputList.append((int)(listVal[18]))
    inputList.append((int)(listVal[19]))
    inputList.append((float)(listVal[20]))
    inputList.append((float)(listVal[21]))
    inputList.append((float)(listVal[22]))
    inputList.append((float)(listVal[23]))
    inputList.append((float)(listVal[24]))
    inputList.append((int)(listVal[25]))
    inputList.append((int)(listVal[26]))
    inputList.append((float)(listVal[27]))
    inputList.append((float)(listVal[28]))
    inputList.append((float)(listVal[29]))
    ans=''
    try:
        ans=network_attack_pipeline(inputList,cnn_model, pmap, fmap, features, scaler)
    except:
        return "Input is not in valid form:"   
    return ans
    
if(st.button('Predict')):
    st.title(predict())


input2 =st.number_input("Enter the input between(0-494021)",step=1)             
def predict1():
    df=pd.read_csv('validation.csv')
    x=df.iloc[input2,:-1]
    lb=df.iloc[input2,-1]
    st.title(f"Labeled as: {lb.upper()}")
    print(df.iloc[input2,-1])
    try:
        res=network_attack_pipeline(x,cnn_model, pmap, fmap, features, scaler)
        st.title(f"Predicted as: {res.upper()}")
    except:
        st.title("There is some essue try on another value...")

if(st.button('Validate')):
    predict1()