"""
Public-facing WastePredictor.

This is the only class consumers need to import:

    from model_lib import WastePredictor
    from model_lib import predict_waste           # convenience function

The model is loaded lazily on first use.  Re-entrant and thread-safe
(the underlying sklearn/torch models are stateless at inference time).
"""
from __future__ import annotations

import threading
from typing import Dict, List, Optional

import pandas as pd

from model_lib._core.loader import load_model_file
from model_lib.types import (
    OUTPUT_COLUMNS,
    PredictionInput,
    PredictionResult,
)


class WastePredictor:
    """
    Production waste predictor (V4 ensemble, R² = 0.98).

    Parameters
    ----------
    model_path:
        Absolute path to ``waste_predictor_v4.pkl``.
        If omitted the loader searches the default locations:
          1. ``$WASTE_PREDICTOR_MODEL_PATH``
          2. ``~/.model_lib/waste_predictor_v4.pkl``
          3. ``<package>/artifacts/waste_predictor_v4.pkl``
          4. ``./waste_predictor_v4.pkl``

    Examples
    --------
    >>> from model_lib import WastePredictor
    >>> predictor = WastePredictor()
    >>> result = predictor.predict(
    ...     production_volume=50_000,
    ...     rain_sum=200,
    ...     temperature_mean=28,
    ...     humidity_mean=85,
    ...     wind_speed_mean=15,
    ...     month=6,
    ... )
    >>> print(result.total_waste_kg)
    """

    OUTPUT_COLS: List[str] = OUTPUT_COLUMNS

    def __init__(self, model_path: Optional[str] = None) -> None:
        self._model_path = model_path
        self._model = None
        self._lock = threading.Lock()

    # ── Lazy loading ──────────────────────────────────────────────────────

    def _ensure_loaded(self) -> None:
        if self._model is None:
            with self._lock:
                if self._model is None:  # double-checked locking
                    self._model = load_model_file(self._model_path)

    @property
    def model(self):
        self._ensure_loaded()
        return self._model

    @property
    def version(self) -> str:
        return str(self.model.test_metrics.get("r2", "unknown"))

    @property
    def metadata(self) -> dict:
        return {
            "model_weights": self.model.weights,
            "output_cols": self.model.output_cols,
            "test_metrics": self.model.test_metrics,
        }

    # ── Model management ──────────────────────────────────────────────────

    def download_and_update_model(
        self,
        s3_uri: str,
        destination: Optional[str] = None,
    ) -> dict:
        """
        Download a new model from S3 and update the current predictor.

        Parameters
        ----------
        s3_uri:
            S3 URI of the model file (e.g., s3://bucket/path/model.pkl)
        destination:
            Optional local path to save the model. If None, saves to the
            default location: ``~/.model_lib/waste_predictor_v4.pkl``

        Returns
        -------
        dict
            Metadata about the newly loaded model including version and metrics.

        Raises
        ------
        ImportError:
            If boto3 is not installed.
        ValueError:
            If the S3 URI is invalid.
        RuntimeError:
            If the download or model loading fails.

        Examples
        --------
        >>> predictor = WastePredictor()
        >>> info = predictor.download_and_update_model(
        ...     s3_uri="s3://my-bucket/models/waste_predictor_v5.pkl"
        ... )
        >>> print(f"Updated to model version: {info['version']}")
        """
        from model_lib._core.loader import download_model_from_s3, load_model_file

        # Download the model
        downloaded_path = download_model_from_s3(s3_uri, destination)

        # Reload the model with thread safety
        with self._lock:
            self._model_path = str(downloaded_path)
            self._model = load_model_file(self._model_path)

        # Return updated metadata
        return {
            "path": str(downloaded_path),
            "version": self.version,
            "metadata": self.metadata,
        }

    # ── Single prediction ─────────────────────────────────────────────────

    def predict(
        self,
        production_volume: float,
        rain_sum: float,
        temperature_mean: float,
        humidity_mean: float,
        wind_speed_mean: float,
        month: int = 6,
        year: int = 2026,
    ) -> PredictionResult:
        """
        Single-row prediction with validated inputs.

        Returns
        -------
        PredictionResult
            Typed wrapper with all eight waste fractions.
        """
        inp = PredictionInput(
            production_volume=production_volume,
            rain_sum=rain_sum,
            temperature_mean=temperature_mean,
            humidity_mean=humidity_mean,
            wind_speed_mean=wind_speed_mean,
            month=month,
            year=year,
        )
        return self.predict_typed(inp)

    def predict_typed(self, inp: PredictionInput) -> PredictionResult:
        """Accept a typed PredictionInput and return a PredictionResult."""
        df = pd.DataFrame(
            [
                {
                    "Year": inp.year,
                    "Month": inp.month,
                    "production_volume": inp.production_volume,
                    "rain_sum": inp.rain_sum,
                    "temperature_mean": inp.temperature_mean,
                    "humidity_mean": inp.humidity_mean,
                    "wind_speed_mean": inp.wind_speed_mean,
                }
            ]
        )
        result_df = self.predict_batch(df)
        row: Dict[str, float] = result_df.iloc[0].to_dict()
        return PredictionResult.from_dict(row, model_version=self.version)

    def predict_dict(self, input_data: Dict) -> Dict[str, float]:
        """
        Accept a plain dict (backward-compatible with predict_api.py).

        Keys: production_volume, rain_sum, temperature_mean,
              humidity_mean, wind_speed_mean, month (optional)
        """
        result = self.predict(
            production_volume=input_data["production_volume"],
            rain_sum=input_data["rain_sum"],
            temperature_mean=input_data["temperature_mean"],
            humidity_mean=input_data["humidity_mean"],
            wind_speed_mean=input_data["wind_speed_mean"],
            month=int(input_data.get("month", 6)),
            year=int(input_data.get("year", 2026)),
        )
        return result.to_dict()

    # ── Batch prediction ──────────────────────────────────────────────────

    def predict_batch(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Batch inference on a DataFrame.

        Required columns: production_volume, rain_sum, temperature_mean,
                          humidity_mean, wind_speed_mean
        Optional columns: Month (int 1-12), Year (int)

        Returns
        -------
        pd.DataFrame with eight waste-fraction columns.
        """
        self._ensure_loaded()
        df = df.copy()
        if "Year" not in df.columns:
            df["Year"] = 2026
        if "Month" not in df.columns:
            df["Month"] = 6

        import numpy as np

        preds = [
            model.predict(df).values
            for model in self.model.models.values()
        ]
        weights = list(self.model.weights.values())
        weighted = np.average(preds, axis=0, weights=weights)
        return pd.DataFrame(weighted, columns=self.OUTPUT_COLS)


# ── Module-level convenience function (drop-in for predict_api.get_waste_prediction) ──

_default_predictor: Optional[WastePredictor] = None
_predictor_lock = threading.Lock()


def predict_waste(
    production_volume: float,
    rain_sum: float,
    temperature_mean: float,
    humidity_mean: float,
    wind_speed_mean: float,
    month: int = 6,
    year: int = 2026,
    model_path: Optional[str] = None,
) -> Dict[str, float]:
    """
    Module-level convenience function.  Identical signature to the original
    ``predict_waste`` in local-module/predict.py.

    Uses a process-wide singleton — safe for multi-threaded servers.
    """
    global _default_predictor
    if _default_predictor is None:
        with _predictor_lock:
            if _default_predictor is None:
                _default_predictor = WastePredictor(model_path)

    return _default_predictor.predict_dict(
        {
            "production_volume": production_volume,
            "rain_sum": rain_sum,
            "temperature_mean": temperature_mean,
            "humidity_mean": humidity_mean,
            "wind_speed_mean": wind_speed_mean,
            "month": month,
            "year": year,
        }
    )
