"""
Example: Censored Energy Window (CEW) Method for Source Detection

This example demonstrates the Censored Energy Window method for optimal ROI
generation and anomaly detection in gamma-ray spectroscopy. The CEW method:

1. Optimizes window selection to maximize signal-to-noise ratio
2. Uses out-of-window (censored) bins to predict in-window background
3. Provides anomaly scores for source detection

Based on: Miller & Dubrawski, "Gamma-Ray Source Detection With Small Sensors",
IEEE Transactions on Nuclear Science, 2018.
"""

import numpy as np
from gammaflow import Spectrum, Spectra
from gammaflow.algorithms import optimize_cew_windows, fit_cew_predictor
from gammaflow.operations import EnergyROI

print("=" * 70)
print("Censored Energy Window (CEW) Method Examples")
print("=" * 70)

# =============================================================================
# Example 1: Basic CEW Workflow
# =============================================================================

print("\n" + "=" * 70)
print("Example 1: Basic CEW Workflow")
print("=" * 70)

# Step 1: Create synthetic data
np.random.seed(42)

# Background training data (source-absent measurements)
print("\nStep 1: Creating background training data...")
bg_spectra = [Spectrum(np.random.poisson(100, 128)) for _ in range(100)]
bg_training = Spectra(bg_spectra)
print(f"  Created {bg_training.n_spectra} background spectra")

# Known or estimated source spectrum
print("\nStep 2: Creating source spectrum...")
source_counts = np.random.poisson(100, 128)
source_counts[50:70] += 800  # Simulated Cs-137 peak
source_counts[90:95] += 400  # Smaller K-40 peak
source_spec = Spectrum(source_counts)
print(f"  Source has peaks in bins 50-70 and 90-95")

# Step 2: Optimize window selection
print("\nStep 3: Optimizing CEW windows...")
window_mask = optimize_cew_windows(
    source_spectrum=source_spec,
    background_spectrum=bg_training.mean_spectrum(),
    method='gradient',
    max_iterations=500,
    verbose=False
)

n_selected = np.sum(window_mask)
print(f"  Optimized windows select {n_selected}/{len(window_mask)} bins")
print(f"  Bins in window 1 (50-70): {np.sum(window_mask[50:70])}")
print(f"  Bins in window 2 (90-95): {np.sum(window_mask[90:95])}")

# Step 3: Fit predictor
print("\nStep 4: Fitting CEW predictor...")
predictor = fit_cew_predictor(
    background_spectra=bg_training,
    window_mask=window_mask,
    alpha=1.0
)
print(f"  {predictor}")

# Step 4: Score test spectra
print("\nStep 5: Scoring test spectra...")

# Background spectrum (should score low)
test_bg = Spectrum(np.random.poisson(100, 128))
bg_score = predictor.score(test_bg)
print(f"  Background spectrum score: {bg_score:.2f}")

# Anomalous spectrum with source (should score high)
test_source_counts = np.random.poisson(100, 128)
test_source_counts[50:70] += 600
test_source_counts[90:95] += 300
test_source = Spectrum(test_source_counts)
source_score = predictor.score(test_source)
print(f"  Source spectrum score: {source_score:.2f}")
print(f"  Score ratio (source/background): {source_score/bg_score:.2f}x")

# =============================================================================
# Example 2: Comparing Optimization Methods
# =============================================================================

print("\n" + "=" * 70)
print("Example 2: Comparing Optimization Methods")
print("=" * 70)

# Same data as Example 1
print("\nTesting different optimization methods...")

methods = ['gradient', 'greedy', 'scipy']
results = {}

for method in methods:
    window_mask = optimize_cew_windows(
        source_spec,
        bg_training.mean_spectrum(),
        method=method,
        n_bins_target=30 if method != 'gradient' else None,
        max_iterations=200 if method == 'gradient' else None,
        verbose=False
    )
    
    # Compute SNR for this window configuration
    numerator = np.dot(window_mask, source_spec.counts)
    denominator = np.sqrt(np.dot(window_mask, bg_training.mean_spectrum().counts))
    snr = numerator / denominator
    
    results[method] = {
        'n_bins': np.sum(window_mask),
        'snr': snr,
    }
    
    print(f"  {method:10s}: {results[method]['n_bins']:3d} bins, SNR={results[method]['snr']:.2f}")

