51def load_trips(
 52    path: str | os.PathLike,
 53    *,
 54    schema: dict[str, pl.DataType] | None = None,
 55) -> pl.DataFrame:
 56    """
 57    Load telematics trip data from CSV or Parquet, validate schema, and
 58    return a normalised Polars DataFrame.
 59
 60    Parameters
 61    ----------
 62    path:
 63        Path to a CSV or Parquet file, or a directory containing Parquet
 64        files (all will be read and concatenated).
 65    schema:
 66        Optional column name overrides for non-standard source files.
 67        Maps source column names to the standard column names. For example:
 68        ``{"gps_speed": "speed_kmh", "accel_x": "acceleration_ms2"}``.
 69
 70    Returns
 71    -------
 72    pl.DataFrame
 73        Normalised DataFrame with standard column names. The ``timestamp``
 74        column is cast to ``Datetime[us, UTC]``. Missing optional columns
 75        are added with null values. Rows are sorted by ``trip_id`` then
 76        ``timestamp``.
 77
 78    Raises
 79    ------
 80    ValueError
 81        If any required column is missing after applying the schema mapping.
 82    FileNotFoundError
 83        If the path does not exist.
 84
 85    Examples
 86    --------
 87    >>> df = load_trips("trips.csv")
 88    >>> df = load_trips("trips.parquet", schema={"gps_speed": "speed_kmh"})
 89    """
 90    path = Path(path)
 91    if not path.exists():
 92        raise FileNotFoundError(f"Path not found: {path}")
 93
 94    if path.is_dir():
 95        parquet_files = list(path.glob("*.parquet"))
 96        if not parquet_files:
 97            raise ValueError(f"No Parquet files found in directory: {path}")
 98        df = pl.concat([pl.read_parquet(f) for f in sorted(parquet_files)])
 99    elif path.suffix.lower() in {".parquet", ".pq"}:
100        df = pl.read_parquet(path)
101    elif path.suffix.lower() in {".csv", ".tsv"}:
102        df = pl.read_csv(path, try_parse_dates=True)
103    else:
104        raise ValueError(
105            f"Unsupported file type '{path.suffix}'. Use CSV or Parquet."
106        )
107
108    # Apply column renames from schema mapping
109    if schema:
110        df = df.rename({src: dst for src, dst in schema.items() if src in df.columns})
111
112    # Validate required columns
113    missing = [c for c in REQUIRED_COLUMNS if c not in df.columns]
114    if missing:
115        raise ValueError(
116            f"Missing required columns: {missing}. "
117            f"Use the 'schema' parameter to map source column names."
118        )
119
120    # Ensure timestamp is datetime
121    if df["timestamp"].dtype == pl.Utf8 or df["timestamp"].dtype == pl.String:
122        df = df.with_columns(
123            pl.col("timestamp").str.to_datetime(
124                format=None, strict=False
125            )
126        )
127
128    if df["timestamp"].dtype not in (pl.Datetime("us", "UTC"), pl.Datetime("ns", "UTC")):
129        # Cast to microsecond UTC if not already timezone-aware
130        if hasattr(df["timestamp"].dtype, "time_zone") and df["timestamp"].dtype.time_zone:
131            df = df.with_columns(
132                pl.col("timestamp").dt.convert_time_zone("UTC").dt.cast_time_unit("us")
133            )
134        else:
135            df = df.with_columns(
136                pl.col("timestamp")
137                .dt.replace_time_zone("UTC")
138                .dt.cast_time_unit("us")
139            )
140
141    # Add missing optional columns as nulls
142    for col in OPTIONAL_COLUMNS:
143        if col not in df.columns:
144            if col == "driver_id":
145                df = df.with_columns(pl.lit("unknown").alias("driver_id"))
146            else:
147                df = df.with_columns(pl.lit(None).cast(pl.Float64).alias(col))
148
149    # Enforce numeric types for key columns
150    numeric_cols = ["latitude", "longitude", "speed_kmh", "acceleration_ms2", "heading_deg"]
151    cast_exprs = [
152        pl.col(c).cast(pl.Float64) for c in numeric_cols if c in df.columns
153    ]
154    if cast_exprs:
155        df = df.with_columns(cast_exprs)
156
157    # Sort for deterministic processing
158    df = df.sort(["trip_id", "timestamp"])
159
160    return df

