import pandas as pd
import numpy as np
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error
import gradio as gr
import warnings
warnings.filterwarnings('ignore')

# --- Načtení a příprava dat ---
df = pd.read_csv('02_simpleml_komplet_pohyb.csv')
zony = pd.read_csv('00_zony.csv')

# Odstranění duplicit a chyb
df = df.dropna()
df = df[df['pocet_cest'] >= 0]
df = df[df['vzdalenost'] >= 0]

# Přidáme vlastnosti cílové zóny
zony_cols = zony[['nazev', 'velikost', 'izolace', 'skola', 'zamestnavatel', 'uzel']].copy()
zony_cols.columns = ['cil', 'cil_velikost', 'cil_izolace', 'cil_skola', 'cil_zamestnavatel', 'cil_uzel']
df = df.merge(zony_cols, on='cil', how='left')

zony_src = zony[['nazev', 'velikost', 'izolace']].copy()
zony_src.columns = ['zdroj', 'zdroj_velikost', 'zdroj_izolace']
df = df.merge(zony_src, on='zdroj', how='left')

df = df.fillna(df.median(numeric_only=True))

# Encoding kategorií
le_dict = {}
cat_cols = ['denni_typ', 'casove_okno', 'vekova_skupina', 'ucel', 'hlavni_mod']
for col in cat_cols:
    le = LabelEncoder()
    df[col + '_enc'] = le.fit_transform(df[col].astype(str))
    le_dict[col] = le

# Features pro model
feature_cols = [
    'rok', 'ctvrtleti', 'denni_typ_enc', 'casove_okno_enc',
    'vekova_skupina_enc', 'ucel_enc', 'hlavni_mod_enc',
    'udalost', 'udalost_velikost', 'podil_spatne_pocasi',
    'vzdalenost', 'cil_velikost', 'cil_izolace', 'cil_skola',
    'cil_zamestnavatel', 'cil_uzel', 'zdroj_velikost', 'zdroj_izolace'
]

X = df[feature_cols]
y = df['pocet_cest']

# Trénink modelu
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
model = GradientBoostingRegressor(n_estimators=200, learning_rate=0.1, max_depth=4, random_state=42)
model.fit(X_train, y_train)

# Baseline (průměr stejného typu dne + časového okna)
baseline_map = df.groupby(['denni_typ', 'casove_okno'])['pocet_cest'].mean().to_dict()

y_pred = model.predict(X_test)
baseline_pred = X_test.apply(
    lambda r: baseline_map.get((
        [k for k,v in le_dict['denni_typ'].classes_ if False] or
        le_dict['denni_typ'].inverse_transform([int(r['denni_typ_enc'])])[0],
        le_dict['casove_okno'].inverse_transform([int(r['casove_okno_enc'])])[0]
    ), y_train.mean()), axis=1
)

mae_model = mean_absolute_error(y_test, y_pred)
mae_baseline = mean_absolute_error(y_test, [y_train.mean()] * len(y_test))

print(f"Model MAE: {mae_model:.2f} | Baseline MAE: {mae_baseline:.2f}")

# Seznamy pro UI
vsechny_zony = sorted(df['zdroj'].unique().tolist())
denny_typy = ['vsedni', 'sobota', 'nedele']
casova_okna = ['rano_spicka', 'dopoledne', 'odpoledne', 'vecer_spicka', 'vecer_pozde']
vekove_skupiny = sorted(df['vekova_skupina'].unique().tolist())
ucely = sorted(df['ucel'].unique().tolist())

