# Paper-scoped implementation; see SOURCE_PROVENANCE.json.
from __future__ import annotations
import pickle
import time
from dataclasses import dataclass, field
from typing import Any, Literal
import numpy as np
from numpy.typing import NDArray
from sklearn.base import BaseEstimator, ClassifierMixin, clone
from sklearn.ensemble import ExtraTreesClassifier, RandomForestClassifier
from sklearn.linear_model import LogisticRegression, RidgeClassifierCV
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.utils.class_weight import compute_sample_weight
from .features import IntervalDistributionSketch, InvariantPhysicsFeatureExtractor, RandomConvSketch, RandomShapeletSketch, WearableFeatureExtractor
from .hmm import TransitionSmoother
from .probability import estimator_proba, global_classes, labels_from_proba

def _softmax(z: NDArray[np.float64]) -> NDArray[np.float64]:
    z = np.asarray(z, dtype=np.float64)
    z = z - np.max(z, axis=1, keepdims=True)
    e = np.exp(z)
    return e / np.sum(e, axis=1, keepdims=True)

def _classifier(kind: str, *, seed: int, n_jobs: int=-1) -> Any:
    if kind == 'extratrees':
        return ExtraTreesClassifier(n_estimators=120, max_features='sqrt', min_samples_leaf=1, class_weight='balanced', random_state=seed, n_jobs=n_jobs)
    if kind == 'rf':
        return RandomForestClassifier(n_estimators=120, max_features='sqrt', min_samples_leaf=1, class_weight='balanced_subsample', random_state=seed, n_jobs=n_jobs)
    if kind == 'logreg':
        return LogisticRegression(C=2.0, max_iter=2000, class_weight='balanced', solver='lbfgs', random_state=seed)
    if kind == 'ridge':
        return RidgeClassifierCV(alphas=np.logspace(-3, 3, 9))
    raise ValueError(f'Unknown classifier kind: {kind}')

class _HMMMixin:
    use_hmm: bool
    smoother_: TransitionSmoother | None
    classes_: NDArray[np.int64] | None
    trained_classes_: NDArray[np.int64] | None

    def _set_classes(self, y: Any, n_classes: int | None) -> None:
        yy = np.asarray(y, dtype=np.int64)
        self.classes_ = global_classes(yy, n_classes)
        self.trained_classes_ = np.unique(yy)

    def _fit_smoother(self, y: NDArray[np.int64], groups: Any | None, time_index: Any | None) -> None:
        self.smoother_ = None
        if self.use_hmm:
            if self.classes_ is None:
                raise RuntimeError('Global class axis is not initialized')
            self.smoother_ = TransitionSmoother().fit(y, groups=groups, time_index=time_index, n_classes=int(self.classes_.size))

    def _maybe_smooth(self, proba: NDArray[np.float64], groups: Any | None, time_index: Any | None) -> NDArray[np.int64]:
        if self.classes_ is None:
            raise RuntimeError('Global class axis is not initialized')
        if getattr(self, 'smoother_', None) is not None and groups is not None:
            return self.smoother_.predict_from_proba(proba, groups=groups, time_index=time_index)
        return labels_from_proba(proba, self.classes_)

