"""
FlexBus — model_train.py
========================
Spusť JEDNOU před nasazením, nebo při aktualizaci dat.

Co dělá:
  1. Natrénuje GradientBoosting model na historických datech (2019–2025)
  2. Uloží model + encodery jako model.pkl  →  app.py ho načte při startu
  3. Natrénuje lineární trend model pro každou oblast → předpovídá vývoj
     poptávky do budoucna (po čtvrtletích)
  4. Uloží trend předpovědi jako trend_forecast.csv
  5. Vypíše přehled výkonu + feature importance

Spuštění:
  python model_train.py

Výstupní soubory (nahraj spolu s app.py na HuggingFace):
  model.pkl            — natrénovaný GBR model + encodery
  trend_forecast.csv   — čtvrtletní předpovědi 2026–2027 pro každou oblast
"""

import warnings
warnings.filterwarnings('ignore')

import pandas as pd
import numpy as np
import pickle
import datetime

from sklearn.ensemble import GradientBoostingRegressor
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import mean_absolute_error, r2_score

# ─────────────────────────────────────────────────────────────
# 1. NAČTENÍ A ČIŠTĚNÍ DAT
# ─────────────────────────────────────────────────────────────
print("=" * 60)
print("FlexBus — trénování modelu")
print("=" * 60)

df   = pd.read_csv('02_simpleml_komplet_pohyb.csv')
zony = pd.read_csv('00_zony.csv')

print(f"\n📂 Načteno {len(df)} záznamů, {df['rok'].min()}–{df['rok'].max()}")

df = df.dropna()
df = df[df['pocet_cest'] >= 0]
df = df[df['vzdalenost'] >= 0]
print(f"   Po čištění: {len(df)} záznamů")

# ─────────────────────────────────────────────────────────────
# 2. FEATURE ENGINEERING
# ─────────────────────────────────────────────────────────────
zony_cil = zony[['nazev','velikost','izolace','skola','zamestnavatel','uzel']].copy()
zony_cil.columns = ['cil','cil_velikost','cil_izolace','cil_skola','cil_zamestnavatel','cil_uzel']
df = df.merge(zony_cil, 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))

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

FEAT = [
    '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[FEAT]
y = df['pocet_cest']

# ─────────────────────────────────────────────────────────────
# 3. TRÉNINK HLAVNÍHO MODELU
# ─────────────────────────────────────────────────────────────
print("\n🔧 Trénuji GradientBoosting model...")

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,
    min_samples_leaf=3,
    subsample=0.85,
    random_state=42
)
model.fit(X_train, y_train)

y_pred     = model.predict(X_test)
mae_model  = mean_absolute_error(y_test, y_pred)
mae_base   = mean_absolute_error(y_test, [y_train.mean()] * len(y_test))
r2         = r2_score(y_test, y_pred)
zlepseni   = round((1 - mae_model / mae_base) * 100, 1)

print(f"\n📊 Výsledky modelu:")
print(f"   MAE (model):    {mae_model:.2f} cestujících")
print(f"   MAE (baseline): {mae_base:.2f} cestujících")
print(f"   Zlepšení:       {zlepseni} %")
print(f"   R² skóre:       {r2:.3f}")

# Cross-validace (5-fold) — pro porotu
cv_scores = cross_val_score(model, X, y, cv=5, scoring='neg_mean_absolute_error')
cv_mae    = -cv_scores.mean()
print(f"   Cross-val MAE:  {cv_mae:.2f} ± {cv_scores.std():.2f} (5-fold)")

# Feature importance
print("\n🔍 Feature importance (top 10):")
feat_imp = sorted(
    zip(FEAT, model.feature_importances_),
    key=lambda x: -x[1]
)
feat_labels = {
    'rok': 'Rok', 'ctvrtleti': 'Čtvrtletí', 'denni_typ_enc': 'Typ dne',
    'casove_okno_enc': 'Čas odjezdu', 'vekova_skupina_enc': 'Věk cestujících',
    'ucel_enc': 'Účel cesty', 'hlavni_mod_enc': 'Dopravní mód',
    'udalost': 'Událost v cíli', 'udalost_velikost': 'Velikost události',
    'podil_spatne_pocasi': 'Počasí', 'vzdalenost': 'Vzdálenost trasy',
    'cil_velikost': 'Velikost cílové obce', 'cil_izolace': 'Izolovanost cíle',
    'cil_skola': 'Škola v cíli', 'cil_zamestnavatel': 'Zaměstnavatel v cíli',
    'cil_uzel': 'Dopravní uzel v cíli', 'zdroj_velikost': 'Velikost výchozí obce',
    'zdroj_izolace': 'Izolovanost výchozí obce'
}
for feat, imp in feat_imp[:10]:
    bar = "█" * int(imp * 200)
    print(f"   {feat_labels.get(feat, feat):<28} {bar:<20} {imp:.4f}")

