from __future__ import annotations

import calendar
import json
import logging
from datetime import datetime, timedelta, date
from pathlib import Path
from typing import Dict, List

import pandas as pd
import requests

from airflow import DAG
from airflow.decorators import task
from airflow.models import Variable
from airflow.operators.empty import EmptyOperator
from airflow.utils.task_group import TaskGroup

DATA_DIR = Path("/opt/airflow/data")
OUT_DIR = Path("/opt/airflow/output/bike_pollution_v2")
OUT_DIR.mkdir(parents=True, exist_ok=True)

default_args = {"owner": "sanzhar", "retries": 2, "retry_delay": timedelta(minutes=2)}


def _safe_read_csv(path: str) -> pd.DataFrame:
    return pd.read_csv(path)


def _safe_write_csv(df: pd.DataFrame, path: Path) -> str:
    df.to_csv(path, index=False)
    return str(path)


def _yyyymmdd(d: date) -> str:
    return d.strftime("%Y%m%d")


def _month_ranges(b: date, e: date):
    cur = date(b.year, b.month, 1)
    while cur <= e:
        last_day = calendar.monthrange(cur.year, cur.month)[1]
        month_end = date(cur.year, cur.month, last_day)
        yield max(b, cur), min(e, month_end)
        if cur.month == 12:
            cur = date(cur.year + 1, 1, 1)
        else:
            cur = date(cur.year, cur.month + 1, 1)


def _build_retry_session() -> requests.Session:
    """Requests session with retries for transient errors / rate limits."""
    from requests.adapters import HTTPAdapter
    from urllib3.util.retry import Retry

    s = requests.Session()
    retry = Retry(
        total=6,
        connect=6,
        read=6,
        backoff_factor=1.0,
        status_forcelist=[429, 500, 502, 503, 504],
        allowed_methods=["GET"],
        raise_on_status=False,
    )
    s.mount("https://", HTTPAdapter(max_retries=retry))
    return s


