Building/convert_teacher_csv.py

from __future__ import annotations

import argparse
from pathlib import Path
import re

import numpy as np
import pandas as pd


VIB_COLUMN_RE = re.compile(r"^(WSMS\d+)\.Acc([XYZ])$", flags=re.IGNORECASE)
SGD_COLUMN_RE = re.compile(r"^(WSGD\d+)\.chdata$", flags=re.IGNORECASE)
OUTPUT_COLUMNS = ["code", "type", "time", "value1", "value2", "value3"]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Convert teacher wide-format CSV into downloads-compatible long-format CSV."
    )
    parser.add_argument("--input", type=str, required=True, help="Input wide-format CSV path.")
    parser.add_argument("--output", type=str, required=True, help="Output long-format CSV path.")
    parser.add_argument(
        "--time-column",
        type=str,
        default="time",
        help="Time column name in the input CSV (default: time).",
    )
    parser.add_argument(
        "--timezone-offset-hours",
        type=float,
        default=8.0,
        help="Offset used when parsing naive datetime strings to epoch seconds (default: 8).",
    )
    return parser.parse_args()


def discover_sensor_columns(columns: list[str]) -> tuple[dict[str, dict[str, str]], dict[str, str]]:
    vibr_map: dict[str, dict[str, str]] = {}
    sgd_map: dict[str, str] = {}
    for column in columns:
        vibr_match = VIB_COLUMN_RE.match(column)
        if vibr_match is not None:
            sensor_code = vibr_match.group(1).upper()
            axis = vibr_match.group(2).upper()
            vibr_map.setdefault(sensor_code, {})[axis] = column
            continue
        sgd_match = SGD_COLUMN_RE.match(column)
        if sgd_match is not None:
            sensor_code = sgd_match.group(1).upper()
            sgd_map[sensor_code] = column
    return vibr_map, sgd_map


def _strip_excel_datetime_wrapper(values: pd.Series) -> pd.Series:
    text = values.astype(str).str.strip()
    text = text.str.replace(r'^="(.*)"$', r"\1", regex=True)
    text = text.str.replace(r'^"(.*)"$', r"\1", regex=True)
    return text


def to_epoch_seconds(time_values: pd.Series, timezone_offset_hours: float) -> pd.Series:
    numeric_time = pd.to_numeric(time_values, errors="coerce")
    if numeric_time.notna().mean() > 0.8:
        return numeric_time.astype("float64")

    text_time = _strip_excel_datetime_wrapper(time_values)
    dt_values = pd.to_datetime(text_time, errors="coerce")
    epoch = (dt_values - pd.Timestamp("1970-01-01")) / pd.Timedelta(seconds=1)
    epoch = epoch - timezone_offset_hours * 3600.0
    epoch = pd.to_numeric(epoch, errors="coerce")

    # Fall back to numeric parsing for rows that were not parsed as datetime.
    fallback_numeric = pd.to_numeric(text_time, errors="coerce")
    epoch = epoch.where(epoch.notna(), fallback_numeric)
    return epoch.astype("float64")


def _empty_output_frame() -> pd.DataFrame:
    return pd.DataFrame(columns=OUTPUT_COLUMNS)


def build_vibr_rows(df: pd.DataFrame, time_column: str, vibr_map: dict[str, dict[str, str]]) -> list[pd.DataFrame]:
    rows: list[pd.DataFrame] = []
    for sensor_code in sorted(vibr_map.keys()):
        axis_columns = vibr_map[sensor_code]
        use_columns = [time_column]
        if "X" in axis_columns:
            use_columns.append(axis_columns["X"])
        if "Y" in axis_columns:
            use_columns.append(axis_columns["Y"])
        if "Z" in axis_columns:
            use_columns.append(axis_columns["Z"])
        if len(use_columns) == 1:
            continue

        sensor_df = df.loc[:, use_columns].copy()
        sensor_df = sensor_df.rename(
            columns={
                axis_columns.get("X", ""): "value1",
                axis_columns.get("Y", ""): "value2",
                axis_columns.get("Z", ""): "value3",
            }
        )
        for column in ["value1", "value2", "value3"]:
            if column not in sensor_df.columns:
                sensor_df[column] = np.nan
        sensor_df["value1"] = pd.to_numeric(sensor_df["value1"], errors="coerce")
        sensor_df["value2"] = pd.to_numeric(sensor_df["value2"], errors="coerce")
        sensor_df["value3"] = pd.to_numeric(sensor_df["value3"], errors="coerce")
        sensor_df = sensor_df[sensor_df[["value1", "value2", "value3"]].notna().any(axis=1)]
        if sensor_df.empty:
            continue

        sensor_df.insert(0, "type", "vibr")
        sensor_df.insert(0, "code", sensor_code)
        rows.append(sensor_df[OUTPUT_COLUMNS])
    return rows


