import gradio as gr
import tensorflow as tf
import numpy as np
import pandas as pd
from datasets import load_dataset
from huggingface_hub import hf_hub_download

# Загрузка файла модели из Hugging Face Hub
model_path = hf_hub_download(repo_id="1sozidatel1/PricePredictorAI_Model", filename="price_model.h5")

# Загрузка модели
model = tf.keras.models.load_model(model_path)

# Карты местоположений и районов
location_mapping = {1: 'Астрахань', 2: 'Волгоград', 3: 'Краснодар', 4: 'Ростов-на-Дону', 5: 'Майкоп', 6: 'Элиста'}
district_mapping = {
    1: [11, 12, 13, 14, 15],
    2: [21, 22, 23, 24, 25, 26, 27, 28],
    3: [31, 32, 33, 34],
    4: [41, 42, 43, 44, 45, 46, 47, 48],
    5: [51],
    6: [61]
}
district_names = {
    11: 'Кировский', 12: 'Ленинский', 13: 'Советский', 14: 'Трусовский', 15: 'Центральный',
    21: 'Ворошиловский', 22: 'Дзержинский', 23: 'Кировский', 24: 'Красноармейский', 25: 'Краснооктябрьский', 26: 'Советский', 27: 'Тракторозаводский', 28: 'Центральный',
    31: 'Западный', 32: 'Карасунский', 33: 'Прикубанский', 34: 'Центральный',
    41: 'Ворошиловский', 42: 'Железнодорожный', 43: 'Кировский', 44: 'Ленинский', 45: 'Октябрьский', 46: 'Первомайский', 47: 'Пролетарский', 48: 'Советский',
    51: 'Центральный',
    61: 'Центральный'
}
author_type_mapping = {0: 'Риелтор', 1: 'Частный владелец'}

# Загрузка данных
dataset = load_dataset("1sozidatel1/sfd_housing_prices_november_2024")
data = pd.DataFrame(dataset['train'])
feature_cols = ['author_type_id', 'location_id', 'district_id', 'floor', 'rooms_count', 'total_meters']

X = data[feature_cols]
X_min = X.min()
X_max = X.max()
y = data['price'].values.reshape(-1, 1)
y_min = np.min(y)
y_max = np.max(y)

def predict_price(author_type, location, district, floor, rooms_count, total_meters):
    author_type_id = list(author_type_mapping.keys())[list(author_type_mapping.values()).index(author_type)]
    location_id = list(location_mapping.keys())[list(location_mapping.values()).index(location)]
    district_id = list(district_names.keys())[list(district_names.values()).index(district)]
    input_data = np.array([[author_type_id, location_id, district_id, floor, rooms_count, total_meters]])
    input_data_normalized = (input_data - X_min.values) / (X_max.values - X_min.values)
    prediction = model.predict(input_data_normalized)
    prediction_unscaled = prediction * (y_max - y_min) + y_min
    return np.round(prediction_unscaled.flatten()[0], 1)

def update_districts(location):
    location_id = list(location_mapping.keys())[list(location_mapping.values()).index(location)]
    districts = [district_names[district_id] for district_id in district_mapping[location_id]]
    return gr.update(choices=districts)

# Создание интерфейса Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Предсказание цен на квартиры в Южном Федеральном Округе")
    gr.Markdown("Этот интерфейс позволяет предсказать цену на недвижимость на основе различных параметров, таких как тип продавца, город, район, этаж, количество комнат и площадь.")
    gr.Markdown("Введите необходимые данные и нажмите кнопку 'Предсказать цену', чтобы получить предсказанную цену.")
    gr.Markdown("❗️Все поля обязательны для заполнения❗️")
    
    with gr.Row():
        with gr.Column():
            author_type = gr.Dropdown(choices=list(author_type_mapping.values()), label="Тип продавца")
            location = gr.Dropdown(choices=list(location_mapping.values()), label="Город")
            district = gr.Dropdown(label="Район")
            floor = gr.Number(label="Этаж", minimum=1, maximum=50)
            rooms_count = gr.Dropdown(choices=[1, 2, 3, 4, 5], label="Количество комнат")
            total_meters = gr.Number(label="Площадь, кв.м", minimum=18, maximum=500)
            predict_btn = gr.Button("Предсказать цену")
        
        with gr.Column():
            price = gr.Number(label="Предсказанная цена, (тыс. руб.)")
    
    location.change(update_districts, inputs=location, outputs=district)
    predict_btn.click(predict_price, inputs=[author_type, location, district, floor, rooms_count, total_meters], outputs=price)

if __name__ == "__main__":
    demo.launch()