import gradio as gr
import pandas as pd
import numpy as np
import pickle

MODEL_PATH = "model.pkl"

with open(MODEL_PATH, 'rb') as f:
    model_package = pickle.load(f)

if isinstance(model_package, dict):
    model = model_package['model']
    scaler = model_package.get('scaler', None)
    features = model_package.get('features', None)
else:
    model = model_package
    scaler = None
    features = None

def create_features(rooms, area, postalcode, pop, pop_dens, frg_pct, emp, tax_income):
    """Create engineered features from input"""
    data = {
        'rooms': rooms,
        'area': area,
        'pop': pop,
        'pop_dens': pop_dens,
        'frg_pct': frg_pct,
        'emp': emp,
        'tax_income': tax_income,

        'rooms_per_sqm': rooms / area,
        'wealth_index': (tax_income / 100000) * (emp / 100000),
        'is_zurich_city': 1 if (postalcode >= 8000 and postalcode < 8100) else 0,
        'pop_emp_ratio': pop / (emp + 1),
        'log_area': np.log1p(area),
        'log_pop': np.log1p(pop),
        'log_tax_income': np.log1p(tax_income)
    }
    return pd.DataFrame([data])

def predict_price(rooms, area, postalcode, pop, pop_dens, frg_pct, emp, tax_income):
    """Predict apartment rental price based on input features"""
    input_data = create_features(rooms, area, postalcode, pop, pop_dens, frg_pct, emp, tax_income)
    
    if features is not None:
        input_data = input_data[features]
    
    if scaler is not None:
        input_data_scaled = scaler.transform(input_data)
        prediction = model.predict(input_data_scaled)[0]
    else:
        prediction = model.predict(input_data)[0]
    
    return f"Geschätzter Mietpreis: CHF {prediction:.2f}/Monat"

with gr.Blocks(title="Zürich Apartment Price Predictor") as demo:
    gr.Markdown("# Zürich Apartment Rent Predictor")
    gr.Markdown("Vorhersage von Mietpreisen für Wohnungen im Kanton Zürich")
    
    with gr.Row():
        with gr.Column():
            rooms = gr.Number(label="Anzahl Zimmer", value=3.5, minimum=1, maximum=10)
            area = gr.Number(label="Wohnfläche (m²)", value=75, minimum=10, maximum=500)
            postalcode = gr.Number(label="Postleitzahl", value=8001, minimum=8000, maximum=8999)
            
        with gr.Column():
            pop = gr.Number(label="Bevölkerung (Gemeinde)", value=420000)
            pop_dens = gr.Number(label="Bevölkerungsdichte", value=5000)
            frg_pct = gr.Number(label="Ausländeranteil (%)", value=30)
            emp = gr.Number(label="Anzahl Arbeitsplätze", value=490000)
            tax_income = gr.Number(label="Durchschn. Steuereinkommen", value=85000)
    
    predict_btn = gr.Button("Preis berechnen", variant="primary")
    output = gr.Textbox(label="Ergebnis", lines=2)
    
    predict_btn.click(
        fn=predict_price,
        inputs=[rooms, area, postalcode, pop, pop_dens, frg_pct, emp, tax_income],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())