import pandas as pd
from pandas import Timestamp
pd.set_option('display.max_rows', 500)
pd.set_option('display.max_columns', 20)


experiment_timestamps = {
    'experiment1': {
        'sharp': [
            Timestamp('2026-09-11 07:38:25.012000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 07:41:36.608000+0000', tz='Europe/Berlin')
        ],
        'blunt': [
            Timestamp('2026-09-11 07:41:46.868000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 07:44:56.526000+0000', tz='Europe/Berlin')
        ]
    },
    'experiment2': {
        'sharp': [
            Timestamp('2026-09-11 07:46:25.223000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 07:49:22.030000+0000', tz='Europe/Berlin')
        ],
        'blunt': [
            Timestamp('2026-09-11 07:49:35.259000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 07:52:29.066000+0000', tz='Europe/Berlin')
        ]
    },
    'experiment3': {
        'sharp': [
            Timestamp('2026-09-11 07:53:22.246000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 07:56:46.017000+0000', tz='Europe/Berlin')
        ],
        'blunt': [
            Timestamp('2026-09-11 07:56:56.503000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:00:20.646000+0000', tz='Europe/Berlin')
        ]
    },
    'experiment4': {
        'sharp': [
            Timestamp('2026-09-11 08:01:07.610000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:04:16.963000+0000', tz='Europe/Berlin')
        ],
        'blunt': [
            Timestamp('2026-09-11 08:04:27.439000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:07:36.995000+0000', tz='Europe/Berlin')
        ]
    },
    'experiment5': {
        'sharp': [
            Timestamp('2026-09-11 08:08:43.364000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:11:52.614000+0000', tz='Europe/Berlin')
        ],
        'blunt': [
            Timestamp('2026-09-11 08:12:02.681000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:15:12.238000+0000', tz='Europe/Berlin')
        ]
    },
    'experiment6': {
        'sharp': [
            Timestamp('2026-09-11 08:16:34.224000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:19:43.475000+0000', tz='Europe/Berlin')
        ],
        'blunt': [
            Timestamp('2026-09-11 08:19:53.749000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:23:03.407000+0000', tz='Europe/Berlin')
        ]
    },
    'experiment7': {
        'sharp': [
            Timestamp('2026-09-11 08:27:06.962000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:30:16.213000+0000', tz='Europe/Berlin')
        ],
        'blunt': [
            Timestamp('2026-09-11 08:30:26.313000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:33:36.075000+0000', tz='Europe/Berlin')
        ]
    },
    'experiment8': {
        'sharp': [
            Timestamp('2026-09-11 08:34:31.407000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:39:16.436000+0000', tz='Europe/Berlin')
        ],
        'blunt': [
            Timestamp('2026-09-11 08:39:26.692000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:44:12.537000+0000', tz='Europe/Berlin')
        ]
    },
    'experiment9': {
        'sharp': [
            Timestamp('2026-09-11 08:45:07.247000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:47:28.660000+0000', tz='Europe/Berlin')
        ],
        'blunt': [
            Timestamp('2026-09-11 08:47:39.048000+0000', tz='Europe/Berlin'),
            Timestamp('2026-09-11 08:50:00.460000+0000', tz='Europe/Berlin')
        ]
    }
}

experiment_configurations = pd.DataFrame(
    [
        [1, "experiment1", 140, 1700, 3, 50],
        [2, "experiment2", 150, 1700, 3, 50],
        [3, "experiment3", 130, 1700, 3, 50],
        [4, "experiment4", 140, 1800, 3, 50],
        [5, "experiment5", 140, 1600, 3, 50],
        [6, "experiment6", 140, 1700, 2, 50],
        [7, "experiment7", 140, 1700, 4, 50],
        [8, "experiment8", 140, 1700, 3, 30],
        [9, "experiment9", 140, 1700, 3, 70],
    ],
    columns=[
        "Experiment id",
        "Experiment name",
        "Feed rate", #  (mm min^-1)
        "Spindle speed", # (min^-1)
        "Axial depth of cut", # (mm)
        "Radial width of cut", # (% of tool diameter)
    ],
)

sensor_properties = {
    "sampling_frequency": 1600,  # Hz
    "window": 512,  # samples 0.32s
    "sensitivity": 2048.0 * 4  # LSB/g
}

classlabels = {
    "sharp": 0, # Good tool type or Desired behavior
    "blunt": 1  # Bad tool type or Anomaly condition
}

# Sensor data with timestamp index and column "x" indicate the vibration data from the smart sensor
sensor_data = pd.read_parquet("mems_sensor_data.parquet")
sensor_data["x"] /= sensor_properties["sensitivity"]

for i in experiment_configurations["Experiment id"].unique():
    experiment_name = f"experiment{i}"
    timestamps = experiment_timestamps[experiment_name]
    for classname in ["sharp", "blunt"]:
        t = timestamps[classname]
        t = [t[0].strftime("%Y-%m-%d %H:%M:%S.%f"), t[1].strftime("%Y-%m-%d %H:%M:%S.%f")]
        mask = (sensor_data.index >= t[0]) & (sensor_data.index <= t[1])
        sensor_data.loc[mask, 'experiment_id'] = i
        sensor_data.loc[mask, 'tool_wear'] = classlabels[classname]

sensor_data = sensor_data.dropna() # Drop NaN values to ensure that only rows with valid experiment_id and tool_wear labels are retained.
sensor_data = sensor_data.join(experiment_configurations.set_index("Experiment id"), on="experiment_id")
# The sensor_data DataFrame now contains the vibration data along with the corresponding experiment configurations and tool wear labels.
print(sensor_data.head(100))