Dataset Creation

#!/usr/bin/env python3

"""Dataset Helper scripts"""

import argparse
import json
import sys

import polars

from infuse_iot.time import InfuseTime
from infuse_iot.util.argparse import ValidFile
from pathlib import Path

class Extract:
    def __init__(self, args):
        self.args = args

    def _windows(self, annotation: polars.DataFrame, labels: list[str]):
        """Return a list of (label, start_unix, end_unix) windows"""
        windows: list[tuple[str, float, float]] = []
        label_indices: dict[str, int] = {}

        if self.args.mode == "window":
            # Labels contain start and end suffixes for their corresponding windows.
            start_suffix = self.args.start_suffix
            end_suffix = self.args.end_suffix
            start_labels = {}
            end_labels = {}
            base_labels = []

            # Determine and match start and end labels from reference labels.
            for label in labels:
                start = label.removesuffix(start_suffix)
                end = label.removesuffix(end_suffix)
                base = label.removesuffix(start_suffix).removesuffix(end_suffix)
                if base == label:
                    # This is not a start/end label
                    continue
                if start != label:
                    start_labels[base] = label
                    label_indices[base] = len(label_indices)
                if end != label:
                    end_labels[base] = label
                base_labels.append(base)

            # Remove duplicates from base_labels
            base_labels = list(dict.fromkeys(base_labels))

            # Ensure labels are matched correctly
            for base in base_labels[:]:
                if base not in start_labels or base not in end_labels:
                    print(f"Warning: {base} is missing start or end labels")
                    sys.exit()

            if not base_labels:
                print("Warning: No start/end labels found")

            # Search for and extract start and end windows from the annotation
            for base in base_labels:
                start_event = start_labels[base]
                end_event = end_labels[base]
                starts = annotation.filter(polars.col("event") == f"b'{start_event}'").sort("unix_time")["unix_time"]
                ends = annotation.filter(polars.col("event") == f"b'{end_event}'").sort("unix_time")["unix_time"]

                if len(starts) != len(ends):
                    print(f"Warning: {base} has {len(starts)} start events but "
                          f"{len(ends)} end events. Check annotation pairs are correct.")
                    sys.exit()

                for start, end in zip(starts, ends, strict=False):
                    windows.append((base, float(start), float(end)))
        else:
            # Search for events and add the before and after padding time.
            for label in labels:
                label_indices[label] = len(label_indices)
                events = annotation.filter(polars.col("event") == f"b'{label}'").sort("unix_time")["unix_time"]
                for event_time in events:
                    windows.append((label, float(event_time) - self.args.before, float(event_time) + self.args.after))

        return label_indices, windows

    def run(self):
        def _parse_time_column(df: polars.DataFrame, column: str = "time") -> polars.Series:
            """Parse a CSV time column (unix seconds or UTC timestamp string) into unix seconds (float)"""
            try:
                return df[column].cast(polars.Float64)
            except polars.exceptions.InvalidOperationError:
                pass

            return (
                df[column]
                .str.to_datetime(time_unit="us", time_zone="UTC", strict=False)
                .dt.timestamp(time_unit="us")
                .cast(polars.Float64)
                / 1_000_000
            )

        with open(self.args.labels, encoding="utf-8") as f:
            labels: list[str] = json.load(f)

        annotation = polars.read_csv(self.args.annotation)
        gps_to_unix = InfuseTime.unix_time_from_gps_seconds
        annotation = annotation.with_columns(
            annotation["timestamp"].map_elements(gps_to_unix, return_dtype=polars.Int64).alias("unix_time")
        )

        label_indices, windows = self._windows(annotation, labels)
        if not windows:
            print("No matching windows found")
            return

        data = polars.read_csv(self.args.data)
        data = data.with_columns(_parse_time_column(data).alias("unix_time"))

        extracted: list[polars.DataFrame] = []
        for window_count, (label, start, end) in enumerate(windows):
            clean_label = label.strip()
            subset = data.filter((polars.col("unix_time") >= start) & (polars.col("unix_time") <= end))
            subset = subset.with_columns(
                polars.lit(clean_label).alias("label"),
                polars.lit(label_indices[label]).alias("label_id"),
                polars.lit(window_count).alias("window_id"),
            )
            extracted.append(subset)

        result = polars.concat(extracted).drop("unix_time")
        result.write_csv(self.args.output)
        print(f"Extracted {len(result)} rows across {len(windows)} windows to {self.args.output}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Use annotation to extract labelled data windows.")
    parser.add_argument("--annotation", "-a", required=True, type=ValidFile, help="Annotation CSV file (time, timestamp, event)")
    parser.add_argument("--data", "-d", required=True, type=ValidFile, help="Data CSV file to extract labelled windows from")
    parser.add_argument("--output", "-o", required=True, type=Path, help="Output CSV file")
    parser.add_argument("--labels", "-l", type=ValidFile, required=True, help="JSON file containing an ordered list of labels")

    mode_subcommands = parser.add_subparsers(dest="mode", required=True, title="extraction modes", metavar="<mode>")

    window_parser = mode_subcommands.add_parser("window", help="Labels mark the start and end of an event window")
    window_parser.add_argument("--start-suffix", default="start", help="Suffix appended to a label for the start of a window")
    window_parser.add_argument("--end-suffix", default="end", help="Suffix appended to a label for the end of a window")

    instant_parser = mode_subcommands.add_parser("instant", help="Labels mark the exact moment of an event")
    instant_parser.add_argument("--before", type=float, default=0.5, help="Seconds before the event to include")
    instant_parser.add_argument("--after", type=float, default=0.5, help="Seconds after the event to include")
    
    args = parser.parse_args()

    extractor = Extract(args)
    extractor.run()

