!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()

!pwd

import os
# print(os.path)

# labcraft_platemap = pd.read_csv("experiment_setup_v2_b/labcraft_results/DP-R3b-20260518_112845/concentration_key.csv")
labcraft_platemap = pd.read_csv("data/concentration_key_corrected_mg2000_24p602_mg250_14p862.csv") #I believe Mg 2000 mM was added second...that means with the larger droplet size. Yikes :'(
labcraft_platemap = labcraft_platemap.rename(columns={'Well ID': 'Well'})
labcraft_platemap = labcraft_platemap.set_index("Well")
labcraft_platemap = labcraft_platemap.fillna(0)

samples = pd.read_csv("data/samples_titration_FINAL.csv", index_col="Well")

labcraft_platemap

# Normalize column headers: '[ATP]_mM' -> '[ATP] mM'
labcraft_platemap.columns = [str(c).replace("_", " ") for c in labcraft_platemap.columns]

labcraft_platemap.columns

# Add base concentration back to titration volumes
from cdk_naming import (
    normalize_reagent_name,
    split_concentration_column_name,
)

base_mm = pd.read_csv(
    "data/base_master_mix_w_smix.csv", index_col="reagent"
)

# normalized reagent key -> (base_conc, unit) lookup
base_lookup = {
    normalize_reagent_name(idx): (row["base_conc_fold"], row["unit"])
    for idx, row in base_mm.iterrows()
}
vol_base_added = 6420

def _reagent_from_col(col):
    """Pull the reagent token from headers like
    '[Magnesium acetate] (mM)', '[ATP]_mM', '[tRNA] ug/ul'.
    """
    # s = str(col).replace("_", " ")
    parts = split_concentration_column_name(col)
    return parts[0] if parts else s  # parts[0] keeps wrappers; normalize strips them

for col in labcraft_platemap.select_dtypes("number").columns:
    key = normalize_reagent_name(_reagent_from_col(col))
    if key in base_lookup:
        base_conc, unit = base_lookup[key]
        labcraft_platemap[col] += base_conc * vol_base_added / labcraft_platemap['actual final volume nL']
        print(f"Added base_conc for '{col}': +{base_conc} {unit}")
    else:
        print(f"  (no base_conc match for '{col}' → key={key!r})")

labcraft_platemap

FACTOR_COLS = labcraft_platemap.columns[1:-1]
# FACTOR_COLS = ['[ATP] mM', '[Creatine phosphate] mM', '[GTP] mM',
#        '[Magnesium acetate] mM', '[PolyP] mM', '[PPK] uM', '[tRNA] ug/ul']
# FACTOR_COLS = ['[ATP]_mM', '[Creatine phosphate]_mM', '[GTP]_mM',
#        '[Magnesium acetate]_mM', '[PolyP]_mM', '[PPK]_uM', '[tRNA]_ug/ul']

# 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

samples.Name.unique()

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

samples.index

# Specify file paths
data_file = "data/20260519-132240-cytation5-pure-timecourse-gfp-DiscoveryPlate_R3b-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

# data.to_csv("samples_final_with_data.csv")

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

data.Name.unique()

data.Generator.unique()

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

g = pr.plot_curves(data=data[data.Generator == "lhs"])

g = pr.plot_curves(data=data[data.Generator == "ges"])

g = pr.plot_curves(data=data[data.Generator.isin(['manual', 'control'])])
data[data.Generator.isin(['manual', 'control'])][FACTOR_COLS]

data.columns

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

kinetics

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

# g = pr.plot_kinetics(data[data.Type=="Control"],  group_by=['Well', 'Name', *FACTOR_COLS])

# Which samples have 0?

# 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, .2, 0.4, 1.0],
  labels=['0-0.05', '0.05-0.2', '0.2-0.4', '>0.4']
)                                                                                                                                                                                                               
# val = "Steady_State_data_normalized"
val = 'steady_state (binned)'
g = sns.pairplot(
  kin[factors + [val]],
  vars=factors,
  hue=val,
  palette ="rocket_r",
  # palette={'0-0.05': '#d9534f', '0.05-0.2': 'orange', '>0.2': 'blue'},
  plot_kws={'alpha': 0.7, 's': 40},
  diag_kind='hist',
)