from __future__ import annotations

import json
import sys
from dataclasses import dataclass
from datetime import date, datetime, time, timedelta
from pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo

import joblib
import numpy as np
import pandas as pd
import yfinance as yf


IST = ZoneInfo("Asia/Kolkata")
YAHOO_NIFTY_SYMBOL = "^NSEI"
BACKEND_ROOT = Path(__file__).resolve().parents[1]
DATA_DIR = BACKEND_ROOT / "data"
MODEL_DIR = BACKEND_ROOT / "models"
OPENING_DATASET_PATH = DATA_DIR / "opening_direction_training_dataset.parquet"
NIFTY_1M_PATH = DATA_DIR / "nifty50_1m.parquet"
NIFTY_1D_PATH = DATA_DIR / "nifty50_1d.parquet"
MODEL_PATH = MODEL_DIR / "nifty_opening_direction_model.joblib"
LATEST_PATH = MODEL_DIR / "latest_prediction.csv"
TEST_PREDICTIONS_PATH = DATA_DIR / "test_predictions.parquet"

DECISION_OVERLAYS = [
    {
        "name": "fifth_minute_momentum_flip",
        "feature": "m5_ret_1m",
        "op": ">=",
        "value": 0.0005085411885759201,
    },
    {
        "name": "vix_stretch_flip",
        "feature": "india_vix_close_vs_sma_20",
        "op": ">=",
        "value": 0.24641908937959742,
    },
]


class ProbabilityBlend:
    def __init__(self, models: list[Any], weights: np.ndarray):
        self.models = models
        self.weights = np.asarray(weights, dtype="float64")
        self.weights = self.weights / self.weights.sum()

    def predict_proba(self, x: pd.DataFrame) -> np.ndarray:
        probs = np.column_stack([predict_proba_up(model, x) for model in self.models])
        prob_up = probs @ self.weights
        return np.column_stack([1.0 - prob_up, prob_up])


@dataclass(frozen=True)
class Prediction:
    input_date: str
    first5_start: str
    first5_end: str
    prediction: str
    prob_up: float
    confidence: float
    threshold: float
    model_name: str

    def to_dict(self) -> dict[str, Any]:
        return {
            "input_date": self.input_date,
            "first5_start": self.first5_start,
            "first5_end": self.first5_end,
            "prediction": self.prediction,
            "prob_up": self.prob_up,
            "confidence": self.confidence,
            "threshold": self.threshold,
            "model_name": self.model_name,
        }


def predict_proba_up(model: Any, x: pd.DataFrame) -> np.ndarray:
    return np.asarray(model.predict_proba(x)[:, 1], dtype="float64")


def safe_div(numer: pd.Series | np.ndarray, denom: pd.Series | np.ndarray) -> pd.Series:
    n = pd.Series(numer, copy=False)
    d = pd.Series(denom, copy=False)
    out = pd.Series(np.nan, index=n.index, dtype="float64")
    mask = d.notna() & np.isfinite(d.to_numpy(dtype="float64")) & (d != 0)
    out.loc[mask] = n.loc[mask].to_numpy(dtype="float64") / d.loc[mask].to_numpy(dtype="float64")
    return out


def load_model() -> dict[str, Any]:
    # Existing artifact was trained as a script, so its custom blend class
    # resolves through __main__ when unpickled.
    sys.modules["__main__"].ProbabilityBlend = ProbabilityBlend
    sys.modules["__main__"].predict_proba_up = predict_proba_up
    payload = joblib.load(MODEL_PATH)
    payload.setdefault("decision_overlays", DECISION_OVERLAYS)
    payload.setdefault("model_name", "nifty_opening_direction_model")
    return payload


def overlay_mask(frame: pd.DataFrame, overlay: dict[str, object]) -> np.ndarray:
    feature = str(overlay["feature"])
    if feature not in frame.columns:
        return np.zeros(len(frame), dtype=bool)
    series = pd.to_numeric(frame[feature], errors="coerce")
    value = float(overlay["value"])
    if overlay["op"] == ">=":
        return (series >= value).fillna(False).to_numpy(dtype=bool)
    if overlay["op"] == "<=":
        return (series <= value).fillna(False).to_numpy(dtype=bool)
    raise ValueError(f"Unsupported overlay op: {overlay['op']}")


