import numpy as np, pandas as pd
from scipy import stats
from scipy.optimize import curve_fit
d = pd.read_csv("viability.csv")
m = d.groupby(["experiment", "compound", "conc_M"], as_index=False).viability.mean()   # average the triplicate wells

def four_pl(logx, bottom, top, log_ic50, hill):     # response falls as concentration rises
    return bottom + (top - bottom) / (1 + 10 ** ((logx - log_ic50) * hill))

rows = []
for (e, comp), g in m.groupby(["experiment", "compound"]):
    p, _ = curve_fit(four_pl, np.log10(g.conc_M), g.viability, p0=[0, 100, -7, 1], maxfev=10000)
    rows.append({"experiment": e, "compound": comp, "pIC50": -p[2]})
fits = pd.DataFrame(rows).pivot(index="experiment", columns="compound", values="pIC50")
print("pIC50 per experiment (-log10 IC50 in mol/L):")
for e, r in fits.iterrows():
    print(f"  experiment {e}: A {r.A:.2f}, B {r.B:.2f}, A - B {r.A - r.B:.2f}")
gm = lambda s: 10 ** (-s.mean()) * 1e9
print(f"Geometric mean IC50: A {gm(fits.A):.0f} nM, B {gm(fits.B):.0f} nM")
diff = fits.A - fits.B                                 # log10 potency difference; > 0 means A is more potent
t = stats.ttest_1samp(diff, 0)
ci = t.confidence_interval()
print(f"Paired (same plates): IC50 ratio B/A = {10 ** diff.mean():.2f} (95% CI {10 ** ci.low:.1f} to {10 ** ci.high:.1f}), "
      f"t = {t.statistic:.2f}, df = {t.df:.0f}, p = {t.pvalue:.3f}")
w = stats.ttest_ind(fits.A, fits.B, equal_var=False)
wci = w.confidence_interval()
print(f"Welch, pairing ignored: IC50 ratio B/A = {10 ** (fits.A.mean() - fits.B.mean()):.2f} (95% CI {10 ** wci.low:.1f} to {10 ** wci.high:.1f}), "
      f"t = {w.statistic:.2f}, df = {w.df:.1f}, p = {w.pvalue:.3f}")

# A common shortcut, shown for comparison
# Shortcut: pool every well from all experiments into one fit per compound and run the extra-sum-of-squares F test
# (separate IC50s vs one shared IC50; bottom, top and Hill slope separate per compound). Wells become the n.
x, y, isB = np.log10(d.conc_M.values), d.viability.values, (d.compound == "B").values
def full(X, bA, tA, iA, hA, bB, tB, iB, hB):
    lx, b = X
    return np.where(b, four_pl(lx, bB, tB, iB, hB), four_pl(lx, bA, tA, iA, hA))
def shared(X, bA, tA, hA, bB, tB, hB, i):
    lx, b = X
    return np.where(b, four_pl(lx, bB, tB, i, hB), four_pl(lx, bA, tA, i, hA))
pf, _ = curve_fit(full, (x, isB), y, p0=[0, 100, -7, 1, 0, 100, -6.6, 1], maxfev=20000)
ps, _ = curve_fit(shared, (x, isB), y, p0=[0, 100, 1, 0, 100, 1, -6.8], maxfev=20000)
rss_f = np.sum((y - full((x, isB), *pf)) ** 2); rss_s = np.sum((y - shared((x, isB), *ps)) ** 2)
df1, df2 = 1, len(y) - 8
F = (rss_s - rss_f) / df1 / (rss_f / df2)
print(f"Shortcut, all {len(y)} wells pooled, F test of one shared IC50: F({df1}, {df2}) = {F:.1f}, p < 1e-15" if stats.f.sf(F, df1, df2) < 1e-15 else f"p = {stats.f.sf(F, df1, df2):.1e}")