import pandas as pd, json
df = pd.read_excel('/Users/maxim.ziatdinov/Code/SciLink/planning_session_20260401_225636/knowledge/PWSdatabase.xlsx')
basins = ['East Texas', 'Arkla', 'Appalachian', 'Williston', 'Southern Oklahoma']
df_sub = df[df['BASIN'].isin(basins)].copy()
for c in ['Li','Sr','Ba','Mg','Rb','TDS']:
    df_sub[c] = pd.to_numeric(df_sub[c], errors='coerce')
stats = df_sub.groupby('BASIN').agg(Li_avg=('Li','mean'),Li_max=('Li','max'),Sr_avg=('Sr','mean'),Sr_max=('Sr','max'),Ba_avg=('Ba','mean'),Ba_max=('Ba','max'),Mg_avg=('Mg','mean'),Mg_max=('Mg','max'),Rb_avg=('Rb','mean'),Rb_max=('Rb','max'),TDS_avg=('TDS','mean'))
answer = {basin: {col: round(val, 2) if pd.notna(val) else None for col, val in row.items()} for basin, row in stats.iterrows()}
summary = "For the top 5 Li basins, average/max concentrations of Sr, Ba, Mg, Rb and average TDS are: " + "; ".join(f"{b}: TDS_avg={answer[b]['TDS_avg']}" for b in basins if b in answer)
print(json.dumps({"answer": answer, "summary": summary}))
