classifier.py

"""Simplified rough sets classifier for fault detection."""

from __future__ import annotations

from datetime import datetime
from typing import TYPE_CHECKING, Optional

import pandas as pd

from config import NON_PZB_WINDOW, PZB
from data_structures import State
from detection_ranges import DEFAULT_INPUT_PARAMETERS, WINDOW_SIZE

if TYPE_CHECKING:
    from data_structures import Channel, DetectionSets
    from data_types import DetectionResults, HistoricalDataPoint


class SimpleRSClassifier:
    """Simplified rough sets fault classifier.

    This classifier performs:
    1. Data smoothing (rolling average)
    2. Differentiation (temporal differences)
    3. Range-based classification (basic rule: if value in range X -> state X)

    """

    calculated_values: dict[str, float]
    data: pd.DataFrame
    detection_sets: DetectionSets
    input_parameters: list[Channel]
    state: State
    window_size: int

    def __init__(
        self,
        detection_sets: DetectionSets,
        input_parameters: Optional[list[Channel]] = None,
        window_size: int = WINDOW_SIZE,
    ) -> None:
        """Initialize the simplified classifier.

        Args:
            detection_sets (DetectionSets): Detection sets containing ranges
            and states
            input_parameters (list[Channel]): List of input sensor channels.
            Defaults to None (DEFAULT_INPUT_PARAMETERS).
            window_size (int): Window size for smoothing and differentiation.
            Defaults to WINDOW_SIZE.

        """
        self.detection_sets = detection_sets
        self.input_parameters = input_parameters or DEFAULT_INPUT_PARAMETERS
        self.window_size = window_size

        # Initialize data storage and reset state
        self.reset()

    def reset(self) -> None:
        """Reset the classifier state and data."""
        self.data = pd.DataFrame()
        self.state = State.NORMAL
        self.calculated_values = {
            f"calculated_value_{i+1}": 0.0 for i in range(5)
        }

    def feed_data_point(
        self, timestamp: datetime, sensor_values: list[float]
    ) -> Optional[State]:
        """Feed a single data point and perform classification.

        Args:
            timestamp (datetime): Timestamp of the measurement
            sensor_values (list[float]): List of sensor values [a1, a2, a3,
            a4, a5]

        Returns:
            Detected state if classification is possible, None otherwise.

        """
        # Create dataframe row
        row_data = {
            param.code_name: value
            for param, value in zip(self.input_parameters, sensor_values)
        }

        new_row = pd.DataFrame([row_data], index=[timestamp])

        # Add to existing data
        if self.data.empty:
            self.data = new_row
        else:
            self.data = pd.concat([self.data, new_row])

        # Sort by timestamp
        self.data.sort_index(inplace=True)

        # Keep only necessary data (sliding window)
        if len(self.data) > self.window_size * 2:
            self.data = self.data.iloc[1:]

        # Calculate detection parameters if we have enough data
        if len(self.data) >= self.window_size * 2:
            self._calculate_detection_parameters()
            return self._classify_current_state()

        return None

    def _calculate_detection_parameters(self) -> None:
        """Calculate smoothed and differentiated values for all parameters."""
        for param in self.input_parameters:
            param_name = param.code_name

            if PZB:
                # Calculate rolling mean (smoothing) and then the difference
                # (differentiation)
                self.data[f"{param_name}_smoothened"] = (
                    self.data[param_name]
                    .rolling(window=self.window_size)
                    .mean()
                )

                self.data[f"{param_name}_differentiated"] = self.data[
                    f"{param_name}_smoothened"
                ].diff(periods=self.window_size)

            else:
                periods = self.window_size if not NON_PZB_WINDOW else 1

                # Calculate differentiation first and then perform smoothing
                self.data[f"{param_name}_differentiated"] = self.data[
                    param_name
                ].diff(periods=periods)
                # ].diff(periods=self.window_size)

                self.data[f"{param_name}_smoothened"] = (
                    self.data[f"{param_name}_differentiated"]
                    .rolling(window=self.window_size)
                    .mean()
                )

    def _classify_current_state(self) -> State:
        """Perform basic range-based classification.

        If calculated value falls in range X, return state X.

        Returns:
            State: Detected state

        """
        # Get the last row with calculated values
        last_row = self.data.iloc[-1]

        # Check if we have valid calculated values
        if pd.isna(last_row.values).any():
            return State.NORMAL

        # Store calculated values for output
        for i, param in enumerate(self.input_parameters):
            param_name = param.code_name

            if PZB:
                diff_col = f"{param_name}_differentiated"
            else:
                diff_col = f"{param_name}_smoothened"

            if diff_col in last_row.index:
                self.calculated_values[f"calculated_value_{i+1}"] = float(
                    last_row[diff_col]
                )
            else:
                self.calculated_values[f"calculated_value_{i+1}"] = 0.0

        # Perform classification for each detection set
        detected_states = []
        detected_state = False

        for detection_set in self.detection_sets.sets:
            param_name = detection_set.channel.code_name

            if PZB:
                diff_col = f"{param_name}_differentiated"
            else:
                diff_col = f"{param_name}_smoothened"

            if diff_col not in last_row.index:
                continue

            differentiated_value = last_row[diff_col]

            # Check each range in the detection set
            for detection_range in detection_set.ranges:
                if detection_range.in_range(differentiated_value):
                    # Basic classification: if value in range -> return state
                    detected_states.append(detection_range.state)
                    detected_state = True
                    break

            # Only detect the state once
            if detected_state:
                break

        # Return the first detected state, or NORMAL if none found
        if detected_states:
            self.state = detected_states[0]
            return self.state

        self.state = State.NORMAL
        return self.state

    def get_calculated_values(self) -> dict[str, float]:
        """Get the current calculated (differentiated) values.

        Returns:
            dict[str, float]: Dictionary of calculated values

        """
        return self.calculated_values.copy()

    def get_current_state(self) -> State:
        """Get the current detected state.

        Returns:
            State: Current detected state

        """
        return self.state

    def process_historical_data(
        self, historical_data: HistoricalDataPoint
    ) -> DetectionResults:
        """Process a batch of historical data.

        Args:
            historical_data (HistoricalDataPoint): List of tuples (timestamp, a1, a2,
            a3, a4, a5, real_state)

        Returns:
            DetectionResults: List of detection results: (timestamp, a1, a2, a3, a4,
            a5, calc_val_1, calc_val_2, calc_val_3, calc_val_4, calc_val_5,
            detected_state, real_state)

        """
        results = []
        self.reset()

        for row in historical_data:
            timestamp = row[0]
            sensor_values = list(row[1:6])  # a1, a2, a3, a4, a5
            real_state = row[6] if len(row) > 6 else 0
            detected_state = self.feed_data_point(timestamp, sensor_values)

            # Only add result if we have a valid detection
            if detected_state is not None:
                calc_values = self.get_calculated_values()

                result = [
                    timestamp,
                    *sensor_values,
                    calc_values["calculated_value_1"],
                    calc_values["calculated_value_2"],
                    calc_values["calculated_value_3"],
                    calc_values["calculated_value_4"],
                    calc_values["calculated_value_5"],
                    float(detected_state.value),
                    float(real_state),
                ]
                results.append(result)

        return results