from __future__ import annotations

import importlib.util
import sys
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

_SCRIPT_PATH = Path(__file__).resolve().parents[2] / "scripts" / "enrich_catalyst_persistence_features.py"
_SPEC = importlib.util.spec_from_file_location("enrich_catalyst_persistence_features", _SCRIPT_PATH)
assert _SPEC and _SPEC.loader
_MODULE = importlib.util.module_from_spec(_SPEC)
sys.modules[_SPEC.name] = _MODULE
_SPEC.loader.exec_module(_MODULE)
compute_catalyst_persistence_features = _MODULE.compute_catalyst_persistence_features


def test_compute_catalyst_persistence_features_counts_only_prior_same_ticker_events() -> None:
    rows = pd.DataFrame(
        [
            {
                "ticker": "AAA",
                "event_date": "2022-01-10",
                "event_type": "earnings_release",
                "__split_name": "train",
                "__split_order": 0,
                "__split_row_idx": 0,
                "__event_date_obj": pd.Timestamp("2022-01-10").date(),
            },
            {
                "ticker": "AAA",
                "event_date": "2022-01-25",
                "event_type": "management_change",
                "__split_name": "train",
                "__split_order": 0,
                "__split_row_idx": 1,
                "__event_date_obj": pd.Timestamp("2022-01-25").date(),
            },
            {
                "ticker": "AAA",
                "event_date": "2022-02-20",
                "event_type": "material_contract",
                "__split_name": "valid",
                "__split_order": 1,
                "__split_row_idx": 0,
                "__event_date_obj": pd.Timestamp("2022-02-20").date(),
            },
            {
                "ticker": "BBB",
                "event_date": "2022-02-20",
                "event_type": "earnings_release",
                "__split_name": "valid",
                "__split_order": 1,
                "__split_row_idx": 1,
                "__event_date_obj": pd.Timestamp("2022-02-20").date(),
            },
        ]
    )

    enriched = compute_catalyst_persistence_features(rows)
    second = enriched.loc[(enriched["ticker"] == "AAA") & (enriched["event_date"] == "2022-01-25")].iloc[0]
    third = enriched.loc[(enriched["ticker"] == "AAA") & (enriched["event_date"] == "2022-02-20")].iloc[0]
    other = enriched.loc[enriched["ticker"] == "BBB"].iloc[0]

    assert second["prior_catalyst_count_20d"] == 1.0
    assert second["prior_catalyst_type_diversity_20d"] == 1.0
    assert third["prior_catalyst_count_20d"] == 0.0
    assert third["prior_catalyst_count_60d"] == 2.0
    assert third["prior_catalyst_type_diversity_60d"] == 2.0
    assert other["prior_catalyst_count_60d"] == 0.0
    assert other["prior_catalyst_type_diversity_60d"] == 0.0