import matplotlib.pyplot as plt
import json
import pandas as pd
import numpy as np

plt.style.use('ggplot')

def plot_overlay_data_from_json(json_file, sensors, use_precise_timestamp=False, slice_select=1):
    # Read the JSON file
    with open(json_file, "r") as f:
        slices = json.load(f)

    # Set up the colormap
    cmap = plt.get_cmap('viridis')

    # Create subplots for each sensor
    fig, axs = plt.subplots(len(sensors), 1, figsize=(12, 2 * len(sensors)), sharex=True)

    for idx, sensor in enumerate(sensors):
        # Plot the overlay of the slices
        for slice_idx, slice_dict in enumerate(slices):
            slice_length = len(slice_dict[sensor])

            # Create timestamp array starting from 0 for each slice
            slice_timestamps = [20 * i for i in range(slice_length)]
            sensor_data = slice_dict[sensor]

            data = pd.DataFrame({sensor: sensor_data}, index=slice_timestamps)
            color = cmap(slice_idx / len(slices))

            axs[idx].plot(data[sensor], color=color, label=f'Slice {slice_idx + 1}')

        axs[idx].set_ylabel(sensor)

    axs[-1].set_xlabel("Timestamp")
    axs[0].legend()

    return fig

plot_overlay_data_from_json('output.json', ["GZ1", "GZ2", "GZ3", "GZ4"], 4)