"""
Museum visitors vs. city population: regression model + FastAPI service.

Run locally (from the repo root):

    uvicorn api.museum_api:app --reload

or run this file directly (e.g. from the IDE):

    python api/museum_api.py
"""
from __future__ import annotations

import os
import sys
from pathlib import Path

if __package__ in (None, ""):
    # Started as a script (python museum_api.py): put the repo root on the
    # path so the `api` package imports below resolve.
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from dataclasses import dataclass, field
from typing import Literal
import numpy as np
import pandas as pd
from app import data_setup, constants
from fastapi import FastAPI, HTTPException, Query
from pydantic import BaseModel, Field
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import KFold, cross_val_score
from dotenv import load_dotenv

load_dotenv()

REQUIRED_COLUMNS = ["Name", "Visitors", "City", "Country", "Population"]


# --------------------------------------------------------------------------- #
# Data + model
# --------------------------------------------------------------------------- #
def clean(df: pd.DataFrame) -> pd.DataFrame:
    """Validate columns, coerce numerics, drop rows that can't be modelled."""
    missing = set(REQUIRED_COLUMNS) - set(df.columns)
    if missing:
        raise ValueError(f"Missing columns: {sorted(missing)}")

    out = df[REQUIRED_COLUMNS].copy()
    for col in ("Visitors", "Population"):
        # Handles strings like "1,234,567"
        out[col] = pd.to_numeric(
            out[col].astype(str).str.replace(r"[,\s]", "", regex=True),
            errors="coerce",
        )
    out = out.dropna(subset=["Visitors", "Population"])
    out = out[(out["Visitors"] > 0) & (out["Population"] > 0)]
    return out.reset_index(drop=True)


@dataclass
class FittedModel:
    reg: LinearRegression
    scale: Literal["log", "linear"]
    data: pd.DataFrame  # the rows actually used for fitting
    n: int
    r2: float
    cv_r2: float | None
    pearson_r: float
    spearman_rho: float
    aggregate_by_city: bool
    predictions: np.ndarray = field(repr=False)

    def _x(self, population: np.ndarray) -> np.ndarray:
        x = np.asarray(population, dtype=float).reshape(-1, 1)
        return np.log10(x) if self.scale == "log" else x

    def predict(self, population) -> np.ndarray:
        y = self.reg.predict(self._x(population))
        return 10 ** y if self.scale == "log" else y

    @property
    def slope(self) -> float:
        return float(self.reg.coef_[0])

    @property
    def intercept(self) -> float:
        return float(self.reg.intercept_)

    def summary(self) -> dict:
        s = {
            "n_samples": self.n,
            "scale": self.scale,
            "aggregate_by_city": self.aggregate_by_city,
            "slope": self.slope,
            "intercept": self.intercept,
            "r2": self.r2,
            "cv_r2_5fold": self.cv_r2,
            "pearson_r": self.pearson_r,
            "spearman_rho": self.spearman_rho,
        }
        if self.scale == "log":
            s["equation"] = (
                f"visitors = {10 ** self.intercept:,.1f} * population^{self.slope:.3f}"
            )
            s["interpretation"] = (
                f"A 10x larger city is associated with "
                f"{10 ** self.slope:.2f}x the visitors."
            )
        else:
            s["equation"] = f"visitors = {self.slope:.4f} * population + {self.intercept:,.0f}"
            s["interpretation"] = (
                f"Each additional 1M residents is associated with "
                f"{self.slope * 1e6:,.0f} more visitors."
            )
        return s


def fit_model(
    df: pd.DataFrame,
    scale: Literal["log", "linear"] = "log",
    aggregate_by_city: bool = False,
) -> FittedModel:
    """
    Fit visitors ~ population.

    scale="log" (default) fits log10(visitors) ~ log10(population). Both
    variables are heavily right-skewed, so a log-log fit is usually far
    better behaved than a linear one and the slope reads as an elasticity.

    aggregate_by_city=True sums visitors per (City, Country) first, so a
    city with many museums counts once rather than once per museum.
    """
    data = clean(df)
    if aggregate_by_city:
        data = (
            data.groupby(["City", "Country"], as_index=False)
            .agg(Visitors=("Visitors", "sum"),
                 Population=("Population", "first"),
                 Name=("Name", lambda s: ", ".join(s)))
        )
    if len(data) < 3:
        raise ValueError(f"Need at least 3 usable rows, got {len(data)}")

    x_raw = data["Population"].to_numpy(float)
    y_raw = data["Visitors"].to_numpy(float)
    X = (np.log10(x_raw) if scale == "log" else x_raw).reshape(-1, 1)
    y = np.log10(y_raw) if scale == "log" else y_raw

    reg = LinearRegression().fit(X, y)

    cv_r2 = None
    if len(data) >= 10:
        cv = KFold(n_splits=5, shuffle=True, random_state=0)
        cv_r2 = float(cross_val_score(LinearRegression(), X, y, cv=cv, scoring="r2").mean())

    xs, ys = pd.Series(X.ravel()), pd.Series(y)
    model = FittedModel(
        reg=reg,
        scale=scale,
        data=data,
        n=len(data),
        r2=float(reg.score(X, y)),
        cv_r2=cv_r2,
        pearson_r=float(xs.corr(ys)),
        # Spearman = Pearson on ranks (avoids a scipy dependency)
        spearman_rho=float(pd.Series(x_raw).rank().corr(pd.Series(y_raw).rank())),
        aggregate_by_city=aggregate_by_city,
        predictions=np.empty(0),
    )
    model.predictions = model.predict(x_raw)
    return model