@dataclass
class HarSculptForestHMM(BaseEstimator, ClassifierMixin, _HMMMixin):
    classifier: Literal['extratrees', 'rf', 'logreg'] = 'extratrees'
    use_hmm: bool = True
    segment_sizes: tuple[int, ...] = (4, 8, 16)
    symbolic_bins: int = 6
    seed: int = 42
    n_jobs: int = -1
    feature_extractor_: WearableFeatureExtractor | None = field(default=None, init=False)
    model_: Any = field(default=None, init=False)
    smoother_: TransitionSmoother | None = field(default=None, init=False)
    classes_: NDArray[np.int64] | None = field(default=None, init=False)
    trained_classes_: NDArray[np.int64] | None = field(default=None, init=False)
    fit_seconds_: float = field(default=0.0, init=False)

    def fit(self, X: Any, y: Any, *, groups: Any | None=None, time_index: Any | None=None, n_classes: int | None=None) -> 'HarSculptForestHMM':
        start = time.perf_counter()
        y = np.asarray(y, dtype=np.int64)
        self._set_classes(y, n_classes)
        self.feature_extractor_ = WearableFeatureExtractor(segment_sizes=self.segment_sizes, symbolic_bins=self.symbolic_bins)
        F = self.feature_extractor_.fit_transform(X, y)
        clf = _classifier(self.classifier, seed=self.seed, n_jobs=self.n_jobs)
        if self.classifier == 'logreg':
            self.model_ = Pipeline([('scale', StandardScaler()), ('clf', clf)])
        else:
            self.model_ = clf
        self.model_.fit(F, y)
        self._fit_smoother(y, groups, time_index)
        self.fit_seconds_ = time.perf_counter() - start
        return self

    def predict_proba(self, X: Any) -> NDArray[np.float64]:
        if self.model_ is None or self.feature_extractor_ is None or self.classes_ is None:
            raise RuntimeError('Model is not fitted')
        F = self.feature_extractor_.transform(X)
        return estimator_proba(self.model_, F, self.classes_)

    def predict(self, X: Any, *, groups: Any | None=None, time_index: Any | None=None) -> NDArray[np.int64]:
        return self._maybe_smooth(self.predict_proba(X), groups, time_index)

    def design_profile(self) -> dict[str, Any]:
        return {'name': 'har_sculpt_forest_hmm', 'route': 'cpu', 'operators': ['zscore', 'magnitude', 'jerk', 'multiscale_interval_stats', 'fft_bandpower', 'autocorr', 'symbolic_transition', 'extratrees', 'hmm' if self.use_hmm else 'argmax'], 'motifs': ['orientation_invariance', 'local_periodicity', 'transition_smoothing', 'tail_class_bias']}

@dataclass
class RocketSketchRidgeHMM(BaseEstimator, ClassifierMixin, _HMMMixin):
    n_kernels: int = 384
    classifier: Literal['ridge', 'logreg'] = 'ridge'
    use_hmm: bool = True
    seed: int = 42
    n_jobs: int = -1
    sketch_: RandomConvSketch | None = field(default=None, init=False)
    model_: Any = field(default=None, init=False)
    smoother_: TransitionSmoother | None = field(default=None, init=False)
    classes_: NDArray[np.int64] | None = field(default=None, init=False)
    trained_classes_: NDArray[np.int64] | None = field(default=None, init=False)
    fit_seconds_: float = field(default=0.0, init=False)

    def fit(self, X: Any, y: Any, *, groups: Any | None=None, time_index: Any | None=None, n_classes: int | None=None) -> 'RocketSketchRidgeHMM':
        start = time.perf_counter()
        y = np.asarray(y, dtype=np.int64)
        self._set_classes(y, n_classes)
        self.sketch_ = RandomConvSketch(n_kernels=self.n_kernels, random_state=self.seed)
        F = self.sketch_.fit_transform(X, y)
        clf = _classifier(self.classifier, seed=self.seed, n_jobs=self.n_jobs)
        self.model_ = Pipeline([('scale', StandardScaler()), ('clf', clf)])
        sample_weight = compute_sample_weight('balanced', y) if self.classifier == 'ridge' else None
        if sample_weight is not None:
            self.model_.fit(F, y, clf__sample_weight=sample_weight)
        else:
            self.model_.fit(F, y)
        self._fit_smoother(y, groups, time_index)
        self.fit_seconds_ = time.perf_counter() - start
        return self

    def predict_proba(self, X: Any) -> NDArray[np.float64]:
        if self.model_ is None or self.sketch_ is None or self.classes_ is None:
            raise RuntimeError('Model is not fitted')
        F = self.sketch_.transform(X)
        return estimator_proba(self.model_, F, self.classes_)

    def predict(self, X: Any, *, groups: Any | None=None, time_index: Any | None=None) -> NDArray[np.int64]:
        return self._maybe_smooth(self.predict_proba(X), groups, time_index)

    def design_profile(self) -> dict[str, Any]:
        return {'name': 'rocket_sketch_ridge_hmm', 'route': 'cpu', 'operators': ['zscore', 'random_conv', 'ppv_pool', 'ridge', 'hmm' if self.use_hmm else 'argmax'], 'motifs': ['phase_tolerant_local_shape', 'fast_cpu', 'transition_smoothing']}

