"""
Censored Energy Window (CEW) method for optimal ROI generation and anomaly detection.

This module implements the Censored Energy Window method from:
Miller & Dubrawski, "Gamma-Ray Source Detection With Small Sensors",
IEEE Transactions on Nuclear Science, 2018.

The CEW method optimizes window selection to maximize signal-to-noise ratio
and uses out-of-window (censored) energy bins to predict and subtract
in-window background contributions.

Typical Workflow:
-----------------
1. Training (uses Spectra - no time ordering needed):
   - Collect background measurements (source-absent)
   - Optimize window selection to maximize SNR
   - Fit predictor to learn background patterns

2. Detection (uses SpectralTimeSeries - time-ordered monitoring):
   - Score time series to detect when source appears
   - Identify time points with anomalously high scores
   - Locate temporal anomalies in continuous monitoring
"""

from typing import Optional, List, Dict, Any, Union, Tuple
import numpy as np
import numpy.typing as npt
from scipy.optimize import minimize
from scipy.linalg import solve

from gammaflow.core.spectrum import Spectrum
from gammaflow.core.spectra import Spectra
from gammaflow.core.time_series import SpectralTimeSeries


def optimize_cew_windows(
    source_spectrum: Spectrum,
    background_spectrum: Spectrum,
    method: str = 'gradient',
    n_bins_target: Optional[int] = None,
    min_bin_width: int = 1,
    max_iterations: int = 1000,
    learning_rate: float = 0.01,
    regularization: float = 0.0,
    convergence_tol: float = 1e-6,
    verbose: bool = False,
) -> np.ndarray:
    """
    Optimize window selection to maximize SNR: w·s / sqrt(w·b)
    
    Uses gradient-based optimization with sigmoid projection for binary constraints.
    
    Parameters
    ----------
    source_spectrum : Spectrum
        Known or estimated source spectrum (signal)
    background_spectrum : Spectrum
        Estimated background from source-absent data
    method : str, optional
        Optimization method: 'gradient', 'greedy', or 'scipy'. Default is 'gradient'.
    n_bins_target : int, optional
        Target number of bins to select. If None, automatically determined.
    min_bin_width : int, optional
        Minimum contiguous bin width for a window region. Default is 1 (no constraint).
    max_iterations : int, optional
        Maximum optimization iterations. Default is 1000.
    learning_rate : float, optional
        Learning rate for gradient descent. Default is 0.01.
    regularization : float, optional
        L1 penalty on window gradient (promotes smooth windows). Default is 0.0.
    convergence_tol : float, optional
        Convergence tolerance for gradient norm. Default is 1e-6.
    verbose : bool, optional
        Print optimization progress. Default is False.
        
    Returns
    -------
    window_mask : np.ndarray
        Binary mask of shape (n_bins,) where 1 = in-window, 0 = censored
        
    Examples
    --------
    >>> source = Spectrum(...)
    >>> background = Spectrum(...)
    >>> window_mask = optimize_cew_windows(source, background, method='gradient')
    >>> n_selected = np.sum(window_mask)
    """
    # Validate inputs
    if source_spectrum.n_bins != background_spectrum.n_bins:
        raise ValueError(
            f"Source and background must have same number of bins. "
            f"Got {source_spectrum.n_bins} and {background_spectrum.n_bins}"
        )
    
    s = source_spectrum.counts
    b = background_spectrum.counts
    n_bins = len(s)
    
    # Ensure background has no zeros (add small regularization)
    b = np.maximum(b, 1e-10)
    
    if method == 'gradient':
        return _gradient_optimize_windows(
            s, b, n_bins_target, max_iterations, learning_rate,
            regularization, convergence_tol, verbose
        )
    elif method == 'greedy':
        return _greedy_optimize_windows(s, b, n_bins_target, verbose)
    elif method == 'scipy':
        return _scipy_optimize_windows(s, b, n_bins_target, verbose)
    else:
        raise ValueError(f"Unknown optimization method: {method}")


