test_did_model.py

import os
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

import pandas as pd

import did_model
import map_did


EXPECTED_TREAT_LIST = [
    "米林县",
    "墨脱县",
    "林芝县",
]

EXPECTED_EXCLUDE_LIST = [
    "波密县",
    "工布江达县",
]

EXPECTED_CONTROL_LIST = [
    "察隅县",
    "边坝县",
    "嘉黎县",
    "洛隆县",
    "八宿县",
    "朗县",
    "比如县",
    "索县",
    "隆子县",
    "丁青县",
    "加查县",
]


class ResearchObjectTests(unittest.TestCase):
    def test_research_groups_match_expected_lists(self):
        self.assertEqual(did_model.ACTIVE_PROFILE_KEY, "project_corridor")
        self.assertEqual(did_model.RESEARCH_OBJECT_LABEL, "工程走廊暴露样本")
        self.assertEqual(did_model.TREAT_LIST, EXPECTED_TREAT_LIST)
        self.assertEqual(did_model.EXCLUDE_LIST, EXPECTED_EXCLUDE_LIST)
        self.assertEqual(did_model.CONTROL_LIST, EXPECTED_CONTROL_LIST)
        self.assertEqual(
            did_model.ALL_COUNTIES,
            EXPECTED_TREAT_LIST + EXPECTED_EXCLUDE_LIST + EXPECTED_CONTROL_LIST,
        )
        self.assertEqual(map_did.MAP_TREAT_LIST, EXPECTED_TREAT_LIST)
        self.assertEqual(map_did.MAP_EXCLUDE_LIST, EXPECTED_EXCLUDE_LIST)
        self.assertEqual(map_did.MAP_CONTROL_LIST, EXPECTED_CONTROL_LIST)
        self.assertEqual(did_model.TREAT_LIST, map_did.MAP_TREAT_LIST)
        self.assertEqual(did_model.CONTROL_LIST, map_did.MAP_CONTROL_LIST)

    def test_policy_shock_date_uses_march_2021_major_project(self):
        self.assertEqual((did_model.POST_YEAR, did_model.POST_MONTH), (2021, 3))

    def test_build_panel_marks_new_treatment_counties(self):
        df = pd.DataFrame(
            {
                "区县": ["林芝县", "察隅县"],
                "ym": pd.PeriodIndex(["2020-01", "2020-01"], freq="M"),
                "is_infra": [0, 0],
            }
        )

        panel = did_model.build_panel(df)

        treat_flag = panel.loc[
            (panel["county"] == "林芝县") & (panel["ym"] == pd.Period("2020-01", freq="M")),
            "treat",
        ].iloc[0]
        control_flag = panel.loc[
            (panel["county"] == "察隅县") & (panel["ym"] == pd.Period("2020-01", freq="M")),
            "treat",
        ].iloc[0]

        self.assertEqual(treat_flag, 1)
        self.assertEqual(control_flag, 0)

    def test_normalize_county_name_handles_common_aliases(self):
        self.assertEqual(did_model.normalize_county_name("米林市"), "米林县")
        self.assertEqual(did_model.normalize_county_name("乃东区"), "乃东县")
        self.assertEqual(did_model.normalize_county_name("巴宜区"), "林芝县")

    def test_is_infrastructure_related_targets_engineering_chain(self):
        infra_record = {
            "行业门类": "租赁和商务服务业",
            "行业大类": "租赁业",
            "行业中类": "机械设备经营租赁",
            "行业小类": "建筑工程机械与设备经营租赁",
            "经营范围": "一般项目：机械设备租赁；建筑工程机械与设备租赁。",
        }
        non_infra_record = {
            "行业门类": "住宿和餐饮业",
            "行业大类": "餐饮业",
            "行业中类": "正餐服务",
            "行业小类": "正餐服务",
            "经营范围": "许可项目：餐饮服务。",
        }

        self.assertTrue(did_model.is_infrastructure_related(infra_record))
        self.assertFalse(did_model.is_infrastructure_related(non_infra_record))

    def test_build_phase_did_terms_split_start_and_followthrough_windows(self):
        panel = pd.DataFrame(
            {
                "year": [2020, 2021, 2022, 2023, 2025],
                "month": [12, 3, 11, 1, 7],
                "treat": [1, 1, 1, 1, 0],
            }
        )

        phased = did_model.build_phase_did_terms(panel)

        self.assertEqual(phased["did_start"].tolist(), [0, 1, 1, 0, 0])
        self.assertEqual(phased["did_followthrough"].tolist(), [0, 0, 0, 1, 0])