32def clean_trips(df: pl.DataFrame) -> pl.DataFrame:
33    """
34    Clean and enrich a raw telematics DataFrame.
35
36    Applies the following steps per trip, in order:
37
38    1. Remove GPS jump rows (speed > 250 km/h).
39    2. Interpolate speed for small gaps (< 5 s within a trip).
40    3. Derive acceleration from speed differences if the
41       ``acceleration_ms2`` column is null or absent.
42    4. Derive jerk (m/s³) as the first difference of acceleration.
43    5. Classify each observation into a road type
44       (``urban`` / ``rural`` / ``motorway``) from speed.
45
46    Parameters
47    ----------
48    df:
49        Raw trip DataFrame as returned by :func:`~insurance_telematics.load_trips`.
50        Must contain ``trip_id``, ``timestamp``, and ``speed_kmh``.
51
52    Returns
53    -------
54    pl.DataFrame
55        Cleaned DataFrame with additional columns:
56        ``acceleration_ms2`` (if derived), ``jerk_ms3``, ``road_type``.
57
58    Notes
59    -----
60    The road type heuristic (urban < 50 km/h, rural 50-100, motorway > 100)
61    is a proxy. It works well in aggregate but misclassifies individual
62    observations — a momentary slow-down on a motorway reads as urban.
63    Trip-level road type statistics (fraction of time in each band) are
64    more robust than per-row labels.
65    """
66    df = _remove_gps_jumps(df)
67    df = _interpolate_speed_gaps(df)
68    df = _derive_acceleration(df)
69    df = _derive_jerk(df)
70    df = _classify_road_type(df)
71    return df

 48def extract_trip_features(df: pl.DataFrame) -> pl.DataFrame:
 49    """
 50    Compute one summary row per trip from cleaned 1Hz observations.
 51
 52    Parameters
 53    ----------
 54    df:
 55        Cleaned trip DataFrame as returned by
 56        :func:`~insurance_telematics.clean_trips`. Must contain:
 57        ``trip_id``, ``timestamp``, ``speed_kmh``, ``acceleration_ms2``,
 58        ``road_type``.
 59
 60    Returns
 61    -------
 62    pl.DataFrame
 63        One row per ``trip_id`` with columns:
 64
 65        ``trip_id``, ``distance_km``, ``duration_min``, ``mean_speed_kmh``,
 66        ``p95_speed_kmh``, ``speed_variation_coeff``, ``harsh_braking_rate``,
 67        ``harsh_accel_rate``, ``harsh_cornering_rate``, ``speeding_fraction``,
 68        ``night_driving_fraction``, ``urban_fraction``, ``driver_id``
 69        (if present in input).
 70
 71    Notes
 72    -----
 73    Distance is estimated as the integral of speed over time at 1Hz:
 74    sum(speed_kmh) / 3600. This is a reasonable proxy when GPS coordinates
 75    are noisy or unavailable for Haversine distance computation.
 76    """
 77    _check_required(df)
 78
 79    has_driver_id = "driver_id" in df.columns
 80    has_heading = "heading_deg" in df.columns and df["heading_deg"].null_count() < len(df)
 81
 82    # Compute per-row flags in a single with_columns call for efficiency
 83    flags = [
 84        # Harsh braking flag
 85        (pl.col("acceleration_ms2") < _HARSH_BRAKE_THRESHOLD)
 86        .cast(pl.Int32)
 87        .alias("_harsh_brake"),
 88        # Harsh acceleration flag
 89        (pl.col("acceleration_ms2") > _HARSH_ACCEL_THRESHOLD)
 90        .cast(pl.Int32)
 91        .alias("_harsh_accel"),
 92        # Night flag
 93        pl.col("timestamp")
 94        .dt.hour()
 95        .is_in(list(_NIGHT_HOURS))
 96        .cast(pl.Int32)
 97        .alias("_night"),
 98        # Urban flag
 99        (pl.col("road_type") == "urban")
100        .cast(pl.Int32)
101        .alias("_urban"),
102        # Speeding flag — road-type-specific threshold
103        pl.when(pl.col("road_type") == "urban")
104        .then(pl.col("speed_kmh") > _SPEED_LIMIT_URBAN)
105        .when(pl.col("road_type") == "rural")
106        .then(pl.col("speed_kmh") > _SPEED_LIMIT_RURAL)
107        .otherwise(pl.col("speed_kmh") > _SPEED_LIMIT_MOTORWAY)
108        .cast(pl.Int32)
109        .alias("_speeding"),
110        # Distance increment: speed_kmh / 3600 = km per second at 1Hz
111        (pl.col("speed_kmh") / 3600.0).alias("_dist_km_increment"),
112    ]
113
114    if has_heading:
115        # Cornering proxy: heading change rate (deg/s) × speed (km/h)
116        # Proportional to lateral acceleration; calibrated empirically.
117        flags.append(
118            (
119                pl.col("heading_deg")
120                .diff()
121                .over("trip_id")
122                .abs()
123                .fill_null(0.0)
124                * pl.col("speed_kmh")
125                * _HARSH_CORNER_DEG_PER_S_KMPH
126                > 1.0
127            )
128            .cast(pl.Int32)
129            .alias("_harsh_corner"),
130        )
131    else:
132        flags.append(pl.lit(0).cast(pl.Int32).alias("_harsh_corner"))
133
134    df = df.with_columns(flags)
135
136    # Aggregate per trip
137    group_cols = ["trip_id"]
138    if has_driver_id:
139        # Carry driver_id through — take first value (it's constant per trip)
140        agg_extra = [pl.col("driver_id").first().alias("driver_id")]
141    else:
142        agg_extra = []
143
144    agg = df.group_by("trip_id").agg(
145        [
146            # Basic trip stats
147            pl.len().alias("n_obs"),
148            (pl.col("_dist_km_increment").sum()).alias("distance_km"),
149            (pl.len() / 60.0).alias("duration_min"),
150            pl.col("speed_kmh").mean().alias("mean_speed_kmh"),
151            pl.col("speed_kmh").quantile(0.95).alias("p95_speed_kmh"),
152            pl.col("speed_kmh").std().alias("_speed_std"),
153            # Event counts (will divide by distance below)
154            pl.col("_harsh_brake").sum().alias("_n_harsh_brake"),
155            pl.col("_harsh_accel").sum().alias("_n_harsh_accel"),
156            pl.col("_harsh_corner").sum().alias("_n_harsh_corner"),
157            # Fraction-based features
158            (pl.col("_speeding").mean()).alias("speeding_fraction"),
159            (pl.col("_night").mean()).alias("night_driving_fraction"),
160            (pl.col("_urban").mean()).alias("urban_fraction"),
161        ]
162        + agg_extra
163    )
164
165    # Derive final rate features
166    agg = agg.with_columns(
167        [
168            # Speed coefficient of variation — std / mean, clamped to avoid div/0
169            (
170                pl.col("_speed_std")
171                / (pl.col("mean_speed_kmh") + 1e-6)
172            ).alias("speed_variation_coeff"),
173            # Event rates per km — clamp distance to 0.01 km minimum
174            (
175                pl.col("_n_harsh_brake")
176                / (pl.col("distance_km").clip(0.01))
177            ).alias("harsh_braking_rate"),
178            (
179                pl.col("_n_harsh_accel")
180                / (pl.col("distance_km").clip(0.01))
181            ).alias("harsh_accel_rate"),
182            (
183                pl.col("_n_harsh_corner")
184                / (pl.col("distance_km").clip(0.01))
185            ).alias("harsh_cornering_rate"),
186        ]
187    )
188
189    # Select and order final output columns
190    output_cols = [
191        "trip_id",
192        "distance_km",
193        "duration_min",
194        "mean_speed_kmh",
195        "p95_speed_kmh",
196        "speed_variation_coeff",
197        "harsh_braking_rate",
198        "harsh_accel_rate",
199        "harsh_cornering_rate",
200        "speeding_fraction",
201        "night_driving_fraction",
202        "urban_fraction",
203    ]
204    if has_driver_id:
205        output_cols.append("driver_id")
206
207    return agg.select(output_cols).sort("trip_id")

 57class DrivingStateHMM:
 58    """
 59    Discrete-time Gaussian HMM for classifying driving trip sequences into
 60    latent driving states (e.g. cautious / normal / aggressive).
 61
 62    Wraps ``hmmlearn.GaussianHMM`` with state-ordering normalisation and
 63    a driver-level feature extraction step.
 64
 65    Parameters
 66    ----------
 67    n_states:
 68        Number of latent states. Default 3. With three states the literature
 69        consistently produces a cautious / normal / aggressive partition.
 70    features:
 71        Trip-level feature columns to use as HMM observations. Defaults to
 72        mean speed, speed variation, harsh braking rate, harsh acceleration.
 73    n_iter:
 74        Maximum EM iterations for hmmlearn (default 200).
 75    covariance_type:
 76        Covariance type passed to GaussianHMM (default "diag").
 77    random_state:
 78        Random seed for reproducibility.
 79
 80    Examples
 81    --------
 82    >>> model = DrivingStateHMM(n_states=3)
 83    >>> model.fit(trip_features_df)
 84    >>> states = model.predict_states(trip_features_df)
 85    >>> driver_features = model.driver_state_features(trip_features_df, states)
 86    """
 87
 88    def __init__(
 89        self,
 90        n_states: int = 3,
 91        features: list[str] | None = None,
 92        n_iter: int = 200,
 93        covariance_type: str = "diag",
 94        random_state: int = 42,
 95    ) -> None:
 96        if not _HMMLEARN_AVAILABLE:
 97            raise ImportError(
 98                "hmmlearn is required for DrivingStateHMM. "
 99                "Install it with: pip install hmmlearn"
100            )
101
102        self.n_states = n_states
103        self.features = features or _DEFAULT_HMM_FEATURES
104        self.n_iter = n_iter
105        self.covariance_type = covariance_type
106        self.random_state = random_state
107
108        self._model: Optional[hmmlearn_hmm.GaussianHMM] = None
109        self._state_order: Optional[np.ndarray] = None  # indices sorted by mean speed
110        self.is_fitted: bool = False
111
112    def fit(self, trip_features: pl.DataFrame) -> "DrivingStateHMM":
113        """
114        Fit the HMM to a DataFrame of trip-level scalar features.
115
116        Each row is treated as one observation. For a per-driver sequence,
117        pass only that driver's trips (ordered by timestamp). For a
118        portfolio-level fit, pass all trips — hmmlearn handles the full
119        sequence as one long chain, which is appropriate for population-level
120        state estimation.
121
122        Parameters
123        ----------
124        trip_features:
125            DataFrame from :func:`~insurance_telematics.extract_trip_features`.
126            Must contain the columns listed in ``self.features``.
127
128        Returns
129        -------
130        self
131        """
132        X = self._to_matrix(trip_features)
133        X = self._standardise(X, fit=True)
134
135        model = hmmlearn_hmm.GaussianHMM(
136            n_components=self.n_states,
137            covariance_type=self.covariance_type,
138            n_iter=self.n_iter,
139            random_state=self.random_state,
140            verbose=False,
141        )
142        with warnings.catch_warnings():
143            warnings.simplefilter("ignore")
144            model.fit(X)
145
146        self._model = model
147
148        # Determine state ordering by the mean speed dimension (feature index 0)
149        # so that state 0 is always the lowest-speed (most cautious) state
150        speed_means = model.means_[:, 0]
151        self._state_order = np.argsort(speed_means)
152        self._state_rank = np.empty_like(self._state_order)
153        self._state_rank[self._state_order] = np.arange(self.n_states)
154
155        self.is_fitted = True
156        return self
157
158    def predict_states(self, trip_features: pl.DataFrame) -> np.ndarray:
159        """
160        Decode the most likely state sequence for a set of trips.
161
162        Parameters
163        ----------
164        trip_features:
165            Same format as used for ``fit()``.
166
167        Returns
168        -------
169        np.ndarray of shape (n_trips,)
170            Integer state labels 0..n_states-1, where 0 = most cautious
171            and n_states-1 = most aggressive.
172        """
173        self._check_fitted()
174        X = self._to_matrix(trip_features)
175        X = self._standardise(X, fit=False)
176        raw_states = self._model.predict(X)
177        # Remap to ordered states
178        return self._state_rank[raw_states]
179
180    def predict_state_probs(self, trip_features: pl.DataFrame) -> np.ndarray:
181        """
182        Compute posterior state probabilities for each trip observation.
183
184        Parameters
185        ----------
186        trip_features:
187            Same format as used for ``fit()``.
188
189        Returns
190        -------
191        np.ndarray of shape (n_trips, n_states)
192            Posterior probabilities, summing to 1.0 across states per row.
193            Column ordering matches the ordered states (0 = most cautious).
194        """
195        self._check_fitted()
196        X = self._to_matrix(trip_features)
197        X = self._standardise(X, fit=False)
198        _, posteriors = self._model.score_samples(X)
199        # Reorder columns to match state ordering
200        return posteriors[:, self._state_order]
201
202    def driver_state_features(
203        self,
204        trip_features: pl.DataFrame,
205        states: np.ndarray,
206    ) -> pl.DataFrame:
207        """
208        Compute per-driver HMM-derived features for use as GLM covariates.
209
210        Parameters
211        ----------
212        trip_features:
213            Trip-level features DataFrame. Must contain ``driver_id``.
214        states:
215            State sequence as returned by :meth:`predict_states`.
216
217        Returns
218        -------
219        pl.DataFrame
220            One row per driver with columns:
221
222            - ``state_{k}_fraction`` for k in 0..n_states-1
223            - ``mean_transition_rate`` — average state transitions per km
224            - ``state_entropy`` — Shannon entropy of state distribution
225            - ``driver_id``
226
227        Notes
228        -----
229        State fractions are the primary actuarial risk features. Following
230        Jiang & Shi (2024), the fraction of time spent in the highest-risk
231        state (state ``n_states - 1``) is the most predictive of claim
232        frequency.
233        """
234        if "driver_id" not in trip_features.columns:
235            raise ValueError(
236                "trip_features must contain 'driver_id' for driver-level aggregation."
237            )
238
239        df = trip_features.with_columns(pl.Series("_state", states.tolist()))
240
241        # Count state occurrences per driver
242        driver_rows = []
243        driver_ids = df["driver_id"].unique().sort().to_list()
244
245        for driver_id in driver_ids:
246            mask = df["driver_id"] == driver_id
247            driver_df = df.filter(mask)
248            driver_states = driver_df["_state"].to_numpy()
249            n_trips = len(driver_states)
250
251            # State fractions
252            state_counts = np.bincount(driver_states, minlength=self.n_states)
253            state_fractions = state_counts / max(n_trips, 1)
254
255            # Transition rate: number of state changes per km
256            n_transitions = np.sum(np.diff(driver_states) != 0)
257            total_km = driver_df["distance_km"].sum() if "distance_km" in driver_df.columns else n_trips
258            transition_rate = n_transitions / max(total_km, 0.01)
259
260            # State entropy (Shannon)
261            state_entropy = float(scipy_entropy(state_fractions))
262
263            row = {"driver_id": driver_id}
264            for k in range(self.n_states):
265                row[f"state_{k}_fraction"] = float(state_fractions[k])
266            row["mean_transition_rate"] = float(transition_rate)
267            row["state_entropy"] = float(state_entropy)
268
269            driver_rows.append(row)
270
271        return pl.DataFrame(driver_rows)
272
273    def _to_matrix(self, trip_features: pl.DataFrame) -> np.ndarray:
274        """Extract feature columns as float64 numpy array."""
275        missing = [c for c in self.features if c not in trip_features.columns]
276        if missing:
277            raise ValueError(f"Feature columns missing from DataFrame: {missing}")
278        return trip_features.select(self.features).to_numpy().astype(np.float64)
279
280    def _standardise(self, X: np.ndarray, *, fit: bool) -> np.ndarray:
281        """Z-score standardise features. Fit mean/std on first call."""
282        if fit:
283            self._mean = np.nanmean(X, axis=0)
284            self._std = np.nanstd(X, axis=0) + 1e-8
285        return (X - self._mean) / self._std
286
287    def _check_fitted(self) -> None:
288        if not self.is_fitted:
289            raise RuntimeError("Call fit() before predict_states().")