def _gradient_optimize_windows(
    s: np.ndarray,
    b: np.ndarray,
    n_bins_target: Optional[int],
    max_iterations: int,
    learning_rate: float,
    regularization: float,
    convergence_tol: float,
    verbose: bool,
) -> np.ndarray:
    """
    Gradient-based optimization with sigmoid projection.
    
    We optimize continuous weights z ∈ ℝⁿ and project to binary via sigmoid:
    w = sigmoid(z) ≈ {0, 1}
    
    Objective: maximize SNR = w·s / sqrt(w·b + ε)
    """
    n_bins = len(s)
    
    # Initialize weights (start with all bins having equal probability)
    if n_bins_target is not None:
        # Initialize to select approximately n_bins_target bins
        init_value = np.log(n_bins_target / (n_bins - n_bins_target))
    else:
        init_value = 0.0
    
    z = np.ones(n_bins) * init_value
    
    best_snr = -np.inf
    best_z = z.copy()
    no_improvement_count = 0
    
    for iteration in range(max_iterations):
        # Project to probabilities using sigmoid
        w = 1.0 / (1.0 + np.exp(-z))
        
        # Compute SNR and its gradient
        numerator = np.dot(w, s)
        denominator = np.sqrt(np.dot(w, b) + 1e-10)
        snr = numerator / denominator
        
        # Gradient of SNR with respect to w
        # d(SNR)/dw = s/sqrt(w·b) - (w·s)/(2*(w·b)^(3/2)) * b
        grad_snr_w = s / denominator - (numerator / (2 * denominator**3)) * b
        
        # Chain rule: d(SNR)/dz = d(SNR)/dw * dw/dz
        # dw/dz = sigmoid(z) * (1 - sigmoid(z))
        sigmoid_derivative = w * (1 - w)
        grad_snr_z = grad_snr_w * sigmoid_derivative
        
        # Add regularization (L1 on gradient of w for smoothness)
        if regularization > 0:
            w_grad = np.gradient(w)
            reg_penalty = regularization * np.sign(w_grad)
            grad_snr_z -= reg_penalty * sigmoid_derivative
        
        # Gradient ascent (maximize SNR)
        z = z + learning_rate * grad_snr_z
        
        # Track best solution
        if snr > best_snr:
            best_snr = snr
            best_z = z.copy()
            no_improvement_count = 0
        else:
            no_improvement_count += 1
        
        # Check convergence
        grad_norm = np.linalg.norm(grad_snr_z)
        if grad_norm < convergence_tol:
            if verbose:
                print(f"Converged at iteration {iteration}, SNR={snr:.4f}")
            break
        
        # Early stopping if no improvement
        if no_improvement_count > 50:
            if verbose:
                print(f"Early stopping at iteration {iteration}, best SNR={best_snr:.4f}")
            break
        
        if verbose and iteration % 100 == 0:
            n_selected = np.sum(w > 0.5)
            print(f"Iter {iteration}: SNR={snr:.4f}, n_bins={n_selected}, grad_norm={grad_norm:.6f}")
    
    # Final projection to binary
    w_final = 1.0 / (1.0 + np.exp(-best_z))
    window_mask = (w_final > 0.5).astype(int)
    
    return window_mask


def _greedy_optimize_windows(
    s: np.ndarray,
    b: np.ndarray,
    n_bins_target: Optional[int],
    verbose: bool,
) -> np.ndarray:
    """
    Greedy optimization: iteratively add bin that maximizes SNR increase.
    """
    n_bins = len(s)
    window_mask = np.zeros(n_bins, dtype=int)
    
    if n_bins_target is None:
        n_bins_target = n_bins // 2  # Default to half
    
    for i in range(n_bins_target):
        best_snr = -np.inf
        best_bin = -1
        
        # Try adding each unselected bin
        for bin_idx in range(n_bins):
            if window_mask[bin_idx] == 1:
                continue  # Already selected
            
            # Temporarily add this bin
            window_mask[bin_idx] = 1
            snr = _compute_snr(window_mask, s, b)
            window_mask[bin_idx] = 0
            
            if snr > best_snr:
                best_snr = snr
                best_bin = bin_idx
        
        # Add the best bin
        if best_bin >= 0:
            window_mask[best_bin] = 1
            if verbose and i % 10 == 0:
                print(f"Added bin {best_bin}, SNR={best_snr:.4f}")
    
    return window_mask


