Untitled

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy import interpolate
import os
import json
import uuid
import h5py
from datetime import datetime

# Get the absolute path of the directory where the script is located
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) if '__file__' in locals() else os.getcwd()

class FPSpectraAugmenter:
    def __init__(self, input_csv, output_dir="augmented_fp_spectra"):
        """
        Initialize the FP spectra augmenter.
        
        Args:
            input_csv: Path to the FPBase spectra CSV file
            output_dir: Directory to save the augmented dataset
        """
        self.input_csv = input_csv
        self.output_dir = output_dir
        self.parent_spectra = {}
        self.family_tree = {}
        
        # Create output directory if it doesn't exist
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            
    def load_emission_spectra(self):
        """Load all emission spectra from the CSV file."""
        print(f"Loading emission spectra from {self.input_csv}...")
        
        # Read the CSV file
        df = pd.read_csv(self.input_csv)
        
        # Extract wavelength column
        wavelengths = df['wavelength'].values
        
        # Get only emission spectrum columns (those with 'em' in the name)
        emission_cols = [col for col in df.columns if ' em' in col]
        
        # Track how many valid spectra we find
        valid_count = 0
        
        # Extract protein names and their emission spectra
        for col in emission_cols:
            protein_name = col.split(' em')[0].strip()
            data = df[col].values
            
            # Find non-NaN data points
            valid_mask = ~np.isnan(data)
            
            if np.sum(valid_mask) > 10:  # Ensure we have enough valid points
                valid_wavelengths = wavelengths[valid_mask]
                valid_data = data[valid_mask]
                
                if np.max(valid_data) > 0:  # Ensure signal exists
                    # Normalize the parent spectrum to max 1
                    valid_data = valid_data / np.max(valid_data)
                    
                    # Generate a unique ID for this parent spectrum
                    parent_id = str(uuid.uuid4())
                    
                    # Store the spectrum
                    self.parent_spectra[parent_id] = {
                        'name': protein_name,
                        'wavelength': valid_wavelengths.tolist(),
                        'emission': valid_data.tolist(),
                        'type': 'parent'
                    }
                    
                    # Initialize family tree entry
                    self.family_tree[parent_id] = {
                        'name': protein_name,
                        'children': []
                    }
                    
                    valid_count += 1
        
        print(f"Loaded {valid_count} valid emission spectra.")
        return valid_count
    
    def generate_children(self, num_children_per_parent=10):
        """
        Generate child spectra for each parent spectrum.
        
        Args:
            num_children_per_parent: Number of children to generate per parent
        """
        print(f"Generating {num_children_per_parent} children for each parent spectrum...")
        
        augmented_spectra = {}
        
        # For each parent spectrum
        for parent_id, parent_data in self.parent_spectra.items():
            parent_name = parent_data['name']
            wavelengths = np.array(parent_data['wavelength'])
            emission = np.array(parent_data['emission'])
            
            print(f"Processing parent: {parent_name}")
            
            # Create interpolation function for this spectrum with smooth edges
            # Add padding points to create smooth transitions to zero
            pad_width = 30  # nm padding on each side
            
            # Get min and max wavelengths
            min_wl = min(wavelengths)
            max_wl = max(wavelengths)
            
            # Create padding wavelengths
            left_pad_wl = np.linspace(min_wl - pad_width, min_wl, 10)
            right_pad_wl = np.linspace(max_wl, max_wl + pad_width, 10)
            
            # Create padding emissions (smoothly decreasing to zero)
            left_pad_em = np.linspace(0, emission[0], 10)
            right_pad_em = np.linspace(emission[-1], 0, 10)
            
            # Combine original and padded data
            padded_wavelengths = np.concatenate([left_pad_wl, wavelengths, right_pad_wl])
            padded_emission = np.concatenate([left_pad_em, emission, right_pad_em])
            
            # Create interpolation function with padded data
            f_interp = interpolate.interp1d(
                padded_wavelengths,
                padded_emission,
                bounds_error=False,
                fill_value=0.0
            )
            
            # Generate children
            for i in range(num_children_per_parent):
                # Generate a unique ID for this child
                child_id = str(uuid.uuid4())
                
                # Choose modification parameters
                amplitude_factor = np.random.uniform(0.3, 3.0)
                wavelength_shift = np.random.uniform(-100, 100)  # nm
                
                # Apply wavelength shift
                shifted_wavelengths = wavelengths + wavelength_shift
                
                # Get expanded wavelength range including padding
                expanded_min = min(min(wavelengths), min(shifted_wavelengths)) - pad_width
                expanded_max = max(max(wavelengths), max(shifted_wavelengths)) + pad_width
                
                # Create common wavelength array with smooth transitions
                common_wavelengths = np.linspace(
                    expanded_min,
                    expanded_max,
                    len(wavelengths) + 20  # Add extra points for smooth transitions
                )
                
                # Interpolate shifted spectrum to common wavelengths
                shifted_emission = f_interp(common_wavelengths - wavelength_shift)
                
                # Apply amplitude scaling
                scaled_emission = shifted_emission * amplitude_factor
                
                # Final emission after modifications (NO NOISE IS ADDED)
                final_emission = scaled_emission
                
                # Ensure no negative values
                final_emission = np.maximum(final_emission, 0)
                
                # Normalize to max 1
                if np.max(final_emission) > 0:
                    final_emission = final_emission / np.max(final_emission)
                
                # Store the child spectrum
                child_name = f"{parent_name}_child_{i+1}"
                
                augmented_spectra[child_id] = {
                    'name': child_name,
                    'wavelength': common_wavelengths.tolist(),
                    'emission': final_emission.tolist(),
                    'parent_id': parent_id,
                    'modifications': {
                        'amplitude_factor': float(amplitude_factor),
                        'wavelength_shift': float(wavelength_shift)
                    },
                    'type': 'child'
                }
                
                # Update family tree
                self.family_tree[parent_id]['children'].append(child_id)
        
        # Merge parent and child spectra
        self.augmented_spectra = {**self.parent_spectra, **augmented_spectra}
        
        print(f"Generated {len(augmented_spectra)} child spectra.")
        return len(augmented_spectra)
    
    def save_database(self):
        """Save the augmented database to files."""
        # Save as HDF5 for efficient storage of numerical data
        h5_file = os.path.join(self.output_dir, "augmented_spectra.h5")
        with h5py.File(h5_file, 'w') as f:
            # Create a group for each spectrum
            for spectrum_id, data in self.augmented_spectra.items():
                group = f.create_group(spectrum_id)
                
                # Store metadata as attributes
                group.attrs['name'] = data['name']
                group.attrs['type'] = data['type']
                
                if data['type'] == 'child':
                    group.attrs['parent_id'] = data['parent_id']
                    
                    # Store modification parameters
                    mods = data['modifications']
                    group.attrs['amplitude_factor'] = mods['amplitude_factor']
                    group.attrs['wavelength_shift'] = mods['wavelength_shift']
                
                # Store the spectral data
                group.create_dataset('wavelength', data=np.array(data['wavelength']))
                group.create_dataset('emission', data=np.array(data['emission']))
        
        # Save family tree as JSON for easy inspection
        json_file = os.path.join(self.output_dir, "family_tree.json")
        with open(json_file, 'w') as f:
            json.dump(self.family_tree, f, indent=2)
        
        # Save metadata about the augmentation
        metadata = {
            'date_created': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            'input_file': self.input_csv, # This will store the full path passed to init
            'num_parents': len(self.parent_spectra),
            'num_children': len(self.augmented_spectra) - len(self.parent_spectra),
            'total_spectra': len(self.augmented_spectra)
        }
        
        meta_file = os.path.join(self.output_dir, "metadata.json")
        with open(meta_file, 'w') as f:
            json.dump(metadata, f, indent=2)
        
        print(f"Saved augmented database to {self.output_dir}:")
        print(f"  - {h5_file} (spectral data)")
        print(f"  - {json_file} (family relationships)")
        print(f"  - {meta_file} (dataset metadata)")
    
    def plot_examples(self, num_examples=5):
        """Plot example families (parent and children)."""
        # Select random parents to plot
        if not self.parent_spectra:
            print("No parent spectra loaded to plot examples.")
            return
            
        parent_ids = list(self.parent_spectra.keys())
        num_to_plot = min(num_examples, len(parent_ids))
        if num_to_plot == 0 :
            print("No parents available to plot.")
            return
        selected_parents = np.random.choice(parent_ids, num_to_plot, replace=False)
        
        for parent_id in selected_parents:
            parent_data = self.augmented_spectra[parent_id]
            parent_name = parent_data['name']
            
            # Get children
            children_ids = self.family_tree[parent_id]['children']
            
            # Create plot
            plt.figure(figsize=(12, 6))
            
            # Plot parent
            plt.plot(
                parent_data['wavelength'], 
                parent_data['emission'], 
                'k-', 
                linewidth=2, 
                label=f"Parent: {parent_name}"
            )
            
            # Plot children (up to 5 to avoid cluttering)
            for i, child_id in enumerate(children_ids[:5]):
                child_data = self.augmented_spectra[child_id]
                plt.plot(
                    child_data['wavelength'],
                    child_data['emission'],
                    '--',
                    alpha=0.7,
                    label=f"Child {i+1}: {child_data['modifications']['wavelength_shift']:.1f}nm shift, "
                          f"{child_data['modifications']['amplitude_factor']:.2f}x amplitude"
                )
            
            plt.xlabel('Wavelength (nm)')
            plt.ylabel('Emission Intensity')
            plt.title(f'Spectrum Family: {parent_name}')
            plt.legend()
            plt.grid(True, alpha=0.3)
            
            # Save figure
            safe_name = parent_name.replace(' ', '_').replace('/', '_')
            plt.savefig(os.path.join(self.output_dir, f"family_{safe_name}.png"))
            plt.close()
        
        print(f"Saved {num_to_plot} example family plots to {self.output_dir}.")

def main():
    # Set parameters
    # Construct paths relative to the script's directory
    input_csv_filename = "fpbase_spectra.csv"
    output_dir_name = "augmented_fp_spectra"
    
    # os.path.join correctly creates paths for the current operating system
    input_csv = os.path.join(SCRIPT_DIR, input_csv_filename)
    output_dir = os.path.join(SCRIPT_DIR, output_dir_name)
    
    children_per_parent = 100
    
    # Check if input file exists
    if not os.path.exists(input_csv):
        print(f"Error: Input CSV file not found at {input_csv}")
        print(f"Please ensure '{input_csv_filename}' is in the same directory as the script.")
        return

    # Initialize augmenter
    augmenter = FPSpectraAugmenter(input_csv, output_dir)
    
    # Process the data
    if augmenter.load_emission_spectra() > 0:
        augmenter.generate_children(children_per_parent)
        augmenter.save_database()
        augmenter.plot_examples(10)
        
        print("\nAugmentation complete!")
        print(f"Created an augmented dataset with {len(augmenter.augmented_spectra)} total spectra")
        print(f"Dataset saved to {output_dir}/")
    else:
        print("No valid spectra loaded. Augmentation cannot proceed.")


if __name__ == "__main__":
    main()