#!/usr/bin/env python3
"""
⁷Be DATA ANALYSIS TOOLKIT
Complete analysis pipeline for half-life measurements
Ready for real experimental data
"""
import numpy as np
from scipy.optimize import curve_fit, minimize
from scipy.stats import chi2, poisson
from scipy.signal import savgol_filter
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
import json
import warnings
warnings.filterwarnings('ignore')

print("=" * 80)
print("⁷Be DATA ANALYSIS TOOLKIT v1.0")
print("Complete pipeline for half-life measurements in heavy fermion materials")
print("=" * 80)

# ============================================================================
# MODULE 1: SPECTRUM ANALYSIS
# ============================================================================
class SpectrumAnalyzer:
    """
    Class for analyzing γ-ray spectra from HPGe detector.
    Includes peak fitting, background subtraction, and efficiency correction.
    """
    
    def __init__(self, energy_calibration=None):
        """
        Initialize spectrum analyzer.
        
        Parameters:
            energy_calibration: tuple (a, b) for E = a*channel + b
        """
        self.energy_cal = energy_calibration or (0.25, 0)  # default: 0.25 keV/channel
        
    def channel_to_energy(self, channel):
        """Convert channel number to energy (keV)."""
        return self.energy_cal[0] * channel + self.energy_cal[1]
    
    def energy_to_channel(self, energy):
        """Convert energy (keV) to channel number."""
        return int((energy - self.energy_cal[1]) / self.energy_cal[0])
    
    @staticmethod
    def gaussian(x, A, mu, sigma):
        """Gaussian peak function."""
        return A * np.exp(-(x - mu)**2 / (2 * sigma**2))
    
    @staticmethod
    def gaussian_with_linear_bg(x, A, mu, sigma, B, C):
        """Gaussian peak with linear background."""
        return A * np.exp(-(x - mu)**2 / (2 * sigma**2)) + B + C * x
    
    def fit_peak(self, channels, counts, peak_center, fit_width=50):
        """
        Fit a Gaussian peak with linear background.
        
        Parameters:
            channels: array of channel numbers
            counts: array of counts per channel
            peak_center: expected peak center (channel)
            fit_width: number of channels to include in fit
            
        Returns:
            dict with fit results
        """
        # Select fit region
        mask = np.abs(channels - peak_center) < fit_width
        x = channels[mask]
        y = counts[mask]
        
        # Initial guesses
        A0 = np.max(y) - np.min(y)
        mu0 = x[np.argmax(y)]
        sigma0 = 3  # typical HPGe resolution
        B0 = np.min(y)
        C0 = 0
        
        try:
            popt, pcov = curve_fit(self.gaussian_with_linear_bg, x, y,
                                   p0=[A0, mu0, sigma0, B0, C0],
                                   sigma=np.sqrt(y + 1),  # Poisson errors
                                   absolute_sigma=True,
                                   maxfev=10000)
            
            A, mu, sigma, B, C = popt
            errors = np.sqrt(np.diag(pcov))
            
            # Calculate net area
            net_area = A * sigma * np.sqrt(2 * np.pi)
            
            # Calculate background under peak
            bg_area = (B + C * mu) * 6 * sigma  # ±3σ range
            
            # Net area error (simplified)
            net_area_err = np.sqrt(net_area + bg_area)
            
            # Chi-squared
            y_fit = self.gaussian_with_linear_bg(x, *popt)
            chi2_val = np.sum((y - y_fit)**2 / (y + 1))
            dof = len(x) - 5
            chi2_red = chi2_val / dof
            
            return {
                'amplitude': A,
                'amplitude_err': errors[0],
                'centroid': mu,
                'centroid_err': errors[1],
                'sigma': abs(sigma),
                'sigma_err': errors[2],
                'fwhm': 2.355 * abs(sigma),
                'net_area': net_area,
                'net_area_err': net_area_err,
                'background': B,
                'bg_slope': C,
                'chi2_red': chi2_red,
                'success': True
            }
            
        except Exception as e:
            return {
                'success': False,
                'error': str(e)
            }
    
    def analyze_spectrum(self, channels, counts, roi_center=478, roi_width=20):
        """
        Complete spectrum analysis for ⁷Be 478 keV peak.
        
        Parameters:
            channels: array of channel numbers
            counts: array of counts
            roi_center: ROI center in keV
            roi_width: ROI half-width in keV
            
        Returns:
            dict with analysis results
        """
        # Convert ROI to channels
        center_ch = self.energy_to_channel(roi_center)
        width_ch = int(roi_width / self.energy_cal[0])
        
        # Fit peak
        result = self.fit_peak(channels, counts, center_ch, fit_width=width_ch)
        
        if result['success']:
            # Add energy values
            result['centroid_keV'] = self.channel_to_energy(result['centroid'])
            result['fwhm_keV'] = result['fwhm'] * self.energy_cal[0]
            
        return result