def _scipy_optimize_windows(
    s: np.ndarray,
    b: np.ndarray,
    n_bins_target: Optional[int],
    verbose: bool,
) -> np.ndarray:
    """
    Use scipy's optimization with continuous relaxation.
    """
    n_bins = len(s)
    
    # Objective: minimize negative SNR
    def objective(w):
        numerator = np.dot(w, s)
        denominator = np.sqrt(np.dot(w, b) + 1e-10)
        snr = numerator / denominator
        return -snr  # Minimize negative SNR
    
    # Gradient
    def gradient(w):
        numerator = np.dot(w, s)
        denominator = np.sqrt(np.dot(w, b) + 1e-10)
        grad = -s / denominator + (numerator / (2 * denominator**3)) * b
        return grad
    
    # Constraints: 0 <= w <= 1 (continuous relaxation)
    bounds = [(0, 1) for _ in range(n_bins)]
    
    # Initial guess
    if n_bins_target is not None:
        w0 = np.zeros(n_bins)
        # Start with top signal bins
        top_indices = np.argsort(s)[-n_bins_target:]
        w0[top_indices] = 1.0
    else:
        w0 = np.ones(n_bins) * 0.5
    
    # Optimize
    result = minimize(
        objective,
        w0,
        method='L-BFGS-B',
        jac=gradient,
        bounds=bounds,
        options={'disp': verbose}
    )
    
    # Threshold to binary
    window_mask = (result.x > 0.5).astype(int)
    
    return window_mask


def _compute_snr(window_mask: np.ndarray, s: np.ndarray, b: np.ndarray) -> float:
    """Compute SNR for given window configuration."""
    numerator = np.dot(window_mask, s)
    denominator = np.sqrt(np.dot(window_mask, b) + 1e-10)
    return numerator / denominator


def fit_cew_predictor(
    background_spectra: Spectra,
    window_mask: np.ndarray,
    alpha: float = 1.0,
    fit_intercept: bool = False,
) -> 'CEWPredictor':
    """
    Fit ridge regression predictor θ to estimate in-window counts from 
    out-of-window counts.
    
    Constraint: θ[i] = 0 for all in-window bins (i where window_mask[i] == 1)
    
    The predictor learns to estimate the in-window background using only
    out-of-window (censored) information, allowing anomaly detection.
    
    Parameters
    ----------
    background_spectra : Spectra
        Collection of source-absent background spectra for training
    window_mask : np.ndarray
        Binary window indicator vector (1 = in-window, 0 = censored)
    alpha : float, optional
        Ridge regression regularization strength. Default is 1.0.
    fit_intercept : bool, optional
        Whether to fit an intercept term. Default is False.
        
    Returns
    -------
    predictor : CEWPredictor
        Fitted predictor with θ coefficients
        
    Examples
    --------
    >>> background_training = Spectra([...])
    >>> window_mask = optimize_cew_windows(source, background.mean_spectrum())
    >>> predictor = fit_cew_predictor(background_training, window_mask, alpha=1.0)
    >>> score = predictor.score(test_spectrum)
    """
    # Validate inputs
    if len(window_mask) != background_spectra.n_bins:
        raise ValueError(
            f"window_mask length ({len(window_mask)}) must match number of bins "
            f"({background_spectra.n_bins})"
        )
    
    # Get counts array: shape (n_spectra, n_bins)
    X = background_spectra.counts  # Full spectra
    n_spectra, n_bins = X.shape
    
    # Identify in-window and out-of-window bins
    in_window_mask = window_mask.astype(bool)
    out_window_mask = ~in_window_mask
    
    if np.sum(in_window_mask) == 0:
        raise ValueError("window_mask has no in-window bins (all zeros)")
    if np.sum(out_window_mask) == 0:
        raise ValueError("window_mask has no out-of-window bins (all ones)")
    
    # Fit ridge regression: predict in-window from out-of-window
    # We want θ such that X @ θ ≈ X @ w (in-window counts)
    # with constraint that θ[in_window] = 0
    
    # Target: in-window counts
    y = X @ window_mask  # shape: (n_spectra,)
    
    # Features: out-of-window counts only
    X_train = X[:, out_window_mask]  # shape: (n_spectra, n_out_window)
    
    # Ridge regression
    theta_out = _fit_ridge_regression(X_train, y, alpha, fit_intercept)
    
    # Construct full θ vector with zeros for in-window bins
    theta = np.zeros(n_bins)
    theta[out_window_mask] = theta_out
    
    # Get energy edges if available
    energy_edges = None
    if background_spectra.is_calibrated:
        if background_spectra.uses_shared_calibration:
            energy_edges = background_spectra.energy_edges.copy()
        else:
            # Use first spectrum's calibration
            energy_edges = background_spectra.spectra[0].energy_edges.copy()
    
    return CEWPredictor(theta, window_mask, alpha, energy_edges)