def apply_decision_overlays(pred: np.ndarray, frame: pd.DataFrame, overlays: list[dict[str, object]]) -> np.ndarray:
    adjusted = np.asarray(pred, dtype="int64").copy()
    for overlay in overlays:
        mask = overlay_mask(frame, overlay)
        adjusted[mask] = 1 - adjusted[mask]
    return adjusted


def directional_confidence(prob_up: np.ndarray, pred: np.ndarray, threshold: float) -> np.ndarray:
    prob_up = np.asarray(prob_up, dtype="float64")
    pred = np.asarray(pred, dtype="int64")
    base_side_prob = np.where(pred == 1, prob_up, 1.0 - prob_up)
    threshold_distance = np.abs(prob_up - float(threshold))
    return np.clip(0.50 + threshold_distance, base_side_prob, 0.99)


def read_training_dataset() -> pd.DataFrame:
    df = pd.read_parquet(OPENING_DATASET_PATH)
    for col in ("date", "first5_start", "first5_end"):
        if col in df.columns:
            df[col] = pd.to_datetime(df[col], errors="coerce")
    return df.sort_values("date").reset_index(drop=True)


def normalize_yahoo_frame(df: pd.DataFrame) -> pd.DataFrame:
    if df.empty:
        return pd.DataFrame(columns=["date", "open", "high", "low", "close", "volume"])
    if isinstance(df.columns, pd.MultiIndex):
        df.columns = [str(c[0]).lower() for c in df.columns]
    else:
        df.columns = [str(c).lower().replace(" ", "_") for c in df.columns]
    df = df.reset_index()
    date_col = next((c for c in df.columns if c.lower() in {"datetime", "date"}), df.columns[0])
    df["date"] = pd.to_datetime(df[date_col], errors="coerce")
    if df["date"].dt.tz is None:
        df["date"] = df["date"].dt.tz_localize("UTC").dt.tz_convert(IST)
    else:
        df["date"] = df["date"].dt.tz_convert(IST)
    rename = {
        "open": "open",
        "high": "high",
        "low": "low",
        "close": "close",
        "adj_close": "close",
        "volume": "volume",
    }
    out = pd.DataFrame({"date": df["date"].dt.tz_localize(None)})
    for src, dst in rename.items():
        if src in df.columns and dst not in out.columns:
            out[dst] = pd.to_numeric(df[src], errors="coerce")
    return out.dropna(subset=["date", "open", "high", "low", "close"]).sort_values("date")


def fetch_yahoo_minutes(period: str = "5d") -> pd.DataFrame:
    raw = yf.download(YAHOO_NIFTY_SYMBOL, period=period, interval="1m", progress=False, prepost=False, auto_adjust=False)
    return normalize_yahoo_frame(raw)


def fetch_yahoo_daily(period: str = "1mo") -> pd.DataFrame:
    raw = yf.download(YAHOO_NIFTY_SYMBOL, period=period, interval="1d", progress=False, prepost=False, auto_adjust=False)
    out = normalize_yahoo_frame(raw)
    out["date"] = pd.to_datetime(out["date"], errors="coerce").dt.normalize()
    return out.drop_duplicates("date", keep="last")


def append_parquet_rows(path: Path, new_rows: pd.DataFrame, subset: list[str]) -> pd.DataFrame:
    if path.exists():
        existing = pd.read_parquet(path)
        combined = pd.concat([existing, new_rows], ignore_index=True)
    else:
        combined = new_rows.copy()
    combined = combined.drop_duplicates(subset=subset, keep="last").sort_values(subset).reset_index(drop=True)
    combined.to_parquet(path, index=False, compression="zstd")
    return combined


