!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')}"

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

from cdk.analysis.cytosol import platereader as pr
# import platereader as pr
# Initialize plotting
pr.plot_setup()

import sys

# !ls bnext/experiments/20251106_labcraft_protein_sweep_01

data = pr.load_platereader_data("data/20251106-191540-cytation5-pure-timecourse-gfp-labcraft-protein-sweep1-biotek-cdk.txt")

platemap = pd.read_csv("data/concentration_key.csv")
# platemap.columns = ["Well", "Mg (mM)", "Ribo (uM)", "T7 (ng/ul)", "Water"]
# platemap["Name"] = "GFP " + platemap["[DNA] (nM)"].astype(str) + " nM"
platemap.rename(columns={'Well ID':"Well"}, inplace=True)
platemap = platemap.fillna(0)
# platemap.columns
sweep_cols= ['Mg_mM', 'Ribo_uM', 'T7poly_ng/ul']
for values, g in platemap.groupby(sweep_cols):
    mg, ribo, T7 = values
    name = f"Mg {mg}; R {ribo}; T7 {T7}"
    platemap.loc[g.index, "Name"] = name

platemap

data = data.merge(platemap)

platemap.head()
data.head()

data

pr.plot_curves(data)

# pr.plot_curves(data, hue=sweep_cols[0], col=sweep_cols[2], palette="viridis")
print(sweep_cols)

sns.relplot(data, x="Time", y="Data", row=sweep_cols[0], col=sweep_cols[2], hue=sweep_cols[1], kind="line") #, hue="Name")

kinetics = pr.kinetic_analysis(data, group_by=["Name"])
kinetics

kinetics.columns
ss = kinetics["Steady State"]["Data"].reset_index().merge(platemap)
# ss = kinetics["Velocity"]["Max"].reset_index().merge(platemap)

# data_name = 'Max' #Data
data_name = 'Data'

pr.plot_kinetics(data, group_by=["Name"])

ss.head()

ss.to_csv('steady_states.csv', index=False)

df = ss

# Average replicates
means = df.groupby(['Mg_mM', 'Ribo_uM', 'T7poly_ng/ul'], as_index=False)[data_name].mean()

# Global min/max for shared color bar
data_min = means[data_name].min()
data_max = means[data_name].max()

t7_values = means['Mg_mM'].unique()
n_plots = len(t7_values)
fig, axes = plt.subplots(1, n_plots, figsize=(7*n_plots, 6), squeeze=False)

for idx, t7 in enumerate(sorted(t7_values)):
    slice_df = means[means['Mg_mM'] == t7]
    pivot = slice_df.pivot(index='Ribo_uM', columns='T7poly_ng/ul', values=data_name)
    ax = axes[0, idx]
    hm = sns.heatmap(
        pivot, annot=True, cmap='viridis',
        vmin=data_min, vmax=data_max,  # shared color scale
        cbar=idx==n_plots-1, cbar_ax=None if idx < n_plots-1 else plt.gcf().add_axes([0.93, 0.3, 0.02, 0.4]),
        ax=ax
    )
    ax.set_title(f'Mg_mM={t7}')
    ax.set_xlabel('T7poly_ng/ul')
    ax.set_ylabel('Ribo_uM')

plt.suptitle("Vmax")
plt.tight_layout(rect=[0, 0, 0.9, 1])  # Leave space for color bar
plt.show()

df = ss
# Average replicates
means = df.groupby(['Mg_mM', 'Ribo_uM', 'T7poly_ng/ul'], as_index=False)[data_name].mean()

# Global min/max for shared color bar
data_min = means[data_name].min()
data_max = means[data_name].max()

t7_values = means['T7poly_ng/ul'].unique()
n_plots = len(t7_values)
fig, axes = plt.subplots(1, n_plots, figsize=(7*n_plots, 6), squeeze=False)

for idx, t7 in enumerate(sorted(t7_values)):
    slice_df = means[means['T7poly_ng/ul'] == t7]
    pivot = slice_df.pivot(index='Ribo_uM', columns='Mg_mM', values=data_name)
    ax = axes[0, idx]
    hm = sns.heatmap(
        pivot, annot=True, cmap='viridis',
        vmin=data_min, vmax=data_max,  # shared color scale
        cbar=idx==n_plots-1, cbar_ax=None if idx < n_plots-1 else plt.gcf().add_axes([0.93, 0.3, 0.02, 0.4]),
        ax=ax
    )
    ax.set_title(f'T7poly_ng/ul={t7}')
    ax.set_xlabel('Mg_mM')
    ax.set_ylabel('Ribo_uM')

plt.suptitle("Vmax")
plt.tight_layout(rect=[0, 0, 0.9, 1])  # Leave space for color bar
plt.show()

# data_name = 'Max' #Data
# Average replicates
means = df.groupby(['Mg_mM', 'Ribo_uM', 'T7poly_ng/ul'], as_index=False)[data_name].mean()

# 3D scatter/Surface for each unique T7poly_ng/ul, or all together using T7poly as color:
fig = plt.figure(figsize=(10, 7))
ax = fig.add_subplot(111, projection='3d')

# Normalize T7poly for color mapping
norm = plt.Normalize(means['T7poly_ng/ul'].min(), means['T7poly_ng/ul'].max())
colors = plt.cm.viridis(norm(means['T7poly_ng/ul']))

