analysis/pump_cycle_detector.py

"""
Testing Cycle Detection Module
Detects testing operation cycles from motor speed time series data,
qualifies them by pressure/flow activity, and ranks by significance.
"""
import pandas as pd
import numpy as np
from typing import List, Dict, Optional
from datetime import datetime
from core import config


class PumpCycleDetector:
    """Detect and qualify testing cycles using motor speed state machine + multi-signal qualification.

    Two-stage pipeline:
      1. Candidate generation: motor speed thresholding finds all motor-on periods
      2. Qualification: filters out brief jogs/false starts using pressure and flow data
    """

    def __init__(self):
        params = config.CYCLE_DETECTION
        self.idle_threshold = params['idle_threshold']
        self.min_cycle_seconds = params['min_cycle_seconds']
        self.smoothing_window = params['smoothing_window']
        self.max_gap_seconds = params['max_gap_seconds']
        # Qualification thresholds
        self.min_peak_pressure = params.get('min_peak_pressure_bar', 50.0)
        self.min_pressure_rise = params.get('min_pressure_rise_bar', 10.0)
        self.min_avg_flow = params.get('min_avg_flow_kg_min', 0.05)
        self.max_cycle_duration_min = params.get('max_cycle_duration_hours', 8.0) * 60

    def detect_cycles(self, motor_speed_df: pd.DataFrame) -> List[Dict]:
        """
        Detect testing cycles from motor speed time series.

        Args:
            motor_speed_df: DataFrame with columns [timestamp, value]
                            where value is motor speed in RPM

        Returns:
            List of cycle dicts with keys:
                cycle_id, start_time, end_time, duration_minutes, peak_speed
        """
        if motor_speed_df.empty or len(motor_speed_df) < 10:
            return []

        df = motor_speed_df.sort_values('timestamp').copy()
        df['speed_smooth'] = self._smooth_signal(df['value'])
        df['is_running'] = df['speed_smooth'] > self.idle_threshold

        # Find transitions
        df['state_change'] = df['is_running'].astype(int).diff().fillna(0)

        # Start events: state_change == 1 (idle -> running)
        starts = df[df['state_change'] == 1]['timestamp'].tolist()
        # Stop events: state_change == -1 (running -> idle)
        stops = df[df['state_change'] == -1]['timestamp'].tolist()

        if not starts and not stops:
            # Check if entire period is running
            if df['is_running'].any():
                return [{
                    'cycle_id': 1,
                    'start_time': df['timestamp'].iloc[0],
                    'end_time': df['timestamp'].iloc[-1],
                    'duration_minutes': (df['timestamp'].iloc[-1] - df['timestamp'].iloc[0]).total_seconds() / 60,
                    'peak_speed': df['value'].max(),
                }]
            return []

        # Pair starts and stops into cycles
        cycles = self._pair_transitions(starts, stops, df)

        # Filter by minimum duration
        cycles = [c for c in cycles if c['duration_seconds'] >= self.min_cycle_seconds]

        # Merge cycles with short gaps between them
        cycles = self._merge_short_gaps(cycles)

        # Compute peak speed for each cycle
        for c in cycles:
            mask = (df['timestamp'] >= c['start_time']) & (df['timestamp'] <= c['end_time'])
            cycle_data = df[mask]
            c['peak_speed'] = cycle_data['value'].max() if not cycle_data.empty else 0

        # Assign sequential IDs
        for i, c in enumerate(cycles, 1):
            c['cycle_id'] = i
            c['duration_minutes'] = c['duration_seconds'] / 60
            del c['duration_seconds']

        return cycles

    def _smooth_signal(self, series: pd.Series) -> pd.Series:
        """Apply rolling median filter to reduce noise"""
        return series.rolling(window=self.smoothing_window, center=True, min_periods=1).median()

    def _pair_transitions(
        self,
        starts: List[datetime],
        stops: List[datetime],
        df: pd.DataFrame,
    ) -> List[Dict]:
        """Pair start/stop transitions into cycles"""
        cycles = []

        # If first data point is already running, prepend a synthetic start
        if df['is_running'].iloc[0]:
            starts = [df['timestamp'].iloc[0]] + starts

        # If last data point is still running, append a synthetic stop
        if df['is_running'].iloc[-1]:
            stops = stops + [df['timestamp'].iloc[-1]]

        # Match each start with the next stop
        stop_idx = 0
        for start in starts:
            # Find the next stop after this start
            while stop_idx < len(stops) and stops[stop_idx] <= start:
                stop_idx += 1

            if stop_idx < len(stops):
                end = stops[stop_idx]
                duration = (end - start).total_seconds()
                cycles.append({
                    'start_time': start,
                    'end_time': end,
                    'duration_seconds': duration,
                })
                stop_idx += 1

        return cycles

    def _merge_short_gaps(self, cycles: List[Dict]) -> List[Dict]:
        """Merge cycles separated by brief motor stops"""
        if len(cycles) <= 1:
            return cycles

        merged = [cycles[0].copy()]
        for c in cycles[1:]:
            gap = (c['start_time'] - merged[-1]['end_time']).total_seconds()
            if gap <= self.max_gap_seconds:
                # Merge: extend the previous cycle
                merged[-1]['end_time'] = c['end_time']
                merged[-1]['duration_seconds'] = (
                    merged[-1]['end_time'] - merged[-1]['start_time']
                ).total_seconds()
            else:
                merged.append(c.copy())

        return merged

    def _qualify_cycle(self, cycle: Dict) -> bool:
        """Check if a cycle is a meaningful testing event.

        A cycle must show evidence of actual compression (pressure) or flow
        to qualify. Motor speed and duration alone are insufficient — they
        only confirm the motor ran, not that meaningful work happened.

        Hard rejects cycles exceeding max_cycle_duration_min (noise spans).
        """
        duration_min = cycle.get('duration_minutes', 0) or 0

        # Hard reject: multi-day noise spans masquerading as cycles
        if duration_min > self.max_cycle_duration_min:
            return False

        peak_pressure = cycle.get('peak_pressure', 0) or 0
        initial_pressure = cycle.get('initial_pressure', 0) or 0
        avg_flow = cycle.get('avg_flow', 0) or 0

        # Criterion 1: Pressure built beyond ambient
        if peak_pressure > self.min_peak_pressure:
            return True

        # Criterion 2: Meaningful pressure rise from initial value
        if initial_pressure > 0 and (peak_pressure - initial_pressure) > self.min_pressure_rise:
            return True

        # Criterion 3: Flow was delivered
        if avg_flow > self.min_avg_flow:
            return True

        return False

    def _sort_and_renumber(self, cycles: List[Dict]) -> List[Dict]:
        """Sort cycles by peak pressure descending (most significant first),
        then by duration as tiebreaker. Re-assign sequential cycle_ids."""
        sorted_cycles = sorted(
            cycles,
            key=lambda c: (-(c.get('peak_pressure', 0) or 0), -(c.get('duration_minutes', 0) or 0)),
        )
        for i, c in enumerate(sorted_cycles, 1):
            c['cycle_id'] = i
        return sorted_cycles

    def enrich_cycle_metadata(
        self,
        cycle: Dict,
        db_connector,
        tags: List[str] = None,
    ) -> Dict:
        """
        Add peak pressure, peak temp, avg flow to a cycle's metadata.

        Args:
            cycle: cycle dict with start_time, end_time
            db_connector: DatabaseConnector instance
            tags: list of tags to query (defaults to CYCLE_DETAIL_TAGS)

        Returns:
            Enriched cycle dict with additional metrics
        """
        if tags is None:
            tags = config.CYCLE_DETAIL_TAGS

        df = db_connector.get_sensor_data(
            tags, cycle['start_time'], cycle['end_time'],
        )

        if df.empty:
            return cycle

        enriched = cycle.copy()

        # Peak discharge pressure
        pt130 = df[df['tag_name'] == 'PT130']
        if not pt130.empty:
            enriched['peak_pressure'] = pt130['value'].max()

        # Peak temperatures
        for tag in ['TT110', 'TT130']:
            data = df[df['tag_name'] == tag]
            if not data.empty:
                enriched[f'peak_{tag}'] = data['value'].max()

        # Average flow
        ft140 = df[df['tag_name'] == 'FT140']
        if not ft140.empty:
            enriched['avg_flow'] = ft140['value'].mean()

        return enriched

    def batch_enrich_cycles(
        self,
        cycles: List[Dict],
        db_connector,
        tags: List[str] = None,
        qualify: bool = True,
    ) -> List[Dict]:
        """
        Enrich ALL cycles with metadata using a single DB query,
        then optionally qualify and rank them.

        Args:
            cycles: list of cycle dicts with start_time, end_time
            db_connector: DatabaseConnector instance
            tags: list of tags to query (defaults to key enrichment tags)
            qualify: if True, filter out non-meaningful cycles and sort by significance

        Returns:
            List of enriched (and optionally qualified + sorted) cycle dicts
        """
        if not cycles:
            return cycles

        if tags is None:
            tags = ['PT130', 'TT110', 'TT130', 'FT140']

        # Single query covering the full time span of all cycles
        earliest = min(c['start_time'] for c in cycles)
        latest = max(c['end_time'] for c in cycles)

        df = db_connector.get_sensor_data(
            tags, earliest, latest,
            table_override='procdatafloattable_utc_15sec',
        )

        if df.empty:
            return cycles

        enriched = []
        for c in cycles:
            ec = c.copy()
            mask = (df['timestamp'] >= c['start_time']) & (df['timestamp'] <= c['end_time'])
            cycle_df = df[mask]

            if cycle_df.empty:
                enriched.append(ec)
                continue

            # Peak discharge pressure
            pt130 = cycle_df[cycle_df['tag_name'] == 'PT130']
            if not pt130.empty:
                pt130_sorted = pt130.sort_values('timestamp')
                ec['peak_pressure'] = pt130_sorted['value'].max()
                ec['initial_pressure'] = pt130_sorted['value'].iloc[0]

            # Peak temperatures
            for tag in ['TT110', 'TT130']:
                data = cycle_df[cycle_df['tag_name'] == tag]
                if not data.empty:
                    ec[f'peak_{tag}'] = data['value'].max()

            # Average flow
            ft140 = cycle_df[cycle_df['tag_name'] == 'FT140']
            if not ft140.empty:
                ec['avg_flow'] = ft140['value'].mean()

            enriched.append(ec)

        # Qualification: filter out non-meaningful cycles
        if qualify:
            qualified = [c for c in enriched if self._qualify_cycle(c)]
            # Sort by significance and re-number
            return self._sort_and_renumber(qualified) if qualified else []

        return enriched