# --------------------------------------------------------------------------- #
# API schemas
# --------------------------------------------------------------------------- #
class MuseumRecord(BaseModel):
    Name: str
    Visitors: float = Field(gt=0)
    City: str
    Country: str
    Population: float = Field(gt=0)


class TrainRequest(BaseModel):
    records: list[MuseumRecord] = Field(min_length=3)
    scale: Literal["log", "linear"] = "log"
    aggregate_by_city: bool = False


class PredictRequest(BaseModel):
    populations: list[float] = Field(min_length=1)


class Prediction(BaseModel):
    population: float
    predicted_visitors: float


# --------------------------------------------------------------------------- #
# App factory
# --------------------------------------------------------------------------- #
def create_app(
    df: pd.DataFrame | None = None,
    scale: Literal["log", "linear"] = "log",
    aggregate_by_city: bool = False,
) -> FastAPI:
    app = FastAPI(title="Museum Visitors Regression", version="1.0")
    app.state.model = fit_model(df, scale, aggregate_by_city) if df is not None else None

    def get_model() -> FittedModel:
        if app.state.model is None:
            raise HTTPException(503, "No model trained yet. POST /train first.")
        return app.state.model

    @app.get("/health")
    def health():
        return {"status": "ok", "model_loaded": app.state.model is not None}

    @app.get("/model")
    def model_summary():
        return get_model().summary()

    @app.post("/train")
    def train(req: TrainRequest):
        df_new = pd.DataFrame([r.model_dump() for r in req.records])
        try:
            app.state.model = fit_model(df_new, req.scale, req.aggregate_by_city)
        except ValueError as e:
            raise HTTPException(422, str(e))
        return app.state.model.summary()

    @app.get("/predict", response_model=Prediction)
    def predict_one(population: float = Query(gt=0)):
        m = get_model()
        return Prediction(population=population,
                          predicted_visitors=float(m.predict([population])[0]))

    @app.post("/predict", response_model=list[Prediction])
    def predict_many(req: PredictRequest):
        if any(p <= 0 for p in req.populations):
            raise HTTPException(422, "Populations must be > 0")
        m = get_model()
        preds = m.predict(req.populations)
        return [Prediction(population=p, predicted_visitors=float(v))
                for p, v in zip(req.populations, preds)]

    @app.get("/data")
    def data():
        """Every row the model was fitted on, with its prediction."""
        m = get_model()
        d = m.data.copy()
        d["Predicted"] = m.predictions
        return d.to_dict(orient="records")

    @app.get("/residuals")
    def residuals(
        top: int = Query(10, ge=1, le=500),
        order: Literal["over", "under"] = "over",
    ):
        """Museums that most over/under-perform what their city size predicts."""
        m = get_model()
        d = m.data.copy()
        d["Predicted"] = m.predictions
        d["Ratio"] = d["Visitors"] / d["Predicted"]  # >1 = outperforms city size
        d = d.sort_values("Ratio", ascending=(order == "under")).head(top)
        return d[["Name", "City", "Country", "Population",
                  "Visitors", "Predicted", "Ratio"]].to_dict(orient="records")

    return app


def _load_data() -> pd.DataFrame | None:
    path = constants.MUSEUM_DATA_FILE

    if not path.exists():
        print("Fetching and conforming museum data and city population...")
        museum_data = data_setup.get_museum_data()
        museum_data.to_csv(path)
    else:
        museum_data = pd.read_csv(path)

    return museum_data


# Module-level app so `uvicorn api.museum_api:app` works.
app = create_app(
    _load_data(),
    scale=os.environ.get("MUSEUM_SCALE", "log"),  # type: ignore[arg-type]
    aggregate_by_city=os.environ.get("MUSEUM_AGG_CITY", "0") == "1",
)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        app,
        host=os.environ.get("HOST", "127.0.0.1"),
        port=int(os.environ.get("PORT", "8000")),
    )