"""Concept drift monitoring utilities.

Provides :class:`DriftMonitor`, which computes Population Stability Index (PSI)
and Jensen-Shannon (JS) divergence between a reference feature distribution
(captured at training time) and a live inference batch.  Optionally integrates
with :class:`~pwml.utilities.neptunehelpers.ExperimentManager` to log drift
metrics to Neptune.
"""

__all__ = ['DriftMonitor']

import logging
import warnings
from typing import Dict, Optional, Tuple
import numpy as np
import pandas as pd
from scipy.spatial.distance import jensenshannon

_log = logging.getLogger(__name__)

# PSI thresholds (industry convention)
_PSI_LOW = 0.1
_PSI_HIGH = 0.2


def _psi(reference: np.ndarray, current: np.ndarray, bins: int = 10) -> float:
    """Compute Population Stability Index between two continuous distributions.

    PSI < 0.1  -> no significant change
    PSI < 0.2  -> moderate change, worth monitoring
    PSI >= 0.2 -> significant change, action recommended

    Args:
        reference: 1-D float array of reference values.
        current:   1-D float array of current (live) values.
        bins:      Number of equal-width bins derived from reference range.

    Returns:
        PSI as a non-negative float.
    """
    if len(reference) == 0 or len(current) == 0:
        return float('nan')

    # Use reference-derived bin edges so both distributions use the same grid.
    min_val = float(np.min(reference))
    max_val = float(np.max(reference))
    if min_val == max_val:
        # Degenerate case: all values identical in reference
        return 0.0

    # Extend edges slightly to avoid boundary zero-count bins
    eps = (max_val - min_val) * 1e-6
    edges = np.linspace(min_val - eps, max_val + eps, bins + 1)

    ref_counts, _ = np.histogram(reference, bins=edges)
    cur_counts, _ = np.histogram(current, bins=edges)

    ref_pct = ref_counts / len(reference)
    cur_pct = cur_counts / len(current)

    # Avoid log(0) by replacing zeros with a small epsilon
    ref_pct = np.where(ref_pct == 0, 1e-9, ref_pct)
    cur_pct = np.where(cur_pct == 0, 1e-9, cur_pct)

    psi = np.sum((cur_pct - ref_pct) * np.log(cur_pct / ref_pct))
    return float(psi)


def _js_divergence(reference: np.ndarray, current: np.ndarray, bins: int = 10) -> float:
    """Jensen-Shannon divergence between two distributions (range [0, 1]).

    A value of 0 means identical distributions; 1 means maximally different.

    For continuous features the distributions are discretised into ``bins``
    equal-width bins derived from the reference range.  For categorical
    features pass integer-encoded arrays.
    """
    if len(reference) == 0 or len(current) == 0:
        return float('nan')

    min_val = float(np.min(reference))
    max_val = float(np.max(reference))
    if min_val == max_val:
        return 0.0

    eps = (max_val - min_val) * 1e-6
    edges = np.linspace(min_val - eps, max_val + eps, bins + 1)

    ref_counts, _ = np.histogram(reference, bins=edges)
    cur_counts, _ = np.histogram(current, bins=edges)

    ref_pct = ref_counts / len(reference)
    cur_pct = cur_counts / len(current)

    # jensenshannon returns the square root of the JS divergence (a distance).
    # Squaring gives the divergence in [0, 1] when base=2 is used.
    js_dist = jensenshannon(ref_pct, cur_pct, base=2.0)
    return float(js_dist ** 2)