with DAG(
    dag_id="bike_pollution_reg_cls_pipeline",
    default_args=default_args,
    start_date=datetime(2025, 1, 1),
    schedule=None,
    catchup=False,
    tags=["etl", "ml", "regression", "classification"],
) as dag:

    @task
    def extract_bike_csv() -> str:
        name = Variable.get("BIKE_CSV_NAME", default_var="bike_hour.csv")
        fp = DATA_DIR / name
        if not fp.exists():
            raise FileNotFoundError(f"CSV not found: {fp}")
        return str(fp)

    @task
    def fetch_air_pollution_daily(bike_csv_path: str) -> str:
        log = logging.getLogger("airflow.task")

        df_bike = pd.read_csv(bike_csv_path)
        df_bike["dteday"] = pd.to_datetime(df_bike["dteday"], errors="coerce").dt.date
        df_bike = df_bike.dropna(subset=["dteday"]).copy()

        start_date = df_bike["dteday"].min()
        end_date = df_bike["dteday"].max()
        if not start_date or not end_date:
            raise RuntimeError("Could not determine date range from bike CSV (dteday missing/empty).")

        email = Variable.get("AQS_EMAIL")
        key = Variable.get("AQS_KEY")

        lat = float(Variable.get("CITY_LAT", default_var="38.9072"))
        lon = float(Variable.get("CITY_LON", default_var="-77.0369"))
        box = float(Variable.get("AQS_BOX_DEG", default_var="0.5"))

        params_map = {
            "pm25": "88101",
            "pm10": "81102",
            "no2": "42602",
            "o3": "44201",
        }

        base_url = "https://aqs.epa.gov/data/api/dailyData/byBox"
        sess = _build_retry_session()

        def _request_month(param_code: str, b: date, e: date) -> pd.DataFrame:
            r = sess.get(
                base_url,
                params={
                    "email": email,
                    "key": key,
                    "param": param_code,
                    "bdate": _yyyymmdd(b),
                    "edate": _yyyymmdd(e),
                    "minlat": lat - box,
                    "maxlat": lat + box,
                    "minlon": lon - box,
                    "maxlon": lon + box,
                },
                timeout=180,
            )

            try:
                js = r.json()
            except Exception:
                raise RuntimeError(f"AQS returned non-JSON response (status={r.status_code}).")

            header = (js.get("Header") or [{}])[0]
            status = header.get("status")
            if r.status_code >= 400 or status == "Failed":
                err = header.get("error", "")
                raise RuntimeError(
                    f"AQS request failed (status_code={r.status_code}). "
                    f"Details: {err}. URL: {header.get('url')}"
                )

            data = js.get("Data", [])
            if not data:
                return pd.DataFrame()

            tmp = pd.DataFrame(data)

            if "date_local" not in tmp.columns or "arithmetic_mean" not in tmp.columns:
                raise RuntimeError(f"AQS response missing expected fields for param {param_code}.")

            tmp["date"] = pd.to_datetime(tmp["date_local"], errors="coerce").dt.date
            tmp["value"] = pd.to_numeric(tmp["arithmetic_mean"], errors="coerce")
            tmp = tmp.dropna(subset=["date", "value"]).copy()

            if tmp.empty:
                return pd.DataFrame()

            daily = tmp.groupby("date", as_index=False)["value"].mean()
            return daily

        def _fetch_param_full(param_code: str, col_name: str) -> pd.DataFrame:
            parts = []
            for mb, me in _month_ranges(start_date, end_date):
                log.info("AQS %s (%s): %s..%s box=%s", col_name, param_code, mb, me, box)
                part = _request_month(param_code, mb, me)
                if not part.empty:
                    parts.append(part)

            if not parts:
                raise RuntimeError(
                    f"AQS returned NO DATA for '{col_name}' (param={param_code}) "
                    f"for {start_date}..{end_date}. "
                    f"Try changing CITY_LAT/CITY_LON or increasing AQS_BOX_DEG."
                )

            dfp = pd.concat(parts, ignore_index=True)
            dfp = dfp.groupby("date", as_index=False)["value"].mean()
            return dfp.rename(columns={"value": col_name})

        pm25 = _fetch_param_full(params_map["pm25"], "pm25")
        pm10 = _fetch_param_full(params_map["pm10"], "pm10")
        no2 = _fetch_param_full(params_map["no2"], "no2")
        o3 = _fetch_param_full(params_map["o3"], "o3")

        dates = pd.date_range(start_date, end_date, freq="D").date
        daily = pd.DataFrame({"date": dates})
        daily = daily.merge(pm25, on="date", how="left")
        daily = daily.merge(pm10, on="date", how="left")
        daily = daily.merge(no2, on="date", how="left")
        daily = daily.merge(o3, on="date", how="left")

        for c in ["pm25", "pm10", "no2", "o3"]:
            daily[c] = pd.to_numeric(daily[c], errors="coerce")

        missing_ratio = daily[["pm25", "pm10", "no2", "o3"]].isna().mean().max()
        if missing_ratio > 0.6:
            raise RuntimeError(
                f"AQS data is too sparse (max missing ratio {missing_ratio:.2f}). "
                f"Increase AQS_BOX_DEG or choose another CITY_LAT/LON."
            )

        daily = daily.sort_values("date").ffill().bfill()
        daily["aq_source"] = "epa_aqs"

        out = OUT_DIR / "air_pollution_daily.csv"
        return _safe_write_csv(daily, out)

    @task
    def build_dataset(bike_csv_path: str, air_daily_path: str) -> str:
        df = pd.read_csv(bike_csv_path)

        for col in ["instant", "casual", "registered"]:
            if col in df.columns:
                df = df.drop(columns=[col])

        df["dteday"] = pd.to_datetime(df["dteday"]).dt.date
        df["hr"] = pd.to_numeric(df["hr"], errors="coerce")
        df["cnt"] = pd.to_numeric(df["cnt"], errors="coerce")
        df = df.dropna(subset=["dteday", "hr", "cnt"]).copy()
        df["hr"] = df["hr"].astype(int)

        df["ts"] = pd.to_datetime(df["dteday"].astype(str)) + pd.to_timedelta(df["hr"], unit="h")
        df = df.sort_values("ts").reset_index(drop=True)

        air = pd.read_csv(air_daily_path)
        air["date"] = pd.to_datetime(air["date"]).dt.date

        merged = df.merge(air, left_on="dteday", right_on="date", how="left").drop(columns=["date"])
        for c in ["pm25", "pm10", "no2", "o3"]:
            merged[c] = pd.to_numeric(merged[c], errors="coerce")

        merged = merged.sort_values("ts").ffill()

        merged["y_next_hour"] = merged["cnt"].shift(-1)
        merged["lag_1"] = merged["cnt"].shift(1)
        merged["lag_24"] = merged["cnt"].shift(24)
        merged["roll_24_mean"] = merged["cnt"].shift(1).rolling(24).mean()
        merged["is_peak"] = merged["hr"].isin([7, 8, 9, 17, 18, 19]).astype(int)
        merged["is_weekend"] = merged["weekday"].isin([0, 6]).astype(int)

        merged = merged.dropna(subset=["y_next_hour", "lag_1", "lag_24", "roll_24_mean"]).copy()

        out = OUT_DIR / "dataset_features.csv"
        return _safe_write_csv(merged, out)

    @task
    def split_train_test(dataset_path: str) -> Dict[str, str]:
        df = pd.read_csv(dataset_path)
        df["ts"] = pd.to_datetime(df["ts"])

        cutoff = df["ts"].max() - pd.Timedelta(days=14)
        train = df[df["ts"] <= cutoff].copy()
        test = df[df["ts"] > cutoff].copy()

        q = float(train["y_next_hour"].quantile(0.75))
        train["y_class"] = (train["y_next_hour"] >= q).astype(int)
        test["y_class"] = (test["y_next_hour"] >= q).astype(int)

        train_path = OUT_DIR / "train.csv"
        test_path = OUT_DIR / "test.csv"
        train.to_csv(train_path, index=False)
        test.to_csv(test_path, index=False)

        meta_path = OUT_DIR / "classification_threshold.json"
        meta_path.write_text(json.dumps({"q75_threshold": q}, indent=2), encoding="utf-8")

        return {"train": str(train_path), "test": str(test_path), "meta": str(meta_path)}

    def _make_X_y_reg(df: pd.DataFrame):
        y = df["y_next_hour"].astype(float)
        X = df.drop(columns=["y_next_hour", "y_class", "ts", "dteday"], errors="ignore")
        return X, y

    def _make_X_y_cls(df: pd.DataFrame):
        y = df["y_class"].astype(int)
        X = df.drop(columns=["y_next_hour", "y_class", "ts", "dteday"], errors="ignore")
        return X, y

    def _build_preprocessor(X: pd.DataFrame):
        from sklearn.compose import ColumnTransformer
        from sklearn.preprocessing import OneHotEncoder, StandardScaler

        cat_candidates = ["season", "weathersit", "holiday", "workingday", "weekday", "mnth", "yr", "aq_source"]
        cat_cols = [c for c in cat_candidates if c in X.columns]
        num_cols = [c for c in X.columns if c not in cat_cols]

        return ColumnTransformer(
            transformers=[
                ("num", StandardScaler(), num_cols),
                ("cat", OneHotEncoder(handle_unknown="ignore"), cat_cols),
            ],
            remainder="drop",
        )

    with TaskGroup(group_id="regression_models") as regression_models:
        reg_start = EmptyOperator(task_id="start")

        @task
        def train_linear_regression(splits: Dict[str, str]) -> str:
            import joblib
            from sklearn.linear_model import LinearRegression
            from sklearn.pipeline import Pipeline
            from sklearn.metrics import mean_absolute_error, mean_squared_error

            train = _safe_read_csv(splits["train"])
            test = _safe_read_csv(splits["test"])

            X_train, y_train = _make_X_y_reg(train)
            X_test, y_test = _make_X_y_reg(test)

            pre = _build_preprocessor(X_train)
            model = Pipeline([("pre", pre), ("model", LinearRegression())])

            model.fit(X_train, y_train)
            preds = model.predict(X_test)

            mae = float(mean_absolute_error(y_test, preds))
            rmse = float(mean_squared_error(y_test, preds) ** 0.5)

            model_dir = OUT_DIR / "models"
            model_dir.mkdir(parents=True, exist_ok=True)
            joblib.dump(model, model_dir / "reg_linear.joblib")

            metrics = {"type": "regression", "model": "linear_regression", "mae": mae, "rmse": rmse}
            mpath = model_dir / "reg_linear_metrics.json"
            mpath.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
            return str(mpath)

        @task
        def train_random_forest_regressor(splits: Dict[str, str]) -> str:
            import joblib
            from sklearn.ensemble import RandomForestRegressor
            from sklearn.pipeline import Pipeline
            from sklearn.metrics import mean_absolute_error, mean_squared_error

            train = _safe_read_csv(splits["train"])
            test = _safe_read_csv(splits["test"])

            X_train, y_train = _make_X_y_reg(train)
            X_test, y_test = _make_X_y_reg(test)

            pre = _build_preprocessor(X_train)
            model = Pipeline([("pre", pre), ("model", RandomForestRegressor(
                n_estimators=250, random_state=42, n_jobs=-1
            ))])

            model.fit(X_train, y_train)
            preds = model.predict(X_test)

            mae = float(mean_absolute_error(y_test, preds))
            rmse = float(mean_squared_error(y_test, preds) ** 0.5)

            model_dir = OUT_DIR / "models"
            model_dir.mkdir(parents=True, exist_ok=True)
            joblib.dump(model, model_dir / "reg_rf.joblib")

            metrics = {"type": "regression", "model": "random_forest_regressor", "mae": mae, "rmse": rmse}
            mpath = model_dir / "reg_rf_metrics.json"
            mpath.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
            return str(mpath)

        @task(trigger_rule="none_failed_min_one_success")
        def select_best_regression(metric_paths: List[str]) -> str:
            best = None
            for p in metric_paths:
                m = json.loads(Path(p).read_text(encoding="utf-8"))
                if best is None or m["mae"] < best["mae"]:
                    best = m
            out = OUT_DIR / "best_regression.json"
            out.write_text(json.dumps(best, indent=2), encoding="utf-8")
            return str(out)

    with TaskGroup(group_id="classification_models") as classification_models:
        cls_start = EmptyOperator(task_id="start")

        @task
        def train_logistic_regression(splits: Dict[str, str]) -> str:
            import joblib
            from sklearn.linear_model import LogisticRegression
            from sklearn.pipeline import Pipeline
            from sklearn.metrics import f1_score, accuracy_score

            train = _safe_read_csv(splits["train"])
            test = _safe_read_csv(splits["test"])

            X_train, y_train = _make_X_y_cls(train)
            X_test, y_test = _make_X_y_cls(test)

            pre = _build_preprocessor(X_train)
            model = Pipeline([("pre", pre), ("model", LogisticRegression(max_iter=2000))])

            model.fit(X_train, y_train)
            preds = model.predict(X_test)

            acc = float(accuracy_score(y_test, preds))
            f1 = float(f1_score(y_test, preds))

            model_dir = OUT_DIR / "models"
            model_dir.mkdir(parents=True, exist_ok=True)
            joblib.dump(model, model_dir / "cls_logreg.joblib")

            metrics = {"type": "classification", "model": "logistic_regression", "accuracy": acc, "f1": f1}
            mpath = model_dir / "cls_logreg_metrics.json"
            mpath.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
            return str(mpath)

        @task
        def train_svm_classifier(splits: Dict[str, str]) -> str:
            import joblib
            from sklearn.svm import SVC
            from sklearn.pipeline import Pipeline
            from sklearn.metrics import f1_score, accuracy_score

            train = _safe_read_csv(splits["train"])
            test = _safe_read_csv(splits["test"])

            X_train, y_train = _make_X_y_cls(train)
            X_test, y_test = _make_X_y_cls(test)

            pre = _build_preprocessor(X_train)
            model = Pipeline([("pre", pre), ("model", SVC(kernel="rbf", C=2.0, gamma="scale"))])

            model.fit(X_train, y_train)
            preds = model.predict(X_test)

            acc = float(accuracy_score(y_test, preds))
            f1 = float(f1_score(y_test, preds))

            model_dir = OUT_DIR / "models"
            model_dir.mkdir(parents=True, exist_ok=True)
            joblib.dump(model, model_dir / "cls_svm.joblib")

            metrics = {"type": "classification", "model": "svm", "accuracy": acc, "f1": f1}
            mpath = model_dir / "cls_svm_metrics.json"
            mpath.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
            return str(mpath)

        @task(trigger_rule="none_failed_min_one_success")
        def select_best_classification(metric_paths: List[str]) -> str:
            best = None
            for p in metric_paths:
                m = json.loads(Path(p).read_text(encoding="utf-8"))
                if best is None or m["f1"] > best["f1"]:
                    best = m
            out = OUT_DIR / "best_classification.json"
            out.write_text(json.dumps(best, indent=2), encoding="utf-8")
            return str(out)

    @task(trigger_rule="none_failed_min_one_success")
    def save_final_summary(best_reg_path: str, best_cls_path: str, splits: Dict[str, str]) -> str:
        best_reg = json.loads(Path(best_reg_path).read_text(encoding="utf-8"))
        best_cls = json.loads(Path(best_cls_path).read_text(encoding="utf-8"))
        threshold = json.loads(Path(splits["meta"]).read_text(encoding="utf-8"))

        summary = {
            "best_regression": best_reg,
            "best_classification": best_cls,
            "classification_threshold": threshold,
        }
        out = OUT_DIR / "final_summary.json"
        out.write_text(json.dumps(summary, indent=2), encoding="utf-8")
        return str(out)

    start = EmptyOperator(task_id="start")
    end = EmptyOperator(task_id="end")

    bike_csv = extract_bike_csv()
    air_daily = fetch_air_pollution_daily(bike_csv)
    dataset = build_dataset(bike_csv, air_daily)
    splits = split_train_test(dataset)

    start >> bike_csv >> air_daily >> dataset >> splits

    reg_m1 = train_linear_regression(splits)
    reg_m2 = train_random_forest_regressor(splits)
    reg_best = select_best_regression([reg_m1, reg_m2])
    regression_models >> reg_best

    cls_m1 = train_logistic_regression(splits)
    cls_m2 = train_svm_classifier(splits)
    cls_best = select_best_classification([cls_m1, cls_m2])
    classification_models >> cls_best

    splits >> regression_models
    splits >> classification_models

    summary = save_final_summary(reg_best, cls_best, splits)
    [reg_best, cls_best] >> summary >> end