!pip install nucleus-cdk==0.5.0rc2 | tail -n2

# Surface a failed install here, rather than as a confusing ModuleNotFoundError
# in the import cell below.
import importlib.metadata as md
assert md.version("nucleus-cdk") == "0.5.0rc2", f"got {md.version('nucleus-cdk')}"

%load_ext autoreload
%autoreload 2
    
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

# Import the cdk platereader module
from cdk.analysis.cytosol import platereader as pr

# Set up plotting
pr.plot_setup()

labcraft_platemap = pd.read_csv("data/concentration_key.csv")
labcraft_platemap = labcraft_platemap.fillna(0)
labcraft_platemap = labcraft_platemap.rename(columns={'Well ID': 'well_id'})
labcraft_platemap = labcraft_platemap.set_index("well_id")
labcraft_platemap = labcraft_platemap.drop(columns=['[PMix]_mg/mL', '[DNA]_nM']) #did not use labcraft here

samples = pd.read_csv("data/samples_titration.csv", index_col="well_id")
samples["Type"] = "Sample"

samples.columns

labcraft_platemap.columns

# samples = samples.reset_index()
samples.head()

# Map concentration_key column names to samples column names
col_rename = {
    '[Creatine phosphate]_mM': '[Creatine phosphate] (mM)',
    '[Magnesium acetate]_mM': '[Magnesium acetate] (mM)',
    '[Potassium glutamate]_mM': '[Potassium glutamate] (mM)',
}
labcraft_platemap = labcraft_platemap.rename(columns=col_rename)


# --- Add base master mix contributions to get true final concentrations ---
base_mm = pd.read_csv("data/base_master_mix.csv", index_col="reagent")
base_mg   = base_mm.loc['magnesium_acetate', 'base_conc']    # 2.469 mM
base_kglu = base_mm.loc['potassium_glutamate', 'base_conc']  # 60.0 mM
# Creatine phosphate not in base_mm → base contribution is 0

labcraft_platemap['[Magnesium acetate] (mM)']   += base_mg
labcraft_platemap['[Potassium glutamate] (mM)'] += base_kglu

!pwd

# Overwrite common concentration columns in samples with labcraft_platemap values
# Use .update() so only matching well_ids are overwritten (no NaNs for missing wells)
common_cols = labcraft_platemap.columns.intersection(samples.columns).tolist()
samples.update(labcraft_platemap[common_cols])
print(f"Updated {len(common_cols)} columns in samples from labcraft_platemap: {common_cols}")
samples.head()

# Add standard wells not in DOE samples
standards = pd.DataFrame([
    {"well_id": "K8", "Type": "Standard", "Name": "Standard"},
    {"well_id": "K9", "Type": "Standard", "Name": "Standard"},
    {"well_id": "K10", "Type": "Standard", "Name": "Standard"},
]).set_index("well_id")
samples = pd.concat([samples, standards])
samples["[Magnesium acetate] (mM)"] = samples["[Magnesium acetate] (mM)"].apply(lambda x: x*95.2/203.3) # Perhaps through artifact miststake, we used the wrong calculation, so inpute Mg might have been wrong
samples["Well"] = samples.index

samples.to_csv("data/samples_final.csv")

samples.Well.unique()

# Specify file paths
data_file = "data/20260309-123116-cytation5-pure-timecourse-gfp-biotek-cdk.txt"
platemap_file = "data/samples_final.csv"


# Load data
data, platemap = pr.load_platereader_data(
    data_file=data_file,
    platemap_file=platemap_file,
    platereader="biotek-cdk"  # Options: "cytation", "envision", "biotek-cdk"
)

# Checkout first few rows
data.head()
# data["conditions"] = data.groupby(factors).ngroup()
# data["Name"] = data["conditions"]

# data = data[~data["Column"].isin([19])]

data.Well.unique()

g = pr.plot_curves(data=data, hue="Well")#[data.Type == "Sample"]) #, hue="conditions")

data.Name.unique()

data = pr.normalize_data_to_controls(data, ctrl_name = 'Standard')

g = pr.plot_curves(data=data[data.Type == "Sample"])

data.columns
FACTOR_COLS = ['[PMix] (mg/mL)','[DNA] (nM)','[Magnesium acetate] (mM)','[Potassium glutamate] (mM)', '[Creatine phosphate] (mM)']

# Perform kinetic analysis using sigmoid_drift model
kinetics = pr.kinetic_analysis(
    data=data,
    group_by=['Well',  *FACTOR_COLS],  # Group by experimental condition
)
# kinetics = kinetics.reset_index()
kinetics.head()

data.head()

# Plot kinetic fits 
g = pr.plot_kinetics(data[data.Well == "I10"],  group_by=['Well', 'well_id', *FACTOR_COLS])

# pr.plot_summary(data)

responses = ["steady_state", "vmax", "lag_time_min"]

kin = kinetics.reset_index()
kin.columns = [
    "_".join([str(x).replace(" ", "_") for x in col if x])
    for col in kin.columns
]
factors = [str(x).replace(" ", "_") for x in FACTOR_COLS]
kin.head()

kin.columns

from plate_qc import (
        qc_summary, plot_replicate_agreement,
        plot_plate_heatmap, plot_qc_overview,
    )
responses = ['Steady_State_data_normalized', 'Velocity_Max']

clean, failed = qc_summary(kin, factors, responses)

kin.head()

import seaborn as sns
kin['steady_state (binned)'] = pd.cut(                                                                                                                                                                                   
  kin['Steady_State_data_normalized'],                                                                                                                                                                        
  bins=[0, 0.05, 0.2, 1.0],                                                                                                                                                                                   
  labels=['0-0.05', '0.05-0.2', '>0.2']                                                                                                                                                                               
)                                                                                                                                                                                                               

g = sns.pairplot(
  kin[factors + ['steady_state (binned)']],
  vars=factors,
  hue='steady_state (binned)',
  palette={'0-0.05': '#d9534f', '0.05-0.2': 'orange', '>0.2': 'green'},
  plot_kws={'alpha': 0.7, 's': 40},
  diag_kind='hist',
)