class FigureBehaviorTests(unittest.TestCase):
    def test_map_extent_focuses_on_sample_bounds(self):
        xlim, ylim = map_did._compute_plot_extent(
            sample_bounds=(90.4677, 27.5554, 98.7531, 30.6686),
            clamp_bounds=(82.0, 27.0, 98.0, 32.0),
        )

        self.assertEqual(xlim, (89.7, 98.0))
        self.assertEqual(ylim, (27.2, 31.1))

    def test_minor_river_files_cover_levels_three_to_five(self):
        self.assertEqual(
            map_did.MINOR_RIVER_FILES,
            ["三级河流.shp", "四级河流.shp", "五级河流.shp"],
        )

    def test_map_selection_uses_explicit_100km_buffer(self):
        self.assertEqual(map_did.CORE_TREAT_KM, 30)
        self.assertEqual(map_did.EXCLUDE_KM, 80)
        self.assertEqual(map_did.CONTROL_KM, 200)

    def test_cut_through_coordinates_match_updated_alignment(self):
        self.assertEqual(
            map_did.CUT_THROUGH_COORDS,
            [(94.85, 29.61), (95.10, 29.20)],
        )

    def test_plot_trend_restores_two_panel_layout(self):
        panel = pd.DataFrame(
            {
                "county": ["A", "A", "A", "A", "B", "B", "B", "B"],
                "ym": pd.PeriodIndex(
                    ["2021-01", "2021-02", "2021-03", "2021-04", "2021-01", "2021-02", "2021-03", "2021-04"],
                    freq="M",
                ),
                "year": [2021] * 8,
                "month": [1, 2, 3, 4, 1, 2, 3, 4],
                "ym_str": ["2021-01", "2021-02", "2021-03", "2021-04"] * 2,
                "treat": [1, 1, 1, 1, 0, 0, 0, 0],
                "post": [0, 0, 1, 1, 0, 0, 1, 1],
                "n_reg": [10.0, 20.0, 50.0, 60.0, 20.0, 20.0, 25.0, 25.0],
            }
        )

        captured = {}

        def fake_savefig(fig, *args, **kwargs):
            captured["figure"] = fig

        with patch("matplotlib.figure.Figure.savefig", new=fake_savefig), patch(
            "matplotlib.pyplot.close", new=lambda fig=None: None
        ):
            did_model.plot_trend(panel)

        titles = [ax.get_title() for ax in captured["figure"].axes]
        self.assertEqual(len(captured["figure"].axes), 2)
        self.assertTrue(any("图2a" in title for title in titles))
        self.assertTrue(any("图2b" in title for title in titles))
        trend_ax = next(ax for ax in captured["figure"].axes if "图2a" in ax.get_title())
        diff_ax = next(ax for ax in captured["figure"].axes if "图2b" in ax.get_title())
        self.assertGreaterEqual(len(trend_ax.lines), 6)
        self.assertGreaterEqual(len(trend_ax.collections), 1)
        self.assertGreaterEqual(len(diff_ax.lines), 2)
        self.assertFalse(any("局部放大" in title for title in titles))
        summary_text = next(
            text for text in trend_ax.texts if "净增差" in text.get_text()
        )
        self.assertGreater(summary_text.get_position()[0], 0.5)

    def test_map_palette_uses_consistent_cool_warm_split(self):
        self.assertEqual(map_did.COLOR_CFG["treat_core"], "#245C5A")
        self.assertEqual(map_did.COLOR_CFG["exclude_ring"], "#CDBE8C")
        self.assertEqual(map_did.COLOR_CFG["control_low"], "#8FB7A2")
        self.assertEqual(map_did.COLOR_CFG["other"], "#D8D3CB")
        self.assertEqual(map_did.COLOR_CFG["project_core_fill"], "#F4CCCC")
        self.assertEqual(map_did.COLOR_CFG["river_main"], "#2F6FA3")
        self.assertEqual(map_did.COLOR_CFG["river_sub"], "#5C97C6")
        self.assertEqual(map_did.COLOR_CFG["river_minor_5"], "#D3E7F3")
        self.assertEqual(map_did.COLOR_CFG["cut"], "#C9252D")

    def test_map_did_supports_custom_output_path(self):
        root = Path(__file__).resolve().parent
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "custom_fig5.png"
            env = os.environ.copy()
            env["MPLBACKEND"] = "Agg"
            env["MAP_OUTPUT_PATH"] = str(output_path)
            env["MAP_SHOW"] = "0"

            result = subprocess.run(
                [sys.executable, "map_did.py"],
                cwd=root,
                env=env,
                capture_output=True,
                text=True,
                timeout=180,
            )

            self.assertEqual(result.returncode, 0, msg=result.stderr)
            self.assertTrue(output_path.exists(), msg=result.stdout)

    def test_prepare_trend_plot_data_normalizes_and_smooths_short_series(self):
        panel = pd.DataFrame(
            {
                "county": ["A", "A", "A", "A", "B", "B", "B", "B"],
                "ym": pd.PeriodIndex(
                    ["2021-01", "2021-02", "2021-03", "2021-04", "2021-01", "2021-02", "2021-03", "2021-04"],
                    freq="M",
                ),
                "year": [2021] * 8,
                "month": [1, 2, 3, 4, 1, 2, 3, 4],
                "ym_str": ["2021-01", "2021-02", "2021-03", "2021-04"] * 2,
                "treat": [1, 1, 1, 1, 0, 0, 0, 0],
                "post": [0, 0, 1, 1, 0, 0, 1, 1],
                "n_reg": [10.0, 20.0, 50.0, 60.0, 20.0, 20.0, 25.0, 25.0],
            }
        )

        series_df, gap_df = did_model.prepare_trend_plot_data(panel, smooth_window=2)

        treat_pre = series_df[(series_df["treat"] == 1) & (series_df["post"] == 0)]
        control_pre = series_df[(series_df["treat"] == 0) & (series_df["post"] == 0)]
        self.assertAlmostEqual(treat_pre["index_raw"].mean(), 100.0, places=6)
        self.assertAlmostEqual(control_pre["index_raw"].mean(), 100.0, places=6)

        treat_feb = series_df[
            (series_df["treat"] == 1) & (series_df["ym"] == pd.Period("2021-02", freq="M"))
        ].iloc[0]
        self.assertAlmostEqual(treat_feb["n_reg_smooth"], 15.0, places=6)

        mar_gap = gap_df.loc[gap_df["ym"] == pd.Period("2021-03", freq="M"), "gap_smooth"].iloc[0]
        treat_mar = series_df[
            (series_df["treat"] == 1) & (series_df["ym"] == pd.Period("2021-03", freq="M"))
        ]["index_smooth"].iloc[0]
        control_mar = series_df[
            (series_df["treat"] == 0) & (series_df["ym"] == pd.Period("2021-03", freq="M"))
        ]["index_smooth"].iloc[0]
        self.assertAlmostEqual(mar_gap, treat_mar - control_mar, places=6)

        pre_cum = gap_df.loc[gap_df["post"] == 0, "cum_excess_growth"]
        self.assertTrue((pre_cum == 0).all())

    def test_seasonal_adjust_series_reduces_month_pattern(self):
        monthly_pattern = [24.0, -18.0, 14.0, -10.0, 6.0, -4.0, 2.0, 5.0, -3.0, 8.0, -6.0, 11.0]
        values = pd.Series(
            [100 + month_idx * 0.8 + monthly_pattern[month_idx % 12] for month_idx in range(36)],
            dtype=float,
        )

        adjusted = did_model.seasonally_adjust_series(values, period=12)
        month_ids = pd.Series([(month_idx % 12) + 1 for month_idx in range(36)])

        before = (values.groupby(month_ids).mean() - values.mean()).abs().mean()
        after = (adjusted.groupby(month_ids).mean() - adjusted.mean()).abs().mean()

        self.assertLess(after, before * 0.4)

    def test_prepare_event_study_panel_reduces_seasonality_within_county(self):
        periods = pd.period_range("2020-01", "2022-12", freq="M")
        pattern_a = [0.42, -0.26, 0.19, -0.14, 0.10, -0.08, 0.07, -0.05, 0.04, -0.03, 0.02, -0.01]
        pattern_b = [-0.30, 0.22, -0.16, 0.13, -0.11, 0.08, -0.07, 0.06, -0.04, 0.03, -0.02, 0.01]
        rows = []
        for county, base, slope, pattern in [
            ("A", 2.0, 0.018, pattern_a),
            ("B", 1.7, 0.015, pattern_b),
        ]:
            for idx, period in enumerate(periods):
                rows.append(
                    {
                        "county": county,
                        "ym": period,
                        "ym_str": str(period),
                        "year": period.year,
                        "month": period.month,
                        "event_time": idx - 18,
                        "ln_reg": base + slope * idx + pattern[period.month - 1],
                    }
                )

        prepared = did_model.prepare_event_study_panel(pd.DataFrame(rows), seasonal_period=12)

        self.assertIn("ln_reg_adj", prepared.columns)
        for county in ["A", "B"]:
            subset = prepared[prepared["county"] == county]
            before = (subset.groupby("month")["ln_reg"].mean() - subset["ln_reg"].mean()).abs().mean()
            after = (subset.groupby("month")["ln_reg_adj"].mean() - subset["ln_reg_adj"].mean()).abs().mean()
            self.assertLess(after, before * 0.65)

    def test_plot_event_study_uses_deseasoned_dynamic_labels(self):
        es_df = pd.DataFrame(
            {
                "et": [-3, -2, -1, 0, 1, 2, 3],
                "beta": [-0.05, 0.02, 0.0, 0.06, 0.18, 0.30, 0.24],
                "ci_lo": [-0.18, -0.11, 0.0, -0.07, 0.03, 0.10, 0.05],
                "ci_hi": [0.08, 0.15, 0.0, 0.19, 0.33, 0.50, 0.43],
            }
        )
        es_df.attrs["pretrend_pval"] = 0.081

        captured = {}

        def fake_savefig(fig, *args, **kwargs):
            captured["figure"] = fig

        with patch("matplotlib.figure.Figure.savefig", new=fake_savefig), patch(
            "matplotlib.pyplot.close", new=lambda fig=None: None
        ):
            did_model.plot_event_study(es_df)

        ax = captured["figure"].axes[0]
        self.assertIn("去季节化", ax.get_title())
        legend_labels = [text.get_text() for text in ax.get_legend().texts]
        self.assertTrue(any("3个月平滑" in label for label in legend_labels))


