import gradio as gr
import pickle
import json
import numpy as np

# Load model and columns
with open("kigali_model.pickle", "rb") as f:
    model = pickle.load(f)

with open("columns.json", "r") as f:
    data_columns = json.load(f)["data_columns"]

# Define the location and property type mappings
location_mapping = {
    'gacuriro': 1,
    'kacyiru': 2,
    'kanombe': 3,
    'kibagabaga': 4,
    'kicukiro': 5,
    'kimironko': 6,
    'nyamirambo': 7,
    'nyarutarama': 8
}

property_type_mapping = {
    'apartment': 1,
    'bungalow': 2,
    'house': 3,
    'villa': 4
}

def transform_data(size_sqm, number_of_bedrooms, number_of_bathrooms, number_of_floors, parking_space, location, property_type):
    # Prepare the input array with zeros
    x = np.zeros(len(data_columns))
    
    # Assign input values to the corresponding columns
    x[0] = size_sqm
    x[1] = number_of_bedrooms
    x[2] = number_of_bathrooms
    x[3] = number_of_floors
    x[5] = parking_space  # Ensure that parking_space aligns with the correct index in your model

    # Apply location mapping
    if location in location_mapping:
        loc_index = data_columns.index(location)
        x[loc_index] = 1

    # Apply property type mapping
    if property_type in property_type_mapping:
        prop_index = data_columns.index(property_type)
        x[prop_index] = 1

    return np.array([x])


def predict(size_sqm, number_of_bedrooms, number_of_bathrooms, number_of_floors, parking_space, location, property_type):
    # Transform input data
    input_data_transformed = transform_data(size_sqm, number_of_bedrooms, number_of_bathrooms, number_of_floors, parking_space, location, property_type)
    
    # Predict using the model
    prediction = model.predict(input_data_transformed)
    return round(prediction[0], 2)  # round prediction for better readability

# Define Gradio interface components
inputs = [
    gr.Number(label="Size (sqm)", value=0),
    gr.Number(label="Number of Bedrooms", value=0),
    gr.Number(label="Number of Bathrooms", value=0),
    gr.Number(label="Number of Floors", value=0),
    gr.Number(label="Parking Space", value=0),
    gr.Dropdown(choices=list(location_mapping.keys()), label="Location"),
    gr.Dropdown(choices=list(property_type_mapping.keys()), label="Property Type"),
    # Add new inputs for other columns, like furnished, proximity, etc.
]


outputs = gr.Textbox(label="Prediction (FRW)")

# Footer content
footer = "Etienne NTAMBARA @AI_Engineer"

# Launch the interface
gr.Interface(
    fn=predict,
    inputs=inputs,
    outputs=outputs,
    title="Property Price Prediction",
    description="Enter property details to get the price prediction.",
    article=footer
).launch(debug=True)