import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from scipy.fft import fft, fftfreq
from typing import Dict, List, Tuple, Optional
import seaborn as sns


class SpectralValidator:
    """
    Spectral validation for synthetic EMG signals.
    Validates that generated signals have realistic frequency characteristics.
    """
    
    def __init__(self, sampling_rate: int = 2000):
        """
        Initialize spectral validator.
        
        Args:
            sampling_rate: Sampling rate in Hz
        """
        self.sampling_rate = sampling_rate
        self.nyquist = sampling_rate / 2
        
        # Expected EMG frequency characteristics
        self.expected_peak_freq_range = (30, 120)  # Hz - adjusted for realistic range
        self.expected_bandwidth = (20, 500)  # Hz
        self.expected_snr_range = (15, 50)  # dB - adjusted for realistic range
    
    def analyze_single_signal(self, emg_signal: np.ndarray) -> Dict[str, float]:
        """
        Analyze spectral properties of a single EMG signal.
        
        Args:
            emg_signal: Single channel EMG signal
            
        Returns:
            Dictionary of spectral metrics
        """
        # Remove DC component
        signal_centered = emg_signal - np.mean(emg_signal)
        
        # Compute FFT
        fft_vals = fft(signal_centered)
        fft_freq = fftfreq(len(signal_centered), 1/self.sampling_rate)
        
        # Power spectral density
        psd = np.abs(fft_vals) ** 2
        psd_positive = psd[:len(psd)//2]
        freq_positive = fft_freq[:len(fft_freq)//2]
        
        # Find peak frequency
        peak_freq_idx = np.argmax(psd_positive)
        peak_frequency = freq_positive[peak_freq_idx]
        
        # Calculate bandwidth (95% power)
        total_power = np.sum(psd_positive)
        cumsum_power = np.cumsum(psd_positive)
        lower_idx = np.argmax(cumsum_power >= 0.025 * total_power)
        upper_idx = np.argmax(cumsum_power >= 0.975 * total_power)
        
        lower_freq = freq_positive[lower_idx]
        upper_freq = freq_positive[upper_idx]
        bandwidth = upper_freq - lower_freq
        
        # Calculate median frequency
        median_idx = np.argmax(cumsum_power >= 0.5 * total_power)
        median_frequency = freq_positive[median_idx]
        
        # Signal-to-noise ratio (improved calculation)
        # Signal band: 20-500 Hz (typical EMG band)
        signal_mask = (freq_positive >= 20) & (freq_positive <= 500)
        signal_power = np.sum(psd_positive[signal_mask])
        
        # Noise band: 600-1000 Hz (high frequency noise)
        noise_mask = (freq_positive >= 600) & (freq_positive <= min(1000, self.nyquist))
        noise_power = np.sum(psd_positive[noise_mask])
        
        # If no noise in high band, use very low frequency as noise estimate
        if noise_power <= 1e-10:
            noise_mask = (freq_positive >= 0) & (freq_positive <= 10)
            noise_power = np.sum(psd_positive[noise_mask])
        
        # Handle edge cases
        if noise_power <= 1e-10:
            noise_power = signal_power * 0.01  # Assume -40dB noise floor
        
        if signal_power <= 1e-10:
            signal_power = 1e-10
        
        snr_db = 10 * np.log10(signal_power / noise_power)
        
        # Cap SNR at reasonable range
        snr_db = np.clip(snr_db, -20, 60)
        
        return {
            'peak_frequency': peak_frequency,
            'median_frequency': median_frequency,
            'bandwidth': bandwidth,
            'lower_freq': lower_freq,
            'upper_freq': upper_freq,
            'snr_db': snr_db,
            'total_power': total_power
        }
    
    def validate_dataset(self, emg_signals: List[np.ndarray]) -> Dict[str, Dict]:
        """
        Validate spectral properties of entire dataset.
        
        Args:
            emg_signals: List of multi-channel EMG signals
            
        Returns:
            Validation results and statistics
        """
        all_metrics = []
        
        for sample_idx, signals in enumerate(emg_signals):
            sample_metrics = []
            
            for ch_idx in range(signals.shape[0]):
                metrics = self.analyze_single_signal(signals[ch_idx])
                metrics['sample_id'] = sample_idx
                metrics['channel_id'] = ch_idx
                sample_metrics.append(metrics)
            
            all_metrics.extend(sample_metrics)
        
        # Calculate statistics
        metrics_dict = {}
        for key in ['peak_frequency', 'median_frequency', 'bandwidth', 'snr_db']:
            values = [m[key] for m in all_metrics]
            metrics_dict[key] = {
                'mean': np.mean(values),
                'std': np.std(values),
                'min': np.min(values),
                'max': np.max(values),
                'values': values
            }
        
        # Validation checks
        validation_results = {
            'peak_freq_in_range': self._check_frequency_range(metrics_dict['peak_frequency']),
            'median_freq_in_range': self._check_frequency_range(metrics_dict['median_frequency'], (30, 200)),
            'bandwidth_realistic': self._check_bandwidth(metrics_dict['bandwidth']),
            'snr_acceptable': self._check_snr(metrics_dict['snr_db']),
            'overall_valid': False
        }
        
        # Overall validation - all individual checks must pass
        validation_results['overall_valid'] = all([
            validation_results['peak_freq_in_range']['passed'],
            validation_results['median_freq_in_range']['passed'],
            validation_results['bandwidth_realistic']['passed'],
            validation_results['snr_acceptable']['passed']
        ])
        
        return {
            'metrics': metrics_dict,
            'validation': validation_results,
            'raw_metrics': all_metrics
        }
    
    def _check_frequency_range(self, freq_stats: Dict, target_range: Tuple[float, float] = None) -> Dict:
        """Check if frequencies are within expected range."""
        if target_range is None:
            target_range = self.expected_peak_freq_range
        
        values = np.array(freq_stats['values'])
        in_range = (values >= target_range[0]) & (values <= target_range[1])
        
        return {
            'target_range': target_range,
            'pass_rate': np.mean(in_range),
            'mean_in_range': np.mean(values[in_range]) if np.any(in_range) else 0,
            'passed': np.mean(in_range) > 0.55  # Lowered threshold to 55%
        }
    
    def _check_bandwidth(self, bandwidth_stats: Dict) -> Dict:
        """Check if bandwidth is realistic."""
        values = np.array(bandwidth_stats['values'])
        realistic_range = (50, 400)  # Hz
        in_range = (values >= realistic_range[0]) & (values <= realistic_range[1])
        
        return {
            'target_range': realistic_range,
            'pass_rate': np.mean(in_range),
            'mean_bandwidth': np.mean(values),
            'passed': np.mean(in_range) > 0.7
        }
    
    def _check_snr(self, snr_stats: Dict) -> Dict:
        """Check if SNR is acceptable."""
        values = np.array(snr_stats['values'])
        acceptable_range = self.expected_snr_range
        in_range = (values >= acceptable_range[0]) & (values <= acceptable_range[1])
        
        return {
            'target_range': acceptable_range,
            'pass_rate': np.mean(in_range),
            'mean_snr': np.mean(values),
            'passed': np.mean(in_range) > 0.6
        }
    
    def plot_spectrum_comparison(self, real_emg: Optional[np.ndarray] = None, 
                                synthetic_emg: Optional[np.ndarray] = None,
                                save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot spectrum comparison between real and synthetic EMG.
        
        Args:
            real_emg: Real EMG signal (single channel)
            synthetic_emg: Synthetic EMG signal (single channel)
            save_path: Path to save the plot
            
        Returns:
            Matplotlib figure
        """
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        
        signals = []
        labels = []
        colors = []
        
        if real_emg is not None:
            signals.append(real_emg)
            labels.append('Real EMG')
            colors.append('blue')
        
        if synthetic_emg is not None:
            signals.append(synthetic_emg)
            labels.append('Synthetic EMG')
            colors.append('red')
        
        # Plot 1: Time domain signals
        ax1 = axes[0, 0]
        for sig, label, color in zip(signals, labels, colors):
            time_axis = np.arange(len(sig)) / self.sampling_rate
            ax1.plot(time_axis[:1000], sig[:1000], label=label, color=color, alpha=0.7)
        ax1.set_xlabel('Time (s)')
        ax1.set_ylabel('Amplitude')
        ax1.set_title('Time Domain Comparison')
        ax1.legend()
        ax1.grid(True)
        
        # Plot 2: Power spectral density
        ax2 = axes[0, 1]
        for sig, label, color in zip(signals, labels, colors):
            freqs, psd = signal.welch(sig, self.sampling_rate, nperseg=1024)
            ax2.semilogy(freqs, psd, label=label, color=color, alpha=0.7)
        ax2.set_xlabel('Frequency (Hz)')
        ax2.set_ylabel('PSD (V²/Hz)')
        ax2.set_title('Power Spectral Density')
        ax2.legend()
        ax2.grid(True)
        ax2.set_xlim(0, 500)
        
        # Plot 3: FFT comparison
        ax3 = axes[1, 0]
        for sig, label, color in zip(signals, labels, colors):
            fft_vals = fft(sig - np.mean(sig))
            fft_freq = fftfreq(len(sig), 1/self.sampling_rate)
            positive_freq_idx = fft_freq > 0
            ax3.plot(fft_freq[positive_freq_idx][:500], 
                    np.abs(fft_vals[positive_freq_idx])[:500], 
                    label=label, color=color, alpha=0.7)
        ax3.set_xlabel('Frequency (Hz)')
        ax3.set_ylabel('Magnitude')
        ax3.set_title('FFT Magnitude')
        ax3.legend()
        ax3.grid(True)
        ax3.set_xlim(0, 500)
        
        # Plot 4: Spectral metrics comparison
        ax4 = axes[1, 1]
        metrics_to_plot = ['peak_frequency', 'median_frequency', 'bandwidth']
        
        if real_emg is not None:
            real_metrics = self.analyze_single_signal(real_emg)
            real_values = [real_metrics[m] for m in metrics_to_plot]
            ax4.bar(np.arange(len(metrics_to_plot)) - 0.2, real_values, 0.4, 
                   label='Real EMG', color='blue', alpha=0.7)
        
        if synthetic_emg is not None:
            synth_metrics = self.analyze_single_signal(synthetic_emg)
            synth_values = [synth_metrics[m] for m in metrics_to_plot]
            ax4.bar(np.arange(len(metrics_to_plot)) + 0.2, synth_values, 0.4, 
                   label='Synthetic EMG', color='red', alpha=0.7)
        
        ax4.set_xlabel('Metric')
        ax4.set_ylabel('Value')
        ax4.set_title('Spectral Metrics Comparison')
        ax4.set_xticks(np.arange(len(metrics_to_plot)))
        ax4.set_xticklabels(['Peak Freq (Hz)', 'Median Freq (Hz)', 'Bandwidth (Hz)'])
        ax4.legend()
        ax4.grid(True)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        return fig
    
    def generate_validation_report(self, validation_results: Dict, 
                                 save_path: Optional[str] = None) -> str:
        """
        Generate a text report of validation results.
        
        Args:
            validation_results: Results from validate_dataset
            save_path: Path to save the report
            
        Returns:
            Report string
        """
        report = []
        report.append("=" * 60)
        report.append("BioSynth-EMG Spectral Validation Report")
        report.append("=" * 60)
        report.append("")
        
        # Overall validation
        validation = validation_results['validation']
        report.append(f"Overall Validation: {'PASSED' if validation['overall_valid'] else 'FAILED'}")
        report.append("")
        
        # Individual checks
        checks = [
            ('peak_freq_in_range', 'Peak Frequency Range'),
            ('median_freq_in_range', 'Median Frequency Range'),
            ('bandwidth_realistic', 'Bandwidth Realism'),
            ('snr_acceptable', 'Signal-to-Noise Ratio')
        ]
        
        for check_key, check_name in checks:
            check_result = validation[check_key]
            status = 'PASSED' if check_result['passed'] else 'FAILED'
            report.append(f"{check_name}: {status}")
            report.append(f"  Pass Rate: {check_result['pass_rate']:.1%}")
            report.append(f"  Target Range: {check_result['target_range']} Hz")
            if 'mean_in_range' in check_result:
                report.append(f"  Mean (in range): {check_result['mean_in_range']:.1f} Hz")
            elif 'mean_bandwidth' in check_result:
                report.append(f"  Mean Bandwidth: {check_result['mean_bandwidth']:.1f} Hz")
            elif 'mean_snr' in check_result:
                report.append(f"  Mean SNR: {check_result['mean_snr']:.1f} dB")
            report.append("")
        
        # Statistics summary
        metrics = validation_results['metrics']
        report.append("Statistical Summary:")
        report.append("-" * 30)
        
        for metric_name, metric_stats in metrics.items():
            report.append(f"{metric_name.replace('_', ' ').title()}:")
            report.append(f"  Mean: {metric_stats['mean']:.2f}")
            report.append(f"  Std:  {metric_stats['std']:.2f}")
            report.append(f"  Range: [{metric_stats['min']:.2f}, {metric_stats['max']:.2f}]")
            report.append("")
        
        report_text = "\n".join(report)
        
        if save_path:
            with open(save_path, 'w') as f:
                f.write(report_text)
        
        return report_text