#!/usr/bin/env python3

import argparse
import pandas as pd


GAP_LEEWAY = 1.15

def split_contiguous_windows(data, sample_rate, min_window, window_column="window_id"):
    """Split source windows wherever the timestamp gap indicates missing samples.

    ``min_window`` is a duration in seconds.  A window is kept when its first to
    last sample duration is at least that long.  The original window id is
    retained for an un-split window; split windows receive ``_0``, ``_1``, ...
    suffixes.
    """
    if sample_rate <= 0:
        raise ValueError("Sample rates must be greater than zero")
    if min_window < 0:
        raise ValueError("--min-window must not be negative")

    sample_period = pd.Timedelta(seconds=1.0 / sample_rate * GAP_LEEWAY)
    result = []
    window_offset = 0
    for original_id, window in data.groupby(window_column, sort=True):
        numeric_id = int(original_id)
        window = window.sort_values("time").copy()
        breaks = window["time"].diff().gt(sample_period).fillna(False)
        groups = breaks.cumsum()
        parts = [part for _, part in window.groupby(groups, sort=False)]
        kept = [part for part in parts
                if part["time"].iloc[-1] - part["time"].iloc[0] >= pd.Timedelta(seconds=min_window)]
        # Do not silently turn a short, unbroken source window into a missing
        # window: min-window is intended to reject fragments created by gaps.
        if len(parts) == 1 and min_window == 0:
            kept = parts
        for number, part in enumerate(kept):
            part = part.copy()
            part[window_column] = numeric_id + window_offset + number
            result.append(part)
        window_offset += max(0, len(kept) - 1)
    if not result:
        return data.iloc[0:0].copy()
    return pd.concat(result, ignore_index=True)


def count_discontinuities(data, sample_rate):
    period = pd.Timedelta(seconds=1 / sample_rate * GAP_LEEWAY)
    return sum(window.sort_values("time")["time"].diff().gt(period).sum()
              for _, window in data.groupby("window_id", sort=False))