def first5_features_from_minutes(minutes: pd.DataFrame, session_date: date | None = None) -> pd.DataFrame:
    if minutes.empty:
        raise RuntimeError("Yahoo returned no minute bars.")
    bars = minutes.copy()
    bars["dt"] = pd.to_datetime(bars["date"], errors="coerce")
    bars["session_date"] = bars["dt"].dt.normalize()
    if session_date is None:
        session_ts = bars["session_date"].max()
    else:
        session_ts = pd.Timestamp(session_date).normalize()
    day = bars[bars["session_date"] == session_ts].sort_values("dt").copy()
    start_dt = pd.Timestamp.combine(session_ts.date(), time(9, 15))
    end_dt = pd.Timestamp.combine(session_ts.date(), time(9, 19))
    first5 = day[(day["dt"] >= start_dt) & (day["dt"] <= end_dt)].head(5).copy()
    if len(first5) < 5:
        raise RuntimeError(f"Need 5 opening bars for {session_ts.date()}, got {len(first5)}.")
    first5["minute_index"] = np.arange(len(first5))
    first5["ret_1m"] = first5["close"].pct_change(fill_method=None)
    first5["range_pct_1m"] = safe_div(first5["high"] - first5["low"], first5["open"])
    first5["body_pct_1m"] = safe_div(first5["close"] - first5["open"], first5["open"])
    row = {
        "date": session_ts,
        "first5_start": first5["dt"].iloc[0],
        "first5_end": first5["dt"].iloc[-1],
        "first5_open": first5["open"].iloc[0],
        "first5_high": first5["high"].max(),
        "first5_low": first5["low"].min(),
        "first5_close": first5["close"].iloc[-1],
        "first5_volume": first5["volume"].sum() if "volume" in first5 else 0.0,
        "first5_bars": len(first5),
        "first5_last_1m_ret": first5["ret_1m"].iloc[-1],
        "first5_ret_std": first5["ret_1m"].std(),
    }
    row["first5_return"] = (row["first5_close"] - row["first5_open"]) / row["first5_open"]
    row["first5_range_pct"] = (row["first5_high"] - row["first5_low"]) / row["first5_open"]
    first5_range = row["first5_high"] - row["first5_low"]
    row["first5_body_to_range"] = (row["first5_close"] - row["first5_open"]) / first5_range if first5_range else np.nan
    row["first5_close_location"] = (row["first5_close"] - row["first5_low"]) / first5_range if first5_range else np.nan
    for idx, (_, candle) in enumerate(first5.iterrows(), start=1):
        for field in ("open", "high", "low", "close", "ret_1m", "range_pct_1m", "body_pct_1m"):
            row[f"m{idx}_{field}"] = candle[field]
        row[f"m{idx}_close_vs_first5_open"] = (candle["close"] - row["first5_open"]) / row["first5_open"]
        row[f"m{idx}_range_share"] = (candle["high"] - candle["low"]) / first5_range if first5_range else np.nan
    row["first5_return_accel"] = row["m5_ret_1m"] - row["m2_ret_1m"]
    row["first5_last2_return"] = (row["m5_close"] - row["m4_open"]) / row["m4_open"]
    row["first5_first2_return"] = (row["m2_close"] - row["m1_open"]) / row["m1_open"]
    row["first5_reversal"] = np.sign(row["first5_first2_return"]) * -np.sign(row["first5_last2_return"])
    row["dow"] = session_ts.dayofweek
    row["dom"] = session_ts.day
    row["month"] = session_ts.month
    return pd.DataFrame([row])


