import pandas as pd, numpy as np
from linearmodels.panel import PanelOLS
import statsmodels.formula.api as smf
def wdi(ind,name):
    d=pd.read_csv(f'/tmp/wdi/indicators/{ind}/data.csv'); d.columns=[c.strip() for c in d.columns]
    d=d.rename(columns={'Country Code':'code','Year':'year','Value':name})
    d['code']=d['code'].str.upper(); d['year']=d['year'].astype(int); return d[['code','year',name]]
mil=wdi('ms.mil.xpnd.gd.zs','mil')
try: debt=wdi('gc.dod.totl.gd.zs','debt')
except Exception as e: debt=None; print("no debt indicator:",e)
try: bal=wdi('gc.nld.totl.gd.zs','bal')
except Exception as e: bal=None
out=[]
def w(s): out.append(s); print(s)
if debt is not None:
    m=mil.merge(debt,on=['code','year'],how='inner').dropna()
    m=m[(m.mil<25)&(m.debt<300)]
    w("MILITARY %GDP vs GOVERNMENT DEBT %GDP")
    w("N=%d, %d countries, %d-%d ; corr=%.3f"%(len(m),m.code.nunique(),m.year.min(),m.year.max(),m.mil.corr(m.debt)))
    fe=PanelOLS.from_formula('debt ~ mil + EntityEffects + TimeEffects',data=m.set_index(['code','year'])).fit(cov_type='clustered',cluster_entity=True)
    w("[Two-way FE] debt ~ military: b=%.3f (se %.3f, p=%.3g)  -> +ve = more defense, more debt"%(fe.params['mil'],fe.std_errors['mil'],fe.pvalues['mil']))
if bal is not None:
    m2=mil.merge(bal,on=['code','year'],how='inner').dropna(); m2=m2[(m2.mil<25)&(m2.bal.abs()<40)]
    fe2=PanelOLS.from_formula('bal ~ mil + EntityEffects + TimeEffects',data=m2.set_index(['code','year'])).fit(cov_type='clustered',cluster_entity=True)
    w("[Two-way FE] fiscal balance ~ military: b=%.3f (p=%.3g)  -> -ve = more defense, worse balance/deficit  (N=%d)"%(fe2.params['mil'],fe2.pvalues['mil'],len(m2)))
open('/tmp/results_mildebt.txt','w').write("\n".join(out)); print("done")