def split_paired_windows(acc_data, gyro_data, acc_rate, gyro_rate, min_window):
    """Split both streams at every gap found in either stream."""
    acc_period = pd.Timedelta(seconds=1 / acc_rate * GAP_LEEWAY)
    gyro_period = pd.Timedelta(seconds=1 / gyro_rate * GAP_LEEWAY)
    acc_result, gyro_result = [], []
    discontinuities = 0
    window_offset = 0
    for window_id in sorted(set(acc_data.window_id) & set(gyro_data.window_id)):
        numeric_id = int(window_id)
        acc = acc_data[acc_data.window_id == window_id].sort_values("time").copy()
        gyro = gyro_data[gyro_data.window_id == window_id].sort_values("time").copy()
        boundaries = set(acc.loc[acc.time.diff().gt(acc_period), "time"])
        boundaries.update(gyro.loc[gyro.time.diff().gt(gyro_period), "time"])
        boundaries = sorted(boundaries)
        discontinuities += len(boundaries)
        starts = [None] + boundaries
        ends = boundaries + [None]
        number = 0
        for start, end in zip(starts, ends):
            acc_part = acc if start is None and end is None else acc[(acc.time >= start if start is not None else True) & (acc.time < end if end is not None else True)]
            gyro_part = gyro if start is None and end is None else gyro[(gyro.time >= start if start is not None else True) & (gyro.time < end if end is not None else True)]
            if acc_part.empty or gyro_part.empty:
                continue
            overlap = min(acc_part.time.iloc[-1], gyro_part.time.iloc[-1]) - max(acc_part.time.iloc[0], gyro_part.time.iloc[0])
            if overlap < pd.Timedelta(seconds=min_window):
                continue
            sub_id = numeric_id + window_offset + number
            acc_part = acc_part.copy()
            gyro_part = gyro_part.copy()
            acc_part["window_id"] = sub_id
            gyro_part["window_id"] = sub_id
            acc_result.append(acc_part)
            gyro_result.append(gyro_part)
            number += 1
        window_offset += max(0, number - 1)
    empty_acc, empty_gyro = acc_data.iloc[0:0].copy(), gyro_data.iloc[0:0].copy()
    return (pd.concat(acc_result, ignore_index=True) if acc_result else empty_acc,
            pd.concat(gyro_result, ignore_index=True) if gyro_result else empty_gyro,
            discontinuities)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Preprocess accelerometer and gyroscope data")
    parser.add_argument("--acc", type=str, required=True, help="Path to the accelerometer CSV file")
    parser.add_argument("--gyro", type=str, help="Path to the gyroscope CSV file")
    parser.add_argument("--output", type=str, required=True, help="Path to the output merged CSV file")
    parser.add_argument("--acc-rate", type=float, required=True, help="Sample rate of the accelerometer data in Hz")
    parser.add_argument("--gyro-rate", type=float, help="Sample rate of the gyroscope data in Hz")
    parser.add_argument("--min-window", type=float, default=0.0,
                        help="Minimum contiguous sub-window duration in seconds (default: 0)")
    args = parser.parse_args()

    # Only require the gyroscope sample rate if gyroscope data is provided
    if args.gyro and not args.gyro_rate:
        raise ValueError("Gyroscope sample rate (--gyro-rate) must be provided if gyroscope data (--gyro) is provided.")

    # Load the accelerometer and gyroscope data
    acc_data = pd.read_csv(args.acc, parse_dates=['time'], date_format="ISO8601", dtype={"sample.x": "int16", "sample.y": "int16", "sample.z": "int16"})
    gyro_data = pd.read_csv(args.gyro, parse_dates=['time'], date_format="ISO8601", dtype={"sample.x": "int16", "sample.y": "int16", "sample.z": "int16"}) if args.gyro else None

    # Rename axes to avoid confusion (sample prefix), but also avoid using '.' for Nordic Edge AI Lab
    acc_data.rename(columns={"sample.x": "acc_x", "sample.y": "acc_y", "sample.z": "acc_z"}, inplace=True)
    if gyro_data is not None:
        gyro_data.rename(columns={"sample.x": "gyro_x", "sample.y": "gyro_y", "sample.z": "gyro_z"}, inplace=True)

    # Find windows in both acc and gyro data with the same window ID (composed of label ID and sample ID columns)
    discontinuities = 0
    if gyro_data is not None:
        common_windows = set(acc_data['window_id']).intersection(set(gyro_data['window_id']))
        acc_data = acc_data[acc_data['window_id'].isin(common_windows)]
        gyro_data = gyro_data[gyro_data['window_id'].isin(common_windows)]
    else:
        discontinuities = count_discontinuities(acc_data, args.acc_rate)
        acc_data = split_contiguous_windows(acc_data, args.acc_rate, args.min_window)
        common_windows = None

    # Within the same window (using Window ID) Find the smallest offset between any two samples, and align the gyroscope data accordingly
    if common_windows is not None and gyro_data is not None:
        for window_id in common_windows:
            acc_window = acc_data[acc_data['window_id'] == window_id].sort_values('time')
            gyro_window = gyro_data[gyro_data['window_id'] == window_id].sort_values('time')
            if not acc_window.empty and not gyro_window.empty:
                time_offset = acc_window['time'].iloc[0] - gyro_window['time'].iloc[0]
                gyro_data.loc[gyro_data['window_id'] == window_id, 'time'] += time_offset
            # Remove samples where either gyro or acc data is outside the start/end of the other
            gyro_data = gyro_data[~((gyro_data['window_id'] == window_id) & 
                                    ((gyro_data['time'] < acc_window['time'].iloc[0]) | 
                                    (gyro_data['time'] > acc_window['time'].iloc[-1])))]
            acc_data = acc_data[~((acc_data['window_id'] == window_id) & 
                                ((acc_data['time'] < gyro_window['time'].iloc[0]) | 
                                    (acc_data['time'] > gyro_window['time'].iloc[-1])))]

        acc_data, gyro_data, discontinuities = split_paired_windows(
            acc_data, gyro_data, args.acc_rate, args.gyro_rate, args.min_window
        )

        # Resample the gyroscope data to match the accelerometer sample rate using Zero Order Hold (ZOH)
        # merge_asof requires its `on` key to be globally sorted, not merely
        # sorted within each window.  `window_id` still keeps matches isolated
        # to their corresponding sub-window via `by` below.
        gyro_data = gyro_data.sort_values(['time', 'window_id'])
        acc_data = acc_data.sort_values(['time', 'window_id'])
        # Match each accelerometer sample with the most recent gyro sample in
        # the same contiguous sub-window (zero-order hold).
        gyro_columns = [column for column in ('gyro_x', 'gyro_y', 'gyro_z')
                        if column in gyro_data.columns]
        gyro_data_resampled = pd.merge_asof(
            acc_data[['window_id', 'time']],
            gyro_data[['window_id', 'time'] + gyro_columns],
            on='time', by='window_id', direction='backward',
        ).drop(columns=['window_id', 'time'])
        gyro_data_resampled.index = acc_data.index

        # Finally, merge the two streams
        merged_data = acc_data.join(gyro_data_resampled)
        merged_data.to_csv(args.output, index=False)
    else:
        merged_data = acc_data

    # Create 1D columns for magnitude and angular velocity (as uint32 to prevent overflow)
    merged_data['acc_magnitude'] = ((merged_data['acc_x'].astype('uint32')**2 + merged_data['acc_y'].astype('uint32')**2 + merged_data['acc_z'].astype('uint32')**2)**0.5).astype('uint32')
    if gyro_data is not None:
        merged_data['gyro_magnitude'] = ((merged_data['gyro_x'].astype('uint32')**2 + merged_data['gyro_y'].astype('uint32')**2 + merged_data['gyro_z'].astype('uint32')**2)**0.5).astype('uint32')

    metadata_columns = [column for column in ('label', 'label_id', 'window_id')
                        if column in merged_data.columns]
    merged_data = merged_data[[column for column in merged_data.columns
                              if column not in metadata_columns] + metadata_columns]
    merged_data.to_csv(args.output, index=False)
    print(f"Non-contiguous gaps detected: {discontinuities}")