@dataclass
class SymbolicIntervalForest(BaseEstimator, ClassifierMixin, _HMMMixin):
    use_hmm: bool = True
    seed: int = 42
    n_jobs: int = -1
    feature_extractor_: WearableFeatureExtractor | None = field(default=None, init=False)
    model_: Any = field(default=None, init=False)
    smoother_: TransitionSmoother | None = field(default=None, init=False)
    classes_: NDArray[np.int64] | None = field(default=None, init=False)
    trained_classes_: NDArray[np.int64] | None = field(default=None, init=False)
    fit_seconds_: float = field(default=0.0, init=False)

    def fit(self, X: Any, y: Any, *, groups: Any | None=None, time_index: Any | None=None, n_classes: int | None=None) -> 'SymbolicIntervalForest':
        start = time.perf_counter()
        y = np.asarray(y, dtype=np.int64)
        self._set_classes(y, n_classes)
        self.feature_extractor_ = WearableFeatureExtractor(include_raw_stats=True, include_magnitude=True, include_jerk=False, include_multiscale=True, segment_sizes=(5, 10, 20), include_spectral=False, include_autocorr=False, include_cross_channel=True, include_symbolic=True, symbolic_bins=8)
        F = self.feature_extractor_.fit_transform(X, y)
        self.model_ = ExtraTreesClassifier(n_estimators=160, max_features='sqrt', class_weight='balanced', random_state=self.seed, n_jobs=self.n_jobs)
        self.model_.fit(F, y)
        self._fit_smoother(y, groups, time_index)
        self.fit_seconds_ = time.perf_counter() - start
        return self

    def predict_proba(self, X: Any) -> NDArray[np.float64]:
        if self.model_ is None or self.feature_extractor_ is None or self.classes_ is None:
            raise RuntimeError('Model is not fitted')
        F = self.feature_extractor_.transform(X)
        return estimator_proba(self.model_, F, self.classes_)

    def predict(self, X: Any, *, groups: Any | None=None, time_index: Any | None=None) -> NDArray[np.int64]:
        return self._maybe_smooth(self.predict_proba(X), groups, time_index)

    def design_profile(self) -> dict[str, Any]:
        return {'name': 'symbolic_interval_forest', 'route': 'cpu', 'operators': ['segment', 'aggregate_stats', 'quantize', 'transition_histogram', 'forest', 'hmm' if self.use_hmm else 'argmax'], 'motifs': ['interpretable_symbols', 'micro_pattern', 'transition_smoothing']}

@dataclass
class HybridCPUEnsemble(BaseEstimator, ClassifierMixin):
    use_hmm: bool = True
    seed: int = 42
    n_jobs: int = -1
    members_: list[Any] = field(default_factory=list, init=False)
    classes_: NDArray[np.int64] | None = field(default=None, init=False)
    trained_classes_: NDArray[np.int64] | None = field(default=None, init=False)
    fit_seconds_: float = field(default=0.0, init=False)

    def fit(self, X: Any, y: Any, *, groups: Any | None=None, time_index: Any | None=None, n_classes: int | None=None) -> 'HybridCPUEnsemble':
        start = time.perf_counter()
        yy = np.asarray(y, dtype=np.int64)
        self.classes_ = global_classes(yy, n_classes)
        self.trained_classes_ = np.unique(yy)
        self.members_ = [HarSculptForestHMM(classifier='extratrees', use_hmm=False, seed=self.seed, n_jobs=self.n_jobs), RocketSketchRidgeHMM(n_kernels=256, classifier='ridge', use_hmm=False, seed=self.seed + 1, n_jobs=self.n_jobs), SymbolicIntervalForest(use_hmm=False, seed=self.seed + 2, n_jobs=self.n_jobs)]
        for m in self.members_:
            m.fit(X, yy, groups=groups, time_index=time_index, n_classes=int(self.classes_.size))
        self.smoother_ = TransitionSmoother().fit(yy, groups=groups, time_index=time_index, n_classes=int(self.classes_.size)) if self.use_hmm else None
        self.fit_seconds_ = time.perf_counter() - start
        return self

    def predict_proba(self, X: Any) -> NDArray[np.float64]:
        if not self.members_ or self.classes_ is None:
            raise RuntimeError('Model is not fitted')
        out = np.mean([m.predict_proba(X) for m in self.members_], axis=0)
        return out / np.maximum(out.sum(axis=1, keepdims=True), 1e-12)

    def predict(self, X: Any, *, groups: Any | None=None, time_index: Any | None=None) -> NDArray[np.int64]:
        P = self.predict_proba(X)
        if getattr(self, 'smoother_', None) is not None and groups is not None:
            return self.smoother_.predict_from_proba(P, groups=groups, time_index=time_index)
        return labels_from_proba(P, self.classes_)

    def design_profile(self) -> dict[str, Any]:
        return {'name': 'hybrid_cpu_ensemble', 'route': 'cpu', 'operators': ['feature_union', 'random_conv', 'symbolic', 'probability_average', 'hmm' if self.use_hmm else 'argmax'], 'motifs': ['multi_view_cpu', 'robustness_by_diversity', 'transition_smoothing']}

