"""
Statistically rigorous weighted calculations for redshift and age estimation in SNID SAGE.

This module implements best-metric weighted estimation methods (preferring RLAP-CCC) for optimal redshift 
and age estimation with full covariance analysis.
"""

import numpy as np
from typing import Union, List, Tuple, Optional
import logging

# Get logger for this module
logger = logging.getLogger(__name__)


def compute_cluster_weights(
    rlap_ccc_values: Union[np.ndarray, List[float]],
    redshift_errors: Union[np.ndarray, List[float]]
) -> np.ndarray:
    """
    Compute canonical cluster weights: w_i = (rlapccc_i)^2 / sigma_z_i^2.

    This is a thin wrapper around calculate_combined_weights for clarity.
    """
    return calculate_combined_weights(rlap_ccc_values, redshift_errors)


def _weighted_mean(values: np.ndarray, weights: np.ndarray) -> float:
    """Compute weighted mean with basic validation; returns NaN if no valid data."""
    if values.size == 0 or weights.size == 0:
        return float('nan')
    sum_w = float(np.sum(weights))
    if sum_w <= 0 or not np.isfinite(sum_w):
        return float('nan')
    return float(np.sum(weights * values) / sum_w)


 


def _weighted_sd_unbiased(values: np.ndarray, weights: np.ndarray) -> float:
    """
    Unbiased weighted standard deviation:
      μ = Σ w_i x_i / Σ w_i
      var_pop = Σ w_i (x_i − μ)^2 / Σ w_i
      N_eff = (Σ w_i)^2 / Σ w_i^2
      var_unbiased = var_pop × N_eff / (N_eff − 1)  (for N_eff > 1; else var_pop)
      sd = sqrt(max(0, var_unbiased))
    """
    if values.size == 0 or weights.size == 0:
        return float('nan')
    valid_mask = (np.isfinite(values) & np.isfinite(weights) & (weights > 0))
    if not np.any(valid_mask):
        return float('nan')
    v = values[valid_mask]
    w = weights[valid_mask]
    sum_w = float(np.sum(w))
    if sum_w <= 0 or not np.isfinite(sum_w):
        return float('nan')
    mean = float(np.sum(w * v) / sum_w)
    var_pop = float(np.sum(w * (v - mean) ** 2) / sum_w)
    if not np.isfinite(var_pop) or var_pop < 0:
        return float('nan')
    sum_w_sq = float(np.sum(w ** 2))
    if sum_w_sq <= 0 or not np.isfinite(sum_w_sq):
        return float('nan')
    n_eff = (sum_w ** 2) / sum_w_sq
    var_unbiased = var_pop * (n_eff / (n_eff - 1.0)) if n_eff > 1 else var_pop
    return float(np.sqrt(max(0.0, var_unbiased)))


def estimate_weighted_redshift(
    redshifts: Union[np.ndarray, List[float]],
    redshift_errors: Union[np.ndarray, List[float]],
    rlap_ccc_values: Union[np.ndarray, List[float]]
) -> float:
    """
    Weighted mean redshift using weights w = (rlapccc)^2 / sigma_z^2.
    """
    z = np.asarray(redshifts, dtype=float)
    sigma = np.asarray(redshift_errors, dtype=float)
    r = np.asarray(rlap_ccc_values, dtype=float)
    if not (len(z) == len(sigma) == len(r)):
        logger.error("Mismatched input lengths for estimate_weighted_redshift")
        return float('nan')
    valid = (np.isfinite(z) & np.isfinite(sigma) & np.isfinite(r) & (sigma > 0))
    if not np.any(valid):
        return float('nan')
    w = compute_cluster_weights(r[valid], sigma[valid])
    return _weighted_mean(z[valid], w)


def estimate_weighted_epoch(
    ages: Union[np.ndarray, List[float]],
    redshift_errors: Union[np.ndarray, List[float]],
    rlap_ccc_values: Union[np.ndarray, List[float]]
) -> float:
    """
    Weighted mean epoch (age) using the same cluster weights as redshift:
    w = (rlapccc)^2 / sigma_z^2.
    """
    t = np.asarray(ages, dtype=float)
    sigma = np.asarray(redshift_errors, dtype=float)
    r = np.asarray(rlap_ccc_values, dtype=float)
    if not (len(t) == len(sigma) == len(r)):
        logger.error("Mismatched input lengths for estimate_weighted_epoch")
        return float('nan')
    valid = (np.isfinite(t) & np.isfinite(sigma) & np.isfinite(r) & (sigma > 0))
    if not np.any(valid):
        return float('nan')
    w = compute_cluster_weights(r[valid], sigma[valid])
    return _weighted_mean(t[valid], w)


def weighted_redshift_error(
    redshifts: Union[np.ndarray, List[float]],
    redshift_errors: Union[np.ndarray, List[float]],
    rlap_ccc_values: Union[np.ndarray, List[float]]
) -> float:
    """
    Uncertainty for redshift reported as unbiased weighted SD within the set.
    Uses weights w = (rlapccc)^2 / sigma_z^2.
    Single-member rule: return that member's redshift_error.
    """
    z = np.asarray(redshifts, dtype=float)
    sigma = np.asarray(redshift_errors, dtype=float)
    r = np.asarray(rlap_ccc_values, dtype=float)
    if not (len(z) == len(sigma) == len(r)):
        logger.error("Mismatched input lengths for weighted_redshift_error")
        return float('nan')
    valid = (np.isfinite(z) & np.isfinite(sigma) & np.isfinite(r) & (sigma > 0) & (r > 0))
    n_valid = int(np.sum(valid))
    if n_valid == 0:
        return float('nan')
    if n_valid == 1:
        return float(sigma[valid][0])
    w = compute_cluster_weights(r[valid], sigma[valid])
    return _weighted_sd_unbiased(z[valid], w)