# ============================================================================
# MODULE 2: DECAY CURVE FITTING
# ============================================================================
class DecayCurveFitter:
    """
    Class for fitting radioactive decay curves and extracting half-lives.
    """
    
    def __init__(self):
        self.results = {}
    
    @staticmethod
    def decay_function(t, A0, t_half):
        """Single exponential decay."""
        return A0 * np.exp(-np.log(2) * t / t_half)
    
    @staticmethod
    def decay_with_background(t, A0, t_half, bg):
        """Decay with constant background."""
        return A0 * np.exp(-np.log(2) * t / t_half) + bg
    
    def fit_halflife(self, times, activities, errors=None, include_background=False):
        """
        Fit decay curve to extract half-life.
        
        Parameters:
            times: array of measurement times (days from start)
            activities: array of measured activities (counts or Bq)
            errors: array of uncertainties (optional, sqrt(N) if None)
            include_background: whether to include background term
            
        Returns:
            dict with fit results
        """
        times = np.array(times)
        activities = np.array(activities)
        
        if errors is None:
            errors = np.sqrt(activities + 1)
        else:
            errors = np.array(errors)
        
        # Initial guesses
        A0_guess = activities[0]
        t_half_guess = 53.22  # ⁷Be half-life
        
        try:
            if include_background:
                popt, pcov = curve_fit(
                    self.decay_with_background, times, activities,
                    p0=[A0_guess, t_half_guess, 0],
                    sigma=errors, absolute_sigma=True,
                    bounds=([0, 1, 0], [np.inf, 200, np.inf])
                )
                A0, t_half, bg = popt
                param_names = ['A0', 't_half', 'background']
            else:
                popt, pcov = curve_fit(
                    self.decay_function, times, activities,
                    p0=[A0_guess, t_half_guess],
                    sigma=errors, absolute_sigma=True,
                    bounds=([0, 1], [np.inf, 200])
                )
                A0, t_half = popt
                bg = 0
                param_names = ['A0', 't_half']
            
            errors_fit = np.sqrt(np.diag(pcov))
            
            # Chi-squared calculation
            if include_background:
                y_fit = self.decay_with_background(times, *popt)
            else:
                y_fit = self.decay_function(times, *popt)
            
            chi2_val = np.sum(((activities - y_fit) / errors)**2)
            dof = len(times) - len(popt)
            chi2_red = chi2_val / dof
            p_value = 1 - chi2.cdf(chi2_val, dof)
            
            # Decay constant
            lambda_decay = np.log(2) / t_half
            lambda_err = np.log(2) / t_half**2 * errors_fit[1]
            
            return {
                'A0': A0,
                'A0_err': errors_fit[0],
                't_half': t_half,
                't_half_err': errors_fit[1],
                'lambda': lambda_decay,
                'lambda_err': lambda_err,
                'background': bg,
                'chi2': chi2_val,
                'chi2_red': chi2_red,
                'dof': dof,
                'p_value': p_value,
                'success': True
            }
            
        except Exception as e:
            return {
                'success': False,
                'error': str(e)
            }
    
    def compare_halflife(self, result1, result2):
        """
        Compare two half-life measurements and calculate significance.
        
        Parameters:
            result1: dict from fit_halflife (reference)
            result2: dict from fit_halflife (test)
            
        Returns:
            dict with comparison results
        """
        t1 = result1['t_half']
        t2 = result2['t_half']
        e1 = result1['t_half_err']
        e2 = result2['t_half_err']
        
        delta = t2 - t1
        delta_err = np.sqrt(e1**2 + e2**2)
        
        # Significance
        if delta_err > 0:
            z_score = abs(delta) / delta_err
        else:
            z_score = np.inf
        
        # P-value (two-tailed)
        from scipy.stats import norm
        p_value = 2 * (1 - norm.cdf(z_score))
        
        return {
            't_half_1': t1,
            't_half_2': t2,
            'delta': delta,
            'delta_err': delta_err,
            'delta_percent': 100 * delta / t1,
            'z_score': z_score,
            'p_value': p_value,
            'significant_3sigma': z_score > 3,
            'significant_5sigma': z_score > 5
        }


