import streamlit as st
from dnafiber.inference import infer
from dnafiber.postprocess.core import refine_segmentation
import numpy as np
from dnafiber.deployment import _get_model
import torch


@st.cache_data
def ui_inference(_model, _image, _device, postprocess=True, id=None):
    return ui_inference_cacheless(
        _model, _image, _device, postprocess=postprocess, id=id
    )


@st.cache_resource
def get_model(model_name):
    model = _get_model(
        device="cuda" if torch.cuda.is_available() else "cpu",
        revision=model_name,
    )
    return model


def ui_inference_cacheless(_model, _image, _device, postprocess=True, id=None):
    """
    A cacheless version of the ui_inference function.
    This function does not use caching and is intended for use in scenarios where caching is not desired.
    """
    h, w = _image.shape[:2]
    with st.spinner("Sliding window segmentation in progress..."):
        if isinstance(_model, list):
            output = None
            for model in _model:
                if isinstance(model, str):
                    model = get_model(model)
                with st.spinner(text="Segmenting with model: {}".format(model)):
                    if output is None:
                        output = infer(
                            model,
                            image=_image,
                            device=_device,
                            scale=st.session_state.get("pixel_size", 0.13),
                            only_probabilities=True,
                        ).cpu()
                    else:
                        output = (
                            output
                            + infer(
                                model,
                                image=_image,
                                device=_device,
                                scale=st.session_state.get("pixel_size", 0.13),
                                only_probabilities=True,
                            ).cpu()
                        )
            output = (output / len(_model)).argmax(1).squeeze().numpy()
        else:
            output = infer(
                _model,
                image=_image,
                device=_device,
                scale=st.session_state.get("pixel_size", 0.13),
            )
    output = output.astype(np.uint8)
    if postprocess:
        with st.spinner("Post-processing segmentation..."):
            output = refine_segmentation(output, fix_junctions=postprocess)
    return output