from edgeeda.viz import export_trials
import pandas as pd, matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt, os, glob, json

out='runs/plots_quick'
os.makedirs(out, exist_ok=True)

# Load trials
df = export_trials('runs/experiment.sqlite')
print('rows:', len(df))
print('columns:', list(df.columns))

# Basic runtime histogram
runtimes = pd.to_numeric(df['runtime_sec'], errors='coerce').dropna()
if not runtimes.empty:
    plt.figure(); runtimes.hist(bins=10)
    plt.xlabel('runtime_sec'); plt.tight_layout(); plt.savefig(os.path.join(out,'runtime_hist.png'), dpi=200); plt.close()
    print('wrote runtime_hist.png')
else:
    print('no runtime data to plot')

# return_code counts
plt.figure(); df['return_code'].value_counts().plot(kind='bar')
plt.xlabel('return_code'); plt.tight_layout(); plt.savefig(os.path.join(out,'return_code_counts.png'), dpi=200); plt.close()
print('wrote return_code_counts.png')

# metadata availability
has_meta = df['metadata_path'].fillna('').apply(lambda x: bool(str(x).strip()))
plt.figure(); has_meta.value_counts().plot(kind='bar'); plt.xticks([0,1],['no metadata','has metadata']); plt.tight_layout(); plt.savefig(os.path.join(out,'metadata_counts.png'), dpi=200); plt.close()
print('wrote metadata_counts.png')

# learning curve from reward, if present
if 'reward' in df.columns:
    r = pd.to_numeric(df['reward'], errors='coerce').dropna()
    if not r.empty:
        df2 = df.copy()
        df2['reward'] = pd.to_numeric(df2['reward'], errors='coerce')
        df2 = df2.dropna(subset=['reward']).sort_values('id')
        best = df2['reward'].cummax()
        plt.figure(); plt.plot(df2['id'].values, best.values)
        plt.xlabel('trial id'); plt.ylabel('best reward so far'); plt.tight_layout(); plt.savefig(os.path.join(out,'learning_curve.png'), dpi=200); plt.close()
        print('wrote learning_curve.png')
    else:
        print('no rewards to plot')
else:
    print('reward column missing')

# area vs wns if metrics present
areas=[]; wnss=[]
for _, r in df.iterrows():
    mj = r.get('metrics') or r.get('metrics_json') or r.get('metrics_json')
    if not mj:
        continue
    if isinstance(mj, str):
        try:
            m = json.loads(mj)
        except Exception:
            continue
    else:
        m = mj
    a = m.get('design__die__area') or m.get('finish__design__die__area')
    w = m.get('timing__setup__wns') or m.get('finish__timing__setup__wns')
    if a is None or w is None:
        continue
    try:
        areas.append(float(a)); wnss.append(float(w))
    except Exception:
        pass
if areas:
    plt.figure(); plt.scatter(areas, wnss); plt.xlabel('die area'); plt.ylabel('WNS'); plt.tight_layout(); plt.savefig(os.path.join(out,'area_vs_wns.png'), dpi=200); plt.close()
    print('wrote area_vs_wns.png')
else:
    print('no area/wns metrics to plot')

print('files:', glob.glob(out+'/*'))