cryosim/data/loader.py

"""DataFrame loaders for CSH2 test campaign data."""

import os
from typing import Optional, List, Dict, Any

import numpy as np
import pandas as pd

from cryosim.data.sensors import TAG_TO_NAME


CSH2_DATA_ROOT = os.path.normpath(os.path.join(
    os.path.dirname(__file__), "..", "..", "..",
    "data_analysis_csh2"
))

_LH2_CLEANED_PATH = os.path.normpath(os.path.join(
    os.path.dirname(__file__), "..", "..", "..",
    "retired_models", "murphy_feb18", "sept24_25_fill_data_1sec.csv"
))

# Default tags: Big 5 + FT140 + AvgPower
_DEFAULT_TAGS = [17, 18, 24, 27, 13, 7, 2]


def load_lh2_fills(path: Optional[str] = None) -> pd.DataFrame:
    """Load pre-pivoted Sept 24-25 LH2 fill data.

    Returns DataFrame indexed by utc_full_timestamp with columns:
    AvgPower, FT140, PT110, PT130, TT110, TT130, M130_Speed
    """
    fpath = path or _LH2_CLEANED_PATH
    if not os.path.exists(fpath):
        raise FileNotFoundError(f"LH2 fill data not found at {fpath}")
    df = pd.read_csv(fpath, parse_dates=["utc_full_timestamp"], index_col="utc_full_timestamp")
    return df


def load_long_format(path: str, tags: Optional[list] = None) -> pd.DataFrame:
    """Load a long-format CSH2 CSV and pivot to wide format."""
    if tags is None:
        tags = [17, 18, 24, 27, 13]
    df = pd.read_csv(path, parse_dates=["utc_full_timestamp"])
    df = df[df["tagindex"].isin(tags)]
    df["sensor"] = df["tagindex"].map(TAG_TO_NAME)
    wide = df.pivot_table(
        index="utc_full_timestamp", columns="sensor", values="val", aggfunc="first"
    )
    wide.index.name = "utc_full_timestamp"
    return wide


def load_fill_metrics(path: Optional[str] = None) -> pd.DataFrame:
    """Load per-fill aggregate metrics."""
    fpath = path or os.path.join(CSH2_DATA_ROOT, "fill_metrics.csv")
    if not os.path.exists(fpath):
        raise FileNotFoundError(f"Fill metrics not found at {fpath}")
    return pd.read_csv(fpath)


# ---------------------------------------------------------------------------
# C8: Calibration data pipeline
# ---------------------------------------------------------------------------

def load_test_data(
    csv_path: str,
    ft140_correction: float = 11.4,
    tags: Optional[List[int]] = None,
    resample: Optional[str] = None,
) -> pd.DataFrame:
    """Load raw test campaign CSV and prepare for calibration.

    Handles:
    - Long-format (tagindex, val) or wide-format (column per sensor)
    - FT140 flow meter correction (raw values are ~11.4x too high for LH2
      due to density calibration mismatch -- divide by ft140_correction)
    - Optional resampling (e.g., '1s', '15s') for large files

    Args:
        csv_path: Path to CSV file.
        ft140_correction: Divisor for FT140 readings. Default 11.4 for LH2.
            Set to 1.0 for LN2 or pre-corrected data.
        tags: Tag indices to load (default: Big 5 + FT140 + Power).
            [17, 18, 24, 27, 13, 7, 2] = PT110, PT130, TT110, TT130,
            MC130_VFD_Speed, FT140, AvgPower
        resample: Resample interval (e.g. '1s', '15s'). None = no resampling.

    Returns:
        Wide-format DataFrame indexed by timestamp with columns like
        PT110, PT130, TT110, TT130, MC130_VFD_Speed, FT140, AvgPower.
        FT140 column is corrected by ft140_correction.
    """
    if tags is None:
        tags = list(_DEFAULT_TAGS)

    df = pd.read_csv(csv_path)

    # --- auto-detect format ------------------------------------------------
    if "tagindex" in df.columns:
        # Long format: pivot to wide
        ts_col = _find_timestamp_col(df)
        df[ts_col] = pd.to_datetime(df[ts_col], format="ISO8601", utc=True)
        df = df[df["tagindex"].isin(tags)]
        df["sensor"] = df["tagindex"].map(TAG_TO_NAME)
        # Drop rows where tagindex wasn't in TAG_TO_NAME
        df = df.dropna(subset=["sensor"])
        wide = df.pivot_table(
            index=ts_col, columns="sensor", values="val", aggfunc="first",
        )
        wide.index.name = "utc_full_timestamp"
    else:
        # Wide format: columns are already sensor names
        ts_col = _find_timestamp_col(df)
        df[ts_col] = pd.to_datetime(df[ts_col], format="ISO8601", utc=True)
        df = df.set_index(ts_col)
        df.index.name = "utc_full_timestamp"
        wide = df

    # Ensure sorted by time
    wide = wide.sort_index()

    # --- FT140 correction --------------------------------------------------
    if "FT140" in wide.columns and ft140_correction != 1.0:
        wide["FT140"] = wide["FT140"] / ft140_correction

    # --- optional resample -------------------------------------------------
    if resample is not None:
        wide = wide.resample(resample).mean()

    return wide


