from __future__ import annotations

import streamlit as st

from cref import source
from cref import views


@st.cache_resource(on_release=source.Source.close)
def data_source(variant: str = "release") -> source.Source:
    return source.Source.from_env(variant)


@st.cache_resource(on_release=source.History.close)
def history(variant: str = "release") -> source.History:
    return source.History.from_env(variant)


@st.cache_resource
def models() -> views.Models:
    return views.load_models()


@st.cache_resource
def tasks():
    return views.load_tasks()


def metric_control() -> str:
    return st.segmented_control(
        "Metric",
        list(views.METRICS),
        default="mase",
        required=True,
        format_func=views.METRICS.get,
        help=views.METRIC_HELP,
    )


def type_controls(models: views.Models) -> set[str]:
    present = [t for t in views.MODEL_TYPES if any(i.get("type") == t for i in models.values())]
    if not present:
        return set(views.MODEL_TYPES)
    picked = st.pills(
        "Model types",
        present,
        selection_mode="multi",
        default=present,
        format_func=lambda t: views.MODEL_TYPES[t].label,
        help="Hides the other models. The seasonal naive baseline is always shown.",
    )
    return set(picked or ())


def board_controls(models: views.Models) -> tuple[str, str, set[str]]:
    variants = data_source("all").get().board is not None
    cells = st.columns([1, 1.5, 3] if variants else [1, 4], vertical_alignment="bottom")
    with cells[0]:
        metric = metric_control()
    variant = "release"
    if variants:
        with cells[1]:
            variant = st.segmented_control(
                "Pretrained samples counted",
                list(views.VARIANTS),
                default="release",
                required=True,
                format_func=views.VARIANTS.get,
                help=views.VARIANT_HELP,
            )
    with cells[-1]:
        types = type_controls(models)
    return metric, variant, types