from __future__ import annotations

import sys
import tempfile
import unittest
from pathlib import Path

import numpy as np
import pandas as pd

TRAINING_ROOT = Path(__file__).resolve().parents[1]
if str(TRAINING_ROOT) not in sys.path:
    sys.path.insert(0, str(TRAINING_ROOT))

from trader_training.onnx_export import LinearHead, export_heads
from trader_training.schemas import FEATURE_ORDER
from trader_training.state_continue_experiment import STATE_FEATURES, _predict_frozen_linear_model, _state_rows_for_age, _train_side_models, _verdict


class StateContinueExperimentTest(unittest.TestCase):
    def test_state_rows_include_required_position_and_frozen_entry_features(self) -> None:
        row = {
            "current_sample_id": "s0",
            "symbol": "BTC-USDT-PERP",
            "current_event_time": pd.Timestamp("2026-01-01T00:05:00Z"),
            "current_open_time_ms": 300_000,
            "side": "LONG",
            "split_id": "fit_inner",
            "walk_forward_fold": 0,
            "time_in_position_minutes": 5,
            "entry_price": 100.0,
            "current_price": 100.1,
            "high_since_entry": 100.2,
            "low_since_entry": 99.95,
            "future_return_bps": 12.0,
            "mae_bps": 3.0,
            "entry_predicted_edge_bps": 8.5,
            "entry_direction_prob": 0.64,
            "add_count": 0.0,
            "minutes_since_last_add": 9999.0,
        }
        for feature_name in FEATURE_ORDER:
            row[feature_name] = 0.0
        frame = pd.DataFrame([row])

        out = _state_rows_for_age(frame, stop_bps=8.0, target_bps=12.0, cost_bps=6.5)

        self.assertEqual(set(STATE_FEATURES), set(STATE_FEATURES).intersection(out.columns))
        self.assertAlmostEqual(5.5, float(out.iloc[0]["expected_continue_edge_bps"]))
        self.assertEqual(1, int(out.iloc[0]["continue_target"]))
        self.assertAlmostEqual(8.5, float(out.iloc[0]["entry_predicted_edge_bps"]))
        self.assertAlmostEqual(0.64, float(out.iloc[0]["entry_direction_prob"]), places=6)
        self.assertAlmostEqual(0.0, float(out.iloc[0]["add_count"]))
        self.assertAlmostEqual(9999.0, float(out.iloc[0]["minutes_since_last_add"]))

    def test_frozen_linear_onnx_weights_are_read_without_row_by_row_runtime(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            model_path = Path(tmp) / "direction.onnx"
            export_heads(
                model_path,
                [
                    LinearHead(
                        "direction",
                        "softmax",
                        np.zeros((len(FEATURE_ORDER), 3), dtype=np.float32),
                        np.array([0.0, 1.0, 2.0], dtype=np.float32),
                    ),
                    LinearHead(
                        "long_expected_net_edge_bps",
                        "identity",
                        np.zeros((len(FEATURE_ORDER), 1), dtype=np.float32),
                        np.array([7.25], dtype=np.float32),
                    ),
                ],
                feature_count=len(FEATURE_ORDER),
            )
            frame = pd.DataFrame({"sample_id": ["s0", "s1"]})
            for feature_name in FEATURE_ORDER:
                frame[feature_name] = 0.0

            out = _predict_frozen_linear_model(
                model_path,
                frame,
                {
                    "direction": ("softmax", ("long_prob", "short_prob", "neutral_prob")),
                    "long_expected_net_edge_bps": ("identity", ("long_edge",)),
                },
            )

            self.assertEqual(["s0", "s1"], out["sample_id"].tolist())
            self.assertTrue(np.allclose(1.0, out[["long_prob", "short_prob", "neutral_prob"]].sum(axis=1)))
            self.assertLess(float(out.iloc[0]["long_prob"]), float(out.iloc[0]["neutral_prob"]))
            self.assertAlmostEqual(7.25, float(out.iloc[0]["long_edge"]), places=6)

    def test_verdict_refuses_state_continue_when_edge_mae_is_not_good_enough(self) -> None:
        results = {}
        for side in ("long", "short"):
            results[f"{side}_market_only"] = {
                "validation_locked": {"continue_auc": 0.61, "edge_mae_vs_constant_ratio": 0.985},
                "latest_stress": {"continue_auc": 0.62, "edge_mae_vs_constant_ratio": 0.984},
                "regressor_converged": True,
            }
            results[f"{side}_market_plus_state"] = {
                "validation_locked": {"continue_auc": 0.63, "edge_mae_vs_constant_ratio": 0.979},
                "latest_stress": {"continue_auc": 0.64, "edge_mae_vs_constant_ratio": 0.978},
                "regressor_converged": True,
            }

        verdict = _verdict(results)

        self.assertEqual("NOT_READY_FOR_FORMAL_CHAIN", verdict["status"])
        self.assertTrue(any("above 0.97" in reason for reason in verdict["reasons"]))

    def test_train_side_models_supports_ridge_regressor_diagnostic(self) -> None:
        rows = []
        for split_id in ("fit_inner", "tune_inner", "validation_locked", "latest_stress"):
            for index, target in enumerate((0, 1)):
                row = {
                    "sample_id": f"{split_id}-{index}",
                    "symbol": "BTC-USDT-PERP",
                    "event_time": pd.Timestamp("2026-01-01T00:00:00Z") + pd.Timedelta(minutes=len(rows)),
                    "split_id": split_id,
                    "position_side": "LONG",
                    "continue_target": target,
                    "expected_continue_edge_bps": -3.0 if target == 0 else 6.0,
                }
                for feature_name in FEATURE_ORDER:
                    row[feature_name] = float(index)
                for feature_name in STATE_FEATURES:
                    row[feature_name] = float(index)
                rows.append(row)
        frame = pd.DataFrame(rows)

        metrics, predictions = _train_side_models(
            frame,
            "LONG",
            [*FEATURE_ORDER, *STATE_FEATURES],
            regressor_kind="ridge",
            ridge_alpha=1.0,
            regression_target_clip_bps=5.0,
        )

        self.assertEqual("ridge", metrics["regressor_kind"])
        self.assertEqual(5.0, metrics["regression_target_clip_bps"])
        self.assertTrue(metrics["regressor_converged"])
        self.assertEqual(8, len(predictions))
        self.assertIn("time_in_position_minutes", predictions.columns)


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