# Reproduction code for 05_marker_validation
# Environment: archaea-bio

import pandas as pd
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# Load feature matrix
fm = pd.read_csv("/Users/ajayvellanki/.claude-science/orgs/a85da623-e9bf-4a84-bae1-c872a64348e0/artifacts/proj_aa6c67020290/bc0e3352-d2ec-4e97-91a6-ae3a4c4a7c27/vd4aee24f_feature_matrix.csv")

markers = ['methanogenesis', 'wood_ljungdahl', 'reverse_gyrase', 'rhodopsin', 'hydrogenase_NiFe', 'nha_antiporter']
hmm = fm.set_index('accession')[markers].astype(bool)

# Load silver hits
sh = pd.read_csv("/Users/ajayvellanki/.claude-science/orgs/a85da623-e9bf-4a84-bae1-c872a64348e0/artifacts/proj_aa6c67020290/4ece5136-d5ca-49a2-b94c-e37040f50441/v03ace195_silver_hits2.tsv", sep='\t', header=None, names=['accession', 'product'], dtype=str)
sh['product'] = sh['product'].fillna('')
p = sh['product'].str.lower()

rules = {
    'methanogenesis'  : p.str.contains('coenzyme m reductase'),
    'wood_ljungdahl'  : p.str.contains(r'acetyl-coa decarbonylase|co dehydrogenase/acetyl-coa|acetyl-coa synthase'),
    'reverse_gyrase'  : p.str.contains('reverse gyrase'),
    'rhodopsin'       : p.str.contains('bacteriorhodopsin|halorhodopsin') | (p.str.contains('rhodopsin') & ~p.str.contains('sensory')),
    'hydrogenase_NiFe': (p.str.contains('nife') | p.str.contains('nickel-dependent hydrogenase') |
                         (p.str.contains('hydrogenase') & p.str.contains('large subunit')) |
                         (p.str.contains('hydrogenase') & ~p.str.contains('dehydrogenase'))),
    'nha_antiporter'  : p.str.contains(r'na\+/h\+ antiporter|sodium/proton antiporter|sodium:proton antiporter|cation:proton antiporter|monovalent cation.proton antiporter|nhac|nhad|nhap|mrp antiporter'),
}

silver = pd.DataFrame({'accession': sh['accession']})
for m in markers:
    silver[m] = rules[m].values
silver_geno = silver.groupby('accession')[markers].any()

allacc = fm['accession']
S = pd.DataFrame(False, index=allacc, columns=markers)
S.update(silver_geno.reindex(allacc).fillna(False))
S = S.loc[allacc].astype(bool)
H = hmm.reindex(allacc).fillna(False).astype(bool)

rows = []
for m in markers:
    y = S[m].values; yh = H[m].values
    tp = int((y & yh).sum()); fp = int((~y & yh).sum()); fn = int((y & ~yh).sum()); tn = int((~y & ~yh).sum())
    prec = tp / (tp + fp) if tp + fp else np.nan
    rec = tp / (tp + fn) if tp + fn else np.nan
    f1 = 2 * prec * rec / (prec + rec) if prec and rec and (prec + rec) > 0 else np.nan
    rows.append(dict(marker=m, silver_pos=int(y.sum()), hmm_pos=int(yh.sum()), TP=tp, FP=fp, FN=fn, TN=tn,
                     precision=round(prec, 4), recall=round(rec, 4), f1=round(f1, 4), accuracy=round((tp + tn) / len(y), 4)))
val = pd.DataFrame(rows)

val_sorted = val.sort_values('f1', ascending=False).reset_index(drop=True)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5.2), gridspec_kw={'width_ratios': [1.05, 1]})
mk = val_sorted['marker'].tolist(); x = np.arange(len(mk)); w = 0.26
ax1.bar(x - w, val_sorted['precision'], w, label='precision', color='#2c7fb8', edgecolor='k', lw=.4)
ax1.bar(x,     val_sorted['recall'],    w, label='recall',    color='#7fcdbb', edgecolor='k', lw=.4)
ax1.bar(x + w, val_sorted['f1'],        w, label='F1',        color='#fdae6b', edgecolor='k', lw=.4)
ax1.set_xticks(x); ax1.set_xticklabels(mk, rotation=30, ha='right', fontsize=8)
ax1.set_ylim(0, 1.05); ax1.set_ylabel('score'); ax1.axhline(1.0, color='grey', lw=.6, ls=':')
ax1.legend(fontsize=8, ncol=3, loc='lower center', framealpha=.9)
ax1.set_title('a  Held-out validation vs RefSeq PGAP annotation (n=2,919)', fontsize=10.5, loc='left', weight='bold')
for sp in ['top', 'right']: ax1.spines[sp].set_visible(False)

ax2.axis('off')
cols = ['marker', 'silver+', 'HMM+', 'TP', 'FP', 'FN', 'prec', 'rec', 'F1']
cell = []
for _, r in val_sorted.iterrows():
    cell.append([r['marker'], int(r['silver_pos']), int(r['hmm_pos']), int(r['TP']), int(r['FP']),
                 int(r['FN']), f"{r['precision']:.2f}", f"{r['recall']:.2f}", f"{r['f1']:.2f}"])
t = ax2.table(cellText=cell, colLabels=cols, loc='center', cellLoc='center')
t.auto_set_font_size(False); t.set_fontsize(7.5); t.scale(1, 1.45)
for j in range(len(cols)): t[0, j].set_facecolor('#e0e0e0'); t[0, j].set_text_props(weight='bold')
ax2.set_title('b  Confusion counts and metrics', fontsize=10.5, loc='left', weight='bold', y=0.86)
fig.text(0.52, 0.02, 'Diagnostic markers (methanogenesis, WL, reverse gyrase, rhodopsin) reach precision = 1.00; '
         'zero halophiles are silver- or HMM-positive for mcr/WL/rgy.', fontsize=7.8, style='italic', ha='center')
fig.tight_layout(rect=[0, 0.04, 1, 1])
fig.savefig('fig_marker_validation.png', dpi=200, bbox_inches='tight')
plt.close(fig)

val_sorted.to_csv('data/marker_validation_heldout.csv', index=False)
print("saved fig_marker_validation.png and marker_validation_heldout.csv")
print(val_sorted[['marker', 'precision', 'recall', 'f1']].to_string(index=False))