292class ContinuousTimeHMM:
293    """
294    Continuous-time HMM for irregularly-sampled telematics observations.
295
296    Uses a generator matrix Q where the transition probability over an
297    interval Δt is P(Δt) = expm(Q × Δt). This handles variable trip
298    lengths and inter-observation time gaps without resampling.
299
300    Implementation uses an EM algorithm:
301    - E-step: forward-backward with time-varying transition matrices P(Δt_i)
302    - M-step: update Q via expected transition counts and holding times
303
304    Parameters
305    ----------
306    n_states:
307        Number of latent states (default 3).
308    features:
309        Feature columns to use as Gaussian emissions per state.
310    n_iter:
311        Number of EM iterations (default 100).
312    tol:
313        Convergence tolerance on log-likelihood (default 1e-4).
314    random_state:
315        Seed for parameter initialisation.
316
317    Notes
318    -----
319    The M-step update for Q follows the standard uniformisation approach:
320    Q_ij (i≠j) ∝ E[N_ij] / E[T_i], where N_ij is the expected number of
321    i→j transitions and T_i is the expected total holding time in state i.
322
323    Numerical stability: scipy.linalg.expm is used for matrix exponential.
324    For typical generator matrices from driving data (rates on the order of
325    0.01-1.0 transitions/minute), expm is stable and accurate.
326
327    Examples
328    --------
329    >>> model = ContinuousTimeHMM(n_states=3)
330    >>> model.fit(trip_features_df, time_deltas=delta_array)
331    >>> states = model.predict_states(trip_features_df, time_deltas=delta_array)
332    """
333
334    def __init__(
335        self,
336        n_states: int = 3,
337        features: list[str] | None = None,
338        n_iter: int = 100,
339        tol: float = 1e-4,
340        random_state: int = 42,
341    ) -> None:
342        if not _SCIPY_AVAILABLE:
343            raise ImportError(
344                "scipy is required for ContinuousTimeHMM. "
345                "Install it with: pip install scipy"
346            )
347
348        self.n_states = n_states
349        self.features = features or _DEFAULT_HMM_FEATURES
350        self.n_iter = n_iter
351        self.tol = tol
352        self.random_state = random_state
353
354        self.Q_: Optional[np.ndarray] = None  # Generator matrix (n_states × n_states)
355        self.means_: Optional[np.ndarray] = None    # Emission means (n_states × n_features)
356        self.covars_: Optional[np.ndarray] = None   # Emission variances (diagonal)
357        self.pi_: Optional[np.ndarray] = None       # Initial state distribution
358        self.is_fitted: bool = False
359        self._state_order: Optional[np.ndarray] = None
360
361    def fit(
362        self,
363        trip_features: pl.DataFrame,
364        time_deltas: np.ndarray | None = None,
365    ) -> "ContinuousTimeHMM":
366        """
367        Fit the CTHMM using the EM algorithm.
368
369        Parameters
370        ----------
371        trip_features:
372            Trip-level features DataFrame (one row per trip, not per second).
373        time_deltas:
374            Array of shape (n_trips,) giving the time interval in minutes
375            between consecutive trips for the same driver. Use a large value
376            (e.g. 1440.0 = 24 hours) between trips from different sessions.
377            If None, all intervals are set to 1.0 (equivalent to discrete HMM
378            with unit time steps).
379
380        Returns
381        -------
382        self
383        """
384        X = self._to_matrix(trip_features)
385        X = self._standardise(X, fit=True)
386        n_obs = X.shape[0]
387
388        if time_deltas is None:
389            time_deltas = np.ones(n_obs, dtype=np.float64)
390        else:
391            time_deltas = np.asarray(time_deltas, dtype=np.float64)
392            if len(time_deltas) != n_obs:
393                raise ValueError(
394                    f"time_deltas length {len(time_deltas)} != n_obs {n_obs}"
395                )
396        # Clamp to small positive to avoid expm(Q * 0)
397        time_deltas = np.clip(time_deltas, 1e-6, None)
398
399        rng = np.random.default_rng(self.random_state)
400
401        # Initialise parameters
402        self.pi_ = np.ones(self.n_states) / self.n_states
403        self.Q_ = self._init_generator(rng)
404        self.means_, self.covars_ = self._init_emissions(X, rng)
405
406        log_likelihood = -np.inf
407        for iteration in range(self.n_iter):
408            # E-step
409            gammas, xis, log_lik_new = self._e_step(X, time_deltas)
410
411            # M-step
412            self._m_step(X, gammas, xis, time_deltas)
413
414            # Convergence check
415            if abs(log_lik_new - log_likelihood) < self.tol:
416                break
417            log_likelihood = log_lik_new
418
419        # Order states by mean of the first feature (speed proxy)
420        self._state_order = np.argsort(self.means_[:, 0])
421        self._state_rank = np.empty_like(self._state_order)
422        self._state_rank[self._state_order] = np.arange(self.n_states)
423
424        self.is_fitted = True
425        return self
426
427    def predict_states(
428        self,
429        trip_features: pl.DataFrame,
430        time_deltas: np.ndarray | None = None,
431    ) -> np.ndarray:
432        """
433        Viterbi decoding of the most likely state sequence.
434
435        Parameters
436        ----------
437        trip_features:
438            Same format as used in ``fit()``.
439        time_deltas:
440            Inter-observation time intervals, same semantics as in ``fit()``.
441
442        Returns
443        -------
444        np.ndarray of shape (n_trips,)
445            State labels 0..n_states-1 with state 0 = most cautious.
446        """
447        self._check_fitted()
448        X = self._to_matrix(trip_features)
449        X = self._standardise(X, fit=False)
450        n_obs = X.shape[0]
451
452        if time_deltas is None:
453            time_deltas = np.ones(n_obs)
454        time_deltas = np.clip(np.asarray(time_deltas, dtype=np.float64), 1e-6, None)
455
456        raw_states = self._viterbi(X, time_deltas)
457        return self._state_rank[raw_states]
458
459    def driver_state_features(
460        self,
461        trip_features: pl.DataFrame,
462        states: np.ndarray,
463    ) -> pl.DataFrame:
464        """
465        Compute driver-level HMM features. Same interface as
466        :meth:`DrivingStateHMM.driver_state_features`.
467        """
468        # Delegate to the same logic — identical to discrete case
469        if "driver_id" not in trip_features.columns:
470            raise ValueError("trip_features must contain 'driver_id'.")
471
472        df = trip_features.with_columns(pl.Series("_state", states.tolist()))
473        driver_rows = []
474        driver_ids = df["driver_id"].unique().sort().to_list()
475
476        for driver_id in driver_ids:
477            driver_df = df.filter(pl.col("driver_id") == driver_id)
478            driver_states = driver_df["_state"].to_numpy()
479            n_trips = len(driver_states)
480
481            state_counts = np.bincount(driver_states, minlength=self.n_states)
482            state_fractions = state_counts / max(n_trips, 1)
483
484            n_transitions = np.sum(np.diff(driver_states) != 0)
485            total_km = (
486                driver_df["distance_km"].sum()
487                if "distance_km" in driver_df.columns
488                else float(n_trips)
489            )
490            transition_rate = n_transitions / max(total_km, 0.01)
491            state_entropy = float(scipy_entropy(state_fractions))
492
493            row = {"driver_id": driver_id}
494            for k in range(self.n_states):
495                row[f"state_{k}_fraction"] = float(state_fractions[k])
496            row["mean_transition_rate"] = float(transition_rate)
497            row["state_entropy"] = float(state_entropy)
498            driver_rows.append(row)
499
500        return pl.DataFrame(driver_rows)
501
502    # ------------------------------------------------------------------
503    # Private EM implementation
504    # ------------------------------------------------------------------
505
506    def _init_generator(self, rng: np.random.Generator) -> np.ndarray:
507        """Initialise a valid generator matrix Q with small random rates."""
508        n = self.n_states
509        Q = rng.uniform(0.01, 0.1, size=(n, n))
510        np.fill_diagonal(Q, 0.0)
511        np.fill_diagonal(Q, -Q.sum(axis=1))
512        return Q
513
514    def _init_emissions(
515        self, X: np.ndarray, rng: np.random.Generator
516    ) -> tuple[np.ndarray, np.ndarray]:
517        """Initialise Gaussian emission parameters by k-means-like split."""
518        n, d = X.shape
519        # Use equally-spaced quantile points as initial means
520        quantiles = np.linspace(0.1, 0.9, self.n_states)
521        means = np.array(
522            [np.quantile(X, q, axis=0) for q in quantiles], dtype=np.float64
523        )
524        # Perturb slightly to break symmetry
525        means += rng.normal(0, 0.05, size=means.shape)
526        covars = np.ones((self.n_states, d), dtype=np.float64)
527        return means, covars
528
529    def _transition_matrix(self, dt: float) -> np.ndarray:
530        """P(dt) = expm(Q * dt), clamped to be a valid stochastic matrix."""
531        P = matrix_expm(self.Q_ * dt)
532        # Numerical correction: clip to [0,1] and renormalise rows
533        P = np.clip(P.real, 0.0, 1.0)
534        row_sums = P.sum(axis=1, keepdims=True)
535        row_sums = np.where(row_sums < 1e-12, 1.0, row_sums)
536        return P / row_sums
537
538    def _emission_log_prob(self, X: np.ndarray) -> np.ndarray:
539        """
540        Diagonal Gaussian log-emission probability.
541
542        Returns array of shape (n_obs, n_states).
543        """
544        n_obs = X.shape[0]
545        log_probs = np.zeros((n_obs, self.n_states))
546        for k in range(self.n_states):
547            diff = X - self.means_[k]
548            var = self.covars_[k]
549            log_probs[:, k] = -0.5 * np.sum(
550                np.log(2 * np.pi * var) + diff ** 2 / var, axis=1
551            )
552        return log_probs
553
554    def _e_step(
555        self, X: np.ndarray, dts: np.ndarray
556    ) -> tuple[np.ndarray, np.ndarray, float]:
557        """
558        Forward-backward algorithm with time-varying transition matrices.
559
560        Returns
561        -------
562        gammas : (n_obs, n_states)
563        xis    : (n_obs-1, n_states, n_states)
564        log_likelihood : float
565        """
566        n_obs, n_states = X.shape[0], self.n_states
567        log_emit = self._emission_log_prob(X)  # (n_obs, n_states)
568
569        # Forward pass (log-scale for numerical stability)
570        log_alpha = np.full((n_obs, n_states), -np.inf)
571        log_alpha[0] = np.log(self.pi_ + 1e-300) + log_emit[0]
572
573        for t in range(1, n_obs):
574            P = self._transition_matrix(dts[t - 1])
575            log_P = np.log(P + 1e-300)
576            # log_alpha[t, j] = log_emit[t, j] + logsumexp_k(log_alpha[t-1, k] + log_P[k, j])
577            for j in range(n_states):
578                log_alpha[t, j] = log_emit[t, j] + _logsumexp(
579                    log_alpha[t - 1] + log_P[:, j]
580                )
581
582        log_likelihood = _logsumexp(log_alpha[-1])
583
584        # Backward pass
585        log_beta = np.zeros((n_obs, n_states))  # log(1) = 0 at final step
586        for t in range(n_obs - 2, -1, -1):
587            P = self._transition_matrix(dts[t])
588            log_P = np.log(P + 1e-300)
589            for i in range(n_states):
590                log_beta[t, i] = _logsumexp(
591                    log_P[i, :] + log_emit[t + 1] + log_beta[t + 1]
592                )
593
594        # Posterior state probabilities gamma
595        # Normalise each row by log P(observations) to get P(z_t | x).
596        log_gamma = log_alpha + log_beta
597        for t in range(n_obs):
598            log_gamma[t] -= _logsumexp(log_gamma[t])
599        gammas = np.exp(log_gamma)
600        gammas = np.clip(gammas, 0.0, None)
601        gammas /= gammas.sum(axis=1, keepdims=True)  # numerical guard
602
603        # Pairwise posteriors xi
604        # Normalise by P(observations) = exp(log_likelihood), not per-timestep.
605        # Per-timestep normalisation is incorrect: it makes each t contribute
606        # equally regardless of probability mass, causing EM to converge to a
607        # spurious fixed point.
608        xis = np.zeros((n_obs - 1, n_states, n_states))
609        for t in range(n_obs - 1):
610            P = self._transition_matrix(dts[t])
611            for i in range(n_states):
612                for j in range(n_states):
613                    xis[t, i, j] = np.exp(
614                        log_alpha[t, i]
615                        + np.log(P[i, j] + 1e-300)
616                        + log_emit[t + 1, j]
617                        + log_beta[t + 1, j]
618                        - log_likelihood
619                    )
620
621        return gammas, xis, log_likelihood
622
623    def _m_step(
624        self,
625        X: np.ndarray,
626        gammas: np.ndarray,
627        xis: np.ndarray,
628        dts: np.ndarray,
629    ) -> None:
630        """Update parameters given posteriors."""
631        n_obs, n_states = gammas.shape
632        n_features = X.shape[1]
633
634        # Update initial distribution
635        self.pi_ = gammas[0] / gammas[0].sum()
636
637        # Update emission parameters
638        gamma_sum = gammas.sum(axis=0)  # (n_states,)
639        for k in range(n_states):
640            w = gammas[:, k]  # (n_obs,)
641            w_sum = w.sum() + 1e-300
642            self.means_[k] = (w[:, None] * X).sum(axis=0) / w_sum
643            diff = X - self.means_[k]
644            self.covars_[k] = (w[:, None] * diff ** 2).sum(axis=0) / w_sum + 1e-6
645
646        # Update generator matrix Q
647        # Expected transition counts: E[N_ij] = sum_t xi[t, i, j]
648        expected_transitions = xis.sum(axis=0)  # (n_states, n_states)
649
650        # Expected holding times: E[T_i] = sum_t gamma[t, i] * dt[t]
651        holding_times = np.zeros(n_states)
652        for t in range(n_obs - 1):
653            holding_times += gammas[t] * dts[t]
654        holding_times = np.clip(holding_times, 1e-6, None)
655
656        # New Q: Q_ij = E[N_ij] / E[T_i] for i≠j
657        Q_new = np.zeros((n_states, n_states))
658        for i in range(n_states):
659            for j in range(n_states):
660                if i != j:
661                    Q_new[i, j] = expected_transitions[i, j] / holding_times[i]
662        np.fill_diagonal(Q_new, -Q_new.sum(axis=1))
663        self.Q_ = Q_new
664
665    def _viterbi(self, X: np.ndarray, dts: np.ndarray) -> np.ndarray:
666        """Viterbi decoding returning most likely state sequence."""
667        n_obs, n_states = X.shape[0], self.n_states
668        log_emit = self._emission_log_prob(X)
669
670        log_delta = np.full((n_obs, n_states), -np.inf)
671        psi = np.zeros((n_obs, n_states), dtype=int)
672        log_delta[0] = np.log(self.pi_ + 1e-300) + log_emit[0]
673
674        for t in range(1, n_obs):
675            P = self._transition_matrix(dts[t - 1])
676            log_P = np.log(P + 1e-300)
677            for j in range(n_states):
678                scores = log_delta[t - 1] + log_P[:, j]
679                best = int(np.argmax(scores))
680                log_delta[t, j] = scores[best] + log_emit[t, j]
681                psi[t, j] = best
682
683        # Backtrack
684        states = np.zeros(n_obs, dtype=int)
685        states[-1] = int(np.argmax(log_delta[-1]))
686        for t in range(n_obs - 2, -1, -1):
687            states[t] = psi[t + 1, states[t + 1]]
688
689        return states
690
691    def _to_matrix(self, trip_features: pl.DataFrame) -> np.ndarray:
692        missing = [c for c in self.features if c not in trip_features.columns]
693        if missing:
694            raise ValueError(f"Feature columns missing: {missing}")
695        return trip_features.select(self.features).to_numpy().astype(np.float64)
696
697    def _standardise(self, X: np.ndarray, *, fit: bool) -> np.ndarray:
698        if fit:
699            self._mean = np.nanmean(X, axis=0)
700            self._std = np.nanstd(X, axis=0) + 1e-8
701        return (X - self._mean) / self._std
702
703    def _check_fitted(self) -> None:
704        if not self.is_fitted:
705            raise RuntimeError("Call fit() before predict_states().")