@dataclass
class InvariantIntervalShapeForest(BaseEstimator, ClassifierMixin, _HMMMixin):
    classifier: Literal['extratrees', 'rf', 'logreg'] = 'extratrees'
    use_hmm: bool = True
    segment_sizes: tuple[int, ...] = (4, 8, 16)
    symbolic_bins: int = 6
    n_shapelets: int = 32
    seed: int = 42
    n_jobs: int = -1
    physics_: InvariantPhysicsFeatureExtractor | None = field(default=None, init=False)
    shapelets_: RandomShapeletSketch | None = field(default=None, init=False)
    model_: Any = field(default=None, init=False)
    smoother_: TransitionSmoother | None = field(default=None, init=False)
    classes_: NDArray[np.int64] | None = field(default=None, init=False)
    trained_classes_: NDArray[np.int64] | None = field(default=None, init=False)
    fit_seconds_: float = field(default=0.0, init=False)

    def fit(self, X: Any, y: Any, *, groups: Any | None=None, time_index: Any | None=None, n_classes: int | None=None) -> 'InvariantIntervalShapeForest':
        start = time.perf_counter()
        y = np.asarray(y, dtype=np.int64)
        self._set_classes(y, n_classes)
        self.physics_ = InvariantPhysicsFeatureExtractor(segment_sizes=self.segment_sizes, symbolic_bins=int(self.symbolic_bins), random_intervals=24, random_state=self.seed, include_axis_features=True).fit(X, y)
        self.shapelets_ = RandomShapeletSketch(n_shapelets=int(self.n_shapelets), random_state=self.seed + 17).fit(X, y)
        F = np.concatenate([self.physics_.transform(X), self.shapelets_.transform(X)], axis=1)
        clf = _classifier(self.classifier, seed=self.seed, n_jobs=self.n_jobs)
        self.model_ = Pipeline([('scale', StandardScaler()), ('clf', clf)]) if self.classifier == 'logreg' else clf
        self.model_.fit(F, y)
        self._fit_smoother(y, groups, time_index)
        self.fit_seconds_ = time.perf_counter() - start
        return self

    def _features(self, X: Any) -> NDArray[np.float64]:
        if self.physics_ is None or self.shapelets_ is None:
            raise RuntimeError('Model is not fitted')
        return np.concatenate([self.physics_.transform(X), self.shapelets_.transform(X)], axis=1)

    def predict_proba(self, X: Any) -> NDArray[np.float64]:
        if self.model_ is None or self.classes_ is None:
            raise RuntimeError('Model is not fitted')
        return estimator_proba(self.model_, self._features(X), self.classes_)

    def predict(self, X: Any, *, groups: Any | None=None, time_index: Any | None=None) -> NDArray[np.int64]:
        return self._maybe_smooth(self.predict_proba(X), groups, time_index)

    def design_profile(self) -> dict[str, Any]:
        return {'name': 'invariant_interval_shape_forest', 'route': 'cpu', 'operators': ['zscore', 'magnitude', 'jerk', 'gram_eig', 'dyadic_intervals', 'quantile_distribution', 'random_shapelet', 'extratrees', 'hmm' if self.use_hmm else 'argmax'], 'motifs': ['orientation_invariance', 'device_shift_robustness', 'interval_distribution', 'local_shape', 'transition_smoothing']}