scat = ax.scatter(means['Mg_mM'], means['Ribo_uM'], means[data_name],
                  c=colors, s=60, edgecolor='k')

# Add colorbar for T7poly_ng/ul
mappable = plt.cm.ScalarMappable(cmap='viridis', norm=norm)
mappable.set_array(means['T7poly_ng/ul'])
fig.colorbar(mappable, ax=ax, label='T7poly_ng/ul')

ax.set_xlabel('Mg_mM')
ax.set_ylabel('Ribo_uM')
ax.set_zlabel('Steady State Value')
ax.set_title('Steady State as a Function of Mg, Ribo, and T7RNAP')

plt.tight_layout()
plt.show()

# Build a continuous palette for Ribo
norm = plt.Normalize(df['Ribo_uM'].min(), df['Ribo_uM'].max())
cmap = plt.cm.viridis

g = sns.relplot(
    data=df,
    x="Mg_mM", y="Data",
    col="T7poly_ng/ul", hue="Ribo_uM",  err_style='band', errorbar='sd',
    kind="line",
    facet_kws={'sharey': True, 'sharex': True},
    hue_norm=(df['Ribo_uM'].min(), df['Ribo_uM'].max())
)

df.columns

def plot_main_effects(data, response_var):
    # Get all columns except the response variable
    factors = [col for col in data.columns if col != response_var]
    
    # Plot each factor's main effect
    for factor in factors:
        plt.figure(figsize=(6, 4))
        sns.pointplot(x=factor, y=response_var, data=data, errorbar="sd", markers="o", linestyles="-")
        plt.title(f"Main Effect of {factor} on {response_var}")
        plt.xlabel(factor)
        plt.ylabel(response_var)
        plt.show()
df["Steady State (AFU)"] = df["Data"]
plot_main_effects(df[['Steady State (AFU)', 'Mg_mM', 'Ribo_uM', 'T7poly_ng/ul']], "Steady State (AFU)")

import math
def plot_main_effects(data, response_var):
    # Get all columns except the response variable
    factors = [col for col in data.columns if col != response_var]
    n = len(factors)

    # Determine subplot grid size
    cols = 3
    rows = math.ceil(n / cols)

    fig, axes = plt.subplots(rows, cols, figsize=(12, 4 * rows))
    axes = axes.flatten()  # flatten in case of 1 row

    for ax, factor in zip(axes, factors):
        sns.pointplot(
            x=factor,
            y=response_var,
            data=data,
            errorbar="sd",
            markers="o",
            linestyles="-",
            ax=ax
        )
        ax.set_title(f"Main Effect of {factor} on {response_var}")
        ax.set_xlabel(factor)
        ax.set_ylabel(response_var)

    # Hide any empty subplots
    for i in range(len(factors), len(axes)):
        axes[i].set_visible(False)

    plt.tight_layout()
    plt.show()

df["Steady State (AFU)"] = df["Data"]
plot_main_effects(df[['Steady State (AFU)', 'Mg_mM', 'Ribo_uM', 'T7poly_ng/ul']], "Steady State (AFU)")

import numpy as np
from itertools import combinations
def plot_all_interactions(data, response_var):
    factors = [col for col in data.columns if col != response_var]
    factor_pairs = list(combinations(factors, 2))
    n = len(factor_pairs)

    cols = 3
    rows = math.ceil(n / cols)

    fig, axes = plt.subplots(rows, cols, figsize=(12, 5 * rows))
    axes = axes.flatten()

    for ax, (factor_x, factor_hue) in zip(axes, factor_pairs):

        levels = np.sort(data[factor_hue].unique())

        # Build a continuous colormap (viridis is great for continuous data)
        cmap = sns.color_palette("viridis", as_cmap=True)

        # Normalize hue values to [0,1] for colormap mapping
        if np.issubdtype(data[factor_hue].dtype, np.number):
            norm_levels = (levels - levels.min()) / (levels.max() - levels.min() + 1e-9)
        else:
            # For categorical: evenly spaced colors
            norm_levels = np.linspace(0, 1, len(levels))

        color_map = {level: cmap(norm) for level, norm in zip(levels, norm_levels)}

        # Plot per hue level
        for level in levels:
            subset = data[data[factor_hue] == level]

            # Scatter points
            sns.scatterplot(
                data=subset,
                x=factor_x,
                y=response_var,
                ax=ax,
                color=color_map[level],
                label=f"{factor_hue}: {level}"
            )

            # CI-smoothed line, same color
            sns.lineplot(
                data=subset,
                x=factor_x,
                y=response_var,
                ax=ax,
                color=color_map[level],
                errorbar="sd"
            )

        ax.set_title(f"Interaction: {factor_x} × {factor_hue}")
        ax.set_xlabel(factor_x)
        ax.set_ylabel(response_var)

    # Hide unused axes
    for i in range(len(factor_pairs), len(axes)):
        axes[i].set_visible(False)

    plt.tight_layout()
    plt.show()

plot_all_interactions(df[['Steady State (AFU)', 'Mg_mM', 'Ribo_uM', 'T7poly_ng/ul']], "Steady State (AFU)")