import argparse
import random
from pathlib import Path

import pandas as pd

REQUIRED_COLUMNS = {"time", "label", "label_id", "window_id"}


def positive_int(value: str) -> int:
    try:
        number = int(value)
    except ValueError as error:
        raise argparse.ArgumentTypeError("must be an integer") from error
    if number <= 0:
        raise argparse.ArgumentTypeError("must be greater than zero")
    return number


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("-i", "--input", required=True, type=Path, help="input CSV file")
    parser.add_argument("-o", "--output", required=True, type=Path, help="output CSV file")
    parser.add_argument(
        "--num-samples",
        required=True,
        type=positive_int,
        help="maximum target number of rows for each regular label_id",
    )
    parser.add_argument(
        "--exception-label-id",
        "--exception-label-ids",
        action="append",
        default=[],
        metavar="LABEL_ID",
        help=(
            "label_id excluded from the common minimum; repeat this option "
            "for multiple exception labels"
        ),
    )
    parser.add_argument(
        "--exception-sample-count",
        type=positive_int,
        metavar="COUNT",
        help=(
            "target row count for exception labels; omit this option to "
            "include all their windows"
        ),
    )
    parser.add_argument(
        "--drop-non-numerical",
        action="store_true",
        help="omit non-numerical columns and time from the output CSV",
    )
    return parser.parse_args()