# ============================================================================
# MODULE 3: EFFICIENCY CALIBRATION
# ============================================================================
class EfficiencyCalibration:
    """
    Class for detector efficiency calibration.
    """
    
    def __init__(self):
        # Standard calibration sources and energies
        self.standard_sources = {
            'Am-241': [(59.54, 0.359)],  # (energy keV, intensity)
            'Ba-133': [(80.99, 0.329), (276.40, 0.071), (302.85, 0.183), 
                       (356.02, 0.621), (383.85, 0.089)],
            'Cs-137': [(661.66, 0.851)],
            'Co-60': [(1173.24, 0.999), (1332.50, 0.9998)],
            'Eu-152': [(121.78, 0.284), (244.70, 0.076), (344.28, 0.265),
                       (778.90, 0.129), (964.08, 0.146), (1408.01, 0.210)]
        }
        
        self.calibration_data = []
        self.fit_params = None
    
    @staticmethod
    def efficiency_function(E, a, b, c, d):
        """
        Empirical efficiency function.
        log(ε) = a + b*log(E) + c*log(E)² + d*log(E)³
        """
        logE = np.log(E)
        return np.exp(a + b*logE + c*logE**2 + d*logE**3)
    
    def add_calibration_point(self, energy, efficiency, error=None):
        """Add a calibration point."""
        self.calibration_data.append({
            'energy': energy,
            'efficiency': efficiency,
            'error': error or 0.05 * efficiency
        })
    
    def fit_calibration(self):
        """Fit efficiency curve to calibration data."""
        if len(self.calibration_data) < 4:
            raise ValueError("Need at least 4 calibration points")
        
        energies = np.array([p['energy'] for p in self.calibration_data])
        efficiencies = np.array([p['efficiency'] for p in self.calibration_data])
        errors = np.array([p['error'] for p in self.calibration_data])
        
        # Initial guess
        p0 = [-5, 0.5, -0.1, 0.01]
        
        popt, pcov = curve_fit(self.efficiency_function, energies, efficiencies,
                               p0=p0, sigma=errors, absolute_sigma=True)
        
        self.fit_params = popt
        self.fit_errors = np.sqrt(np.diag(pcov))
        
        return popt
    
    def get_efficiency(self, energy):
        """Get efficiency at given energy."""
        if self.fit_params is None:
            raise ValueError("Run fit_calibration first")
        return self.efficiency_function(energy, *self.fit_params)
    
    def get_efficiency_at_478keV(self):
        """Get efficiency at ⁷Be γ energy."""
        return self.get_efficiency(477.6)