# --- Predikční funkce ---
def predikuj(zdroj, cil, denni_typ, casove_okno, vekova_skupina, ucel, udalost, udalost_velikost, pocasi, rok=2025, ctvrtleti=2):
    try:
        row = {
            'rok': rok,
            'ctvrtleti': ctvrtleti,
            'denni_typ_enc': le_dict['denni_typ'].transform([denni_typ])[0] if denni_typ in le_dict['denni_typ'].classes_ else 0,
            'casove_okno_enc': le_dict['casove_okno'].transform([casove_okno])[0] if casove_okno in le_dict['casove_okno'].classes_ else 0,
            'vekova_skupina_enc': le_dict['vekova_skupina'].transform([vekova_skupina])[0] if vekova_skupina in le_dict['vekova_skupina'].classes_ else 0,
            'ucel_enc': le_dict['ucel'].transform([ucel])[0] if ucel in le_dict['ucel'].classes_ else 0,
            'hlavni_mod_enc': le_dict['hlavni_mod'].transform(['autobus'])[0],
            'udalost': 1 if udalost else 0,
            'udalost_velikost': float(udalost_velikost),
            'podil_spatne_pocasi': float(pocasi),
            'vzdalenost': float(df[((df['zdroj']==zdroj) & (df['cil']==cil))]['vzdalenost'].mean() or df['vzdalenost'].mean()),
        }

        zony_cil = zony[zony['nazev'] == cil]
        zony_src = zony[zony['nazev'] == zdroj]

        row['cil_velikost'] = float(zony_cil['velikost'].values[0]) if len(zony_cil) else 0.5
        row['cil_izolace'] = float(zony_cil['izolace'].values[0]) if len(zony_cil) else 0.5
        row['cil_skola'] = int(zony_cil['skola'].values[0]) if len(zony_cil) else 0
        row['cil_zamestnavatel'] = int(zony_cil['zamestnavatel'].values[0]) if len(zony_cil) else 0
        row['cil_uzel'] = int(zony_cil['uzel'].values[0]) if len(zony_cil) else 0
        row['zdroj_velikost'] = float(zony_src['velikost'].values[0]) if len(zony_src) else 0.5
        row['zdroj_izolace'] = float(zony_src['izolace'].values[0]) if len(zony_src) else 0.5

        X_input = pd.DataFrame([row])[feature_cols]
        pocet = float(model.predict(X_input)[0])
        pocet = max(0, round(pocet, 1))

        if pocet >= 8:
            doporuceni = "🟢 VYSLAT — dostatečná poptávka"
            flexbus = "Pevný spoj"
        elif pocet >= 4:
            doporuceni = "🟡 ČEKAT NA OBJEDNÁVKY — on-demand mikrobus"
            flexbus = "FlexBus (objednat 1 hod předem)"
        else:
            doporuceni = "🔴 NEVYSÍLAT — příliš nízká poptávka"
            flexbus = "Nevysílat (možnost taxi záchrana)"

        nejistota = "±2–4 cestující (model pracuje s trendem, ne s jistotou)"

        return (
            f"**Předpověď:** {pocet} cestujících",
            f"**Doporučení:** {doporuceni}",
            f"**Režim:** {flexbus}",
            f"**Přesnost modelu:** MAE = {mae_model:.2f} vs baseline {mae_baseline:.2f}",
            f"**Nejistota:** {nejistota}"
        )
    except Exception as e:
        return f"Chyba: {e}", "", "", "", ""

# --- Gradio UI ---
with gr.Blocks(title="FlexBus — Chytrá mobilita Plzeňský kraj") as demo:
    gr.Markdown("# 🚌 FlexBus Prediktor\n### AI doporučení pro dispečera mikrobusu — Plzeňský kraj")
    gr.Markdown(f"*Model trénován na reálných datech | MAE modelu: {mae_model:.2f} | Baseline MAE: {mae_baseline:.2f}*")

    with gr.Row():
        with gr.Column():
            zdroj_input = gr.Dropdown(vsechny_zony, label="Odkud (zdroj)", value=vsechny_zony[0])
            cil_input = gr.Dropdown(vsechny_zony, label="Kam (cíl)", value=vsechny_zony[1])
            denni_typ_input = gr.Dropdown(denny_typy, label="Typ dne", value="sobota")
            casove_okno_input = gr.Dropdown(casova_okna, label="Čas odjezdu", value="vecer_spicka")
            vekova_input = gr.Dropdown(vekove_skupiny, label="Věková skupina", value=vekove_skupiny[0])
            ucel_input = gr.Dropdown(ucely, label="Účel cesty", value=ucely[0])

        with gr.Column():
            udalost_input = gr.Checkbox(label="Koná se událost v cíli? (koncert, zápas...)")
            udalost_vel_input = gr.Slider(0, 5, value=0, step=1, label="Velikost události (0=žádná, 5=velká)")
            pocasi_input = gr.Slider(0.0, 1.0, value=0.3, step=0.1, label="Podíl špatného počasí (0=hezky, 1=hodně špatně)")

            btn = gr.Button("🔮 Predikovat poptávku", variant="primary")

            out1 = gr.Markdown()
            out2 = gr.Markdown()
            out3 = gr.Markdown()
            out4 = gr.Markdown()
            out5 = gr.Markdown()

    btn.click(
        predikuj,
        inputs=[zdroj_input, cil_input, denni_typ_input, casove_okno_input,
                vekova_input, ucel_input, udalost_input, udalost_vel_input, pocasi_input],
        outputs=[out1, out2, out3, out4, out5]
    )

    gr.Markdown("""
    ---
    ### Jak číst výsledek
    - 🟢 **8+ cestujících** → pevný spoj jede automaticky
    - 🟡 **4–7 cestujících** → FlexBus čeká na objednávky (cestující objednají 1 hod předem)
    - 🔴 **0–3 cestující** → spoj nejede, k dispozici záchranný taxi tarif

    *AI etika: pracujeme pouze s agregovanými počty, žádné sledování jednotlivců.*
    """)

demo.launch()