def _fit_ridge_regression(
    X: np.ndarray,
    y: np.ndarray,
    alpha: float,
    fit_intercept: bool,
) -> np.ndarray:
    """
    Fit ridge regression: minimize ||Xθ - y||² + alpha||θ||²
    
    Closed-form solution: θ = (X'X + alpha*I)^(-1) X'y
    """
    n_samples, n_features = X.shape
    
    if fit_intercept:
        # Add intercept column
        X = np.column_stack([np.ones(n_samples), X])
    
    # Closed-form solution
    XtX = X.T @ X
    Xty = X.T @ y
    
    # Add ridge penalty
    ridge_matrix = XtX + alpha * np.eye(X.shape[1])
    
    # Solve
    theta = solve(ridge_matrix, Xty, assume_a='pos')
    
    if fit_intercept:
        # Return coefficients without intercept for consistency
        # (intercept handling can be added later if needed)
        return theta[1:]
    
    return theta


class CEWPredictor:
    """
    Censored Energy Window predictor for anomaly scoring.
    
    Trained on background spectra (Spectra) and used to score time series
    (SpectralTimeSeries) for anomaly detection.
    
    Scoring formula: (w - θ)·x
    where:
    - w = window mask (in-window = 1, censored = 0)
    - θ = regression coefficients (predicts in-window from out-of-window)
    - x = observed spectrum
    
    Higher scores indicate presence of source (anomaly).
    
    Workflow:
    1. Train on background using Spectra (no time ordering needed)
    2. Score time series using SpectralTimeSeries (detects when source appears)
    
    Attributes
    ----------
    theta : np.ndarray
        Regression coefficients (shape: n_bins). Zero for in-window bins.
    window_mask : np.ndarray
        Binary window indicator (shape: n_bins). 1=in-window, 0=censored.
    alpha : float
        Ridge regularization used in fitting.
    energy_edges : np.ndarray or None
        Energy calibration for ROI conversion.
    n_bins : int
        Number of energy bins.
    scoring_vector : np.ndarray
        Pre-computed (w - θ) for efficient scoring.
        
    Examples
    --------
    >>> # Training: Use Spectra (background samples)
    >>> bg_spectra = Spectra([...])
    >>> predictor = fit_cew_predictor(bg_spectra, window_mask)
    >>> 
    >>> # Detection: Use SpectralTimeSeries (time-ordered monitoring)
    >>> time_series = SpectralTimeSeries([...])
    >>> scores = predictor.score_time_series(time_series)
    >>> anomalies = scores > threshold
    """
    
    def __init__(
        self,
        theta: np.ndarray,
        window_mask: np.ndarray,
        alpha: float,
        energy_edges: Optional[np.ndarray] = None,
    ):
        self.theta = np.asarray(theta)
        self.window_mask = np.asarray(window_mask, dtype=int)
        self.alpha = alpha
        self.energy_edges = energy_edges.copy() if energy_edges is not None else None
        self.n_bins = len(self.theta)
        
        # Pre-compute scoring vector for efficiency
        self.scoring_vector = self.window_mask - self.theta
        
        # Validate
        if len(self.window_mask) != self.n_bins:
            raise ValueError("theta and window_mask must have same length")
        if self.energy_edges is not None and len(self.energy_edges) != self.n_bins + 1:
            raise ValueError("energy_edges must have length n_bins + 1")
    
    def score(self, spectrum: Spectrum) -> float:
        """
        Score a single spectrum.
        
        Parameters
        ----------
        spectrum : Spectrum
            Spectrum to score
            
        Returns
        -------
        float
            CEW score. Higher values indicate source presence.
        """
        if spectrum.n_bins != self.n_bins:
            raise ValueError(
                f"Spectrum has {spectrum.n_bins} bins, expected {self.n_bins}"
            )
        
        return float(np.dot(self.scoring_vector, spectrum.counts))
    
    def score_spectra(self, spectra: Spectra) -> np.ndarray:
        """
        Vectorized scoring for multiple spectra.
        
        Parameters
        ----------
        spectra : Spectra
            Collection of spectra to score
            
        Returns
        -------
        np.ndarray
            Array of scores with shape (n_spectra,)
        """
        if spectra.n_bins != self.n_bins:
            raise ValueError(
                f"Spectra have {spectra.n_bins} bins, expected {self.n_bins}"
            )
        
        # Matrix multiplication: (1, n_bins) @ (n_spectra, n_bins).T = (n_spectra,)
        return self.scoring_vector @ spectra.counts.T
    
    def score_time_series(self, time_series: SpectralTimeSeries) -> np.ndarray:
        """
        Score time series for anomaly detection.
        
        This is specifically for time-ordered anomaly detection where you want
        to detect when a source appears in a time series of measurements.
        
        Parameters
        ----------
        time_series : SpectralTimeSeries
            Time-ordered series of spectra to score for anomalies
            
        Returns
        -------
        np.ndarray
            Array of anomaly scores with shape (n_spectra,). Higher scores
            indicate potential source presence at that time point.
            
        Examples
        --------
        >>> # Train on background
        >>> bg_spectra = Spectra([...])  # Collection of background measurements
        >>> predictor = fit_cew_predictor(bg_spectra, window_mask)
        >>> 
        >>> # Score time series for anomaly detection
        >>> time_series = SpectralTimeSeries([...])  # Time-ordered measurements
        >>> scores = predictor.score_time_series(time_series)
        >>> 
        >>> # Detect anomalies (times when source appeared)
        >>> threshold = np.percentile(scores, 99)
        >>> anomaly_times = time_series.timestamps[scores > threshold]
        """
        return self.score_spectra(time_series)
    
    def to_rois(self, labels: Optional[List[str]] = None) -> List:
        """
        Convert window mask to list of EnergyROI objects.
        
        Finds contiguous regions in window_mask and creates an ROI for each.
        Stores predictor in metadata for later use.
        
        Parameters
        ----------
        labels : list of str, optional
            Labels for each ROI. If None, auto-generated as "CEW Window 1", etc.
            
        Returns
        -------
        list of EnergyROI
            List of ROI objects
        """
        from gammaflow.operations.roi import EnergyROI
        
        if self.energy_edges is None:
            raise ValueError("Cannot create ROIs: energy calibration not available")
        
        # Find contiguous regions
        regions = []
        in_region = False
        start_idx = None
        
        for i, val in enumerate(self.window_mask):
            if val == 1 and not in_region:
                # Start of new region
                start_idx = i
                in_region = True
            elif val == 0 and in_region:
                # End of region
                regions.append((start_idx, i))
                in_region = False
        
        # Handle case where last bin is in-window
        if in_region:
            regions.append((start_idx, len(self.window_mask)))
        
        # Create ROIs
        rois = []
        for idx, (start, end) in enumerate(regions):
            e_min = self.energy_edges[start]
            e_max = self.energy_edges[end]
            
            if labels is not None and idx < len(labels):
                label = labels[idx]
            else:
                label = f"CEW Window {idx + 1}"
            
            roi = EnergyROI(
                e_min=e_min,
                e_max=e_max,
                label=label,
                method="Censored Energy Windows",
                metadata={
                    'cew_predictor': self,
                    'window_indices': (start, end),
                    'alpha': self.alpha,
                }
            )
            rois.append(roi)
        
        return rois
    
    def to_dict(self) -> Dict[str, Any]:
        """
        Serialize predictor to dictionary.
        
        Returns
        -------
        dict
            Dictionary containing all predictor parameters
        """
        return {
            'theta': self.theta.tolist(),
            'window_mask': self.window_mask.tolist(),
            'alpha': self.alpha,
            'energy_edges': self.energy_edges.tolist() if self.energy_edges is not None else None,
        }
    
    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> 'CEWPredictor':
        """
        Deserialize predictor from dictionary.
        
        Parameters
        ----------
        d : dict
            Dictionary from to_dict()
            
        Returns
        -------
        CEWPredictor
            Reconstructed predictor
        """
        theta = np.array(d['theta'])
        window_mask = np.array(d['window_mask'], dtype=int)
        alpha = d['alpha']
        energy_edges = np.array(d['energy_edges']) if d['energy_edges'] is not None else None
        
        return cls(theta, window_mask, alpha, energy_edges)
    
    def __repr__(self) -> str:
        n_windows = np.sum(self.window_mask)
        return (
            f"CEWPredictor(n_bins={self.n_bins}, n_in_window={n_windows}, "
            f"alpha={self.alpha})"
        )

