app.py

import io
from pathlib import Path
from typing import Callable, Optional, cast
from urllib.parse import parse_qsl

import altair as alt
import polars as pl
import reacton.core
import solara
import solara.lab
from cmap import Catalog, Colormap
from ipyvuetify.extra import FileInput

from make_link import encode_url
from viewer import AxisProperties, ColorTransform, ProteinView, RoutedView

DEFAULT_CMAP = "tol:rainbow_PuRd"
NORM_CATEGORIES = ["linear", "diverging", "categorical"]
VMIN_DEFAULT = 0.0
VMAX_DEFAULT = 1.0
HIGHLIGHT_COLOR = "#e933f8"
MISSING_DATA_COLOR = "#8c8c8c"
CMAP_OPTIONS = list(Catalog().namespaced_keys())

# BASE_URL = "http://localhost:8765" # local testing
BASE_URL = "https://huggingface.co/spaces/Jhsmit/ipymolstar-annotate-colors"

pth = Path(__file__).parent


@solara.component
def FileInputComponent(
    on_file: Callable[[solara.components.file_drop.FileInfo | None], None],
):
    """Adaptation of _FileDrop."""

    file_info, set_file_info = solara.use_state(None)
    wired_files, set_wired_files = solara.use_state(
        cast(Optional[list[solara.components.file_drop.FileInfo]], None)
    )

    file_drop = FileInput.element(on_file_info=set_file_info, multiple=False)  # type: ignore

    def wire_files():
        if not file_info:
            set_wired_files([])
            return

        real = cast(FileInput, solara.get_widget(file_drop))

        # workaround for @observe being cleared
        real.version += 1
        real.reset_stats()

        set_wired_files(
            cast(list[solara.components.file_drop.FileInfo], real.get_files())
        )

    solara.use_effect(wire_files, [file_info])

    def handle_file():
        if not wired_files:
            on_file(None)
            return
        if on_file:
            f = wired_files[0].copy()
            f["data"] = None
            on_file(f)

    solara.lab.use_task(handle_file, dependencies=[wired_files])

    return file_drop


@solara.component
def ColorPickerMenuButton(title: str, color: solara.Reactive[str]):
    local_color = solara.use_reactive(color.value)

    def on_open(value: bool):
        if not value:
            color.set(local_color.value)

    btn = solara.Button(title, color=local_color.value)
    with solara.lab.Menu(
        activator=btn, close_on_content_click=False, on_open_value=on_open
    ):
        solara.v.ColorPicker(
            v_model=local_color.value,
            on_v_model=local_color.set,
        )


empty_frame = pl.DataFrame()
R_DEFAULT = ""
V_DEFAULT = ""

# %%