# ─────────────────────────────────────────────────────────────
# 4. TREND FORECASTING — předpověď do budoucna po čtvrtletích
# ─────────────────────────────────────────────────────────────
print("\n📈 Trénuji trend model (lineární regrese per oblast)...")

# Agreguj historická data na čtvrtletní úrovni per oblast
ts = df.groupby(['rok', 'ctvrtleti', 'zdroj'])['pocet_cest'].sum().reset_index()
ts['cas_idx'] = (ts['rok'] - ts['rok'].min()) * 4 + (ts['ctvrtleti'] - 1)

forecast_rows = []
oblasti = ts['zdroj'].unique()

# Budoucí čtvrtletí: 2026 Q1-Q4 + 2027 Q1-Q4
future_quarters = []
for rok in [2026, 2027]:
    for q in [1, 2, 3, 4]:
        future_quarters.append({'rok': rok, 'ctvrtleti': q})

cas_min = ts['rok'].min()

for oblast in oblasti:
    sub = ts[ts['zdroj'] == oblast].sort_values('cas_idx')
    if len(sub) < 3:
        continue  # příliš málo dat pro trend

    X_t = sub['cas_idx'].values.reshape(-1, 1)
    y_t = sub['pocet_cest'].values

    # Lineární trend
    lr = LinearRegression()
    lr.fit(X_t, y_t)

    trend_r2 = r2_score(y_t, lr.predict(X_t))

    # Sezónní korekce — průměr dle čtvrtletí
    sezon = sub.groupby('ctvrtleti')['pocet_cest'].mean()
    global_mean = sub['pocet_cest'].mean()
    sezon_koef = {q: (sezon.get(q, global_mean) / global_mean if global_mean > 0 else 1.0)
                  for q in [1, 2, 3, 4]}

    # Predikce budoucích čtvrtletí
    for fq in future_quarters:
        cas_idx = (fq['rok'] - cas_min) * 4 + (fq['ctvrtleti'] - 1)
        pred_base = float(lr.predict([[cas_idx]])[0])
        pred_sez  = pred_base * sezon_koef.get(fq['ctvrtleti'], 1.0)
        pred_final = max(0.0, round(pred_sez, 1))

        # Směr trendu
        if lr.coef_[0] > 0.5:
            trend_dir = "↑ Rostoucí"
        elif lr.coef_[0] < -0.5:
            trend_dir = "↓ Klesající"
        else:
            trend_dir = "→ Stabilní"

        forecast_rows.append({
            'oblast':         oblast,
            'rok':            fq['rok'],
            'ctvrtleti':      fq['ctvrtleti'],
            'predikce_cest':  pred_final,
            'trend':          trend_dir,
            'trend_sklon':    round(float(lr.coef_[0]), 3),
            'trend_r2':       round(trend_r2, 3),
        })

forecast_df = pd.DataFrame(forecast_rows)
forecast_df.to_csv('trend_forecast.csv', index=False)

print(f"   Zpracováno {len(oblasti)} oblastí → trend_forecast.csv")
print("\n   Ukázka předpovědí:")
print(forecast_df[forecast_df['rok'] == 2026][
    ['oblast', 'ctvrtleti', 'predikce_cest', 'trend']
].to_string(index=False))

# ─────────────────────────────────────────────────────────────
# 5. ULOŽENÍ MODELU
# ─────────────────────────────────────────────────────────────
payload = {
    'model':      model,
    'le_dict':    le_dict,
    'mae_model':  mae_model,
    'mae_base':   mae_base,
    'zlepseni':   zlepseni,
    'r2':         r2,
    'cv_mae':     cv_mae,
    'feat_imp':   feat_imp,
    'feat_labels': feat_labels,
    'features':   FEAT,
}

with open('model.pkl', 'wb') as f:
    pickle.dump(payload, f)

print(f"\n✅ model.pkl uložen")
print(f"✅ trend_forecast.csv uložen")
print(f"\nNahraj oba soubory na HuggingFace Space spolu s app.py.")
print("=" * 60)