# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/51_deploy_app.ipynb (unless otherwise specified).

__all__ = ['deploy_app', 'StreamlitApp']

# Cell
import json
import requests
from dataclasses import dataclass, field
from pathlib import Path
from textwrap import dedent
from typing import Union

from jinja2 import Template, DebugUndefined


# Internal Cell
PathStr = Union[Path, str]

# Cell
def deploy_app(app="app.py", from_jupyter=True):
    """Deploy a streamlit app using ngrok"""
    if not Path(app).is_file():
        print(f"ERROR: the app {app} does not exist!")
        return

    if from_jupyter:
        get_ipython().system_raw('ngrok http 8501 &')
        resp = requests.get("http://localhost:4040/api/tunnels")
        tunnel = json.loads(resp.content)["tunnels"][0]
        local = tunnel["config"]["addr"]
        public = tunnel["public_url"]
    else:
        raise NotImplementedError(f"Deployment outside jupyter currently not supported")

    print(dedent(f"""\
        Let's create a tunnel to port {local.split(":")[-1]}!
        1. Create in a cell : !nohup streamlit run {app}
        2. Run that cell
        3. Click on this link: {public}

        Note: we recommend to put the the code `!nohup streamlit ...`
        after this line of code (instead of a new dedicated cell).
    """))

# Internal Cell
TEMPLATE_FLOW = """
from pathlib import Path

import streamlit as st

{{ specific_import }}

st.set_page_config(page_title="ML deployment, by unpackAI", page_icon="🚀")
st.image("https://unpackai.github.io/unpackai_logo.svg")
st.title("{{ title }}")
st.write("*by {{ author }}*")
st.write("---")

{{ load_model }}

{{ display_prediction }}

{% if (input or "").endswith("s") -%}
select = st.sidebar.radio("How to load {{ input }}?", ["from file{% if input.endswith("s") %}s{% endif %}", "from URL"])
{% else -%}
select = st.sidebar.radio("How to load {{ input }}?", ["from file", "from URL"])
{% endif %}
st.sidebar.write("---")

if select == "from URL":
    url = st.sidebar.text_input("url")
    if url:
        display_prediction(url)

else:
    {% if (input or "").endswith("s")   -%}
    files = st.sidebar.file_uploader("Choose {{ input }}, accept_multiple_files=True)
    for file in files:  # type:ignore # this is an iterable
        display_prediction(file)
    {% else -%}
    file = st.sidebar.file_uploader("Choose {{ input }}")
    if file:
        display_prediction(file)
    {% endif %}
"""

@dataclass
class TemplateCode:
    specific_import:str = ""
    load_model:str = ""
    display_predict:str = ""


# Internal Cell

# NOTE: We need re-implementation of elements defined during loading

TEMPLATE_CV_FASTAI = TemplateCode(
    specific_import="""
        from {# Go around nbdev #}unpackai.deploy.cv import get_image, get_learner, PathStr, dummy_function
    """
        ,
    load_model="""\
        {{ implem_4_model }}

        learn = get_learner(Path(__file__).parent / "{{ model }}")
        vocab = learn.dls.vocab
    """,
    display_predict="""
        def display_prediction(pic):
            img = get_image(pic)
            with learn.no_bar():
                prediction, idx, probabilities = learn.predict(img)
            col_img, col_pred = st.columns(2)
            col_img.image(img, caption=getattr(pic, "name", None))
            col_pred.write(f"### {prediction}")
            col_pred.metric(f"Probability", f"{probabilities[idx].item()*100:.2f}%")
    """,
)


# Internal Cell

# NOTE: Loading model and prediction depends on the module (e.g. "regression")

TEMPLATE_TABULAR_PYCARET = TemplateCode(
    specific_import="""
        import pandas as pd
        from pycaret.{{ module }} import load_model, predict_model
    """,
    load_model="""model = load_model("{{ model }}")""",
    display_predict="""
        def display_prediction(csv):
            df = pd.read_csv(csv)
            predictions = predict_model(model, data = df)
            st.dataframe(predictions)
    """,
)


# Internal Cell
TEMPLATE_NLP_HF = TemplateCode(
    specific_import="""
        # TODO
    """,
    load_model="""
        # TODO
    """,
    display_predict="""
        def display_prediction(csv:PathStr):
            pass # TODO
    """,
)


# Internal Cell

LIST_TEMPLATES = {
    "CV": {"fastai": TEMPLATE_CV_FASTAI},
    "Tabular": {"pycaret": TEMPLATE_TABULAR_PYCARET},
    "NLP": {"hugging_face": TEMPLATE_NLP_HF},
}


# Cell
@dataclass
class StreamlitApp:
    """Class to generate Streamlit App for different applications

    Available applications: "CV", "Tabular", "NLP"
    """

    application: str
    framework: str = field(init=False, default="")
    content: str = field(init=False, default="")
    _title: str = field(init=False, default="Streamlit App")

    def __post_init__(self):
        applications = list(LIST_TEMPLATES)
        if self.application not in applications:
            raise AttributeError(
                f"ERROR: Application {self.application} not supported. "
                f"Available applications: {applications}"
            )

    @property
    def template(self) -> Template:
        """Get the template"""
        base_template = Template(TEMPLATE_FLOW, undefined=DebugUndefined)
        try:
            template_framework: TemplateCode = LIST_TEMPLATES[self.application][
                self.framework
            ]
        except KeyError as missing:
            raise AttributeError(f"Element {missing} not found among list of templates")

        # Some elements can only be rendered later
        return Template(
            base_template.render(
                specific_import=dedent(template_framework.specific_import),
                display_prediction=dedent(template_framework.display_predict),
                load_model=dedent(template_framework.load_model),
            )
        )

    def _render(
        self, framework: str, title: str, author: str, model: PathStr, **kwargs
    ):
        """Generic function for rendering"""
        self.framework = framework
        self._title = title
        input_types = {
            "CV": "images",
            "Tabular": "CSV",
            "NLP": "Text",
        }
        self.content = self.template.render(
            title=title,
            author=author,
            model=model,
            input=input_types.get(self.application, "input"),
            **kwargs,
        )
        return self

    def render_fastai(
        self, title: str, author: str, model: PathStr, implem_4_model: str
    ) -> "StreamlitApp":
        """Render an app based on template

        Args:
            title: title of the App
            author: author of the App
            model: path of .pkl model to load (exported with `learn.export(...)`)
            implem_4_model: extra implementations needed to load the model
                (e.g. function used for labelling)
        """
        return self._render(
            "fastai", title, author, model, implem_4_model=dedent(implem_4_model)
        )

    def render_pycaret(
        self, title: str, author: str, model: PathStr, module: str
    ) -> "StreamlitApp":
        """Render an app based on template

        Args:
            title: title of the App
            author: author of the App
            model: path of .pkl model to load (exported with `learn.export(...)`)
            module: regression, classification, etc.
        """
        return self._render(
            "pycaret", title, author, model.replace(".pkl", ""), module=module
        )

    def append(self, content: str) -> "StreamlitApp":
        """Add additional content to the app"""
        self.content += dedent(content)
        return self

    def save(self, dest: PathStr, show=False):
        """Write the app to a file"""
        Path(dest).write_text(self.content, encoding="utf-8")
        print(f"Saved app '{self._title}' to '{dest}'")
        if show:
            print("-" * 20)
            print(self.content)
