from __future__ import annotations

import sys
import importlib.util
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_peer_surprise_features.py"
_SPEC = importlib.util.spec_from_file_location("enrich_peer_surprise_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_peer_features = _MODULE.compute_peer_features


def test_compute_peer_features_uses_prior_same_sector_peers_only() -> None:
    rows = pd.DataFrame(
        [
            {
                "ticker": "AAA",
                "event_date": "2022-01-10",
                "event_type": "earnings_release",
                "sector": "Technology",
                "earnings_surprise_pct": 10.0,
                "sue_hist_mean_4q": 5.0,
                "sue_hist_pos_rate_4q": 0.75,
                "__split_name": "train",
                "__split_order": 0,
                "__split_row_idx": 0,
                "__event_date_obj": pd.Timestamp("2022-01-10").date(),
            },
            {
                "ticker": "BBB",
                "event_date": "2022-02-10",
                "event_type": "earnings_release",
                "sector": "Technology",
                "earnings_surprise_pct": 20.0,
                "sue_hist_mean_4q": 8.0,
                "sue_hist_pos_rate_4q": 1.0,
                "__split_name": "train",
                "__split_order": 0,
                "__split_row_idx": 1,
                "__event_date_obj": pd.Timestamp("2022-02-10").date(),
            },
            {
                "ticker": "CCC",
                "event_date": "2022-03-10",
                "event_type": "earnings_release",
                "sector": "Technology",
                "earnings_surprise_pct": 25.0,
                "sue_hist_mean_4q": 6.0,
                "sue_hist_pos_rate_4q": 0.5,
                "__split_name": "valid",
                "__split_order": 1,
                "__split_row_idx": 0,
                "__event_date_obj": pd.Timestamp("2022-03-10").date(),
            },
            {
                "ticker": "AAA",
                "event_date": "2022-04-10",
                "event_type": "earnings_release",
                "sector": "Technology",
                "earnings_surprise_pct": 15.0,
                "sue_hist_mean_4q": 7.0,
                "sue_hist_pos_rate_4q": 0.75,
                "__split_name": "test",
                "__split_order": 2,
                "__split_row_idx": 0,
                "__event_date_obj": pd.Timestamp("2022-04-10").date(),
            },
        ]
    )

    enriched = compute_peer_features(rows)

    ccc = enriched.loc[enriched["ticker"] == "CCC"].iloc[0]
    assert ccc["peer_sector_event_count_365d"] == 2.0
    assert ccc["peer_sector_surprise_median_365d"] == 15.0
    assert ccc["peer_relative_surprise_pct_365d"] == 10.0
    assert ccc["peer_sector_sue_hist_mean_4q_median_365d"] == 6.5
    assert ccc["peer_relative_sue_hist_mean_4q_365d"] == -0.5

    aaa_late = enriched.loc[(enriched["ticker"] == "AAA") & (enriched["event_date"] == "2022-04-10")].iloc[0]
    assert aaa_late["peer_sector_event_count_365d"] == 2.0
    assert aaa_late["peer_sector_surprise_median_365d"] == 22.5
    assert aaa_late["peer_relative_surprise_pct_365d"] == -7.5


def test_compute_peer_features_ignores_same_day_and_other_sectors() -> None:
    rows = pd.DataFrame(
        [
            {
                "ticker": "AAA",
                "event_date": "2022-01-10",
                "event_type": "earnings_release",
                "sector": "Technology",
                "earnings_surprise_pct": 10.0,
                "sue_hist_mean_4q": 4.0,
                "sue_hist_pos_rate_4q": 0.5,
                "__split_name": "train",
                "__split_order": 0,
                "__split_row_idx": 0,
                "__event_date_obj": pd.Timestamp("2022-01-10").date(),
            },
            {
                "ticker": "BBB",
                "event_date": "2022-01-10",
                "event_type": "earnings_release",
                "sector": "Technology",
                "earnings_surprise_pct": 12.0,
                "sue_hist_mean_4q": 6.0,
                "sue_hist_pos_rate_4q": 0.75,
                "__split_name": "train",
                "__split_order": 0,
                "__split_row_idx": 1,
                "__event_date_obj": pd.Timestamp("2022-01-10").date(),
            },
            {
                "ticker": "CCC",
                "event_date": "2022-01-20",
                "event_type": "earnings_release",
                "sector": "Healthcare",
                "earnings_surprise_pct": 15.0,
                "sue_hist_mean_4q": 8.0,
                "sue_hist_pos_rate_4q": 1.0,
                "__split_name": "valid",
                "__split_order": 1,
                "__split_row_idx": 0,
                "__event_date_obj": pd.Timestamp("2022-01-20").date(),
            },
        ]
    )

    enriched = compute_peer_features(rows)
    same_day_bbb = enriched.loc[enriched["ticker"] == "BBB"].iloc[0]
    health_ccc = enriched.loc[enriched["ticker"] == "CCC"].iloc[0]

    assert pd.isna(same_day_bbb["peer_sector_surprise_median_365d"])
    assert same_day_bbb["peer_sector_event_count_365d"] == 0.0
    assert pd.isna(health_ccc["peer_sector_surprise_median_365d"])
    assert health_ccc["peer_sector_event_count_365d"] == 0.0