import numpy as np
import plotly.graph_objs as go
from plotly.subplots import make_subplots
from IPython.display import Image, display
from radarsimpy import Radar, Transmitter, Receiver

# Set to True for interactive plots; False renders a static JPEG (e.g. for HTML export)
INTERACTIVE = False


def show(fig):
    if INTERACTIVE:
        fig.show()
    else:
        display(Image(fig.to_image(format="jpg", scale=2)))

# Define transmitter channel at the origin
tx_channel = dict(
    location=(0, 0, 0),  # Initial position (will move according to motion plan)
)

# Create FMCW transmitter
tx = Transmitter(
    f=[77e9 - 50e6, 77e9 + 50e6],  # 76.95 to 77.05 GHz (100 MHz bandwidth)
    t=[0, 80e-6],  # 80 microsecond chirp duration
    tx_power=25,  # 25 dBm transmit power (~316 mW)
    prp=100e-6,  # 100 microsecond pulse repetition period
    pulses=256,  # 256 pulses for Doppler processing
    channels=[tx_channel],
)

# Define receiver channel co-located with transmitter
rx_channel = dict(
    location=(0, 0, 0),  # Monostatic configuration (TX and RX at same location)
)

# Create receiver with specified noise and gain parameters
rx = Receiver(
    fs=2e6,  # 2 MHz sampling rate
    noise_figure=8,  # 8 dB noise figure
    rf_gain=20,  # 20 dB RF gain
    load_resistor=500,  # 500 Ohm load resistor
    baseband_gain=30,  # 30 dB baseband gain
    channels=[rx_channel],
)

# Create radar system from transmitter and receiver
radar = Radar(transmitter=tx, receiver=rx)

# The timestamp array contains the time value for EACH SAMPLE (not just each pulse)
# This allows high-fidelity motion simulation where position is updated at every ADC sample
# You can use this to define ANY arbitrary motion trajectory as a function of time
timestamps = radar.time_prop["timestamp"]

# Note: timestamps.shape = (frames × channels, pulses, samples_per_pulse)
# Motion is evaluated at each of these timestamps for accurate simulation

# ============================================================================
# EXAMPLE MOTION PATTERNS - Uncomment any to try different trajectories
# ============================================================================

# Option 1: Linear motion at constant velocity (ACTIVE)
# Simple constant velocity motion along x-axis
location_x = -12 * timestamps  # Moving at -12 m/s
location_y = 0 * timestamps  # No y motion
location_z = 0 * timestamps  # No z motion

# Option 2: Sinusoidal motion (UNCOMMENT TO TRY)
# Oscillating motion - useful for vibration/jitter analysis
# location_x = 10 * np.sin(2 * np.pi * timestamps)  # Oscillate with 1 Hz
# location_y = 0 * timestamps
# location_z = 0 * timestamps

# Option 3: Circular motion (UNCOMMENT TO TRY)
# Radar moving in a circle - simulates rotating platform
# radius = 20
# angular_freq = 2 * np.pi * 0.5  # 0.5 Hz rotation
# location_x = 30 + radius * np.cos(angular_freq * timestamps)
# location_y = radius * np.sin(angular_freq * timestamps)
# location_z = 0 * timestamps

# Option 4: Accelerated motion (UNCOMMENT TO TRY)
# Non-uniform velocity - simulates vehicle acceleration
# acceleration = 5  # m/s^2
# location_x = 0.5 * acceleration * timestamps**2
# location_y = 0 * timestamps
# location_z = 0 * timestamps

# Option 5: Complex trajectory (UNCOMMENT TO TRY)
# Combine multiple components for realistic vehicle motion
# location_x = -12 * timestamps + 2 * np.sin(4 * np.pi * timestamps)  # Linear + oscillation
# location_y = 3 * np.sin(2 * np.pi * timestamps)  # Side-to-side motion
# location_z = 0 * timestamps

# Option 6: Motion with rotation (UNCOMMENT TO TRY)
# Radar translating and rotating - simulates vehicle turning
# location_x = -12 * timestamps
# location_y = 0 * timestamps
# location_z = 0 * timestamps
# rotation_yaw = 0.5 * timestamps  # Yaw: rotating at 0.5 rad/s
# rotation_pitch = 0 * timestamps  # Pitch: no pitching
# rotation_roll = 0 * timestamps   # Roll: no rolling
# rotation = (rotation_yaw, rotation_pitch, rotation_roll)  # Order: (yaw, pitch, roll)

# ============================================================================

# Apply motion plan to radar
# IMPORTANT: When using motion planning (time-varying location/rotation arrays),
# both speed and rotation_rate MUST be set to (0, 0, 0)
# All motion should be encoded in the location and rotation arrays
radar.set_motion(
    location=(location_x, location_y, location_z),
    speed=(0, 0, 0),  # MUST be (0,0,0) when using motion planning arrays
    # rotation=(rotation_yaw, rotation_pitch, rotation_roll),  # Optional: add rotation (yaw, pitch, roll)
    # rotation_rate=(0, 0, 0),  # MUST be (0,0,0) when using motion planning arrays
)

# Note: With linear motion at -12 m/s in the negative x direction:
# - Radar starts at x=0 and moves in negative x direction (away from target at x=50m)
# - Target is stationary at x=50m (positive x direction)
# - This creates opening geometry, producing positive Doppler shift

# Define stationary corner reflector target
target_1 = {
    "model": "../models/cr.stl",  # 3D mesh model of trihedral corner reflector
    "unit": "m",  # Units for the model
    "location": (50, 0, 0),  # Fixed position at x=50m, y=0, z=0
    "speed": (0, 0, 0),  # Stationary target (no motion)
}

targets = [target_1]

# Geometry: Radar starts at origin (0,0,0) and moves in -x direction (away from target)
#           Target stays at (50,0,0) in the positive x direction
# Result: Radar moves away from the target, creating positive range rate (opening)

from scipy import signal
import radarsimpy.processing as proc

# Design windowing functions for sidelobe suppression
range_window = signal.windows.chebwin(radar.sample_prop["samples_per_pulse"], at=60)
doppler_window = signal.windows.chebwin(
    radar.radar_prop["transmitter"].waveform_prop["pulses"], at=60
)

# Perform 2D FFT to generate Range-Doppler map
# rwin: window applied along range (fast-time) dimension
# dwin: window applied along Doppler (slow-time/pulse) dimension
range_doppler = proc.range_doppler_fft(baseband, rwin=range_window, dwin=doppler_window)