# Reproduction code for 02_salinity_ph_metadata
# Environment: archaea-bio

import pandas as pd
import numpy as np
import urllib.request

# 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")

tm = fm[['accession','organism_name','phylum','class','order','family','genus','species',
         'checkm_completeness','checkm_contamination','hq_genome',
         'ogt_measured','ogt_provenance','lifestyle_tax']].copy()

# ---------------- THERMOPHILY: measured OGT >= 60C ----------------
tm['thermophily'] = np.where(tm['ogt_measured']>=60, 1,
                     np.where(tm['ogt_measured'].notna(), 0, np.nan))
tm['thermophily_evidence'] = np.where(tm['ogt_measured'].notna(),'measured_OGT','missing')

# ---------------- METHANOGENESIS: recognized methanogenic orders ----------------
methanogen_orders = {
 'Methanobacteriales','Methanococcales','Methanomicrobiales','Methanosarcinales',
 'Methanocellales','Methanopyrales','Methanomassiliicoccales','Methanotrichales',
 'Methanonatronarchaeales','Methanofastidiosales','Methanophagales',
 'Methanoliparales','Methanomethyliales'
}
tm['methanogenesis'] = tm['order'].isin(methanogen_orders).astype(int)
tm.loc[tm['organism_name'].str.contains('Methanodesulfokora', na=False), 'methanogenesis']=1
tm['methanogenesis_evidence'] = np.where(tm['methanogenesis']==1,'methanogenic_lineage','non_methanogen_taxon')

# ---------------- HALOPHILY: taxonomy seed ----------------
halo_orders = {'Halobacteriales','Haloferacales','Natrialbales'}
halo_genera = {'Haloarcula','Halobacterium','Haloferax','Halorubrum','Natronomonas',
               'Haloquadratum','Natrialba','Halomicrobium','Salinibacter'}
tm['halophily'] = ((tm['class']=='Halobacteria') |
                   (tm['order'].isin(halo_orders)) |
                   (tm['genus'].isin(halo_genera))).astype(int)
tm.loc[tm['order'].isin({'Methanonatronarchaeales'}),'halophily']=1
tm['halophily_evidence']='taxonomy_seed'

# ---------------- ACIDOPHILY: taxonomy seed ----------------
acid_orders = {'Sulfolobales','Acidilobales','Thermoplasmatales','Nitrosotaleales'}
acid_genera = {'Picrophilus','Ferroplasma','Acidiplasma','Thermoplasma','Sulfolobus',
               'Acidianus','Metallosphaera','Sulfuracidifontes','Acidilobus','Picrophilaceae'}
tm['acidophily'] = ((tm['order'].isin(acid_orders)) |
                    (tm['genus'].isin(acid_genera))).astype(int)
tm['acidophily_evidence']='taxonomy_seed'

# Download Madin et al. condensed traits DB
url = "https://raw.githubusercontent.com/bacteria-archaea-traits/bacteria-archaea-traits/master/output/condensed_traits_NCBI.csv"
urllib.request.urlretrieve(url, "condensed_traits_NCBI.csv")
md = pd.read_csv('condensed_traits_NCBI.csv', low_memory=False)

arch = md[md['superkingdom']=='Archaea'] if 'superkingdom' in md.columns else md
arch2 = arch.copy()

sal_map = {'extreme-halophilic':'extreme','moderate-halophilic':'moderate','non-halophilic':'non'}
arch2['sal_class'] = arch2['range_salinity'].map(sal_map)

sp_sal = arch2.dropna(subset=['sal_class']).groupby('species')['sal_class'].agg(lambda x: x.mode().iloc[0])
gen_sal = arch2.dropna(subset=['sal_class']).groupby('genus')['sal_class'].agg(lambda x: x.mode().iloc[0])
sp_ph  = arch2.dropna(subset=['optimum_ph']).groupby('species')['optimum_ph'].median()
gen_ph = arch2.dropna(subset=['optimum_ph']).groupby('genus')['optimum_ph'].median()

tm2 = tm.copy()

def match(series_sp, series_gen, sp, gen):
    if pd.notna(sp) and sp in series_sp.index: return series_sp[sp],'species'
    if pd.notna(gen) and gen in series_gen.index: return series_gen[gen],'genus'
    return np.nan,'none'

sal_vals=[]; sal_prov=[]; ph_vals=[]; ph_prov=[]
for _,r in tm2.iterrows():
    sv,sp_= match(sp_sal,gen_sal,r['species'],r['genus']); sal_vals.append(sv); sal_prov.append(sp_)
    pv,pp_= match(sp_ph,gen_ph,r['species'],r['genus']);  ph_vals.append(pv); ph_prov.append(pp_)
tm2['measured_salinity_class']=sal_vals; tm2['salinity_provenance']=sal_prov
tm2['measured_optimum_ph']=ph_vals;      tm2['ph_provenance']=ph_prov

# Refine halophily with measured salinity
halo_refined = tm2['halophily'].copy()
halo_ev = tm2['halophily_evidence'].copy()
mask_meas_pos = tm2['measured_salinity_class'].isin(['extreme','moderate'])
mask_meas_neg = tm2['measured_salinity_class']=='non'
halo_refined[mask_meas_pos]=1; halo_ev[mask_meas_pos]='measured_salinity'
halo_refined[mask_meas_neg]=0; halo_ev[mask_meas_neg]='measured_salinity'
tm2['halophily']=halo_refined; tm2['halophily_evidence']=halo_ev

# Refine acidophily with measured pH
acid_refined = tm2['acidophily'].copy(); acid_ev=tm2['acidophily_evidence'].copy()
mask_acid_pos = tm2['measured_optimum_ph']<=5.0
mask_acid_neg = tm2['measured_optimum_ph']>=6.0
acid_refined[mask_acid_pos]=1; acid_ev[mask_acid_pos]='measured_pH'
acid_refined[mask_acid_neg]=0; acid_ev[mask_acid_neg]='measured_pH'
tm2['acidophily']=acid_refined; tm2['acidophily_evidence']=acid_ev

traits=['halophily','methanogenesis','thermophily','acidophily']
tm2['n_traits']=tm2[traits].fillna(0).astype(int).sum(axis=1)

tm2[['accession','organism_name','genus','species',
     'measured_salinity_class','salinity_provenance',
     'measured_optimum_ph','ph_provenance',
     'ogt_measured','ogt_provenance']].to_csv('salinity_ph_metadata.csv', index=False)
print("salinity_ph_metadata.csv written")
print("coverage — salinity:", tm2['measured_salinity_class'].notna().sum(),
      "pH:", tm2['measured_optimum_ph'].notna().sum())