@dataclass
class StateRocketIntervalHMM(BaseEstimator, ClassifierMixin, _HMMMixin):
    n_kernels: int = 384
    n_intervals: int = 24
    classifier: Literal['ridge', 'logreg'] = 'ridge'
    use_hmm: bool = True
    seed: int = 42
    n_jobs: int = -1
    rocket_: RandomConvSketch | None = field(default=None, init=False)
    interval_: IntervalDistributionSketch | None = field(default=None, init=False)
    state_: WearableFeatureExtractor | None = field(default=None, init=False)
    model_: Any = field(default=None, init=False)
    smoother_: TransitionSmoother | None = field(default=None, init=False)
    classes_: NDArray[np.int64] | None = field(default=None, init=False)
    trained_classes_: NDArray[np.int64] | None = field(default=None, init=False)
    fit_seconds_: float = field(default=0.0, init=False)

    def fit(self, X: Any, y: Any, *, groups: Any | None=None, time_index: Any | None=None, n_classes: int | None=None) -> 'StateRocketIntervalHMM':
        start = time.perf_counter()
        y = np.asarray(y, dtype=np.int64)
        self._set_classes(y, n_classes)
        self.rocket_ = RandomConvSketch(n_kernels=int(self.n_kernels), random_state=self.seed).fit(X, y)
        self.interval_ = IntervalDistributionSketch(n_random_intervals=int(self.n_intervals), random_state=self.seed + 29).fit(X, y)
        self.state_ = WearableFeatureExtractor(include_raw_stats=False, include_magnitude=True, include_jerk=True, include_multiscale=False, include_spectral=True, include_autocorr=True, include_cross_channel=True, include_symbolic=True, symbolic_bins=6).fit(X, y)
        F = self._features(X)
        clf = _classifier(self.classifier, seed=self.seed, n_jobs=self.n_jobs)
        self.model_ = Pipeline([('scale', StandardScaler()), ('clf', clf)])
        if self.classifier == 'ridge':
            self.model_.fit(F, y, clf__sample_weight=compute_sample_weight('balanced', y))
        else:
            self.model_.fit(F, y)
        self._fit_smoother(y, groups, time_index)
        self.fit_seconds_ = time.perf_counter() - start
        return self

    def _features(self, X: Any) -> NDArray[np.float64]:
        if self.rocket_ is None or self.interval_ is None or self.state_ is None:
            raise RuntimeError('Model is not fitted')
        return np.concatenate([self.rocket_.transform(X), self.interval_.transform(X), self.state_.transform(X)], axis=1)

    def predict_proba(self, X: Any) -> NDArray[np.float64]:
        if self.model_ is None or self.classes_ is None:
            raise RuntimeError('Model is not fitted')
        return estimator_proba(self.model_, self._features(X), self.classes_)

    def predict(self, X: Any, *, groups: Any | None=None, time_index: Any | None=None) -> NDArray[np.int64]:
        return self._maybe_smooth(self.predict_proba(X), groups, time_index)

    def design_profile(self) -> dict[str, Any]:
        return {'name': 'state_rocket_interval_hmm', 'route': 'cpu', 'operators': ['zscore', 'random_conv', 'ppv_pool', 'dyadic_intervals', 'quantile_distribution', 'symbolic_transition', 'ridge', 'hmm' if self.use_hmm else 'argmax'], 'motifs': ['phase_tolerant_local_shape', 'interval_distribution', 'state_dynamics', 'fast_cpu', 'transition_smoothing']}

