# %%
import numpy as np
import sympy as sp
from scipy.optimize import root_scalar
import ultraplot as uplt
from smitfit.symbol import Symbols
from smitfit.model import Model
import polars as pl
# %%

s = Symbols("TF1, TF2, TFR, R, T_TF, T_R, kD_MD, kD_MR", positive=True)
# %%
mb_ribosome = s.TFR + s.R - s.T_R  # type: ignore
mb_TF = s.TF1 + 2 * s.TF2 + s.TFR - s.T_TF  # type: ignore
eq_MD = s.TF1**2 - s.kD_MD * s.TF2  # type: ignore
eq_MR = (s.TF1 * s.R) - s.kD_MR * s.TFR  # type: ignore

#
knowns = ["T_TF", "T_R", "kD_MD", "kD_MR"]

# solve for: TF1
# take Monomer-dimer equillibrium, put it in mass balance TF to eliminate TF2
sub_TF2 = (s.TF2, sp.solve(eq_MD, s.TF2)[0])

# same for monomer dimer, eliminate TF_R
sub_mb = (s.R, sp.solve(mb_ribosome, s.R)[0])
sub_TFR = (s.TFR, sp.solve(eq_MR.subs(*sub_mb), s.TFR)[0])

# we know have an expr to find free TF monomer
eq_TF1 = mb_TF.subs([sub_TF2, sub_TFR])
eq_TF1


# %%

d = {
    s.TF2: sp.solve(eq_MD, s.TF2)[0],
    s.TFR: sp.solve(mb_TF, s.TFR)[0],
    s.R: sp.solve(mb_ribosome, s.R)[0],
}
m = Model(d)

# %%

ld = sp.lambdify([s.TF1] + [s[k] for k in knowns], eq_TF1)

# %%


def solve_system(params: dict) -> dict:
    args = tuple(params[k] for k in knowns)
    # root find TF1
    sol = root_scalar(ld, bracket=(0, params["T_TF"]), args=args)
    # calculate the others
    ans = m(**params, TF1=sol.root)

    return {"TF1": sol.root, **ans}


def make_df(records: list[dict]) -> pl.DataFrame:
    df = pl.DataFrame(records)
    cols = [
        (pl.col("TF1") + pl.col("TF2") + pl.col("TFR")).alias("total TF"),
        (pl.col("TFR") + pl.col("R")).alias("total R"),
    ]
    df = df.with_columns(cols)
    return df


# %%
# The concentration of TF exceeds that of ribosomes (∼50 μM and ∼30 μM, respectively)[12,13]
# TF binds free ribosomal 50S subunits with a Kd of ∼1 μM
# Purified TF forms dimers with a Kd of 1–2 μM (ref. 14).
# " Real-time observation of trigger factor function on translating ribosomes", https://doi.org/10.1038/nature05225

# "the cytosol contains 2.6 moles of trigger factor per mole of ribosomes. "
# The “trigger factor cycle” includes ribosomes, presecretory proteins, and the plasma membrane
#

ecoli_params = {
    "T_TF": 50,
    "T_R": 30,
    "kD_MD": 1,
    "kD_MR": 2,
}

solve_system(ecoli_params)

# %%

vmin, vmax = 1e-6, 1000e-6
total_tf_protomer = np.logspace(
    0, 4, endpoint=True, num=100
)  # total TF protomer from 1 uM to 1

input_params = [ecoli_params | {"T_TF": v} for v in total_tf_protomer]
input_params

# %%
output_records = [solve_system(params) for params in input_params]
df = make_df(output_records)


# %%
species = ["TF1", "TF2", "TFR"]
cycle = iter(uplt.Cycle("default"))
fig, axes = uplt.subplots(nrows=2, aspect=2.5, axwidth="120mm")
for s in species:
    axes[0].plot(
        total_tf_protomer,
        df.select(pl.col(s) / pl.col("total TF")),
        label=s,
        **next(cycle),
    )

axes[0].format(title="TF")
species = ["TFR", "R"]
for s in species:
    axes[1].plot(
        total_tf_protomer,
        df.select(pl.col(s) / pl.col("total R")),
        label=s,
        **next(cycle),
    )
axes[1].format(title="ribosome")

axes.format(
    ylim=(0, 1),
    xscale="log",
    xformatter="sci",
    ylabel="Fractional population",
    xlabel="Total protomer TF (uM)",
)
axes.legend(loc="r", ncols=1)

# %%
# lets repeat for ribosome titration, fixed 100 nM TF protomer concentration

microscopy_params = {
    "T_TF": 1,  # 100 nM
    "kD_MD": 1,
    "kD_MR": 1,
}

# %%

total_ribosome = np.linspace(0, 5, endpoint=True)
input_params = [microscopy_params | {"T_R": v} for v in total_ribosome]

# %%
output_records = [solve_system(params) for params in input_params]
df = make_df(output_records)

# %%
species = ["TF1", "TF2", "TFR"]
cycle = iter(uplt.Cycle("default"))

fig, axes = uplt.subplots(nrows=2, aspect=2.5, axwidth="120mm")
for s in species:
    axes[0].plot(
        total_ribosome,
        df.select(pl.col(s) / pl.col("total TF")),
        label=s,
        **next(cycle),
    )

axes[0].format(title="TF")

species = ["TFR", "R"]
for s in species:
    axes[1].plot(
        total_ribosome,
        df.select(pl.col(s) / pl.col("total R")),
        label=s,
        **next(cycle),
    )
axes[1].format(title="ribosome")

axes.format(
    ylim=(0, 1),
    # xscale="log",
    xformatter="sci",
    ylabel="Fractional population",
    xlabel="Total protomer TF (uM)",
)
axes.legend(loc="r", ncols=1)

# %%