import pandas as pd
import pickle
import streamlit as st

@st.cache_resource
def load_model():
    with open('model_rf.pkl', 'rb') as file:
        model = pickle.load(file)
    return model

cols = ['type', 'amount', 'amount', 'old_balance_ori', 'new_balance_ori', 'old_balance_dest']


def run():
    st.title('Transaction Fraud Prediction')
    st.write('This is a simple web app to predict transaction is a fraud or not using random forest.')
    st.write('Please fill in the form below to get the prediction.')

    amount = st.number_input('Transfer Amount', min_value=0.0)
    old_balance_ori = st.number_input('Old Balance Origin', min_value=0.0)
    new_balance_ori = st.number_input('New Balance Origin', min_value=0.0)
    old_balance_dest = st.number_input('Old Balance Destination', min_value=0.0)
    new_balance_dest = st.number_input('New Balance Destination', min_value=0.0)
    type = st.selectbox('Transaction type', ['CASH_OUT', 'TRANSFER', 'DEBIT', 'CASH_IN', 'PAYMENT'])

    if st.button("Predict"):
        model = load_model()
        data = {'type': type, 'amount': amount, 'old_balance_ori': old_balance_ori, 'new_balance_ori': new_balance_ori,
                'old_balance_dest': old_balance_dest, 'new_balance_dest': new_balance_dest}
        features = pd.DataFrame(data, index=[0])
        prediction = model.predict(features)
        if prediction == 0:
            st.success('The model predicts that the transaction is not a fraud.')
        else:
            st.error('The model predicts that the transaction is a fraud.')


if __name__ == '__main__':
    run()