# ============================================================================
# MODULE 4: DATA I/O
# ============================================================================
class DataManager:
    """
    Class for managing experimental data files.
    """
    
    def __init__(self, data_dir='.'):
        self.data_dir = data_dir
        self.measurements = []
    
    def load_spectrum(self, filename):
        """
        Load spectrum from file.
        Supports: .spe, .csv, .txt formats
        """
        # Determine format
        if filename.endswith('.spe'):
            return self._load_spe(filename)
        elif filename.endswith('.csv'):
            return self._load_csv(filename)
        else:
            return self._load_txt(filename)
    
    def _load_csv(self, filename):
        """Load CSV spectrum file."""
        data = np.genfromtxt(filename, delimiter=',', skip_header=1)
        return {
            'channels': data[:, 0].astype(int),
            'counts': data[:, 1].astype(float),
            'filename': filename
        }
    
    def _load_txt(self, filename):
        """Load text spectrum file."""
        data = np.genfromtxt(filename, skip_header=1)
        return {
            'channels': data[:, 0].astype(int),
            'counts': data[:, 1].astype(float),
            'filename': filename
        }
    
    def _load_spe(self, filename):
        """Load Ortec .spe format."""
        counts = []
        with open(filename, 'r') as f:
            lines = f.readlines()
            in_data = False
            for line in lines:
                if '$DATA:' in line:
                    in_data = True
                    continue
                if in_data:
                    if line.startswith('$'):
                        break
                    try:
                        counts.append(int(line.strip()))
                    except:
                        pass
        
        return {
            'channels': np.arange(len(counts)),
            'counts': np.array(counts, dtype=float),
            'filename': filename
        }
    
    def add_measurement(self, sample_name, timestamp, spectrum_file, 
                        live_time, real_time):
        """Add a measurement to the dataset."""
        spectrum = self.load_spectrum(spectrum_file)
        self.measurements.append({
            'sample': sample_name,
            'timestamp': timestamp,
            'spectrum': spectrum,
            'live_time': live_time,
            'real_time': real_time,
            'dead_time_fraction': 1 - live_time/real_time
        })
    
    def save_results(self, results, filename):
        """Save analysis results to JSON."""
        with open(filename, 'w') as f:
            json.dump(results, f, indent=2, default=str)
    
    def load_results(self, filename):
        """Load analysis results from JSON."""
        with open(filename, 'r') as f:
            return json.load(f)


# ============================================================================
# MODULE 5: COMPLETE ANALYSIS PIPELINE
# ============================================================================
class AnalysisPipeline:
    """
    Complete analysis pipeline for ⁷Be half-life measurements.
    """
    
    def __init__(self):
        self.spectrum_analyzer = SpectrumAnalyzer()
        self.decay_fitter = DecayCurveFitter()
        self.data_manager = DataManager()
        self.efficiency_cal = EfficiencyCalibration()
        
        self.samples = {}
        self.results = {}
    
    def analyze_sample(self, sample_name, times, peak_areas, errors=None):
        """
        Analyze a single sample dataset.
        
        Parameters:
            sample_name: name of the sample
            times: array of times (days from start)
            peak_areas: array of peak areas (net counts)
            errors: array of uncertainties
        """
        # Fit decay curve
        fit_result = self.decay_fitter.fit_halflife(times, peak_areas, errors)
        
        if fit_result['success']:
            self.results[sample_name] = fit_result
            print(f"\n{sample_name}:")
            print(f"  t₁/₂ = {fit_result['t_half']:.2f} ± {fit_result['t_half_err']:.2f} days")
            print(f"  χ²/dof = {fit_result['chi2_red']:.2f}")
        else:
            print(f"\n{sample_name}: Fit failed - {fit_result['error']}")
        
        return fit_result
    
    def compare_samples(self, ref_name, test_name):
        """
        Compare half-lives between reference and test samples.
        """
        if ref_name not in self.results or test_name not in self.results:
            raise ValueError("Both samples must be analyzed first")
        
        comparison = self.decay_fitter.compare_halflife(
            self.results[ref_name], 
            self.results[test_name]
        )
        
        print(f"\n{'='*60}")
        print(f"COMPARISON: {test_name} vs {ref_name}")
        print(f"{'='*60}")
        print(f"  t₁/₂({ref_name}) = {comparison['t_half_1']:.2f} days")
        print(f"  t₁/₂({test_name}) = {comparison['t_half_2']:.2f} days")
        print(f"  Δt₁/₂ = {comparison['delta']:.2f} ± {comparison['delta_err']:.2f} days")
        print(f"  Δt₁/₂ = {comparison['delta_percent']:.2f}%")
        print(f"  Significance: {comparison['z_score']:.1f}σ")
        print(f"  p-value: {comparison['p_value']:.2e}")
        
        if comparison['significant_5sigma']:
            print(f"  ⭐ 5σ DISCOVERY! ⭐")
        elif comparison['significant_3sigma']:
            print(f"  ✓ 3σ evidence")
        else:
            print(f"  ✗ Not significant")
        
        return comparison
    
    def generate_report(self, output_file='analysis_report.txt'):
        """Generate a complete analysis report."""
        with open(output_file, 'w') as f:
            f.write("=" * 80 + "\n")
            f.write("⁷Be ELECTRON CAPTURE ANALYSIS REPORT\n")
            f.write(f"Generated: {datetime.now()}\n")
            f.write("=" * 80 + "\n\n")
            
            f.write("HALF-LIFE MEASUREMENTS\n")
            f.write("-" * 40 + "\n")
            for name, result in self.results.items():
                if result['success']:
                    f.write(f"{name}:\n")
                    f.write(f"  t₁/₂ = {result['t_half']:.3f} ± {result['t_half_err']:.3f} days\n")
                    f.write(f"  χ²/dof = {result['chi2_red']:.3f}\n")
                    f.write(f"  p-value = {result['p_value']:.4f}\n\n")
            
        print(f"Report saved to {output_file}")


