tests/test_predict_response_denormalized.py

import tempfile
import unittest
import json
from pathlib import Path

import numpy as np
from unittest.mock import patch

from bucklake_ai.config import Settings
from bucklake_ai.schemas import PredictRequest
from bucklake_ai.service import InferenceService


class PredictResponseDenormalizedTests(unittest.TestCase):
    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        root_dir = Path(self.temp_dir.name)
        data_dir = root_dir / "data"
        data_dir.mkdir(parents=True, exist_ok=True)
        (data_dir / "bl_symbol.csv").write_text(
            "id,symbol,market,company_name,sector,industry,country,currency,ipo_date,market_cap\n",
            encoding="utf-8",
        )
        self.scaler_path = root_dir / "scalers_v21_rolling_7d.json"
        self.scaler_path.write_text(
            json.dumps(
                {
                    "index": {
                        "mean": [11630.00801763615, 11632.439510618715, 11699.49961104118, 11556.167112385792, 1855572862.157339, 5344228574915.616],
                        "scale": [10451.767239272729, 10453.620729932496, 10507.837853482002, 10393.161849307528, 1700927598.7351134, 8624911182529.906],
                    },
                    "stock": {
                        "mean": [69.22109597071513, 69.27929937205207, 70.12196645745344, 68.34393667735944, 13.586773233899697, 0.0],
                        "scale": [611.5967864850386, 613.200961031293, 621.5358163904571, 602.4584422610227, 2.2932616337759755, 1.0],
                    },
                }
            ),
            encoding="utf-8",
        )
        self.settings = Settings(
            app_name="BuckLakeAI",
            app_version="2.0.0",
            model_version="anchored-path-v1",
            input_contract_version="v2.1.0",
            root_dir=root_dir,
            hf_repo_id="parkerjj/BuckLake-Stock-Model",
            hf_model_filename="stock_prediction_model_anchored-path-v1.keras",
            hf_scaler_filename="scalers_v21_rolling_7d.json",
            hf_cache_dir=root_dir / ".cache" / "huggingface",
            hf_token=None,
            model_weights_path=root_dir / ".cache" / "huggingface" / "stock_prediction_model_anchored-path-v1.keras",
            text_encoder_model="test-encoder",
            text_encoder_device="cpu",
            text_encoder_max_seq_length=512,
            enable_finbert_sentiment=False,
            finbert_model_name="ProsusAI/finbert",
            preload_model=False,
            preload_text_encoder=False,
            scaler_artifact_path=self.scaler_path,
        )

    def tearDown(self):
        self.temp_dir.cleanup()

    def _request(self) -> PredictRequest:
        bars = [
            {
                "open": 125.0 + i,
                "close": 126.0 + i,
                "high": 127.0 + i,
                "low": 124.0 + i,
                "volume": 1000.0 + i,
                "amount": 0.0,
            }
            for i in range(30)
        ]
        return PredictRequest.model_validate(
            {
                "symbol": "BX",
                "published_at": "2026-05-11T07:41:38Z",
                "market_bars": {
                    "stock": bars,
                    "inx": bars,
                    "dj": bars,
                    "ixic": bars,
                    "ndx": bars,
                },
            }
        )

    def _v2_request(self) -> PredictRequest:
        bars = [
            {
                "open": 100.0 + i,
                "close": 101.0 + i,
                "high": 102.0 + i,
                "low": 99.0 + i,
                "volume": 1000.0 + i,
                "amount": 0.0,
            }
            for i in range(30)
        ]
        return PredictRequest.model_validate(
            {
                "symbol": "BX",
                "published_at": "2026-05-11T07:41:38Z",
                "text_session": "post_market",
                "market_bars": {
                    "stock": bars,
                    "inx": bars,
                    "dj": bars,
                    "ixic": bars,
                    "ndx": bars,
                },
            }
        )

    def test_build_response_renders_anchored_path_prices_from_raw_request_close(self):
        settings = self.settings.__class__(
            **{**self.settings.__dict__, "model_version": "anchored_path_v1"}
        )
        service = InferenceService(settings)
        request = self._request()
        zero_1 = np.zeros((3, 1), dtype=np.float32)
        predictions = [
            zero_1,
            zero_1,
            np.full((3, 1), 0.01, dtype=np.float32),
            np.full((3, 1), 0.02, dtype=np.float32),
            np.zeros((3, 4), dtype=np.float32),
            np.array([[0.7], [0.6], [0.55]], dtype=np.float32),
            np.array([[0.1], [0.2], [0.3]], dtype=np.float32),
            np.array([[0.5], [0.5], [0.5]], dtype=np.float32),
        ]

        response = service._build_response(request, predictions, elapsed_ms=1.0)

        last_close = request.market_bars.stock[-1].close
        self.assertAlmostEqual(response.derived.predicted_stock_close[0], last_close, places=5)
        self.assertGreater(response.derived.predicted_stock_high[0], last_close)
        self.assertLess(response.derived.predicted_stock_low[0], last_close)
        self.assertAlmostEqual(response.derived.predicted_direction_probability[0], 0.7, places=6)
        self.assertAlmostEqual(response.derived.predicted_direction_probability[1], 0.6, places=6)
        self.assertAlmostEqual(response.derived.predicted_direction_probability[2], 0.55, places=6)

    def test_anchored_path_raw_low_high_match_rendered_stock_prices(self):
        settings = self.settings.__class__(
            **{**self.settings.__dict__, "model_version": "anchored_path_v1"}
        )
        service = InferenceService(settings)
        request = self._request()
        zero_1 = np.zeros((3, 1), dtype=np.float32)
        predictions = [
            zero_1,
            zero_1,
            np.full((3, 1), 0.01, dtype=np.float32),
            np.full((3, 1), 0.02, dtype=np.float32),
            np.zeros((3, 4), dtype=np.float32),
            np.array([[0.7], [0.6], [0.55]], dtype=np.float32),
            np.array([[0.1], [0.2], [0.3]], dtype=np.float32),
            np.array([[0.5], [0.5], [0.5]], dtype=np.float32),
        ]

        response = service._build_response(request, predictions, elapsed_ms=1.0)

        self.assertEqual(
            response.outputs.stock_low,
            [[row[3]] for row in response.outputs.stock_price_output],
        )
        self.assertEqual(
            response.outputs.stock_high,
            [[row[2]] for row in response.outputs.stock_price_output],
        )

    def test_build_response_uses_output_names_when_predictions_are_mapped(self):
        settings = self.settings.__class__(
            **{**self.settings.__dict__, "model_version": "anchored_path_v1"}
        )
        service = InferenceService(settings)
        request = self._request()
        zero_1 = np.zeros((3, 1), dtype=np.float32)
        predictions = {
            "stock_direction": np.array([[0.7], [0.6], [0.55]], dtype=np.float32),
            "text_gate": np.array([[0.5], [0.6], [0.7]], dtype=np.float32),
            "risk_score": np.array([[0.1], [0.2], [0.3]], dtype=np.float32),
            "rendered_ohlc_prices": np.zeros((3, 4), dtype=np.float32),
            "lower_spans": np.full((3, 1), 0.02, dtype=np.float32),
            "upper_spans": np.full((3, 1), 0.01, dtype=np.float32),
            "open_log_gaps": zero_1,
            "cum_close_log_returns": zero_1,
        }

        response = service._build_response(request, predictions, elapsed_ms=1.0)

        self.assertAlmostEqual(response.derived.predicted_direction_probability[0], 0.7, places=6)
        np.testing.assert_allclose(
            response.outputs.risk_score,
            [[0.1], [0.2], [0.3]],
            rtol=1e-6,
        )
        np.testing.assert_allclose(
            response.outputs.text_gate,
            [[0.5], [0.6], [0.7]],
            rtol=1e-6,
        )

    def test_anchored_path_exposes_only_real_raw_heads(self):
        settings = self.settings.__class__(
            **{**self.settings.__dict__, "model_version": "anchored_path_v1"}
        )
        service = InferenceService(settings)
        request = self._request()
        zero_1 = np.zeros((3, 1), dtype=np.float32)
        risk_score = np.array([[0.1], [0.2], [0.3]], dtype=np.float32)
        text_gate = np.array([[0.5], [0.6], [0.7]], dtype=np.float32)
        predictions = [
            zero_1,
            zero_1,
            np.full((3, 1), 0.01, dtype=np.float32),
            np.full((3, 1), 0.02, dtype=np.float32),
            np.zeros((3, 4), dtype=np.float32),
            np.array([[0.7], [0.6], [0.55]], dtype=np.float32),
            risk_score,
            text_gate,
        ]

        response = service._build_response(request, predictions, elapsed_ms=1.0)

        self.assertEqual(response.outputs.stock_volume_output, [])
        np.testing.assert_allclose(
            response.outputs.risk_score,
            [[0.1], [0.2], [0.3]],
            rtol=1e-6,
        )
        np.testing.assert_allclose(
            response.outputs.text_gate,
            [[0.5], [0.6], [0.7]],
            rtol=1e-6,
        )

    def test_anchored_path_v2_renders_index_path_outputs_from_request_anchors(self):
        settings = self.settings.__class__(
            **{**self.settings.__dict__, "model_version": "anchored-path-v2"}
        )
        service = InferenceService(settings)
        request = self._v2_request()
        zero_1 = np.zeros((3, 1), dtype=np.float32)
        index_path = np.array(
            [
                [np.log(1.01), np.log(1.02), np.log(1.03), np.log(1.04)],
                [0.0, 0.0, 0.0, 0.0],
                [np.log(0.99), np.log(0.98), np.log(1.01), np.log(1.02)],
            ],
            dtype=np.float32,
        )
        predictions = {
            "cum_close_log_returns": zero_1,
            "open_log_gaps": zero_1,
            "upper_spans": np.full((3, 1), 0.01, dtype=np.float32),
            "lower_spans": np.full((3, 1), 0.02, dtype=np.float32),
            "rendered_ohlc_prices": np.zeros((3, 4), dtype=np.float32),
            "stock_direction": np.array([[0.7], [0.6], [0.55]], dtype=np.float32),
            "risk_score": np.array([[0.1], [0.2], [0.3]], dtype=np.float32),
            "text_gate": np.array([[0.5], [0.6], [0.7]], dtype=np.float32),
            "index_inx_path_output": index_path,
            "index_dj_path_output": index_path,
            "index_ixic_path_output": index_path,
            "index_ndx_path_output": index_path,
        }

        response = service._build_response(request, predictions, elapsed_ms=1.0)

        anchor = request.market_bars.inx[-1].close
        expected_open = anchor * np.exp(index_path[0, 0])
        expected_close = anchor * np.exp(index_path[0, 1])
        expected_high = max(expected_open, expected_close) * np.exp(index_path[0, 2])
        expected_low = min(expected_open, expected_close) * np.exp(-index_path[0, 3])
        np.testing.assert_allclose(
            response.outputs.index_inx_output[0],
            [expected_open, expected_close, expected_high, expected_low],
            rtol=1e-6,
        )
        self.assertEqual(len(response.outputs.index_dj_output), 3)
        self.assertEqual(len(response.outputs.index_ixic_output[0]), 4)
        self.assertEqual(len(response.outputs.index_ndx_output), 3)

    def test_anchored_path_v2_uses_previous_close_for_pre_market_index_output_anchor(self):
        settings = self.settings.__class__(
            **{**self.settings.__dict__, "model_version": "anchored-path-v2"}
        )
        service = InferenceService(settings)
        request = PredictRequest.model_validate(
            {
                **self._v2_request().model_dump(),
                "text_session": "pre_market",
            }
        )
        zero_1 = np.zeros((3, 1), dtype=np.float32)
        index_path = np.zeros((3, 4), dtype=np.float32)
        predictions = {
            "cum_close_log_returns": zero_1,
            "open_log_gaps": zero_1,
            "upper_spans": np.full((3, 1), 0.01, dtype=np.float32),
            "lower_spans": np.full((3, 1), 0.02, dtype=np.float32),
            "rendered_ohlc_prices": np.zeros((3, 4), dtype=np.float32),
            "stock_direction": np.array([[0.7], [0.6], [0.55]], dtype=np.float32),
            "risk_score": np.array([[0.1], [0.2], [0.3]], dtype=np.float32),
            "text_gate": np.array([[0.5], [0.6], [0.7]], dtype=np.float32),
            "index_inx_path_output": index_path,
            "index_dj_path_output": index_path,
            "index_ixic_path_output": index_path,
            "index_ndx_path_output": index_path,
        }

        response = service._build_response(request, predictions, elapsed_ms=1.0)

        self.assertAlmostEqual(
            response.outputs.index_inx_output[0][1],
            request.market_bars.inx[-2].close,
            places=6,
        )

    def test_anchored_path_v1_keeps_empty_index_outputs(self):
        settings = self.settings.__class__(
            **{**self.settings.__dict__, "model_version": "anchored_path_v1"}
        )
        service = InferenceService(settings)
        request = self._request()
        zero_1 = np.zeros((3, 1), dtype=np.float32)
        predictions = [
            zero_1,
            zero_1,
            np.full((3, 1), 0.01, dtype=np.float32),
            np.full((3, 1), 0.02, dtype=np.float32),
            np.zeros((3, 4), dtype=np.float32),
            np.array([[0.7], [0.6], [0.55]], dtype=np.float32),
            np.array([[0.1], [0.2], [0.3]], dtype=np.float32),
            np.array([[0.5], [0.5], [0.5]], dtype=np.float32),
        ]

        response = service._build_response(request, predictions, elapsed_ms=1.0)

        self.assertEqual(response.outputs.index_inx_output, [])
        self.assertEqual(response.outputs.index_dj_output, [])
        self.assertEqual(response.outputs.index_ixic_output, [])
        self.assertEqual(response.outputs.index_ndx_output, [])

    def test_prediction_summary_uses_service_logger(self):
        settings = self.settings.__class__(
            **{**self.settings.__dict__, "model_version": "anchored_path_v1"}
        )
        service = InferenceService(settings)
        request = self._request()
        zero_1 = np.zeros((3, 1), dtype=np.float32)
        predictions = [
            zero_1,
            zero_1,
            np.full((3, 1), 0.01, dtype=np.float32),
            np.full((3, 1), 0.02, dtype=np.float32),
            np.zeros((3, 4), dtype=np.float32),
            np.array([[0.7], [0.6], [0.55]], dtype=np.float32),
            np.array([[0.1], [0.2], [0.3]], dtype=np.float32),
            np.array([[0.5], [0.5], [0.5]], dtype=np.float32),
        ]

        with patch("bucklake_ai.service.logger.info") as info:
            service._build_response(request, predictions, elapsed_ms=1.234)

        info.assert_called_once()
        self.assertEqual(
            info.call_args.args[0],
            "Prediction for symbol %s completed in %.3f ms | base %.2f | %s",
        )
        self.assertEqual(info.call_args.args[1], "BX")
        self.assertEqual(info.call_args.args[3], 155.0)
        self.assertEqual(
            info.call_args.args[4],
            "day1 155.00 (+0.00%) day2 155.00 (+0.00%) day3 155.00 (+0.00%)",
        )

    def test_build_response_logs_anchored_path_prediction_summary(self):
        settings = self.settings.__class__(
            **{**self.settings.__dict__, "model_version": "anchored_path_v1"}
        )
        service = InferenceService(settings)
        request = self._request()
        zero_1 = np.zeros((3, 1), dtype=np.float32)
        predictions = [
            zero_1,
            zero_1,
            np.full((3, 1), 0.01, dtype=np.float32),
            np.full((3, 1), 0.02, dtype=np.float32),
            np.zeros((3, 4), dtype=np.float32),
            np.array([[0.7], [0.6], [0.55]], dtype=np.float32),
            np.array([[0.1], [0.2], [0.3]], dtype=np.float32),
            np.array([[0.5], [0.5], [0.5]], dtype=np.float32),
        ]

        with self.assertLogs(level="INFO") as logs:
            service._build_response(request, predictions, elapsed_ms=1.234)

        self.assertIn(
            "Prediction for symbol BX completed in 1.234 ms | base 155.00 | "
            "day1 155.00 (+0.00%) day2 155.00 (+0.00%) day3 155.00 (+0.00%)",
            "\n".join(logs.output),
        )


if __name__ == "__main__":
    unittest.main()