def build_model_row(first5_row: pd.DataFrame) -> pd.DataFrame:
    dataset = read_training_dataset()
    latest_context = dataset.iloc[[-1]].copy()
    output = latest_context.copy()
    for col in first5_row.columns:
        output[col] = first5_row[col].iloc[0]
    if {"first5_open", "nifty_close"}.issubset(output.columns):
        output["first5_gap_from_prev_close"] = (output["first5_open"] - output["nifty_close"]) / output["nifty_close"]
        output["first5_close_vs_prev_close"] = (output["first5_close"] - output["nifty_close"]) / output["nifty_close"]
    if {"first5_range_pct", "nifty_range_pct"}.issubset(output.columns):
        output["first5_range_vs_prev_range"] = output["first5_range_pct"] / output["nifty_range_pct"]
    if {"first5_return", "nifty_ret_1"}.issubset(output.columns):
        output["first5_return_x_prev_ret"] = output["first5_return"] * output["nifty_ret_1"]
        output["gap_x_prev_ret"] = output["first5_gap_from_prev_close"] * output["nifty_ret_1"]
    if {"first5_return", "banknifty_ret_1"}.issubset(output.columns):
        output["first5_return_x_bank_ret_1"] = output["first5_return"] * output["banknifty_ret_1"]
    if {"first5_range_pct", "india_vix_ret_1"}.issubset(output.columns):
        output["first5_range_x_vix_ret_1"] = output["first5_range_pct"] * output["india_vix_ret_1"]
    output["target"] = np.nan
    output["day_return"] = np.nan
    return output


def predict_row(row: pd.DataFrame) -> Prediction:
    payload = load_model()
    model = payload["model"]
    features = payload["features"]
    threshold = float(payload["threshold"])
    missing = [c for c in features if c not in row.columns]
    if missing:
        raise RuntimeError(f"Feature row is missing {len(missing)} features; first missing: {missing[:5]}")
    prob_up = predict_proba_up(model, row[features])
    raw_pred = (prob_up >= threshold).astype("int64")
    pred = apply_decision_overlays(raw_pred, row, payload.get("decision_overlays", DECISION_OVERLAYS))
    confidence = directional_confidence(prob_up, pred, threshold)
    prediction = Prediction(
        input_date=pd.to_datetime(row["date"].iloc[0]).date().isoformat(),
        first5_start=str(pd.to_datetime(row["first5_start"].iloc[0])),
        first5_end=str(pd.to_datetime(row["first5_end"].iloc[0])),
        prediction="UP" if int(pred[0]) == 1 else "DOWN",
        prob_up=float(prob_up[0]),
        confidence=float(confidence[0]),
        threshold=threshold,
        model_name=str(payload.get("model_name", "nifty_opening_direction_model")),
    )
    pd.DataFrame([prediction.to_dict()]).to_csv(LATEST_PATH, index=False)
    return prediction


def latest_saved_prediction() -> dict[str, Any]:
    if LATEST_PATH.exists():
        return pd.read_csv(LATEST_PATH).iloc[-1].to_dict()
    summary_path = MODEL_DIR / "summary.json"
    if summary_path.exists():
        return json.loads(summary_path.read_text(encoding="utf-8"))
    raise FileNotFoundError("No latest prediction is available yet.")


def _json_ready_frame(df: pd.DataFrame, limit: int | None = None) -> list[dict[str, Any]]:
    out = df.copy()
    if limit is not None:
        out = out.tail(limit)
    for col in out.columns:
        if pd.api.types.is_datetime64_any_dtype(out[col]):
            out[col] = out[col].dt.strftime("%Y-%m-%d %H:%M:%S")
    out = out.replace({np.nan: None})
    return out.to_dict(orient="records")


def load_model_summary() -> dict[str, Any]:
    summary_path = MODEL_DIR / "summary.json"
    if not summary_path.exists():
        return {}
    return json.loads(summary_path.read_text(encoding="utf-8"))


def load_candidate_results() -> list[dict[str, Any]]:
    path = MODEL_DIR / "candidate_results.csv"
    if not path.exists():
        return []
    return _json_ready_frame(pd.read_csv(path).head(12))


def load_test_predictions() -> pd.DataFrame:
    if not TEST_PREDICTIONS_PATH.exists():
        return pd.DataFrame()
    df = pd.read_parquet(TEST_PREDICTIONS_PATH)
    df["date"] = pd.to_datetime(df["date"], errors="coerce")
    return df.sort_values("date").reset_index(drop=True)