@dataclass
class PosteriorCPUEnsemble(BaseEstimator, ClassifierMixin):
    use_hmm: bool = True
    seed: int = 42
    n_jobs: int = -1
    weights: tuple[float, ...] = (0.34, 0.33, 0.33)
    members_: list[Any] = field(default_factory=list, init=False)
    smoother_: TransitionSmoother | None = field(default=None, init=False)
    classes_: NDArray[np.int64] | None = field(default=None, init=False)
    trained_classes_: NDArray[np.int64] | None = field(default=None, init=False)
    fit_seconds_: float = field(default=0.0, init=False)

    def fit(self, X: Any, y: Any, *, groups: Any | None=None, time_index: Any | None=None, n_classes: int | None=None) -> 'PosteriorCPUEnsemble':
        start = time.perf_counter()
        yy = np.asarray(y, dtype=np.int64)
        self.classes_ = global_classes(yy, n_classes)
        self.trained_classes_ = np.unique(yy)
        self.members_ = [StateRocketIntervalHMM(n_kernels=256, n_intervals=16, use_hmm=False, seed=self.seed, n_jobs=self.n_jobs), InvariantIntervalShapeForest(n_shapelets=24, use_hmm=False, seed=self.seed + 1, n_jobs=self.n_jobs), SymbolicIntervalForest(use_hmm=False, seed=self.seed + 2, n_jobs=self.n_jobs)]
        for member in self.members_:
            member.fit(X, yy, groups=groups, time_index=time_index, n_classes=int(self.classes_.size))
        self.smoother_ = TransitionSmoother().fit(yy, groups=groups, time_index=time_index, n_classes=int(self.classes_.size)) if self.use_hmm else None
        self.fit_seconds_ = time.perf_counter() - start
        return self

    def predict_proba(self, X: Any) -> NDArray[np.float64]:
        if not self.members_ or self.classes_ is None:
            raise RuntimeError('Model is not fitted')
        probs = [m.predict_proba(X) for m in self.members_]
        w = np.asarray(self.weights, dtype=np.float64)
        if w.size != len(probs):
            w = np.ones((len(probs),), dtype=np.float64)
        w = w / np.maximum(w.sum(), 1e-12)
        out = sum((float(wi) * pi for wi, pi in zip(w, probs)))
        return out / np.maximum(out.sum(axis=1, keepdims=True), 1e-12)

    def predict(self, X: Any, *, groups: Any | None=None, time_index: Any | None=None) -> NDArray[np.int64]:
        P = self.predict_proba(X)
        if self.smoother_ is not None and groups is not None:
            return self.smoother_.predict_from_proba(P, groups=groups, time_index=time_index)
        return labels_from_proba(P, self.classes_)

    def design_profile(self) -> dict[str, Any]:
        return {'name': 'posterior_cpu_ensemble', 'route': 'cpu', 'operators': ['posterior_weighting', 'random_conv', 'invariant_physics', 'symbolic_transition', 'probability_ensemble', 'hmm' if self.use_hmm else 'argmax'], 'motifs': ['multi_view_cpu', 'posterior_over_motifs', 'state_dynamics', 'orientation_invariance', 'transition_smoothing']}

def build_cpu_estimator(spec: dict[str, Any] | str) -> Any:
    if isinstance(spec, str):
        name = spec
        params: dict[str, Any] = {}
    else:
        name = str(spec.get('name', spec.get('algorithm', 'har_sculpt_forest_hmm')))
        params = dict(spec.get('params', {}))
    if name == 'har_sculpt_forest_hmm':
        return HarSculptForestHMM(**params)
    if name == 'rocket_sketch_ridge_hmm':
        return RocketSketchRidgeHMM(**params)
    if name == 'symbolic_interval_forest':
        return SymbolicIntervalForest(**params)
    if name == 'hybrid_cpu_ensemble':
        return HybridCPUEnsemble(**params)
    if name == 'invariant_interval_shape_forest':
        return InvariantIntervalShapeForest(**params)
    if name == 'state_rocket_interval_hmm':
        return StateRocketIntervalHMM(**params)
    if name == 'posterior_cpu_ensemble':
        return PosteriorCPUEnsemble(**params)
    raise ValueError(f'Unknown CPU estimator: {name}')

def estimate_pickle_size_mb(model: Any) -> float:
    return len(pickle.dumps(model)) / (1024.0 * 1024.0)