class DriftMonitor:
    """Monitor concept drift in model inputs by comparing distributions.

    Usage::

        monitor = DriftMonitor(numeric_features=['price', 'weight'],
                               category_features=['brand'])
        monitor.fit(reference_df)           # capture reference distribution

        # Later, during inference:
        report = monitor.compute(live_df)
        # report: {'price': {'psi': 0.03, 'js': 0.01, 'drifted': False}, ...}

    Args:
        numeric_features (list[str]): Names of numeric input columns to monitor.
        category_features (list[str]): Names of categorical input columns to monitor.
        bins (int): Number of histogram bins used for numeric PSI/JS computation.
            Default 10.
        psi_threshold (float): PSI threshold above which a feature is flagged as
            drifted.  Default 0.2 (industry convention).
        experiment_manager: Optional
            :class:`~pwml.utilities.neptunehelpers.ExperimentManager` instance.
            When provided, ``compute()`` logs each drift metric to Neptune via
            ``set_experiment_property``.
    """

    def __init__(
        self,
        numeric_features=None,
        category_features=None,
        bins: int = 10,
        psi_threshold: float = _PSI_HIGH,
        experiment_manager=None,
    ):
        self.numeric_features = list(numeric_features or [])
        self.category_features = list(category_features or [])
        self.bins = bins
        self.psi_threshold = psi_threshold
        self.experiment_manager = experiment_manager

        self._reference: dict = {}  # feature_name -> np.ndarray of reference values
        self._is_fitted = False

    def fit(self, reference: pd.DataFrame) -> 'DriftMonitor':
        """Capture the reference distribution from a training/baseline DataFrame.

        Args:
            reference: DataFrame containing at least the columns declared in
                ``numeric_features`` and ``category_features``.

        Returns:
            self (for method chaining).
        """
        self._reference = {}

        for col in self.numeric_features:
            if col not in reference.columns:
                warnings.warn(
                    "DriftMonitor.fit: numeric feature '{0}' not found in "
                    "reference DataFrame.".format(col),
                    UserWarning, stacklevel=2)
                continue
            self._reference[col] = reference[col].dropna().to_numpy(dtype=np.float64)

        for col in self.category_features:
            if col not in reference.columns:
                warnings.warn(
                    "DriftMonitor.fit: category feature '{0}' not found in "
                    "reference DataFrame.".format(col),
                    UserWarning, stacklevel=2)
                continue
            # Encode categories as integer indices
            cats = pd.Categorical(reference[col].dropna())
            self._reference[col] = (cats.codes.astype(np.float64), cats.categories)

        self._is_fitted = True
        _log.info('DriftMonitor fitted on %d rows.', len(reference))
        return self

    def compute(self, live: pd.DataFrame) -> dict:
        """Compute PSI and JS divergence between the reference and a live batch.

        Args:
            live: DataFrame of current inference inputs.

        Returns:
            dict mapping feature name to a sub-dict::

                {
                    'psi':     float,  # Population Stability Index
                    'js':      float,  # Jensen-Shannon divergence in [0, 1]
                    'drifted': bool,   # True if PSI >= psi_threshold
                }

        Raises:
            RuntimeError: if ``fit()`` has not been called first.
        """
        if not self._is_fitted:
            raise RuntimeError(
                'DriftMonitor.compute() called before fit(). '
                'Call fit(reference_df) first.')

        report = {}

        for col in self.numeric_features:
            if col not in self._reference:
                continue
            if col not in live.columns:
                _log.warning("DriftMonitor.compute: column '%s' missing from live batch.", col)
                continue

            ref_arr = self._reference[col]
            cur_arr = live[col].dropna().to_numpy(dtype=np.float64)

            psi_val = _psi(ref_arr, cur_arr, bins=self.bins)
            js_val = _js_divergence(ref_arr, cur_arr, bins=self.bins)
            drifted = (not np.isnan(psi_val)) and (psi_val >= self.psi_threshold)

            report[col] = {'psi': psi_val, 'js': js_val, 'drifted': drifted}

            if drifted:
                _log.warning(
                    "Drift detected in feature '%s': PSI=%.4f (threshold=%.2f).",
                    col, psi_val, self.psi_threshold)

        for col in self.category_features:
            if col not in self._reference:
                continue
            if col not in live.columns:
                _log.warning("DriftMonitor.compute: column '%s' missing from live batch.", col)
                continue

            ref_codes, ref_cats = self._reference[col]
            # Encode live data using the same category set
            cur_cats = pd.Categorical(live[col].dropna(), categories=ref_cats)
            cur_codes = cur_cats.codes.astype(np.float64)
            # Unknown categories get code -1; treat as out-of-vocabulary
            cur_codes = cur_codes[cur_codes >= 0]

            psi_val = _psi(ref_codes, cur_codes, bins=len(ref_cats))
            js_val = _js_divergence(ref_codes, cur_codes, bins=len(ref_cats))
            drifted = (not np.isnan(psi_val)) and (psi_val >= self.psi_threshold)

            report[col] = {'psi': psi_val, 'js': js_val, 'drifted': drifted}

            if drifted:
                _log.warning(
                    "Drift detected in feature '%s': PSI=%.4f (threshold=%.2f).",
                    col, psi_val, self.psi_threshold)

        # Log to Neptune if experiment_manager is provided
        if self.experiment_manager is not None:
            for col, metrics in report.items():
                self.experiment_manager.set_experiment_property(
                    'drift/{0}/psi'.format(col), metrics['psi'])
                self.experiment_manager.set_experiment_property(
                    'drift/{0}/js'.format(col), metrics['js'])
                self.experiment_manager.set_experiment_property(
                    'drift/{0}/drifted'.format(col), metrics['drifted'])

        return report
