import numpy as np
import polars as pl
import matplotlib.pyplot as plt
import stumpy

def generate_data():
    np.random.seed(42)
    timestamps = np.linspace(0, 10, 500)
    base_signal = np.sin(2 * np.pi * timestamps)

    shifts = [0, 0.5, 1, -0.3]
    sensor_ids = [f"sensor_{i}" for i in range(len(shifts))]
    data = []
    for shift, sensor_id in zip(shifts, sensor_ids):
      shift_samples = int(shift * 60)
      signal = np.roll(base_signal, shift_samples) + np.random.normal(0, 0.1, len(base_signal))
      for ts, val in zip(timestamps, signal):
        data.append({"timestamp": ts, "value": val, "sensor_id": sensor_id})

  df = pl.DataFrame(data)
  return df, timestamps, sensor_ids

def plot_signals(df, title):
    plt.figure(figsize=(10, 6))
    sensor_ids = df["sensor_id"].unique().to_list()
    for sensor_id in sensor_ids:
        sensor_data = df.filter(pl.col("sensor_id") == sensor_id)
        plt.plot(sensor_data["timestamp"].to_numpy(), sensor_data["value"].to_numpy(), label=sensor_id)
    plt.title(title)
    plt.xlabel("Timestamp")
    plt.ylabel("Value")
    plt.legend()
    plt.show()

df, timestamps, sensor_ids = generate_data()
plot_signals(df, "Original Time Series")

def shift_signal(signal, shift_amount):
    """
    Shift a signal to the left or right without wrapping.
    If shifted left and points run out, the series ends early.
    If shifted right, the beginning is filled with NaNs.
    """
    if shift_amount > 0:  # Shift right
        return np.concatenate((np.full(shift_amount, np.nan), signal[:-shift_amount]))
    elif shift_amount < 0:  # Shift left
        shift_amount = abs(shift_amount)
        return signal[shift_amount:]
    return signal  # No shift


def align_and_facet_plot_all(reference_sensor_id, df, timestamps, window_size=50):
    # Prepare signals and sensor IDs
    sensor_ids = df["sensor_id"].unique().to_list()
    reference_signal = df.filter(pl.col("sensor_id") == reference_sensor_id)["value"].to_numpy()

    # Define a consistent style for each sensor
    styles = {sensor_id: {"color": plt.cm.tab10(i)} for i, sensor_id in enumerate(sensor_ids)}

    fig, axes = plt.subplots(2, 1, figsize=(12, 10), sharex=True)

    # Plot all original signals on the top facet
    for sensor_id in sensor_ids:
        sensor_signal = df.filter(pl.col("sensor_id") == sensor_id)["value"].to_numpy()
        axes[0].plot(
            timestamps[:len(sensor_signal)],
            sensor_signal,
            label=f"Original: {sensor_id}",
            **styles[sensor_id]
        )

    axes[0].set_title("Original Signals")
    axes[0].set_ylabel("Value")
    axes[0].legend()

    # Plot the reference signal and aligned signals on the bottom facet
    axes[1].plot(
        timestamps[:len(reference_signal)],
        reference_signal,
        label=f"Reference: {reference_sensor_id}",
        linewidth=2,
        **styles[reference_sensor_id]
    )

    for sensor_id in sensor_ids:
        if sensor_id == reference_sensor_id:
            continue

        target_signal = df.filter(pl.col("sensor_id") == sensor_id)["value"].to_numpy()

        # Compute distance profile and find optimal shift
        distance_profile = stumpy.mass(reference_signal[:window_size], target_signal)
        optimal_shift_index = np.argmin(distance_profile)

        # Align the target signal using clipping instead of wrapping
        aligned_signal = shift_signal(target_signal, -optimal_shift_index)

        # Plot the aligned signal
        axes[1].plot(
            timestamps[:len(aligned_signal)],
            aligned_signal,
            label=f"Aligned: {sensor_id}",
            **styles[sensor_id]
        )

    axes[1].set_title("Aligned Signals (Shifted)")
    axes[1].set_xlabel("Timestamp")
    axes[1].set_ylabel("Value")
    axes[1].legend()

    plt.tight_layout()
    plt.show()

# Example: Align and facet plot for sensor_0 with all other sensors
align_and_facet_plot_all("sensor_0", df, timestamps, window_size=450)