app.py

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import pickle
from datasets import load_dataset

# Cargar dataset desde Hugging Face
dataset = load_dataset("psabino/precios_de_viviendas_madrid", split="train")

# Convertir el dataset a un DataFrame de pandas
df = pd.DataFrame(dataset)

# Crear columna 'tipo_vivienda' a partir del título
def extraer_tipo(titulo):
    titulo = titulo.lower()
    if "chalet" in titulo:
        return "chalet"
    elif "ático" in titulo or "atico" in titulo:
        return "ático"
    elif "piso" in titulo:
        return "piso"
    else:
        return "otro"

df["tipo_vivienda"] = df["titulo"].apply(extraer_tipo)

# Convertir tipo de vivienda en variables dummy
df = pd.get_dummies(df, columns=["tipo_vivienda"])

# Asegurar que las columnas dummy existan aunque no estén presentes
for col in ["tipo_vivienda_ático", "tipo_vivienda_chalet", "tipo_vivienda_piso"]:
    if col not in df.columns:
        df[col] = 0

# Selección de columnas
columnas_modelo = ["metros_cuadrados", "habitaciones", "baños",
                   "tipo_vivienda_ático", "tipo_vivienda_chalet", "tipo_vivienda_piso"]
X = df[columnas_modelo]
y = df["precio"]

# Entrenar modelo
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
model = LinearRegression()
model.fit(X_train, y_train)

# Guardar modelo entrenado y columnas
modelo_datos = {
    "modelo": model,
    "columnas": columnas_modelo
}

with open("modelo_vivienda_madrid.pkl", "wb") as f:
    pickle.dump(modelo_datos, f)

print("✅ Modelo entrenado y guardado como 'modelo_vivienda_madrid.pkl'")