"""
Advanced ML Model for Mumbai Local Train Crowd Density Prediction
Uses RandomForest with enhanced feature engineering
"""

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import cross_val_score
from sklearn.metrics import accuracy_score, classification_report
import joblib
import os

LABEL_MAP = {"Low": 0, "Medium": 1, "Extreme": 2}
REVERSE_LABEL_MAP = {0: "Low", 1: "Medium", 2: "Extreme"}

DAY_ORDER = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]
DAY_MAP = {d: i for i, d in enumerate(DAY_ORDER)}

HOUR_WEIGHTS = {
    5: 0.3, 6: 0.8, 7: 1.9, 8: 2.0, 9: 1.7,
    10: 1.1, 11: 1.0, 12: 1.0, 13: 1.0, 14: 1.0,
    15: 1.1, 16: 1.3, 17: 1.8, 18: 2.0, 19: 1.9,
    20: 1.4, 21: 0.8, 22: 0.5
}

LINE_MAP = {"Western": 0, "Central": 1, "Harbour": 2}
DIR_MAP = {"Northbound": 0, "Southbound": 1}

MODEL_PATH = os.path.join(os.path.dirname(__file__), "../data/model.pkl")
ENCODERS_PATH = os.path.join(os.path.dirname(__file__), "../data/encoders.pkl")


def engineer_features(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["day_num"] = df["day_of_week"].map(DAY_MAP)
    df["is_weekday"] = (df["day_num"] < 5).astype(int)
    df["is_weekend"] = (df["day_num"] >= 5).astype(int)
    df["hour_weight"] = df["hour"].map(HOUR_WEIGHTS).fillna(0.5)
    df["line_num"] = df["line"].map(LINE_MAP).fillna(0)
    df["dir_num"] = df["direction"].map(DIR_MAP).fillna(0)

    # Peak hour flags
    df["is_morning_peak"] = ((df["hour"] >= 7) & (df["hour"] <= 9)).astype(int)
    df["is_evening_peak"] = ((df["hour"] >= 17) & (df["hour"] <= 19)).astype(int)
    df["is_peak"] = ((df["is_morning_peak"] == 1) | (df["is_evening_peak"] == 1)).astype(int)
    df["is_offpeak"] = ((df["hour"] < 6) | (df["hour"] > 21)).astype(int)

    # Directional rush interactions
    df["morning_southbound"] = ((df["is_morning_peak"] == 1) & (df["dir_num"] == 1)).astype(int)
    df["evening_northbound"] = ((df["is_evening_peak"] == 1) & (df["dir_num"] == 0)).astype(int)
    df["rush_dir_interaction"] = df["morning_southbound"] + df["evening_northbound"]

    # Weekend peak dampening
    df["effective_weight"] = df["hour_weight"] * (1.2 if 1 else 1.0)
    df["weekday_weight"] = df["hour_weight"] * df["is_weekday"] * 1.2 + df["hour_weight"] * df["is_weekend"] * 0.8

    # Station encoding (label encoding)
    le_station = LabelEncoder()
    le_dest = LabelEncoder()
    all_stations = list(pd.concat([df["station"], df["destination"]]).unique())
    le_station.fit(all_stations)
    le_dest.fit(all_stations)
    df["station_enc"] = le_station.transform(df["station"])
    df["dest_enc"] = le_dest.transform(df["destination"])

    return df, le_station, le_dest


FEATURE_COLS = [
    "hour", "day_num", "is_weekday", "is_weekend",
    "hour_weight", "line_num", "dir_num",
    "is_morning_peak", "is_evening_peak", "is_peak", "is_offpeak",
    "morning_southbound", "evening_northbound", "rush_dir_interaction",
    "weekday_weight", "station_enc", "dest_enc"
]


def train_model(df: pd.DataFrame):
    df_feat, le_station, le_dest = engineer_features(df)
    df_feat["crowd_label"] = df_feat["crowd_level"].map(LABEL_MAP)

    X = df_feat[FEATURE_COLS]
    y = df_feat["crowd_label"]

    model = RandomForestClassifier(
        n_estimators=200,
        max_depth=12,
        min_samples_split=4,
        random_state=42,
        n_jobs=-1,
        class_weight="balanced"
    )
    model.fit(X, y)

    preds = model.predict(X)
    acc = accuracy_score(y, preds)
    report = classification_report(y, preds, target_names=["Low", "Medium", "Extreme"])

    # Save model
    os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
    joblib.dump(model, MODEL_PATH)
    joblib.dump({"le_station": le_station, "le_dest": le_dest}, ENCODERS_PATH)

    return model, acc, report, le_station, le_dest


def load_model():
    if os.path.exists(MODEL_PATH) and os.path.exists(ENCODERS_PATH):
        model = joblib.load(MODEL_PATH)
        encoders = joblib.load(ENCODERS_PATH)
        return model, encoders["le_station"], encoders["le_dest"]
    return None, None, None


def predict_crowd(
    model, le_station, le_dest,
    line: str, station: str, destination: str,
    direction: str, hour: int, day_of_week: str
) -> dict:
    try:
        station_enc = le_station.transform([station])[0]
        dest_enc = le_dest.transform([destination])[0]
    except ValueError:
        station_enc = 0
        dest_enc = 0

    day_num = DAY_MAP.get(day_of_week, 0)
    is_weekday = int(day_num < 5)
    is_weekend = int(day_num >= 5)
    hour_weight = HOUR_WEIGHTS.get(hour, 0.5)
    line_num = LINE_MAP.get(line, 0)
    dir_num = DIR_MAP.get(direction, 0)
    is_morning_peak = int(7 <= hour <= 9)
    is_evening_peak = int(17 <= hour <= 19)
    is_peak = int(is_morning_peak or is_evening_peak)
    is_offpeak = int(hour < 6 or hour > 21)
    morning_southbound = int(is_morning_peak and dir_num == 1)
    evening_northbound = int(is_evening_peak and dir_num == 0)
    rush_dir = morning_southbound + evening_northbound
    weekday_weight = hour_weight * 1.2 if is_weekday else hour_weight * 0.8

    features = [[
        hour, day_num, is_weekday, is_weekend,
        hour_weight, line_num, dir_num,
        is_morning_peak, is_evening_peak, is_peak, is_offpeak,
        morning_southbound, evening_northbound, rush_dir,
        weekday_weight, station_enc, dest_enc
    ]]

    pred = model.predict(features)[0]
    proba = model.predict_proba(features)[0]

    label = REVERSE_LABEL_MAP[pred]
    confidence = float(max(proba)) * 100

    return {
        "crowd_level": label,
        "crowd_code": int(pred),
        "confidence": round(confidence, 1),
        "probabilities": {
            "Low": round(float(proba[0]) * 100, 1),
            "Medium": round(float(proba[1]) * 100, 1),
            "Extreme": round(float(proba[2]) * 100, 1),
        }
    }