# ============================================================================
# DEMONSTRATION
# ============================================================================
print("\n" + "=" * 80)
print("DEMONSTRATION: Simulated Data Analysis")
print("=" * 80)

# Create synthetic data
np.random.seed(42)

def generate_synthetic_data(t_half_true, A0=10000, n_points=100, duration=150):
    """Generate synthetic decay curve data."""
    times = np.linspace(0, duration, n_points)
    true_activity = A0 * np.exp(-np.log(2) * times / t_half_true)
    observed = np.random.poisson(true_activity)
    errors = np.sqrt(observed + 1)
    return times, observed, errors

# Generate data for three samples
print("\nGenerating synthetic data...")
samples_data = {
    'BeO (reference)': generate_synthetic_data(53.22),
    'Pd (control)': generate_synthetic_data(53.22),
    'CePd3 (test)': generate_synthetic_data(47.0),  # HF effect
}

# Run analysis
pipeline = AnalysisPipeline()

print("\nFitting decay curves...")
for name, (times, activity, errors) in samples_data.items():
    pipeline.analyze_sample(name, times, activity, errors)

# Compare samples
print("\nComparing samples...")
result_comparison = pipeline.compare_samples('BeO (reference)', 'CePd3 (test)')

# Generate visualization
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
fig.suptitle('⁷Be Decay Analysis - Demonstration', fontsize=14, fontweight='bold')

# Plot 1: Raw decay curves
ax1 = axes[0, 0]
colors = {'BeO (reference)': 'blue', 'Pd (control)': 'green', 'CePd3 (test)': 'red'}
for name, (times, activity, errors) in samples_data.items():
    ax1.errorbar(times, activity, yerr=errors, fmt='o', markersize=2, 
                 label=name, color=colors[name], alpha=0.5)
ax1.set_xlabel('Time (days)', fontsize=11)
ax1.set_ylabel('Activity (counts)', fontsize=11)
ax1.set_title('Decay Curves', fontsize=11)
ax1.legend()
ax1.grid(True, alpha=0.3)
ax1.set_yscale('log')

# Plot 2: Fitted curves
ax2 = axes[0, 1]
for name, result in pipeline.results.items():
    if result['success']:
        times = samples_data[name][0]
        A0 = result['A0']
        t_half = result['t_half']
        y_fit = A0 * np.exp(-np.log(2) * times / t_half)
        ax2.plot(times, y_fit, '-', linewidth=2, label=f"{name}: t₁/₂={t_half:.1f}d", 
                 color=colors[name])
ax2.set_xlabel('Time (days)', fontsize=11)
ax2.set_ylabel('Activity (counts)', fontsize=11)
ax2.set_title('Fitted Decay Curves', fontsize=11)
ax2.legend()
ax2.grid(True, alpha=0.3)
ax2.set_yscale('log')

# Plot 3: Residuals
ax3 = axes[1, 0]
for name, result in pipeline.results.items():
    if result['success']:
        times, activity, errors = samples_data[name]
        A0 = result['A0']
        t_half = result['t_half']
        y_fit = A0 * np.exp(-np.log(2) * times / t_half)
        residuals = (activity - y_fit) / errors
        ax3.scatter(times, residuals, s=10, label=name, color=colors[name], alpha=0.5)