class FigureNarrativeTests(unittest.TestCase):
    def test_build_event_study_summary_reports_pretrend_and_peak(self):
        es_df = pd.DataFrame(
            {
                "et": [-2, -1, 0, 1, 2, 3],
                "beta": [0.03, 0.0, 0.08, 0.22, 0.35, 0.28],
                "ci_lo": [-0.12, 0.0, -0.10, 0.04, 0.12, 0.05],
                "ci_hi": [0.18, 0.0, 0.26, 0.40, 0.58, 0.51],
            }
        )
        es_df.attrs["pretrend_pval"] = 0.0741

        summary = did_model.build_event_study_summary(es_df)

        self.assertIn("p=0.074", summary["pretrend"])
        self.assertIn("不能拒绝平行趋势", summary["pretrend"])
        self.assertIn("t+2", summary["post"])
        self.assertIn("3个", summary["post"])

    def test_build_forest_takeaways_summarizes_consistency_and_infra_focus(self):
        rob_res = [
            {"label": "基准DID(Post=2021-03)", "beta": 0.29, "pval": 0.014},
            {"label": "Post=2024-12(核准)", "beta": 0.15, "pval": 0.097},
            {"label": "Post=2025-07(开工)", "beta": 0.33, "pval": 0.041},
            {"label": "安慰剂(虚拟2022-06)", "beta": 0.11, "pval": 0.100},
        ]
        het_res = {
            "林芝核心段": {"label": "林芝核心段", "beta": 0.23, "pval": 0.033},
            "基建关联行业": {"label": "基建关联行业", "beta": 0.33, "pval": 0.045},
            "非基建关联": {"label": "非基建关联", "beta": 0.27, "pval": 0.010},
        }

        notes = did_model.build_forest_takeaways(rob_res, het_res)

        self.assertTrue(any("基准/核准/开工" in note for note in notes))
        self.assertTrue(any("安慰剂" in note and "未达5%显著" in note for note in notes))
        self.assertTrue(any("基建关联行业" in note for note in notes))

    def test_build_heterogeneity_notes_mentions_phase_gap_when_available(self):
        het = {
            "基建关联行业": {"label": "基建关联行业", "beta": 0.32, "pval": 0.030, "ci_lo": 0.1, "ci_hi": 0.5},
            "非基建关联": {"label": "非基建关联", "beta": 0.21, "pval": 0.080, "ci_lo": -0.02, "ci_hi": 0.4},
            "重大工程启动期": {
                "label": "重大工程启动期",
                "beta": 0.28,
                "pval": 0.021,
                "ci_lo": 0.05,
                "ci_hi": 0.51,
                "phase_diff_pval": 0.034,
            },
            "持续推进期": {
                "label": "持续推进期",
                "beta": 0.10,
                "pval": 0.290,
                "ci_lo": -0.08,
                "ci_hi": 0.28,
                "phase_diff_pval": 0.034,
            },
        }

        notes = did_model.build_heterogeneity_notes(het, baseline_beta=0.25)

        self.assertIn("阶段差异", notes["时间异质性\n（启动期 vs 持续推进期）"])
        self.assertIn("p=0.034", notes["时间异质性\n（启动期 vs 持续推进期）"])

    def test_selection_rationale_text_describes_corridor_groups(self):
        stats = {
            "treat_core_count": 3,
            "exclude_ring_count": 2,
            "control_low_count": 11,
        }

        note = map_did.build_selection_rationale_text(stats)

        self.assertIn("工程走廊", note)
        self.assertIn("0–30km定义为核心处理县，共3县", note)
        self.assertIn("30–80km定义为潜在外溢排除县，共2县", note)
        self.assertIn("80–200km定义为低暴露对照县，共11县", note)