def dashboard_payload() -> dict[str, Any]:
    summary = load_model_summary()
    latest = latest_saved_prediction()
    test = load_test_predictions()
    daily = pd.read_parquet(NIFTY_1D_PATH)
    daily["date"] = pd.to_datetime(daily["date"], errors="coerce")
    daily = daily.sort_values("date").tail(180)
    dataset = read_training_dataset()
    opening = dataset[["date", "first5_return", "first5_range_pct", "first5_close_location"]].tail(120).copy()

    if not test.empty:
        recent_predictions = test.tail(40).copy()
        recent_accuracy = float(recent_predictions["correct"].mean())
        direction_mix = test.groupby("prediction")["correct"].agg(["count", "mean"]).reset_index()
        monthly = (
            test.assign(month=test["date"].dt.strftime("%Y-%m"))
            .groupby("month", as_index=False)["correct"]
            .mean()
            .rename(columns={"correct": "accuracy"})
        )
    else:
        recent_predictions = pd.DataFrame()
        recent_accuracy = None
        direction_mix = pd.DataFrame()
        monthly = pd.DataFrame()

    metrics = {
        "validation_accuracy": summary.get("validation_accuracy"),
        "test_accuracy": summary.get("test_accuracy"),
        "baseline_test_accuracy": summary.get("baseline_test_accuracy"),
        "validation_auc": summary.get("validation_auc"),
        "test_auc": summary.get("test_auc"),
        "test_brier": summary.get("test_brier"),
        "feature_count": summary.get("feature_count"),
        "recent_accuracy": recent_accuracy,
        "recent_accuracy_days": int(len(recent_predictions)) if not recent_predictions.empty else 0,
        "total_test_days": int(len(test)) if not test.empty else 0,
    }
    return {
        "latest": latest,
        "metrics": metrics,
        "summary": summary,
        "candidates": load_candidate_results(),
        "charts": {
            "daily_close": _json_ready_frame(daily[["date", "open", "high", "low", "close"]]),
            "opening_features": _json_ready_frame(opening),
            "monthly_accuracy": _json_ready_frame(monthly),
            "direction_mix": _json_ready_frame(direction_mix),
            "recent_predictions": _json_ready_frame(recent_predictions),
        },
        "data_status": {
            "nifty_1m_rows": int(len(pd.read_parquet(NIFTY_1M_PATH, columns=["date"]))),
            "nifty_1d_rows": int(len(pd.read_parquet(NIFTY_1D_PATH, columns=["date"]))),
            "training_rows": int(len(dataset)),
            "test_prediction_rows": int(len(test)),
            "latest_daily_date": pd.to_datetime(daily["date"]).max().date().isoformat(),
        },
    }


def refresh_first5_prediction(session_date: date | None = None) -> Prediction:
    minutes = fetch_yahoo_minutes(period="5d")
    append_parquet_rows(NIFTY_1M_PATH, minutes, ["date"])
    first5 = first5_features_from_minutes(minutes, session_date=session_date)
    row = build_model_row(first5)
    dataset = read_training_dataset()
    merged = pd.concat([dataset, row], ignore_index=True)
    merged = merged.drop_duplicates(subset=["date"], keep="last").sort_values("date").reset_index(drop=True)
    merged.to_parquet(OPENING_DATASET_PATH, index=False, compression="zstd")
    return predict_row(row)


def refresh_daily_data() -> dict[str, Any]:
    daily = fetch_yahoo_daily(period="1mo")
    combined = append_parquet_rows(NIFTY_1D_PATH, daily, ["date"])
    return {
        "rows": int(len(combined)),
        "latest_date": pd.to_datetime(combined["date"]).max().date().isoformat(),
        "path": str(NIFTY_1D_PATH),
    }


def seconds_until_next_ist_run(run_time: time = time(9, 20)) -> float:
    now = datetime.now(IST)
    target = datetime.combine(now.date(), run_time, tzinfo=IST)
    if now >= target:
        target += timedelta(days=1)
    return max(1.0, (target - now).total_seconds())