"""
Tests for Censored Energy Window (CEW) algorithm.
"""

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


class TestCEWWindowOptimization:
    """Test window optimization functions."""
    
    def test_optimize_gradient_basic(self):
        """Test gradient-based optimization."""
        # Create synthetic source and background
        source = Spectrum(np.concatenate([np.ones(32) * 10, np.ones(32)]))
        background = Spectrum(np.ones(64))
        
        window_mask = optimize_cew_windows(
            source, background,
            method='gradient',
            max_iterations=100,
        )
        
        assert window_mask.shape == (64,)
        assert np.all((window_mask == 0) | (window_mask == 1))
        assert np.sum(window_mask) > 0  # Should select some bins
    
    def test_optimize_greedy_basic(self):
        """Test greedy optimization."""
        source = Spectrum(np.concatenate([np.ones(32) * 10, np.ones(32)]))
        background = Spectrum(np.ones(64))
        
        window_mask = optimize_cew_windows(
            source, background,
            method='greedy',
            n_bins_target=20,
        )
        
        assert np.sum(window_mask) == 20
    
    def test_optimize_scipy_basic(self):
        """Test scipy optimization."""
        source = Spectrum(np.concatenate([np.ones(32) * 10, np.ones(32)]))
        background = Spectrum(np.ones(64))
        
        window_mask = optimize_cew_windows(
            source, background,
            method='scipy',
            n_bins_target=20,
        )
        
        assert window_mask.shape == (64,)
        assert np.all((window_mask == 0) | (window_mask == 1))
    
    def test_optimize_selects_high_snr_bins(self):
        """Test that optimization selects high SNR bins."""
        # Source has peak in first half, background uniform
        source_counts = np.concatenate([
            np.random.poisson(1000, 32),  # High signal
            np.random.poisson(10, 32),    # Low signal
        ])
        background_counts = np.ones(64) * 10
        
        source = Spectrum(source_counts)
        background = Spectrum(background_counts)
        
        window_mask = optimize_cew_windows(
            source, background,
            method='greedy',
            n_bins_target=20,
        )
        
        # Should select more from first half (higher SNR)
        first_half_selected = np.sum(window_mask[:32])
        second_half_selected = np.sum(window_mask[32:])
        
        assert first_half_selected > second_half_selected
    
    def test_optimize_incompatible_bins_raises(self):
        """Test error when source and background have different bins."""
        source = Spectrum(np.ones(64))
        background = Spectrum(np.ones(32))
        
        with pytest.raises(ValueError, match="same number of bins"):
            optimize_cew_windows(source, background)
    
    def test_optimize_unknown_method_raises(self):
        """Test error for unknown optimization method."""
        source = Spectrum(np.ones(64))
        background = Spectrum(np.ones(64))
        
        with pytest.raises(ValueError, match="Unknown optimization method"):
            optimize_cew_windows(source, background, method='invalid')


