"""
Mumbai Local Train Crowd Density Predictor
FastAPI Backend - Advanced Edition
"""

import os
import sys
import json
from typing import Optional
from pathlib import Path

import pandas as pd
import numpy as np
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel

# Fix path
sys.path.insert(0, str(Path(__file__).parent))

from src.model import (
    train_model, load_model, predict_crowd,
    LINE_MAP, DIR_MAP, DAY_MAP, HOUR_WEIGHTS
)
from src.risk_engine import (
    get_risk_info, build_reasoning, get_best_hours,
    get_travel_tip, RISK_CONFIG,
    WESTERN_STATIONS_ORDER, CENTRAL_STATIONS_ORDER, HARBOUR_STATIONS_ORDER
)

# ─────────────────────────────────────────
# App Setup
# ─────────────────────────────────────────
app = FastAPI(
    title="Mumbai Train Crowd Predictor",
    description="AI-powered crowd density prediction for Mumbai Local Trains",
    version="2.0.0"
)

BASE_DIR = Path(__file__).parent
DATA_PATH = BASE_DIR / "data" / "raw" / "mumbai_local_crowd_data.csv"

app.mount("/static", StaticFiles(directory=str(BASE_DIR / "static")), name="static")
templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))

# ─────────────────────────────────────────
# Global State
# ─────────────────────────────────────────
MODEL = None
LE_STATION = None
LE_DEST = None
DF = None
MODEL_ACCURACY = None
MODEL_REPORT = None

STATIONS_BY_LINE = {
    "Western": WESTERN_STATIONS_ORDER,
    "Central": CENTRAL_STATIONS_ORDER,
    "Harbour": HARBOUR_STATIONS_ORDER
}


@app.on_event("startup")
async def startup_event():
    global MODEL, LE_STATION, LE_DEST, DF, MODEL_ACCURACY, MODEL_REPORT

    print("📂 Loading dataset...")
    DF = pd.read_csv(DATA_PATH)

    # Try loading cached model first
    MODEL, LE_STATION, LE_DEST = load_model()
    if MODEL is None:
        print("🧠 Training model...")
        MODEL, MODEL_ACCURACY, MODEL_REPORT, LE_STATION, LE_DEST = train_model(DF)
        print(f"✅ Model trained! Accuracy: {MODEL_ACCURACY:.2%}")
    else:
        # Quick retrain to get accuracy
        _, MODEL_ACCURACY, MODEL_REPORT, _, _ = train_model(DF)
        print(f"✅ Model loaded! Accuracy: {MODEL_ACCURACY:.2%}")


# ─────────────────────────────────────────
# Schemas
# ─────────────────────────────────────────
class PredictRequest(BaseModel):
    line: str
    station: str
    destination: str
    direction: str
    hour: int
    day_of_week: str


class HourlyRequest(BaseModel):
    line: str
    station: str
    destination: str
    direction: str
    day_of_week: str


# ─────────────────────────────────────────
# Routes
# ─────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})


@app.get("/api/stations/{line}")
async def get_stations(line: str):
    if line not in STATIONS_BY_LINE:
        raise HTTPException(status_code=404, detail="Line not found")
    return {"stations": STATIONS_BY_LINE[line]}


@app.get("/api/lines")
async def get_lines():
    return {"lines": list(STATIONS_BY_LINE.keys())}


@app.get("/api/days")
async def get_days():
    return {"days": list(DAY_MAP.keys())}


@app.get("/api/hours")
async def get_hours():
    return {"hours": list(HOUR_WEIGHTS.keys())}


@app.post("/api/predict")
async def predict(req: PredictRequest):
    if MODEL is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")

    result = predict_crowd(
        MODEL, LE_STATION, LE_DEST,
        req.line, req.station, req.destination,
        req.direction, req.hour, req.day_of_week
    )

    risk_info = get_risk_info(result["crowd_level"])
    reasons = build_reasoning(
        req.line, req.station, req.destination,
        req.direction, req.hour, req.day_of_week,
        result["crowd_level"]
    )
    tip = get_travel_tip(result["crowd_level"], req.direction, req.hour)

    return {
        "crowd_level": result["crowd_level"],
        "crowd_code": result["crowd_code"],
        "confidence": result["confidence"],
        "probabilities": result["probabilities"],
        "risk_score": risk_info["score"],
        "risk_color": risk_info["color"],
        "risk_bg": risk_info["bg"],
        "risk_emoji": risk_info["emoji"],
        "risk_badge": risk_info["badge"],
        "advice": risk_info["advice"],
        "travel_tip": tip,
        "reasons": reasons,
        "input": {
            "line": req.line,
            "station": req.station,
            "destination": req.destination,
            "direction": req.direction,
            "hour": req.hour,
            "day": req.day_of_week
        }
    }


@app.post("/api/hourly-forecast")
async def hourly_forecast(req: HourlyRequest):
    if MODEL is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")

    results = get_best_hours(
        MODEL, LE_STATION, LE_DEST,
        req.line, req.station, req.destination,
        req.direction, req.day_of_week,
        predict_crowd
    )

    # Enrich with colors
    for r in results:
        info = RISK_CONFIG[r["crowd_level"]]
        r["color"] = info["color"]
        r["bg"] = info["bg"]
        r["emoji"] = info["emoji"]
        r["score"] = info["score"]

    best = sorted(results, key=lambda x: x["risk_score"])[:3]

    return {
        "hourly": results,
        "best_hours": best,
        "peak_hour": max(results, key=lambda x: x["risk_score"])
    }


@app.get("/api/stats")
async def get_stats():
    if DF is None:
        raise HTTPException(status_code=503, detail="Data not loaded")

    total = len(DF)
    by_line = DF.groupby("line").size().to_dict()
    by_crowd = DF.groupby("crowd_level").size().to_dict()

    return {
        "total_records": total,
        "by_line": by_line,
        "by_crowd": by_crowd,
        "model_accuracy": round(MODEL_ACCURACY * 100, 2) if MODEL_ACCURACY else None,
        "lines": list(STATIONS_BY_LINE.keys()),
        "total_stations": sum(len(v) for v in STATIONS_BY_LINE.values())
    }


@app.get("/api/heatmap/{line}")
async def get_heatmap(line: str, direction: str = "Southbound", day: str = "Monday"):
    if MODEL is None or line not in STATIONS_BY_LINE:
        raise HTTPException(status_code=404)

    stations = STATIONS_BY_LINE[line]
    hours = list(range(5, 23))
    heatmap = []

    for station in stations[:6]:  # Limit for performance
        row = {"station": station, "hours": {}}
        for h in hours:
            dest = stations[-1] if station != stations[-1] else stations[0]
            try:
                pred = predict_crowd(MODEL, LE_STATION, LE_DEST, line, station, dest, direction, h, day)
                row["hours"][h] = {
                    "level": pred["crowd_level"],
                    "score": RISK_CONFIG[pred["crowd_level"]]["score"],
                    "color": RISK_CONFIG[pred["crowd_level"]]["color"]
                }
            except Exception:
                row["hours"][h] = {"level": "Low", "score": 15, "color": "#10b981"}
        heatmap.append(row)

    return {"heatmap": heatmap, "hours": hours, "stations": stations[:6]}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)