from pathlib import Path

import altair as alt
import numpy as np
import pandas as pd
import solara
import sympy as sp

#
P1, P2, PT, k_on, k_off, kD = sp.symbols("P_1 P_2 P_T k_on k_off k_D", positive=True)

sol = sp.solve(
    [
        -2 * k_on * P1 * P1 + 2 * k_off * P2,
        P1 + 2 * P2 - PT,
        (k_off / k_on) - kD,
    ],
    [P1, P2, k_on, k_off],
    dict=True,
)

solve_for = [P1, P2]
inputs = [PT, kD]

lambdas = {s: sp.lambdify(inputs, sol[0][s]) for s in solve_for}
ld_total = sp.lambdify(inputs, sol[0][P1] + sol[0][P2])


def make_chart(df: pd.DataFrame, dark: bool) -> alt.Chart:
    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="Ratio PT/kD"),
            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
    rule_color = "white" if dark else "black"
    rules = (
        alt.Chart(source)
        .transform_pivot("species", value="y", groupby=["PT"])
        .mark_rule(color=rule_color)
        .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")

    PT = solara.use_reactive(10.0)
    kD = solara.use_reactive(1.0)

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

    ans = {k: ld(PT.value, kD.value) for k, ld in lambdas.items()}

    solara.Title("Dimerization Kinetics")
    with solara.Card("Calculate concentrations from kD"):
        solara.Markdown(md)
        with solara.GridFixed(columns=2):
            with solara.Tooltip("Total protomer concentration"):
                solara.InputFloat("PT", value=PT)
            with solara.Tooltip("Dissociation constant"):
                solara.InputFloat("kD", value=kD)
            solara.Markdown(f"### Concentration monomer: {ans[P1]:.2f}")
            solara.Markdown(f"### Concentration dimer: {ans[P2]:.2f}")

    # create a vector of PT values ranging from 0.1 times kD to 1000 times kD
    def update():
        PT_values = np.logspace(vmin.value, vmax.value, endpoint=True, num=100)
        ans = {
            k: ld(PT_values, 1) / ld_total(PT_values, 1) for k, ld in lambdas.items()
        }

        # put the results in a dataframe, together with input PT values
        df = pd.DataFrame(dict(PT=PT_values) | {k.name: v for k, v in ans.items()})
        return make_chart(df, dark_effective)

    chart = solara.use_memo(update, [vmin.value, vmax.value])

    with solara.Card("Fraction monomer/dimer vs ratio over kD"):
        with solara.Row():
            with solara.Tooltip("X axis lower limit (log10)"):
                solara.InputFloat("xmin", value=vmin)
            with solara.Tooltip("X axis upper limit (log10)"):
                solara.InputFloat("xmax", value=vmax)
        solara.HTML(tag="div", style="height: 10px")
        solara.FigureAltair(chart)