from __future__ import annotations

from functools import lru_cache
from pathlib import Path

import numpy as np
import pandas as pd

TARGETS = ["mill_power_mw", "feed_rate_tph", "target_impurity_pct"]
TARGET_LABELS = {
    "mill_power_mw": "Mill Power (MW)",
    "feed_rate_tph": "Feed Rate (tph)",
    "target_impurity_pct": "Impurity (%)",
}
CONTEXT_WINDOWS = {
    "48 hours": 48,
    "96 hours (4 days)": 96,
    "168 hours (1 week)": 168,
}
HORIZONS = {
    "12 hours": 12,
    "24 hours": 24,
    "48 hours": 48,
}
RANDOM_SEED = 42
ASSETS_DIR = Path(__file__).resolve().parent.parent / "assets"
TELEMETRY_FILE = "telemetry_synthetic.parquet"


def generate_mineral_processing_telemetry(n_timestamps: int = 1000) -> pd.DataFrame:
    """Synthetic comminution/flotation plant telemetry."""
    rng = np.random.default_rng(RANDOM_SEED)
    timestamps = pd.date_range("2024-01-01", periods=n_timestamps, freq="h", tz="UTC")

    hour = np.arange(n_timestamps) % 24
    day = np.arange(n_timestamps) // 24

    shift_cycle = 0.6 * np.sin(2 * np.pi * hour / 24) + 0.25 * np.sin(2 * np.pi * hour / 12)
    hardness_drift = 0.15 * np.sin(2 * np.pi * day / 14)

    feed_rate_tph = 285 + 55 * shift_cycle + 20 * hardness_drift + rng.normal(0, 8, n_timestamps)
    feed_rate_tph = np.clip(feed_rate_tph, 180, 420)

    mill_power_mw = (
        7.5
        + 0.018 * feed_rate_tph
        + 2.2 * hardness_drift
        + 0.35 * shift_cycle
        + rng.normal(0, 0.25, n_timestamps)
    )
    mill_power_mw = np.clip(mill_power_mw, 5.0, 16.0)

    impurity_driver = (
        2.4
        - 0.004 * feed_rate_tph
        + 0.35 * hardness_drift
        + rng.normal(0, 0.08, n_timestamps)
    )
    target_impurity_pct = np.zeros(n_timestamps)
    target_impurity_pct[0] = impurity_driver[0]
    for t in range(1, n_timestamps):
        target_impurity_pct[t] = 0.82 * target_impurity_pct[t - 1] + 0.18 * impurity_driver[t]
    target_impurity_pct = np.clip(target_impurity_pct, 0.4, 4.5)

    return pd.DataFrame(
        {
            "timestamp": timestamps,
            "mill_power_mw": mill_power_mw.round(3),
            "feed_rate_tph": feed_rate_tph.round(2),
            "target_impurity_pct": target_impurity_pct.round(4),
        }
    )


def ensure_telemetry_asset() -> Path:
    """Persist synthetic telemetry to assets/ for fast Space startup."""
    ASSETS_DIR.mkdir(parents=True, exist_ok=True)
    path = ASSETS_DIR / TELEMETRY_FILE
    if not path.exists():
        generate_mineral_processing_telemetry().to_parquet(path, index=False)
    return path


@lru_cache(maxsize=1)
def load_telemetry() -> pd.DataFrame:
    path = ensure_telemetry_asset()
    frame = pd.read_parquet(path)
    frame["timestamp"] = pd.to_datetime(frame["timestamp"], utc=True)
    return frame.sort_values("timestamp").reset_index(drop=True)


def telemetry_summary(telemetry: pd.DataFrame) -> pd.DataFrame:
    rows = []
    for col in TARGETS:
        rows.append(
            {
                "Signal": TARGET_LABELS[col],
                "Rows": len(telemetry),
                "Start (UTC)": telemetry["timestamp"].min(),
                "End (UTC)": telemetry["timestamp"].max(),
                "Min": round(telemetry[col].min(), 4),
                "Max": round(telemetry[col].max(), 4),
                "Mean": round(telemetry[col].mean(), 4),
                "Std": round(telemetry[col].std(), 4),
            }
        )
    return pd.DataFrame(rows)


def correlation_matrix(telemetry: pd.DataFrame) -> pd.DataFrame:
    return telemetry[TARGETS].corr().round(3)


def build_context_frame(
    telemetry: pd.DataFrame,
    target: str,
    anchor_time: pd.Timestamp,
    context_steps: int,
) -> pd.DataFrame:
    history = telemetry.loc[telemetry["timestamp"] <= anchor_time].tail(context_steps)
    if len(history) < max(24, context_steps // 4):
        raise ValueError("Not enough history before the selected anchor time.")
    return pd.DataFrame(
        {
            "item_id": target,
            "timestamp": history["timestamp"],
            "target": history[target].astype(float),
        }
    )


def holdout_actuals(
    telemetry: pd.DataFrame,
    target: str,
    anchor_time: pd.Timestamp,
    horizon: int,
) -> pd.DataFrame:
    future = telemetry.loc[
        (telemetry["timestamp"] > anchor_time)
        & (telemetry["timestamp"] <= anchor_time + pd.Timedelta(hours=horizon))
    ]
    if future.empty:
        return pd.DataFrame(columns=["timestamp", "actual"])
    return future[["timestamp", target]].rename(columns={target: "actual"})


def compute_metrics(actual: pd.Series, predicted: pd.Series) -> dict[str, float]:
    diff = actual.values - predicted.values
    mae = float(np.mean(np.abs(diff)))
    rmse = float(np.sqrt(np.mean(diff**2)))
    mape = float(np.mean(np.abs(diff / np.clip(actual.values, 1e-6, None))) * 100)
    return {"mae": mae, "rmse": rmse, "mape_pct": mape}