@solara.component
def MainApp():
    dark_effective = solara.lab.use_dark_effective()
    title = solara.use_reactive("My annotated protein view")
    description = solara.use_reactive("")
    molecule_id = solara.use_reactive("1QYN")
    data = solara.use_reactive(empty_frame)
    warning_text = solara.use_reactive("")

    residue_column = solara.use_reactive(R_DEFAULT)
    color_column = solara.use_reactive(V_DEFAULT)

    label = solara.use_reactive("value")
    unit = solara.use_reactive("au")

    highlight_color = solara.use_reactive(HIGHLIGHT_COLOR)
    missing_data_color = solara.use_reactive(MISSING_DATA_COLOR)
    autoscale_y = solara.use_reactive(True)

    cmap_name = solara.use_reactive(DEFAULT_CMAP)
    reverse = solara.use_reactive(False)
    full_cmap_name = cmap_name.value + "_r" if reverse.value else cmap_name.value
    cmap = Colormap(full_cmap_name)

    vmin = solara.use_reactive(VMIN_DEFAULT)
    vmax = solara.use_reactive(VMAX_DEFAULT)
    norm_type = solara.use_reactive(NORM_CATEGORIES[0])

    rc = reacton.core.get_render_context()

    def on_file(file_info: solara.components.file_drop.FileInfo | None):
        if not file_info:
            data.set(pl.DataFrame())
            return

        try:
            df = pl.read_csv(file_info["file_obj"])
        except Exception as e:
            warning_text.set(str(e))
            return
        if len(df.columns) < 2:
            warning_text.set(f"Expected at least 2 columns, got {len(df.columns)}")
            data.set(pl.DataFrame())
            return

        warning_text.set("")

        # order matters!
        # with solara.batch_update():
        # https://github.com/widgetti/solara/issues/637
        with rc:
            residue_column.set(df.columns[0])
            color_column.set(df.columns[1])
            data.set(df)

    colors = ColorTransform(
        name=full_cmap_name,
        norm_type=norm_type.value,
        vmin=vmin.value,
        vmax=vmax.value,
        missing_data_color=missing_data_color.value,
        highlight_color=highlight_color.value,
    )

    axis_properties = AxisProperties(
        label=label.value,
        unit=unit.value,
        autoscale_y=autoscale_y.value,
    )

    if data.value.is_empty():
        data_view = pl.DataFrame({"residue_number": [], "value": []})
    else:
        data_view = pl.DataFrame(
            {
                "residue_number": data.value[residue_column.value],
                "value": data.value[color_column.value],
            }
        )

    def load_example_data():
        bio = io.BytesIO(Path("example_data.csv").read_bytes())
        bio.seek(0)
        file_info = solara.components.file_drop.FileInfo(
            name="example_data.csv",
            size=Path("example_data.csv").stat().st_size,
            file_obj=bio,
            data=None,
        )

        # with solara.batch_update():
        # https://github.com/widgetti/solara/issues/637
        with rc:
            molecule_id.set("6GOX")
            title.set("SecA HDX-MS protein local flexibility")
            on_file(file_info)
            vmin.set(4e4)
            vmax.set(1e4)
            autoscale_y.set(False)
            description.set(Path("example_data_description.md").read_text())

    with solara.AppBar():
        with solara.Tooltip("Load example data and settings"):
            solara.Button(
                icon_name="mdi-test-tube",
                icon=True,
                on_click=load_example_data,
            )

    query_string = encode_url(
        title=title.value,
        molecule_id=molecule_id.value,
        colors=colors,
        axis_properties=axis_properties,
        data=data_view,
        description=description.value,
    )

    ProteinView(
        title.value,
        molecule_id=molecule_id.value,
        data=data_view,
        colors=colors,
        axis_properties=axis_properties,
        dark_effective=dark_effective,
        description=description.value,
    )

    with solara.Sidebar():
        with solara.Card("Settings"):
            solara.InputText(label="Title", value=title)
            solara.InputText(label="PDB ID", value=molecule_id)
            solara.Text("Choose .csv data file:")
            FileInputComponent(on_file)

            if warning_text.value:
                solara.Warning(warning_text.value)

            if not data.value.is_empty():
                with solara.Row():
                    solara.Select(
                        label="Residue Column",
                        value=residue_column,
                        values=list(data.value.columns),
                    )
                    solara.Select(
                        label="Color Column",
                        value=color_column,
                        values=list(data.value.columns),
                    )

            with solara.Row():
                solara.InputText(label="Label", value=label)
                solara.InputText(label="Unit", value=unit)

            solara.Text("Colors")
            with solara.Row(gap="10px", justify="space-around"):
                ColorPickerMenuButton("Highlight", highlight_color)
                ColorPickerMenuButton("Missing data", missing_data_color)

            def set_cmap_name(name: str):
                try:
                    Colormap(name)
                    cmap_name.set(name)
                except TypeError:
                    pass

            # with solara.Row():
            solara.v.Autocomplete(
                v_model=cmap_name.value,
                on_v_model=set_cmap_name,
                items=CMAP_OPTIONS,
            )

            solara.Select(
                label="Normalization type", value=norm_type, values=NORM_CATEGORIES
            )

            with solara.Row():

                def set_vmin(value: float):
                    if norm_type.value == "diverging":
                        vmin.set(value)
                        vmax.set(-value)
                    else:
                        vmin.set(value)

                solara.InputFloat(
                    label="vmin",
                    value=vmin.value,
                    on_value=set_vmin,
                    disabled=norm_type.value == "categorical",
                )
                solara.InputFloat(
                    label="vmax",
                    value=vmax,
                    disabled=norm_type.value in ["diverging", "categorical"],
                )
            with solara.GridFixed(columns=2):
                with solara.Tooltip("Reverses the color map"):
                    solara.Checkbox(label="Reverse", value=reverse)

                with solara.Tooltip(
                    "Uncheck to use color range as the scatterplot y scale"
                ):
                    solara.Checkbox(label="Autoscale Y", value=autoscale_y)

            solara.Image(cmap._repr_png_(height=24), width="100%")
            solara.InputTextArea(
                label="Description", value=description, continuous_update=True
            )
            solara.Div(style={"height": "10px"})

            solara.Button(
                label="Open view in new tab",
                attributes={"href": BASE_URL + "?" + query_string, "target": "_blank"},
                block=True,
            )


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

    dark_effective = solara.lab.use_dark_effective()
    dark_previous = solara.use_previous(dark_effective)
    if dark_previous != dark_effective:
        if dark_effective:
            alt.themes.enable("dark")
        else:
            alt.themes.enable("default")

    # todo: aways true, check valid
    if route.search:
        query_dict = {k: v for k, v in parse_qsl(route.search)}
        # needs more keys but if this is there then at least there has been an attempt
        if "molecule_id" in query_dict:
            RoutedView()

    else:
        MainApp()


# %%