def _find_timestamp_col(df: pd.DataFrame) -> str:
    """Find the timestamp column in a DataFrame."""
    for candidate in ("utc_full_timestamp", "timestamp", "time", "datetime"):
        if candidate in df.columns:
            return candidate
    # Fall back to first column that contains 'time'
    for col in df.columns:
        if "time" in col.lower():
            return col
    raise ValueError(
        f"Cannot find timestamp column. Columns: {list(df.columns)}"
    )


def extract_fills(
    df: pd.DataFrame,
    pressure_col: str = "PT130",
    speed_col: str = "MC130_VFD_Speed",
    min_pressure_rise: float = 50.0,
    min_duration_s: float = 30.0,
) -> List[Dict[str, Any]]:
    """Extract individual fill periods from a continuous test dataset.

    Detects fills by finding periods where:
    1. VFD speed > 5% (pump is running)
    2. Discharge pressure rises by at least *min_pressure_rise* bar
    3. Duration exceeds *min_duration_s*

    Args:
        df: Wide-format DataFrame (output of :func:`load_test_data`).
        pressure_col: Column name for discharge pressure.
        speed_col: Column name for VFD speed.
        min_pressure_rise: Minimum pressure rise to qualify as a fill [bar].
        min_duration_s: Minimum duration to qualify as a fill [s].

    Returns:
        List of dicts, each with:
        - ``start``: timestamp
        - ``end``: timestamp
        - ``duration_s``: float
        - ``P_start``: float (bar)
        - ``P_peak``: float (bar)
        - ``speed_mean``: float (%)
        - ``FT140_mean``: float (kg/min, corrected) -- NaN if FT140 absent
        - ``data``: DataFrame slice
    """
    if pressure_col not in df.columns:
        raise KeyError(f"Pressure column '{pressure_col}' not in DataFrame")
    if speed_col not in df.columns:
        raise KeyError(f"Speed column '{speed_col}' not in DataFrame")

    speed = df[speed_col].fillna(0)
    running = (speed > 5.0).astype(int)

    # Find contiguous runs where the pump is running
    transitions = running.diff().fillna(0)
    starts = df.index[transitions == 1]
    stops = df.index[transitions == -1]

    # Handle edge cases: running at start/end of data
    if running.iloc[0] > 0:
        starts = starts.insert(0, df.index[0])
    if running.iloc[-1] > 0:
        stops = stops.append(pd.DatetimeIndex([df.index[-1]]))

    fills: List[Dict[str, Any]] = []
    for s, e in zip(starts, stops):
        segment = df.loc[s:e]
        if len(segment) < 2:
            continue

        duration_s = (segment.index[-1] - segment.index[0]).total_seconds()
        if duration_s < min_duration_s:
            continue

        p_start = segment[pressure_col].iloc[0]
        p_peak = segment[pressure_col].max()
        if (p_peak - p_start) < min_pressure_rise:
            continue

        ft140_mean = float("nan")
        if "FT140" in segment.columns:
            ft140_mean = segment["FT140"].mean()

        fills.append({
            "start": segment.index[0],
            "end": segment.index[-1],
            "duration_s": duration_s,
            "P_start": float(p_start),
            "P_peak": float(p_peak),
            "speed_mean": float(segment[speed_col].mean()),
            "FT140_mean": float(ft140_mean),
            "data": segment,
        })

    return fills


def fills_to_calibration_input(
    fills: List[Dict[str, Any]],
    speed_scale: float = 0.664,
) -> List[Dict[str, float]]:
    """Convert extracted fills to calibration-ready format.

    Maps from test data format to the dict format expected by
    :func:`cryosim.calibrate`::

        [{'Pexit_barg': ..., 'speed_f': ..., 'Ptank_barg': ...,
          'Psat_barg': ..., 'measured_mdot_kgpm': ...}, ...]

    Args:
        fills: Output of :func:`extract_fills`.
        speed_scale: VFD% to speed fraction mapping. Default 0.664 for LH2,
            5.0 for LN2.

    Returns:
        List of calibration input dicts.
    """
    cal_fills: List[Dict[str, float]] = []
    for f in fills:
        speed_pct = f["speed_mean"]
        speed_f = (speed_pct / 100.0) * speed_scale

        # Use mean FT140 as measured mass flow; skip if unavailable
        mdot = f.get("FT140_mean", float("nan"))
        if np.isnan(mdot):
            continue

        cal_fills.append({
            "Pexit_barg": f["P_peak"],
            "speed_f": speed_f,
            "Ptank_barg": 7.0,   # default inlet tank pressure
            "Psat_barg": 2.0,    # default saturation pressure
            "measured_mdot_kgpm": mdot,
        })
    return cal_fills