import glob
import os
import numpy as np
import matplotlib.pyplot as plt

def load_spectrum_data(filename):
    """
    Load spectrum data from a text file in the provided format.
    
    This function:
      - Reads the file line by line.
      - Finds the "$DATA:" section.
      - Skips the first line after "$DATA:" (assumed to be the channel range header).
      - Collects all subsequent lines (until the next section indicated by a line starting with '$').
      - Parses these lines into a NumPy array.
      
    Parameters:
      filename (str): Path to the spectrum file.
      
    Returns:
      numpy.ndarray: Array of count values.
    """
    with open(filename, 'r') as f:
        lines = f.readlines()

    data_lines = []
    in_data_section = False
    skip_first_data_line = True  # flag to skip the channel-range header

    for line in lines:
        # Look for the beginning of the data section
        if line.strip().startswith("$DATA:"):
            in_data_section = True
            continue
        if in_data_section:
            # If we hit a new section header, exit the data section
            if line.strip().startswith("$"):
                break
            # Skip the first line after "$DATA:" if it contains exactly two numbers (channel range header)
            if skip_first_data_line:
                parts = line.strip().split()
                if len(parts) == 2:
                    skip_first_data_line = False
                    continue
                skip_first_data_line = False  # even if not two numbers, do not skip subsequent lines
            # Append non-empty lines
            if line.strip():
                data_lines.append(line.strip())

    # Combine the collected lines into one string and parse numbers
    data_str = " ".join(data_lines)
    # Convert the string of numbers into a NumPy array
    data = np.fromstring(data_str, sep=' ')
    return data

# Get a list of all .Spe files in the current directory
spe_files = glob.glob("*.Spe")

# To store data for the combined plot
combined_data = []

for file in spe_files:
    # Load the spectrum data
    spectrum_data = load_spectrum_data(file)
    # Create an array for channel numbers (one per data point)
    channels = np.arange(len(spectrum_data))
    
    # Store the data for the combined plot (use file name without extension for legend)
    base_name = os.path.splitext(os.path.basename(file))[0]
    combined_data.append((base_name, channels, spectrum_data))
    
    # Plot individual spectrum
    plt.figure(figsize=(10, 6))
    plt.step(channels, spectrum_data, where='mid', color='blue')
    plt.xlabel('Channel')
    plt.ylabel('Counts')
    plt.title(f"Spectrum: {base_name}")
    plt.grid(True)
    plt.tight_layout()
    
    # Save individual figure as PNG
    output_filename = base_name + '.png'
    plt.savefig(output_filename)
    plt.close()  # Close the figure to free memory

# Now create a combined plot with all spectra
plt.figure(figsize=(12, 8))
for name, channels, spectrum_data in combined_data:
    # Do not specify a color so that the default cycle gives different colors for each curve
    plt.step(channels, spectrum_data, where='mid', label=name)
    
plt.xlabel('Channel')
plt.ylabel('Counts')
plt.title("Combined Spectrum Data")
plt.grid(True)
plt.legend()
plt.tight_layout()

# Save and display the combined plot
plt.savefig('combined_spectrum.png')
plt.show()