class TestCEWPredictorFitting:
    """Test fitting CEW predictors."""
    
    def test_fit_predictor_basic(self):
        """Test basic predictor fitting."""
        # Create background training data
        bg_spectra = [Spectrum(np.random.poisson(100, 64)) for _ in range(20)]
        bg_training = Spectra(bg_spectra)
        
        # Define window mask
        window_mask = np.zeros(64, dtype=int)
        window_mask[20:40] = 1  # Select middle bins
        
        # Fit predictor
        predictor = fit_cew_predictor(bg_training, window_mask, alpha=1.0)
        
        assert isinstance(predictor, CEWPredictor)
        assert predictor.n_bins == 64
        assert len(predictor.theta) == 64
        assert np.all(predictor.theta[window_mask == 1] == 0)  # In-window coefficients are zero
    
    def test_predictor_score_spectrum(self):
        """Test scoring a single spectrum."""
        # Create simple predictor
        theta = np.ones(64) * 0.5
        window_mask = np.zeros(64, dtype=int)
        window_mask[20:40] = 1
        theta[window_mask == 1] = 0  # Zero in-window
        
        predictor = CEWPredictor(theta, window_mask, alpha=1.0)
        
        # Score a spectrum
        spec = Spectrum(np.ones(64) * 100)
        score = predictor.score(spec)
        
        assert isinstance(score, float)
        # Score should be positive (in-window counts minus predicted background)
        assert score != 0
    
    def test_predictor_score_spectra_collection(self):
        """Test scoring multiple spectra."""
        theta = np.ones(64) * 0.5
        window_mask = np.zeros(64, dtype=int)
        window_mask[20:40] = 1
        theta[window_mask == 1] = 0
        
        predictor = CEWPredictor(theta, window_mask, alpha=1.0)
        
        # Score multiple spectra
        spectra_list = [Spectrum(np.random.poisson(100, 64)) for _ in range(10)]
        spectra = Spectra(spectra_list)
        
        scores = predictor.score_spectra(spectra)
        
        assert scores.shape == (10,)
        assert np.all(np.isfinite(scores))
    
    def test_predictor_detects_anomaly(self):
        """Test that predictor gives higher scores for anomalous spectra."""
        # Create background training
        bg_spectra = [Spectrum(np.random.poisson(100, 64)) for _ in range(50)]
        bg_training = Spectra(bg_spectra)
        
        # Define window (middle bins)
        window_mask = np.zeros(64, dtype=int)
        window_mask[25:35] = 1
        
        # Fit predictor
        predictor = fit_cew_predictor(bg_training, window_mask, alpha=1.0)
        
        # Score background spectrum (should be low)
        bg_spec = Spectrum(np.random.poisson(100, 64))
        bg_score = predictor.score(bg_spec)
        
        # Score anomalous spectrum (high counts in window)
        anom_counts = np.random.poisson(100, 64)
        anom_counts[25:35] += 500  # Add signal in window
        anom_spec = Spectrum(anom_counts)
        anom_score = predictor.score(anom_spec)
        
        # Anomalous spectrum should have higher score
        assert anom_score > bg_score
    
    def test_fit_window_mask_validation(self):
        """Test validation of window_mask."""
        bg_training = Spectra([Spectrum(np.ones(64)) for _ in range(10)])
        
        # Wrong length
        with pytest.raises(ValueError, match="must match number of bins"):
            fit_cew_predictor(bg_training, np.ones(32, dtype=int), alpha=1.0)
        
        # All zeros
        with pytest.raises(ValueError, match="no in-window bins"):
            fit_cew_predictor(bg_training, np.zeros(64, dtype=int), alpha=1.0)
        
        # All ones
        with pytest.raises(ValueError, match="no out-of-window bins"):
            fit_cew_predictor(bg_training, np.ones(64, dtype=int), alpha=1.0)


