from pathlib import Path
import numpy as np
import pandas as pd
import statsmodels.formula.api as smf

source = Path("inclusive_finance_did_panel.csv")
df = pd.read_csv(source)

# 根据财政部等五部门公布的 2019 年试点名单建立城市层面的 treatment。
pilot_roots = ["北京市", "天津市", "承德市", "廊坊市", "大同市", "长治市", "包头市", "鄂尔多斯市", "鞍山市", "辽阳市", "大连市", "松原市", "双鸭山市", "齐齐哈尔市", "上海市", "常州市", "泰州市", "温州市", "台州市", "宁波市", "芜湖市", "合肥市", "泉州市", "厦门市", "上饶市", "萍乡市", "济南市", "烟台市", "青岛市", "焦作市", "洛阳市", "宜昌市", "襄阳市", "株洲市", "浏阳市", "佛山市", "东莞市", "玉林市", "来宾市", "三亚市", "重庆市", "宜宾市", "广安市", "铜仁市", "玉溪市", "安康市", "榆林市", "酒泉市", "海东市", "吴忠市", "乌鲁木齐市", "日喀则市"]
df["treatment"] = df["city"].apply(lambda x: int(any(root in x or x in root for root in pilot_roots)))
df["post"] = (df["year"] >= 2020).astype(int)
df["did"] = df["treatment"] * df["post"]
df["analysis_sample"] = df["year"].ne(2019)
sample = df.loc[df["analysis_sample"]].copy()

numeric = ["inclusive_finance_index", "treatment", "post", "did"]
print("\nDESCRIPTIVE STATISTICS")
print(df[numeric].describe().T.round(3))
print("\nCORRELATIONS")
print(df[numeric].corr().round(3))

m0 = smf.ols("inclusive_finance_index ~ did", data=sample).fit(cov_type="HC1")
m1 = smf.ols("inclusive_finance_index ~ did + C(city_code) + C(year)", data=sample).fit(cov_type="cluster", cov_kwds={"groups": sample["city_code"]})
check = df.copy()
check["post_2019"] = (check["year"] >= 2019).astype(int)
check["did_2019"] = check["treatment"] * check["post_2019"]
m2 = smf.ols("inclusive_finance_index ~ did_2019 + C(city_code) + C(year)", data=check).fit(cov_type="cluster", cov_kwds={"groups": check["city_code"]})

print("\nREGRESSION TABLE")
for name, model, term in [("M0 raw", m0, "did"), ("M1 city/year FE", m1, "did"), ("M2 post=2019", m2, "did_2019")]:
    print(f"{name:18s} coef={model.params[term]:9.4f} se={model.bse[term]:9.4f} p={model.pvalues[term]:.4f} N={int(model.nobs)} R2={model.rsquared:.4f}")

print("\nMAIN MODEL SUMMARY")
print(m1.summary())