# =============================================================================
# Example 3: Time Series Anomaly Detection
# =============================================================================

print("\n" + "=" * 70)
print("Example 3: Time Series Anomaly Detection")
print("=" * 70)

# Create time series with injected anomaly
print("\nCreating time series with source appearing at time 40-50...")
from gammaflow import SpectralTimeSeries

spectra_list = []
for i in range(100):
    counts = np.random.poisson(100, 128)
    
    # Inject source in middle of time series
    if 40 <= i < 50:
        counts[50:70] += 500
        counts[90:95] += 250
    
    spectra_list.append(Spectrum(counts, timestamp=float(i), real_time=1.0))

# Use SpectralTimeSeries for time-ordered anomaly detection
test_series = SpectralTimeSeries(spectra_list)

# Use score_time_series() for time-ordered anomaly detection
print("\nScoring time series for anomaly detection...")
scores = predictor.score_time_series(test_series)

print(f"  Mean score (background): {np.mean(scores[:40]):.2f}")
print(f"  Mean score (source present): {np.mean(scores[40:50]):.2f}")
print(f"  Mean score (background): {np.mean(scores[50:]):.2f}")

# Detect anomaly using simple threshold
threshold = np.mean(scores[:40]) + 3 * np.std(scores[:40])
detections = scores > threshold
print(f"\nDetections (score > threshold):")
print(f"  False alarms (time 0-39): {np.sum(detections[:40])}")
print(f"  True detections (time 40-49): {np.sum(detections[40:50])}")
print(f"  False alarms (time 50-99): {np.sum(detections[50:])}")

# =============================================================================
# Example 4: CEW with Energy Calibration and ROIs
# =============================================================================

print("\n" + "=" * 70)
print("Example 4: CEW with Energy Calibration and ROIs")
print("=" * 70)

# Create calibrated spectra
print("\nCreating calibrated spectra...")
energy_edges = np.linspace(0, 3000, 129)

bg_cal_spectra = [
    Spectrum(np.random.poisson(100, 128), energy_edges=energy_edges)
    for _ in range(50)
]
bg_cal_training = Spectra(bg_cal_spectra)

source_cal_counts = np.random.poisson(100, 128)
source_cal_counts[35:45] += 1000  # Peak at ~820 keV (Cs-137)
source_cal_counts[75:80] += 600   # Peak at ~1460 keV (K-40)
source_cal = Spectrum(source_cal_counts, energy_edges=energy_edges)

# Optimize windows
print("\nOptimizing windows...")
window_mask = optimize_cew_windows(
    source_cal,
    bg_cal_training.mean_spectrum(),
    method='greedy',
    n_bins_target=25
)

# Fit predictor
predictor_cal = fit_cew_predictor(bg_cal_training, window_mask, alpha=1.0)

# Convert to ROIs
print("\nConverting CEW windows to ROIs...")
rois = EnergyROI.from_cew_predictor(
    predictor_cal,
    labels=["Cs-137 Region", "K-40 Region"]  # Custom labels
)

print(f"\nCreated {len(rois)} ROIs:")
for roi in rois:
    print(f"  {roi.label:20s}: [{roi.e_min:7.1f}, {roi.e_max:7.1f}] keV, width={roi.width:.1f} keV")
    print(f"    Method: {roi.method}")
    print(f"    CEW predictor stored in metadata: {'cew_predictor' in roi.metadata}")

# =============================================================================
# Example 5: Adaptive Window Selection
# =============================================================================

print("\n" + "=" * 70)
print("Example 5: Adaptive Window Selection (Multiple n_bins)")
print("=" * 70)

# Test different numbers of bins
print("\nTesting different window sizes...")
n_bins_values = [10, 20, 30, 40, 50]

best_snr = -np.inf
best_n_bins = None
best_window = None

