class ToyLinearModel:
    def __init__(self, weights):
        self.weights = weights

    def predict(self, X):
        return X @ self.weights


if __name__ == "__main__":
    rng = np.random.default_rng(0)
    n, d = 500, 3
    X_train = rng.normal(size=(n, d))
    true_weights = np.array([2.0, -1.0, 0.5])
    y_train = X_train @ true_weights + rng.normal(0, 0.1, n)

    model = ToyLinearModel(true_weights)
    registry = ModelRegistry()
    metadata = ModelMetadata(
        model_name="price_predictor",
        version="v1",
        training_data_hash=hash_dataset(X_train, y_train),
        hyperparameters={"model_type": "linear"},
        eval_metrics={"mse": 0.01},
    )
    registry.register(model, metadata)
    print("Registered versions:",
          registry.list_versions("price_predictor"))

    service = PredictionService(model, expected_n_features=d)
    print(service.handle_request({"features": [1.0, 2.0, 3.0]}))
    print(service.handle_request({"features": [1.0, 2.0]}))  # bad input
    print("Health report:", service.health_report())

    # simulate distribution shift: live feature 0 has drifted to
    # a much wider range than the training distribution.
    live_feature_0_shifted = rng.normal(3, 2, size=200)
    live_feature_0_stable = rng.normal(0, 1, size=200)

    psi_shifted = detect_feature_drift(X_train[:, 0], live_feature_0_shifted)
    psi_stable = detect_feature_drift(X_train[:, 0], live_feature_0_stable)
    print(f"\nPSI (shifted feature): {psi_shifted:.3f}")
    print(f"PSI (stable feature):  {psi_stable:.3f}")
    print("(PSI > 0.2 is a commonly used 'investigate' threshold; "
          "PSI > 0.25 is a commonly used 'significant shift' threshold.)")
