import os
import pandas as pd
import numpy as np
from scipy.signal import stft, hann
import glob

# Define the base directory where DAMAGE_X folders are located
damage_base_path = 'D:/thesis/data/converted/raw/'

# Define sensor directories
sensor_dirs = {
    'sensor1': os.path.join(damage_base_path, 'sensor1'),
    'sensor2': os.path.join(damage_base_path, 'sensor2')
}

# Define STFT parameters
window_size = 1024
hop_size = 512
window = hann(window_size)
Fs = 1024

def verify_stft(damage_num, test_num, sensor_num):
    """
    Verifies the STFT of an individual test run against the aggregated STFT data.

    Parameters:
    - damage_num (int): Damage case number.
    - test_num (int): Test run number.
    - sensor_num (int): Sensor number (1 or 2).
    """
    # Mapping sensor number to suffix
    sensor_suffix = f'_0{sensor_num}'
    
    # Construct the file name for the individual test run
    individual_file_name = f'DAMAGE_{damage_num}_TEST{test_num}{sensor_suffix}.csv'
    individual_file_path = os.path.join(damage_base_path, f'DAMAGE_{damage_num}', individual_file_name)
    
    # Check if the individual file exists
    if not os.path.isfile(individual_file_path):
        print(f"File {individual_file_path} does not exist. Skipping verification for this test run.")
        return
    
    # Read the individual test run CSV
    try:
        df_individual = pd.read_csv(individual_file_path)
    except Exception as e:
        print(f"Error reading {individual_file_path}: {e}. Skipping verification for this test run.")
        return
    
    # Ensure the CSV has exactly two columns: 'Timestamp (s)' and 'Sensor X'
    if df_individual.shape[1] != 2:
        print(f"Unexpected number of columns in {individual_file_path}. Expected 2, got {df_individual.shape[1]}. Skipping.")
        return
    
    # Extract vibration data
    vibration_data = df_individual.iloc[:, 1].values
    
    # Perform STFT
    frequencies, times, Zxx = stft(
        vibration_data, 
        fs=Fs, 
        window=window, 
        nperseg=window_size, 
        noverlap=window_size - hop_size
    )
    
    # Compute magnitude and transpose
    stft_magnitude = np.abs(Zxx).T  # Shape: (513, 513)
    
    # Select random row indices to verify (e.g., 3 random rows)
    np.random.seed(42)  # For reproducibility
    sample_row_indices = np.random.choice(stft_magnitude.shape[0], size=3, replace=False)
    
    # Read the aggregated STFT CSV
    aggregated_file_name = f'stft_data{sensor_num}_{damage_num}.csv'
    aggregated_file_path = os.path.join(sensor_dirs[f'sensor{sensor_num}'], aggregated_file_name)
    
    if not os.path.isfile(aggregated_file_path):
        print(f"Aggregated file {aggregated_file_path} does not exist. Skipping verification for this test run.")
        return
    
    try:
        df_aggregated = pd.read_csv(aggregated_file_path)
    except Exception as e:
        print(f"Error reading {aggregated_file_path}: {e}. Skipping verification for this test run.")
        return
    
    # Calculate the starting row index in the aggregated CSV
    # Each test run contributes 513 rows
    start_row = (test_num - 1) * 513
    end_row = start_row + 513  # Exclusive
    
    # Ensure the aggregated CSV has enough rows
    if df_aggregated.shape[0] < end_row:
        print(f"Aggregated file {aggregated_file_path} does not have enough rows for Test {test_num}. Skipping.")
        return
    
    # Extract the corresponding STFT block from the aggregated CSV
    df_aggregated_block = df_aggregated.iloc[start_row:end_row].values  # Shape: (513, 513)
    
    # Compare selected rows
    all_match = True
    for row_idx in sample_row_indices:
        individual_row = stft_magnitude[row_idx]
        aggregated_row = df_aggregated_block[row_idx]
        
        # Check if the rows are almost equal within a tolerance
        if np.allclose(individual_row, aggregated_row, atol=1e-6):
            verification_status = "MATCH"
        else:
            verification_status = "MISMATCH"
            all_match = False
        
        # Print the comparison details
        print(f"Comparing Damage {damage_num}, Test {test_num}, Sensor {sensor_num}, Row {row_idx}: {verification_status}")
        print(f"Individual STFT Row {row_idx}: {individual_row[:5]} ... {individual_row[-5:]}")
        print(f"Aggregated STFT Row {row_idx + start_row}: {aggregated_row[:5]} ... {aggregated_row[-5:]}\n")
    
    # If all sampled rows match, print a verification success message
    if all_match:
        print(f"STFT of DAMAGE_{damage_num}_TEST{test_num}_{sensor_num}.csv is verified. On `stft_data{sensor_num}_{damage_num}.csv` start at rows {start_row} to {end_row} with 513 rows.\n")
    else:
        print(f"STFT of DAMAGE_{damage_num}_TEST{test_num}_{sensor_num}.csv has discrepancies in `stft_data{sensor_num}_{damage_num}.csv` start at rows {start_row} to {end_row} with 513 rows.\n")

# Define the number of damage cases and test runs
num_damage_cases = 6  # Adjust to 30 as per your dataset
num_test_runs = 5

# Iterate through all damage cases, test runs, and sensors
for damage_num in range(1, num_damage_cases + 1):
    for test_num in range(1, num_test_runs + 1):
        for sensor_num in [1, 2]:
            verify_stft(damage_num, test_num, sensor_num)