#!/usr/bin/env python
"""Summary figures for the UC HII accretion test."""
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from astropy.table import Table
LOG=open('make_figures.log','w')
def p(*a):
    s=' '.join(str(x) for x in a);print(s);print(s,file=LOG);LOG.flush()

m=Table.read('master_phase1_classified.fits')
ag=Table.read('uchii_atlasgal.fits')
co=Table.read('uchii_cores.fits')
cen=Table.read('alma_census.csv',format='csv')

# ---- Fig 1: classification + ALMA coverage ----
fig,axes=plt.subplots(1,3,figsize=(15,4.2))
rc=np.array([str(x) for x in m['radio_class']])
styp=np.array([str(x) for x in m['simple_type']])
hii=styp=='HII region'
counts=[ (rc[hii]=='UCHII').sum(),(rc[hii]=='compact').sum(),(rc[hii]=='classical').sum(),
         (hii&~np.isin(rc,['UCHII','compact','classical'])).sum() ]
ax=axes[0]
ax.bar(['UC HII\n(<=0.1pc)','compact\n(0.1-0.5)','classical\n(>0.5)','no radio\nsize'],counts,
       color=['#c0392b','#e67e22','#f1c40f','#bdc3c7'])
for i,c in enumerate(counts): ax.text(i,c+3,str(c),ha='center')
ax.set_ylabel('# RMS HII regions'); ax.set_title('(a) Radio-size classification (CORNISH)')

# ALMA coverage by class
ax=axes[1]
classes=['UCHII','compact','classical']
crc=np.array([str(x) for x in cen['radio_class']])
nobs=np.array(cen['n_obs'],float)
obsfrac=[ (nobs[(crc==k)]>0).mean()*100 for k in classes ]
ax.bar(['UC HII','compact','classical'],obsfrac,color=['#c0392b','#e67e22','#f1c40f'])
for i,c in enumerate(obsfrac): ax.text(i,c+1,f'{c:.0f}%',ha='center')
ax.set_ylabel('% with ALMA archival data'); ax.set_ylim(0,80); ax.set_title('(b) ALMA coverage')

# ---- Fig 1c: mass vs scale ----
ax=axes[2]
Mclump=np.array(ag['ag_Mclump'],float); Mclump=Mclump[np.isfinite(Mclump)]
Mcore_near=np.array(co['core_nearest_Mcore'],float); Mcore_near=Mcore_near[np.isfinite(Mcore_near)]
Mdust=np.array(co['core_Mdust_thin'],float); Mdust=Mdust[np.isfinite(Mdust)]
bins=np.logspace(-2,4.5,28)
ax.hist(Mclump,bins=bins,alpha=0.6,label=f'clump ~0.5-1pc (ATLASGAL), med={np.median(Mclump):.0f}',color='#2980b9')
ax.hist(Mcore_near,bins=bins,alpha=0.6,label=f'core at star (ALMAGAL), med={np.median(Mcore_near):.1f}',color='#27ae60')
if len(Mdust): ax.hist(Mdust[Mdust>0],bins=bins,alpha=0.7,label=f'dust-only (ff-subtr), >0 only',color='#c0392b')
ax.set_xscale('log'); ax.set_xlabel('mass [Msun]'); ax.set_ylabel('# UC HII')
ax.axvline(8,ls='--',color='k',lw=0.8); ax.text(8.5,ax.get_ylim()[1]*0.8,'8 Msun',fontsize=8)
ax.legend(fontsize=7.5); ax.set_title('(c) Reservoir mass vs spatial scale')
plt.tight_layout(); plt.savefig('fig_summary.png',dpi=130); p('wrote fig_summary.png')

# ---- Fig 2: the accreting-fraction funnel ----
fig,ax=plt.subplots(figsize=(7.5,5))
inal=np.array(co['in_almagal'])
N=int(inal.sum())
sep_pc=np.array(co['core_nearest_sep_pc'],float)
Mnear=np.array(co['core_nearest_Mcore'],float)
Mdust=np.array(co['core_Mdust_thin'],float)
stages=['UC HII\nw/ ALMAGAL\ncore','core coincident\n(<=0.02pc)','+ core >=5 Msun','+ dust-dominated\n(ff-subtracted\n>=5 Msun)']
v_coinc=int((inal&(sep_pc<=0.02)).sum())
v_m5=int((inal&(sep_pc<=0.02)&(Mnear>=5)).sum())
v_dust=int((inal&(sep_pc<=0.02)&(Mdust>=5)).sum())
vals=[N,v_coinc,v_m5,v_dust]
cols=['#34495e','#2980b9','#e67e22','#c0392b']
ax.bar(range(4),vals,color=cols)
for i,v in enumerate(vals):
    ax.text(i,v+0.4,f'{v}\n({100*v/N:.0f}%)',ha='center',fontsize=10)
ax.set_xticks(range(4)); ax.set_xticklabels(stages,fontsize=8.5)
ax.set_ylabel('# UC HII'); ax.set_ylim(0,N+4)
ax.set_title(f'Accreting-fraction funnel (N={N} UC HII with ALMAGAL cores)')
plt.tight_layout(); plt.savefig('fig_accreting_funnel.png',dpi=130); p('wrote fig_accreting_funnel.png')
p('DONE')
LOG.close()