def build_sgd_rows(df: pd.DataFrame, time_column: str, sgd_map: dict[str, str]) -> list[pd.DataFrame]:
    rows: list[pd.DataFrame] = []
    for sensor_code in sorted(sgd_map.keys()):
        sensor_column = sgd_map[sensor_code]
        sensor_df = df.loc[:, [time_column, sensor_column]].copy()
        sensor_df = sensor_df.rename(columns={sensor_column: "value1"})
        sensor_df["value1"] = pd.to_numeric(sensor_df["value1"], errors="coerce")
        sensor_df = sensor_df[sensor_df["value1"].notna()]
        if sensor_df.empty:
            continue

        sensor_df["value2"] = np.nan
        sensor_df["value3"] = np.nan
        sensor_df.insert(0, "type", "sgd")
        sensor_df.insert(0, "code", sensor_code)
        rows.append(sensor_df[OUTPUT_COLUMNS])
    return rows


def convert_dataframe(input_df: pd.DataFrame, time_column: str, timezone_offset_hours: float) -> pd.DataFrame:
    if time_column not in input_df.columns:
        raise ValueError(f"Time column '{time_column}' not found in input CSV.")

    vibr_map, sgd_map = discover_sensor_columns(input_df.columns.tolist())
    if not vibr_map and not sgd_map:
        raise ValueError("No sensor columns detected. Expected WSMSxxxx.AccX/Y/Z or WSGDxxxx.chdata columns.")

    df = input_df.copy()
    df[time_column] = to_epoch_seconds(df[time_column], timezone_offset_hours=timezone_offset_hours)
    df = df[df[time_column].notna()]
    if df.empty:
        return _empty_output_frame()

    rows: list[pd.DataFrame] = []
    rows.extend(build_vibr_rows(df, time_column=time_column, vibr_map=vibr_map))
    rows.extend(build_sgd_rows(df, time_column=time_column, sgd_map=sgd_map))
    if not rows:
        return _empty_output_frame()

    output_df = pd.concat(rows, ignore_index=True)
    output_df["time"] = pd.to_numeric(output_df["time"], errors="coerce").astype("float64")
    output_df = output_df[output_df["time"].notna()]
    output_df = output_df.sort_values(["time", "code"], kind="mergesort").reset_index(drop=True)
    return output_df.loc[:, OUTPUT_COLUMNS]


def print_summary(df: pd.DataFrame, output_path: Path) -> None:
    print(f"Output saved to: {output_path}")
    print(f"Rows: {len(df)}")
    if df.empty:
        print("No data rows were generated.")
        return
    print(f"Sensors: {df['code'].nunique()}")
    print("Rows by sensor:")
    counts = df.groupby(["code", "type"]).size().sort_index()
    for (code, sensor_type), count in counts.items():
        print(f"  - {code} ({sensor_type}): {int(count)}")


def main() -> None:
    args = parse_args()
    input_path = Path(args.input).resolve()
    output_path = Path(args.output).resolve()

    if not input_path.exists():
        raise FileNotFoundError(f"Input file not found: {input_path}")

    input_df = pd.read_csv(input_path)
    output_df = convert_dataframe(
        input_df=input_df,
        time_column=args.time_column,
        timezone_offset_hours=float(args.timezone_offset_hours),
    )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_df.to_csv(output_path, index=False)
    print_summary(output_df, output_path=output_path)


if __name__ == "__main__":
    main()