# imports
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import sigmf.sigmffile as sigmf

# Function to load SigMF files and extract samples
def load_sigmf_file(file_path):
    print(f"Loading SigMF file: {file_path}")
    signal = sigmf.fromfile(file_path,skip_checksum=True)
    return signal

#paths
dataset_path = '<path to folder containing dataset.csv>'        # should not include '/dataset.csv'
sigmf_path = '<path to folder containing the sigmfs folder>'   # should not include '/sigmfs'


# Load the dataset
dataset = pd.read_csv(dataset_path+'/dataset.csv')

# Select a frame with high SNR for better visualization
frame_entry = dataset[dataset['snr'] > 20].iloc[0]
print(f"Frame informations: \n{frame_entry}")

# Load the corresponding SigMF file
filename = frame_entry['sigmf_file']
file_path = filename.replace('.',sigmf_path)
signal = load_sigmf_file(file_path)

# Get frame start and end sample indices
frame_start = int(frame_entry['sigmf_file_offset'])
frame_end = frame_start + int(frame_entry['sigmf_file_n_samples'])

# Read samples
samples = signal.read_samples()

# Plot the samples
plt.figure(figsize=(10, 13))
plt.subplot(4,1,1)
plt.subplots_adjust(hspace=0.5)
# whole signal with frame start and end
plt.plot(np.real(samples), label='In-Phase')
# plt.plot(np.imag(samples), label='Quadrature')
plt.axvline(x=frame_start, color='r', linestyle='--', label='Frame Start')
plt.axvline(x=frame_end, color='r', linestyle='--', label='Frame End')
plt.title('Signal Samples (Real)')
plt.xlabel('Sample Index')
plt.ylabel('Amplitude')
plt.grid()
plt.subplot(4,1,2)
# zoomed in view of the frame
plt.plot(np.real(samples), label='In-Phase')
plt.axvline(x=frame_start, color='r', linestyle='--', label='Frame Start')
plt.xlim(frame_start-5000, frame_start+5000)
plt.title('Zoomed Frame Start Samples (Real)')
plt.xlabel('Sample Index')
plt.ylabel('Amplitude')
plt.grid()

plt.subplot(4,1,3)
#instantaneous frequency
instantaneous_phase = np.unwrap(np.angle(samples))
instantaneous_frequency = np.diff(instantaneous_phase) / (2.0*np.pi)  # Normalized frequency
plt.plot(instantaneous_frequency, label='Instantaneous Frequency')
plt.axvline(x=frame_start, color='r', linestyle='--', label='Frame Start')
plt.axvline(x=frame_end, color='r', linestyle='--', label='Frame End')
plt.title('Instantaneous Frequency')
plt.xlabel('Sample Index')
plt.ylabel('Normalized Frequency')
plt.grid()

plt.subplot(4,1,4)
# zoom in view of instantaneous frequency at frame start
plt.plot(instantaneous_frequency, label='Instantaneous Frequency')
plt.axvline(x=frame_start, color='r', linestyle='--', label='Frame Start')
plt.xlim(frame_start-5000, frame_start+5000)
plt.title('Zoomed Instantaneous Frequency at Frame Start')
plt.xlabel('Sample Index')
plt.ylabel('Normalized Frequency')
plt.grid()
plt.show()