class TestCEWPredictor:
    """Test CEWPredictor class."""
    
    def test_predictor_init(self):
        """Test predictor initialization."""
        theta = np.ones(64) * 0.5
        window_mask = np.zeros(64, dtype=int)
        window_mask[20:40] = 1
        
        predictor = CEWPredictor(theta, window_mask, alpha=1.0)
        
        assert predictor.n_bins == 64
        assert predictor.alpha == 1.0
        assert len(predictor.scoring_vector) == 64
    
    def test_predictor_to_rois_without_calibration(self):
        """Test that to_rois raises without calibration."""
        theta = np.ones(64) * 0.5
        window_mask = np.zeros(64, dtype=int)
        window_mask[20:40] = 1
        
        predictor = CEWPredictor(theta, window_mask, alpha=1.0, energy_edges=None)
        
        with pytest.raises(ValueError, match="energy calibration not available"):
            predictor.to_rois()
    
    def test_predictor_to_rois_with_calibration(self):
        """Test converting predictor to ROIs."""
        theta = np.ones(64) * 0.5
        window_mask = np.zeros(64, dtype=int)
        window_mask[10:20] = 1
        window_mask[40:50] = 1  # Two separate regions
        theta[window_mask == 1] = 0
        
        energy_edges = np.linspace(0, 3000, 65)
        predictor = CEWPredictor(theta, window_mask, alpha=1.0, energy_edges=energy_edges)
        
        rois = predictor.to_rois()
        
        assert len(rois) == 2
        assert all(isinstance(roi, EnergyROI) for roi in rois)
        assert rois[0].method == "Censored Energy Windows"
        assert "cew_predictor" in rois[0].metadata
    
    def test_predictor_serialization(self):
        """Test to_dict and from_dict."""
        theta = np.ones(64) * 0.5
        window_mask = np.zeros(64, dtype=int)
        window_mask[20:40] = 1
        energy_edges = np.linspace(0, 3000, 65)
        
        predictor = CEWPredictor(theta, window_mask, alpha=1.0, energy_edges=energy_edges)
        
        # Serialize
        d = predictor.to_dict()
        
        assert 'theta' in d
        assert 'window_mask' in d
        assert 'alpha' in d
        assert 'energy_edges' in d
        
        # Deserialize
        predictor2 = CEWPredictor.from_dict(d)
        
        assert np.allclose(predictor2.theta, predictor.theta)
        assert np.array_equal(predictor2.window_mask, predictor.window_mask)
        assert predictor2.alpha == predictor.alpha
    
    def test_predictor_repr(self):
        """Test __repr__."""
        theta = np.ones(64) * 0.5
        window_mask = np.zeros(64, dtype=int)
        window_mask[20:40] = 1
        
        predictor = CEWPredictor(theta, window_mask, alpha=1.0)
        
        repr_str = repr(predictor)
        assert "CEWPredictor" in repr_str
        assert "n_bins=64" in repr_str
        assert "n_in_window=20" in repr_str


class TestCEWIntegration:
    """Test end-to-end CEW workflow."""
    
    def test_complete_workflow(self):
        """Test complete CEW workflow."""
        np.random.seed(42)
        
        # 1. Create synthetic data
        # Background: uniform
        bg_spectra = [Spectrum(np.random.poisson(100, 128)) for _ in range(50)]
        bg_training = Spectra(bg_spectra)
        bg_mean = bg_training.mean_spectrum()
        
        # Source: has peak in bins 60-80
        source_counts = np.random.poisson(100, 128)
        source_counts[60:80] += 500
        source_spec = Spectrum(source_counts)
        
        # 2. Optimize windows
        window_mask = optimize_cew_windows(
            source_spec,
            bg_mean,
            method='greedy',
            n_bins_target=30,
        )
        
        # Windows should include the peak region
        assert np.sum(window_mask[60:80]) > 0
        
        # 3. Fit predictor
        predictor = fit_cew_predictor(bg_training, window_mask, alpha=1.0)
        
        # 4. Score test spectra
        # Background spectrum (low score expected)
        test_bg = Spectrum(np.random.poisson(100, 128))
        bg_score = predictor.score(test_bg)
        
        # Anomalous spectrum (high score expected)
        test_anom_counts = np.random.poisson(100, 128)
        test_anom_counts[60:80] += 300
        test_anom = Spectrum(test_anom_counts)
        anom_score = predictor.score(test_anom)
        
        # Anomalous should score higher
        assert anom_score > bg_score
    
    def test_roi_integration(self):
        """Test integration with ROI system."""
        # Create predictor with calibrated spectra
        edges = np.linspace(0, 3000, 65)
        bg_spectra = [
            Spectrum(np.random.poisson(100, 64), energy_edges=edges)
            for _ in range(20)
        ]
        bg_training = Spectra(bg_spectra)
        
        window_mask = np.zeros(64, dtype=int)
        window_mask[20:40] = 1
        
        predictor = fit_cew_predictor(bg_training, window_mask, alpha=1.0)
        
        # Convert to ROIs
        rois = EnergyROI.from_cew_predictor(predictor)
        
        assert len(rois) > 0
        assert all(isinstance(roi, EnergyROI) for roi in rois)
        assert all(roi.method == "Censored Energy Windows" for roi in rois)