class DiagnosticTableTests(unittest.TestCase):
    def test_build_diagnostic_tables_returns_factor_and_time_tables(self):
        raw = pd.DataFrame(
            {
                "公司名称": ["t1", "t2", "t3", "t4", "c1", "c2", "c3", "c4"],
                "区县": ["米林县", "米林县", "米林县", "米林县", "察隅县", "察隅县", "察隅县", "察隅县"],
                "行业门类": ["租赁和商务服务业", "建筑业", "租赁和商务服务业", "建筑业", "批发和零售业", "住宿和餐饮业", "批发和零售业", "住宿和餐饮业"],
                "行业大类": ["商务服务业", "房屋建筑业", "商务服务业", "房屋建筑业", "零售业", "餐饮业", "零售业", "餐饮业"],
                "year": [2021, 2021, 2023, 2024, 2021, 2021, 2023, 2024],
                "month": [1, 6, 1, 1, 1, 6, 1, 1],
                "ym": pd.PeriodIndex(["2021-01", "2021-06", "2023-01", "2024-01", "2021-01", "2021-06", "2023-01", "2024-01"], freq="M"),
                "is_infra": [0, 1, 1, 0, 0, 0, 0, 0],
            }
        )
        panel = pd.DataFrame(
            {
                "county": ["米林县", "米林县", "米林县", "米林县", "察隅县", "察隅县", "察隅县", "察隅县"],
                "ym": pd.PeriodIndex(["2021-01", "2021-06", "2023-01", "2024-01", "2021-01", "2021-06", "2023-01", "2024-01"], freq="M"),
                "ym_str": ["2021-01", "2021-06", "2023-01", "2024-01", "2021-01", "2021-06", "2023-01", "2024-01"],
                "year": [2021, 2021, 2023, 2024, 2021, 2021, 2023, 2024],
                "month": [1, 6, 1, 1, 1, 6, 1, 1],
                "sample_group": ["treat_core", "treat_core", "treat_core", "treat_core", "control_low", "control_low", "control_low", "control_low"],
                "treat": [1, 1, 1, 1, 0, 0, 0, 0],
                "post": [0, 1, 1, 1, 0, 1, 1, 1],
                "n_reg": [1.0, 2.0, 2.0, 3.0, 1.0, 1.0, 1.0, 2.0],
                "n_infra": [0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                "ln_reg": pd.Series([1.0, 2.0, 2.0, 3.0, 1.0, 1.0, 1.0, 2.0]).map(lambda v: __import__("math").log1p(v)),
                "ln_infra": pd.Series([0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0]).map(lambda v: __import__("math").log1p(v)),
                "did": [0, 1, 1, 1, 0, 0, 0, 0],
            }
        )

        tables = did_model.build_diagnostic_tables(raw, panel)

        self.assertIn("group_summary", tables)
        self.assertIn("annual_gap", tables)
        self.assertIn("county_contributions", tables)
        self.assertIn("sector_factors", tables)
        self.assertIn("sector_big_factors", tables)
        self.assertIn("time_window_factors", tables)
        self.assertIn("did_diff", tables["sector_factors"].columns)
        self.assertIn("window_label", tables["time_window_factors"].columns)
        self.assertTrue((tables["time_window_factors"]["window_label"] == "启动期").any())

    def test_build_trend_change_summary_highlights_treat_gain(self):
        panel = pd.DataFrame(
            {
                "county": ["米林县", "米林县", "察隅县", "察隅县"],
                "treat": [1, 1, 0, 0],
                "post": [0, 1, 0, 1],
                "n_reg": [80.0, 136.0, 40.0, 56.0],
            }
        )

        summary = did_model.build_trend_change_summary(panel)

        self.assertAlmostEqual(summary["treat"]["pre"], 80.0)
        self.assertAlmostEqual(summary["treat"]["post"], 136.0)
        self.assertAlmostEqual(summary["treat"]["diff"], 56.0)
        self.assertAlmostEqual(summary["control"]["diff"], 16.0)
        self.assertAlmostEqual(summary["gap_diff"], 40.0)
        self.assertIn("+70.0%", summary["treat"]["label"])

    def test_compute_trend_axis_limits_downweights_raw_spikes(self):
        series_df = pd.DataFrame(
            {
                "treat": [1, 1, 1, 0, 0, 0],
                "index_raw": [100.0, 120.0, 260.0, 100.0, 118.0, 280.0],
                "index_smooth": [102.0, 118.0, 176.0, 98.0, 112.0, 165.0],
            }
        )

        lower, upper = did_model.compute_trend_axis_limits(series_df)

        self.assertLess(upper, 240.0)
        self.assertGreater(upper, 176.0)
        self.assertLess(lower, 98.0)


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