def choose_window_positions(
    window_positions: list[list[int]],
    target_rows: int | None,
) -> set[int]:
    """Return positions from whole windows until the target is reached."""

    windows = list(window_positions)
    if target_rows is None:
        return {position for window in windows for position in window}

    random.shuffle(windows)
    selected: set[int] = set()
    for window in windows:
        selected.update(window)
        if len(selected) >= target_rows:
            break
    return selected


def balance_dataframe(
    data: pd.DataFrame,
    num_samples: int,
    exception_label_ids: set[str],
    exception_sample_count: int | None,
) -> tuple[pd.DataFrame, dict[str, int], int]:
    """Select complete windows while retaining the input row order."""

    data = data.reset_index(drop=True)
    label_counts = data.groupby("label_id", sort=False).size()
    regular_counts = label_counts[~label_counts.index.isin(exception_label_ids)]
    common_target = (
        min(num_samples, int(regular_counts.min()))
        if not regular_counts.empty
        else num_samples
    )

    grouped_windows = data.groupby(["label_id", "window_id"], sort=False).indices
    windows_by_label: dict[str, list[list[int]]] = {}
    for (label_id, _window_id), positions in grouped_windows.items():
        windows_by_label.setdefault(label_id, []).append(list(positions))

    selected_positions: set[int] = set()
    selected_counts: dict[str, int] = {}
    for label_id, windows in windows_by_label.items():
        target = (
            exception_sample_count
            if label_id in exception_label_ids
            else common_target
        )
        selected = choose_window_positions(windows, target)
        selected_positions.update(selected)
        selected_counts[label_id] = len(selected)

    positions = sorted(selected_positions)
    return data.iloc[positions], selected_counts, common_target


def main() -> None:
    args = parse_args()
    exception_label_ids = set(args.exception_label_id)
    if args.exception_sample_count is not None and not exception_label_ids:
        raise SystemExit("--exception-sample-count requires --exception-label-id")

    try:
        # dtype=str plus keep_default_na=False prevents numeric coercion and
        # prevents values such as "NA" from becoming missing values.
        data = pd.read_csv(args.input, dtype=str, keep_default_na=False)
        missing = REQUIRED_COLUMNS - set(data.columns)
        if missing:
            names = ", ".join(sorted(missing))
            raise ValueError(f"{args.input} is missing required column(s): {names}")

        selected, counts, common_target = balance_dataframe(
            data,
            args.num_samples,
            exception_label_ids,
            args.exception_sample_count,
        )
        if args.drop_non_numerical:
            non_numerical = [
                column
                for column in selected.columns
                if not pd.to_numeric(selected[column], errors="coerce").notna().all()
            ]
            selected = selected.drop(columns=[*non_numerical, "time"], errors="ignore")
        selected.to_csv(args.output, index=False)
    except (OSError, pd.errors.ParserError, ValueError) as error:
        raise SystemExit(str(error)) from error

    print("Input rows by label_id:")
    for label_id, count in data.groupby("label_id", sort=False).size().items():
        print(f"label_id={label_id!r}: {count} row(s)")
    print(f"Common target: {common_target} row(s)")
    for label_id in sorted(counts):
        print(f"label_id={label_id!r}: selected {counts[label_id]} row(s)")
    print(f"Wrote {len(selected)} row(s) to {args.output}")

if __name__ == "__main__":
    main()