for n_bins in n_bins_values:
    window_mask = optimize_cew_windows(
        source_spec,
        bg_training.mean_spectrum(),
        method='greedy',
        n_bins_target=n_bins,
        verbose=False
    )
    
    predictor_temp = fit_cew_predictor(bg_training, window_mask, alpha=1.0)
    
    # Evaluate on test set
    test_scores_bg = [predictor_temp.score(Spectrum(np.random.poisson(100, 128))) for _ in range(20)]
    test_scores_src = [predictor_temp.score(test_source) for _ in range(20)]
    
    # Separation metric (higher is better)
    separation = (np.mean(test_scores_src) - np.mean(test_scores_bg)) / np.std(test_scores_bg)
    
    print(f"  n_bins={n_bins:2d}: separation={separation:.2f}")
    
    if separation > best_snr:
        best_snr = separation
        best_n_bins = n_bins
        best_window = window_mask

print(f"\nBest configuration: n_bins={best_n_bins}, separation={best_snr:.2f}")

# =============================================================================
# Example 6: Predictor Serialization
# =============================================================================

print("\n" + "=" * 70)
print("Example 6: Predictor Serialization")
print("=" * 70)

# Serialize predictor to dictionary
print("\nSerializing predictor...")
predictor_dict = predictor.to_dict()
print(f"  Serialized keys: {list(predictor_dict.keys())}")

# Deserialize
print("\nDeserializing predictor...")
predictor_loaded = predictor.__class__.from_dict(predictor_dict)
print(f"  {predictor_loaded}")

# Verify it works
test_spec = Spectrum(np.random.poisson(100, 128))
score_original = predictor.score(test_spec)
score_loaded = predictor_loaded.score(test_spec)

print(f"\nVerification:")
print(f"  Original predictor score: {score_original:.4f}")
print(f"  Loaded predictor score: {score_loaded:.4f}")
print(f"  Match: {np.isclose(score_original, score_loaded)}")

# =============================================================================
# Example 7: Real-World Scenario (Background Subtraction)
# =============================================================================

print("\n" + "=" * 70)
print("Example 7: Background Subtraction with CEW")
print("=" * 70)

# Simulate real measurement scenario
print("\nSimulating measurement scenario...")

# Collect background
print("  Phase 1: Collecting background (no source)...")
bg_measurements = [Spectrum(np.random.poisson(120, 128)) for _ in range(100)]
bg_data = Spectra(bg_measurements)
print(f"    Collected {bg_data.n_spectra} background spectra")

# Estimate source spectrum from a short measurement
print("  Phase 2: Short source measurement...")
gross_counts = np.random.poisson(120, 128)  # Background
gross_counts[50:70] += 400  # Source contribution
gross_spec = Spectrum(gross_counts)

# Estimate net source by subtracting background
net_source = Spectrum(gross_spec.counts - bg_data.mean_spectrum().counts)
print(f"    Net source spectrum created")

# Train CEW with background
print("  Phase 3: Training CEW detector...")
window_mask = optimize_cew_windows(
    net_source,
    bg_data.mean_spectrum(),
    method='gradient',
    max_iterations=300
)
detector = fit_cew_predictor(bg_data, window_mask, alpha=0.5)
print(f"    Trained: {detector}")

# Test on new measurements
print("  Phase 4: Testing detector...")
# Clean background
clean_bg = Spectrum(np.random.poisson(120, 128))
clean_score = detector.score(clean_bg)

# With source
with_source_counts = np.random.poisson(120, 128)
with_source_counts[50:70] += 350
with_source = Spectrum(with_source_counts)
source_score = detector.score(with_source)

print(f"\nResults:")
print(f"  Clean background score: {clean_score:.2f}")
print(f"  Measurement with source score: {source_score:.2f}")
print(f"  Detection ratio: {source_score/clean_score:.2f}x")

print("\n" + "=" * 70)
print("CEW Examples Complete!")
print("=" * 70)
print("\nKey Takeaways:")
print("  1. CEW optimizes window selection for maximum SNR")
print("  2. Predictor uses censored (out-of-window) bins for background estimation")
print("  3. Works with both uncalibrated and calibrated spectra")
print("  4. Integrates seamlessly with ROI system")
print("  5. Effective for time series anomaly detection")
print("  6. Multiple optimization methods available (gradient, greedy, scipy)")
print("  7. Predictors are serializable for deployment")

