import streamlit as st
from fengxai.render import setModel, getModel, updateModel, applyBindingController
from fengxai.inference import imagePrediction
from PIL import Image


## model
setModel("distance", 30.07)                                                         # (英尺/米)到物体的距离
setModel("sensorHeight", 24)                                                        # (毫米)传感器高度
setModel("imageObjectHeight", 2724)                                                 # (像素)图像中物体的高度
setModel("lensFocalLength", 60)                                                     # (像素)镜头焦距
setModel("imageHeight", 4912)                                                       # (像素)图片的高度
## and controller
updateModel(                                                                        # 物体高度
    "objectHeight", 
    str(getModel("sensorHeight") * getModel("imageObjectHeight") * getModel("distance") / (getModel("lensFocalLength") * getModel("imageHeight")))
)
## (default)set image info
setModel("imageInfo")

def changeUploadFile():
    updateModel("imageInfo", {
        "imageFile": getModel("imageUploadFile"),
        "isUploadImageFile": True                                                   # 图片上传开关，控制ai模型在正确的事件下识别图像物体
    })
    
def changeDelete():
    updateModel("imageInfo", "")

## view: side 左侧边栏
def render():
    st.set_page_config(layout="wide") 

    with st.sidebar:
        st.title("example")
        st.image("assets/e2.jpg")
        if st.button('use'):
            updateModel("imageInfo", {
                "imageFile": "assets/e2.jpg",
                "isUploadImageFile": True
            })

    # page
    st.title("Rjxai v4_1: :blue[Get object height]")
    with st.container():
        st.title("Step 1: Rjxai Image Identification")
        st.file_uploader(
            label = "Upload your image here: png, jpg, jpeg", 
            type=['png','jpg','jpeg'],
            on_change=changeUploadFile,
            key="imageUploadFile",
        )
        if getModel('imageInfo') != "":
            def callback(output):
                print("output=", output)
                updateModel("imageHeight", output["imageHeight"])
                updateModel("imageObjectHeight", output["objectHeight"])
                c1, c2 = st.columns([0.7, 0.3])
                with c1.container():
                    imageInput = Image.open(output["imageOutput"])
                    st.image(imageInput)
                with c2.container():
                    st.button("delete image", on_click=changeDelete)

            if getModel("imageInfo")["isUploadImageFile"]:
                def tmpCallback(output):
                    updateModel("imageInfo", {
                        "imageFile": getModel("imageInfo")["imageFile"],
                        "tmpUploadImageInfo": output,                               # 存放图片的临时信息
                        "isUploadImageFile": False
                    })
                    print("output=", output)
                    callback(output)
                imagePrediction(getModel('imageInfo')["imageFile"], tmpCallback)
            else:
                callback(getModel("imageInfo")["tmpUploadImageInfo"])

    with st.container():
        st.title("Step 2: Calculate real object height")

    col1, col2 = st.columns(2)
    with col1.container():
        st.subheader("Input:")
        st.number_input('(pixels) image height', key="imageHeight")
        st.number_input('(pixels) image object height', key="imageObjectHeight")
        st.number_input('(mm) sensor height', key="sensorHeight")
        st.number_input('(mm) lens focal length', key="lensFocalLength")
        st.number_input('(feet or meters) distance', key="distance")
        
    with col2:
        st.subheader("Output:")
        st.text_area("(feet or meters) object height:", key="objectHeight", disabled=True)

# control 绑定
applyBindingController(render)