def weighted_epoch_error(
    ages: Union[np.ndarray, List[float]],
    redshift_errors: Union[np.ndarray, List[float]],
    rlap_ccc_values: Union[np.ndarray, List[float]]
) -> float:
    """
    Uncertainty for age reported as unbiased weighted SD within the set.
    Uses redshift-based weights w = (rlapccc)^2 / sigma_z^2.
    Single-member rule: return NaN (cannot estimate SD from one point).
    """
    t = np.asarray(ages, dtype=float)
    sigma = np.asarray(redshift_errors, dtype=float)
    r = np.asarray(rlap_ccc_values, dtype=float)
    if not (len(t) == len(sigma) == len(r)):
        logger.error("Mismatched input lengths for weighted_epoch_error")
        return float('nan')
    valid = (np.isfinite(t) & np.isfinite(sigma) & np.isfinite(r) & (sigma > 0) & (r > 0))
    n_valid = int(np.sum(valid))
    if n_valid == 0:
        return float('nan')
    if n_valid == 1:
        return float('nan')
    w = compute_cluster_weights(r[valid], sigma[valid])
    return _weighted_sd_unbiased(t[valid], w)


 

def calculate_combined_weights(
    rlap_ccc_values: Union[np.ndarray, List[float]],
    uncertainties: Union[np.ndarray, List[float]]
) -> np.ndarray:
    """
    Calculate combined weights using both best-metric quality and individual uncertainties.
    
    This implements the statistically correct approach for weighted averaging when
    both quality indicators (best metric) and individual uncertainties are available.
    
    Parameters
    ----------
    rlap_ccc_values : array-like
        Best metric quality scores (prefer RLAP-CCC; fallback to RLAP)
    uncertainties : array-like
        Individual uncertainty estimates for each template (e.g., redshift errors)
        
    Returns
    -------
    np.ndarray
        Combined weights = (metric)² / σ²
        
    Notes
    -----
    Statistical Formulation:
    - Quality weight: q_i = (metric_i)²
    - Precision weight: p_i = 1/σ²_i
    - Combined weight: w_i = q_i × p_i = (metric_i)² / σ²_i
    
    This gives high-quality templates with low uncertainty the highest influence,
    which is statistically optimal for uncertainty propagation.
    """
    rlap_ccc_values = np.asarray(rlap_ccc_values, dtype=float)
    uncertainties = np.asarray(uncertainties, dtype=float)
    
    # Validate inputs
    if len(rlap_ccc_values) != len(uncertainties):
        raise ValueError("Metric values and uncertainties must have same length")
    
    if len(rlap_ccc_values) == 0:
        return np.array([])
    
    # Handle zero uncertainties (perfect measurements) by using a small floor value
    # This prevents infinite weights while preserving the relative ordering
    min_uncertainty = np.min(uncertainties[uncertainties > 0]) if np.any(uncertainties > 0) else 1e-6
    uncertainty_floor = min_uncertainty * 0.1
    safe_uncertainties = np.maximum(uncertainties, uncertainty_floor)
    
    # Calculate combined weights using squared best-metric values (RLAP-CCC preferred)
    quality_weights = rlap_ccc_values ** 2
    precision_weights = 1.0 / (safe_uncertainties ** 2)  # Inverse variance weighting
    combined_weights = quality_weights * precision_weights
    
    logger.debug(f"Combined weighting: Best-metric [{rlap_ccc_values.min():.2f}, {rlap_ccc_values.max():.2f}], "
                f"uncertainties [{uncertainties.min():.4f}, {uncertainties.max():.4f}], "
                f"weights [{combined_weights.min():.2e}, {combined_weights.max():.2e}]")
    
    return combined_weights


def apply_exponential_weighting(rlap_ccc_values: Union[np.ndarray, List[float]]) -> np.ndarray:
    """
    Apply squared-metric weighting to RLAP-CCC/RLAP values for template prioritization.
    
    This helper now implements w = (metric)² to match the main pipeline's
    weighting policy when per-template σ is unavailable (quality-only case).
    
    Parameters
    ----------
    rlap_ccc_values : array-like
        Raw best-metric values from template matching
        
    Returns
    -------
    np.ndarray
        Squared metric weights
        
    Notes
    -----
    Transformation: w = (metric)²
    """
    rlap_ccc_values = np.asarray(rlap_ccc_values, dtype=float)
    
    # Handle empty input
    if len(rlap_ccc_values) == 0:
        return np.array([])
    
    # Apply squared weighting: w = x^2
    exponential_weights = rlap_ccc_values ** 2
    
    # Log the transformation for debugging
    if len(rlap_ccc_values) > 0:
        logger.debug(f"Squared weighting: Best-metric range [{rlap_ccc_values.min():.2f}, {rlap_ccc_values.max():.2f}] "
                    f"→ weight range [{exponential_weights.min():.2e}, {exponential_weights.max():.2e}]")
    
    return exponential_weights




# Exports
__all__ = [
    'calculate_combined_weights',
    'apply_exponential_weighting',
    'compute_cluster_weights',
    'estimate_weighted_redshift',
    'estimate_weighted_epoch',
    'weighted_redshift_error',
    'weighted_epoch_error'
]