# %%

from pathlib import Path

import altair as alt
import numpy as np
import pandas as pd
import solara
import solara.lab
import sympy as sp
from scipy.optimize import root_scalar

# %%
P1, P2, P4, PT, kD1, kD2 = sp.symbols("P_1 P_2 P_4 P_T k_D1 k_D2", positive=True)

# %%
sub_p1_p2 = (P1, sp.solve(kD1 * (P2 / P1**2) - 1, P1)[0])
sub_p2_p4 = (P2, sp.solve(kD2 * (P4 / P2**2) - 1, P2)[0])
sub_p4_p2 = (P4, sp.solve(kD2 * (P4 / P2**2) - 1, P4)[0])

# %%
mass_balance = P1 + 2 * P2 + 4 * P4 - PT

eq_p4 = mass_balance.subs([sub_p1_p2, sub_p2_p4])
eq_p2 = mass_balance.subs([sub_p1_p2, sub_p4_p2])

# %%


def make_df(vmin: float, vmax: float, kD_1_v: float, kD2_v: float) -> pd.DataFrame:
    PT_values = np.logspace(np.log10(vmin), np.log10(vmax), endpoint=True, num=100)

    kd_subs = [(kD1, kD_1_v), (kD2, kD2_v)]

    ld = sp.lambdify([P4, PT], eq_p4.subs(kd_subs))
    P4_values = np.array(
        [root_scalar(ld, bracket=(0, PT_v), args=(PT_v,)).root for PT_v in PT_values]
    )

    ld = sp.lambdify([P2, PT], eq_p2.subs(kd_subs))
    P2_values = np.array(
        [root_scalar(ld, bracket=(0, PT_v), args=(PT_v,)).root for PT_v in PT_values]
    )

    P1_values = PT_values - 2 * P2_values - 4 * P4_values

    columns = {"P1": P1_values, "P2": P2_values, "P4": P4_values}
    total = np.sum(list(columns.values()), axis=0)

    df = pd.DataFrame(dict(PT=PT_values) | {k: v / total for k, v in columns.items()})

    return df


def make_chart(df: pd.DataFrame) -> alt.LayerChart:
    source = df.melt("PT", var_name="species", value_name="y")

    # Create a selection that chooses the nearest point & selects based on x-value
    nearest = alt.selection_point(
        nearest=True, on="pointerover", fields=["PT"], empty=False
    )

    # The basic line
    line = (
        alt.Chart(source)
        .mark_line(interpolate="basis")
        .encode(
            x=alt.X(
                "PT:Q",
                scale=alt.Scale(type="log"),
                title="Total protomer concentration",
            ),
            y=alt.Y("y:Q", title="Fraction of total"),
            color="species:N",
        )
        .properties(width="container")
    )

    # Draw points on the line, and highlight based on selection
    points = (
        line.mark_point()
        .encode(opacity=alt.condition(nearest, alt.value(1), alt.value(0)))
        .properties(width="container")
    )

    # Draw a rule at the location of the selection
    rules = (
        alt.Chart(source)
        .transform_pivot("species", value="y", groupby=["PT"])
        .mark_rule(color="black")
        .encode(
            x="PT:Q",
            opacity=alt.condition(nearest, alt.value(0.3), alt.value(0)),
            tooltip=[
                alt.Tooltip(c, type="quantitative", format=".2f") for c in df.columns
            ],
        )
        .add_params(nearest)
        .properties(width="container")
    )

    # Put the five layers into a chart and bind the data
    chart = (
        alt.layer(line, points, rules)
        .properties(height=300)
        .configure(autosize="fit-x")
    )

    return chart


md = """
This app calculates monomer and dimer concentrations given a total amount of protomer PT and the 
dissociation constant KD. More info on how and why can be found [HuggingFace](https://huggingface.co/spaces/Jhsmit/binding-kinetics) (right click, open new tab).
"""


@solara.component
def Page():
    solara.Style(Path("style.css"))

    dark_effective = solara.lab.use_dark_effective()
    if dark_effective is True:
        alt.themes.enable("dark")

    elif dark_effective is False:
        alt.themes.enable("default")

    kD1 = solara.use_reactive(1.0)
    kD2 = solara.use_reactive(100)

    vmin = solara.use_reactive(1e-3)
    vmax = solara.use_reactive(1e3)

    async def update():
        df = make_df(vmin.value, vmax.value, kD1.value, kD2.value)
        chart = make_chart(df)

        return chart

    task: solara.lab.Task = solara.lab.use_task(
        update, dependencies=[kD1.value, kD2.value, vmin.value, vmax.value]
    )

    solara.Title("Tetramerization Kinetics")

    with solara.Card("Fraction monomer/dimer/tetramer"):
        with solara.GridFixed(columns=2):
            with solara.Tooltip("Dissociation constant monomer/dimer"):
                solara.InputFloat("kD1", value=kD1)
            with solara.Tooltip("Dissociation constant dimer/tetramer"):
                solara.InputFloat("kD2", value=kD2)
            with solara.Tooltip("X axis lower limit"):
                solara.InputFloat("xmin", value=vmin)
            with solara.Tooltip("X axis upper limit"):
                solara.InputFloat("xmax", value=vmax)
        solara.HTML(tag="div", style="height: 10px")

        if task.finished:
            solara.FigureAltair(task.value)


# %%