#!/usr/bin/env python3
"""
spectrum_waterfall.py - Capture a waterfall spectrogram from a GRX receiver.

This example retrieves the FFT parameters and then requests a waterfall
JPEG image from the Spectrum API. The resulting image shows signal
activity over time across the monitored frequency band.

Usage:
    python spectrum_waterfall.py [HOST] [CHANNEL]

Arguments:
    HOST      IP address or hostname of the GRX (default: YOUR_GRX_IP)
    CHANNEL   RX channel index, starting from 0 (default: 0)

Prerequisites:
    pip install grpcio grpcio-tools Pillow
    Run gen_pb.sh to generate the Python bindings.
"""

import datetime
import io
import sys

import grpc
import PIL.Image

# Import the generated gRPC stubs for the Spectrum API
import Spectrumd_pb2 as spectrumd_pb2
import Spectrumd_pb2_grpc

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

host = sys.argv[1] if len(sys.argv) >= 2 else "YOUR_GRX_IP"
rx_channel = int(sys.argv[2]) if len(sys.argv) >= 3 else 0

# The Spectrum API listens on TCP port 5306
SPECTRUMD_PORT = 5306

# Duration of the waterfall capture in seconds
CAPTURE_DURATION = 10

# ---------------------------------------------------------------------------
# Connect to the GRX receiver
# ---------------------------------------------------------------------------

print(f"Connecting to {host}:{SPECTRUMD_PORT}...")
channel = grpc.insecure_channel(f"{host}:{SPECTRUMD_PORT}")
spectrumd = Spectrumd_pb2_grpc.SpectrumdStub(channel)

# ---------------------------------------------------------------------------
# Query FFT parameters
# ---------------------------------------------------------------------------

# Retrieve the static FFT aggregation parameters for this channel.
# These tell us the center frequency, sample rate, FFT size, and how
# many FFT blocks are aggregated into a single line of the waterfall.
params = spectrumd.GetAggregatedFFTProperties(
    spectrumd_pb2.AggregatedFFTRequest(rx_channel_index=rx_channel)
)

# Duration of a single waterfall line
line_duration = params.aggregation_factor * params.fft_size / params.sample_rate
num_lines = int(CAPTURE_DURATION / line_duration)

print(f"\nFFT Parameters:")
print(f"  Center frequency:    {params.center_frequency / 1e6:.3f} MHz")
print(f"  Sample rate:         {params.sample_rate / 1e6:.3f} MSps")
print(f"  FFT size:            {params.fft_size}")
print(f"  Aggregation factor:  {params.aggregation_factor}")
print(f"  Line duration:       {line_duration * 1000:.1f} ms")
print(f"\nCapturing {num_lines} lines ({num_lines * line_duration:.1f} s)...")

# ---------------------------------------------------------------------------
# Request waterfall JPEG
# ---------------------------------------------------------------------------

# Build the request with reasonable display settings.
# min_level / max_level define the color scale range in dBm.
# Adjust these values depending on your signal environment.
request = spectrumd_pb2.GetWaterfallJPEGRequest(
    num_lines=num_lines,
    min_level=-180.0,
    max_level=-90.0,
    jpeg_quality=90,
    aggregation_type=spectrumd_pb2.GetWaterfallJPEGRequest.AVERAGE,
    rx_channel_index=rx_channel,
)

# This call blocks until all lines have been collected
result = spectrumd.GetWaterfallJPEG(request)

# ---------------------------------------------------------------------------
# Display the image
# ---------------------------------------------------------------------------

image = PIL.Image.open(io.BytesIO(result.image))

start_time = datetime.datetime.fromtimestamp(result.timestamp - num_lines * line_duration)
end_time = datetime.datetime.fromtimestamp(result.timestamp)

print(f"\nWaterfall captured: {start_time} to {end_time}")
print(f"Image size: {image.size[0]} x {image.size[1]} pixels")

# Show the image using the system's default image viewer
image.show(title=f"GRX Waterfall - {params.center_frequency / 1e6:.1f} MHz")

# ---------------------------------------------------------------------------
# Clean up
# ---------------------------------------------------------------------------

channel.close()