#!/bin/python3

import argparse, os, sys
from datetime import datetime, timezone
import scipy.signal
import numpy as np
import matplotlib.pyplot as plt


def readBin(file_path, start_time, duration, sample_rate_numerator, sample_rate_denominator):
    # Calculate the number of samples to read based on the duration and sample rate
    sample_rate = sample_rate_numerator / sample_rate_denominator
    num_samples_to_read = int(duration * sample_rate)

    # Calculate the byte offset to start reading from based on the start time
    byte_offset = int(start_time * sample_rate * 4)  # 4 bytes per sample (sc16 format)

    # Read the specified portion of the file
    with open(file_path, 'rb') as f:
        f.seek(byte_offset)
        data = f.read(num_samples_to_read * 4)  # Read the required number of samples

    return data

def plot_spectrum(rawIQ, sample_rate_numerator, sample_rate_denominator, center_frequency, fft_size, output_file):
    # Convert raw IQ data to numpy array of complex numbers
    iq_data = np.frombuffer(rawIQ, dtype=np.int16).astype(np.float32)
    iq_data = iq_data[0::2] + 1j * iq_data[1::2]  # Combine I and Q components

    # Calculate the sample rate
    sample_rate = sample_rate_numerator / sample_rate_denominator

    # Compute the power spectral density
    freqs, psd = scipy.signal.welch(iq_data, 
                                    fs=sample_rate,
                                    window=('kaiser', 3), 
                                    nperseg=fft_size, 
                                    noverlap=fft_size // 2,
                                    nfft=fft_size,
                                    detrend=False,
                                    return_onesided=False, 
                                    scaling='density')
    freqs = np.fft.fftshift(freqs)  
    psd = np.fft.fftshift(psd) 

    # Shift the frequencies to be centered around the center frequency
    freqs += center_frequency

    # Plot the spectrum
    plt.figure(figsize=(10, 6))
    plt.plot(freqs, 10 * np.log10(psd))
    plt.title('Spectrum of OneWeb IQ Data')
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('Power Spectral Density (dB/Hz)')
    plt.grid()
    plt.xlim(center_frequency - sample_rate / 2, center_frequency + sample_rate / 2)
    plt.savefig(output_file)
    plt.close()


if __name__ == "__main__":
    
    parser = argparse.ArgumentParser(description='Plots spectrum of OneWeb IQ file.')
    parser.add_argument('--file', '-f', metavar='file_path', type=str, help='OneWeb IQ file to plot', required=True)
    parser.add_argument('--start', '-s', metavar='start_time',default=0.0, type=float, help='Start time for plotting', required=False)
    parser.add_argument('--duration', '-d', metavar='duration', default=10e-3, type=float, help='Duration for plotting', required=False)
    parser.add_argument('--fft_size', '-n', metavar='fft_size', default=1024, type=int, help='FFT size for spectrum plot', required=False)
    parser.add_argument('--output', '-o', metavar='output_file', type=str, help='Output file for spectrum plot (e.g., spectrum.png)', required=True)
    
    args = parser.parse_args()
    file_path = os.path.abspath(args.file)
    start_time = args.start
    duration = args.duration
    fft_size = args.fft_size
    output_file = args.output

    # -- Parameters for the OneWeb recording --
    end_unix_timestamp = 1718571050 # Time of the end of the recording in Unix timestamp
    sample_rate_numerator = 491520000
    sample_rate_denominator = 1
    center_frequency = 11574999999.999992
    tle1 = '1 50474U 21132F   24168.55889395  .00000036  00000+0  57407-4 0  9997'
    tle2  = '2 50474  87.9139 120.9025 0001975  84.3721 275.7636 13.18688156125309'
    sat_name = 'ONEWEB-0394'
    sample_size_bytes = 4 # Each sample is 4 bytes (2 bytes for I and 2 bytes for Q)

    # -- Dependent parameters --
    file_duration_seconds = os.path.getsize(file_path) / (sample_rate_numerator / sample_rate_denominator * sample_size_bytes)
    start_unix_timestamp = end_unix_timestamp - file_duration_seconds
    start_time_dt = datetime.fromtimestamp(start_unix_timestamp, tz=timezone.utc)

    # -- Read in data
    rawIQ = readBin(file_path, start_time, duration, sample_rate_numerator, sample_rate_denominator)

    # -- Plot spectrum
    plot_spectrum(rawIQ, sample_rate_numerator, sample_rate_denominator, center_frequency, fft_size, output_file)