ax3.axhline(y=0, color='black', linestyle='--')
ax3.axhline(y=2, color='gray', linestyle=':', alpha=0.5)
ax3.axhline(y=-2, color='gray', linestyle=':', alpha=0.5)
ax3.set_xlabel('Time (days)', fontsize=11)
ax3.set_ylabel('Normalized residuals', fontsize=11)
ax3.set_title('Fit Residuals', fontsize=11)
ax3.legend()
ax3.grid(True, alpha=0.3)

# Plot 4: Half-life comparison
ax4 = axes[1, 1]
sample_names = list(pipeline.results.keys())
t_halfs = [pipeline.results[n]['t_half'] for n in sample_names]
t_half_errs = [pipeline.results[n]['t_half_err'] for n in sample_names]
x = np.arange(len(sample_names))
bars = ax4.bar(x, t_halfs, yerr=t_half_errs, capsize=5, 
               color=[colors[n] for n in sample_names], edgecolor='black')
ax4.axhline(y=53.22, color='black', linestyle='--', label='Standard ⁷Be')
ax4.set_xticks(x)
ax4.set_xticklabels([n.split()[0] for n in sample_names])
ax4.set_ylabel('Half-life (days)', fontsize=11)
ax4.set_title('Measured Half-lives', fontsize=11)
ax4.legend()
ax4.grid(True, alpha=0.3, axis='y')

plt.tight_layout()
plt.savefig('/mnt/user-data/outputs/Be7_Data_Analysis_Demo.png', dpi=150, bbox_inches='tight')
print("\n✅ Analysis demonstration figure saved!")

# ============================================================================
# USER GUIDE
# ============================================================================
print("\n" + "=" * 80)
print("USER GUIDE: How to use this toolkit with real data")
print("=" * 80)

print("""
STEP 1: PREPARE YOUR DATA
═════════════════════════

Your data should be in one of these formats:
  - CSV: channel,counts (with header)
  - TXT: channel counts (whitespace separated)
  - SPE: Ortec format

Example CSV:
  channel,counts
  0,12
  1,15
  2,23
  ...

STEP 2: ANALYZE SPECTRA
═══════════════════════

>>> from be7_data_analysis_toolkit import SpectrumAnalyzer
>>> analyzer = SpectrumAnalyzer(energy_calibration=(0.25, 0))
>>> result = analyzer.analyze_spectrum(channels, counts, roi_center=478)
>>> print(f"Net area: {result['net_area']} ± {result['net_area_err']}")

STEP 3: FIT DECAY CURVES
════════════════════════

>>> from be7_data_analysis_toolkit import DecayCurveFitter
>>> fitter = DecayCurveFitter()
>>> 
>>> # Your measurement data
>>> times = [0, 1, 2, 3, ...]  # days
>>> activities = [1000, 987, 974, ...]  # counts
>>> 
>>> result = fitter.fit_halflife(times, activities)
>>> print(f"t₁/₂ = {result['t_half']:.2f} ± {result['t_half_err']:.2f} days")

STEP 4: COMPARE SAMPLES
═══════════════════════

>>> result_ref = fitter.fit_halflife(times_ref, activities_ref)
>>> result_test = fitter.fit_halflife(times_test, activities_test)
>>> 
>>> comparison = fitter.compare_halflife(result_ref, result_test)
>>> print(f"Δt₁/₂ = {comparison['delta']:.2f} days ({comparison['z_score']:.1f}σ)")

STEP 5: FULL PIPELINE
═════════════════════

>>> from be7_data_analysis_toolkit import AnalysisPipeline
>>> pipeline = AnalysisPipeline()
>>> 
>>> pipeline.analyze_sample('BeO', times_BeO, activities_BeO)
>>> pipeline.analyze_sample('CePd3', times_CePd3, activities_CePd3)
>>> 
>>> pipeline.compare_samples('BeO', 'CePd3')
>>> pipeline.generate_report('my_analysis.txt')
""")

print("\n" + "=" * 80)
print("TOOLKIT READY FOR USE")
print("=" * 80)

plt.close()