import glob
import os
import time

import numpy as np
import streamlit as st
import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
from SOURCE.siamese.model import SiameseDropout, SiameseNetwork

#from SOURCE.yolo_files import detect

MEDIA_ROOT = 'media/signatures/'
SIGNATURE_ROOT = 'media/UserSignaturesSquare/'
YOLO_RESULT = 'results/yolov5/'
YOLO_OP = 'crops/DLSignature/'


st.set_page_config(page_title="Signature Detection", page_icon="📈")

st.markdown("# Detect Signatures")
st.sidebar.header("Plotting Demo")
st.write(
    """Detect signatures of the given image!"""
)

if "predict" not in st.session_state:
    # Initialize session state.
    st.session_state.update({
        # Default page.
        "predict": False,
    })




def select_document():
    '''
        Selects the document from the dropdown menu and displays the image.
        Returns an integer represeting the id of the document selected.
    '''
    left, right = st.columns(2) # Create two columns
    # dropdown box in left column
    files = ["01_039_merged","01_068_merged", "02_042", "02_068","06_018_merged", "01_042_merged", "02_039", "02_043_merged", "03_018", "10_043"]
    selection_left = str(left.selectbox('Select document to run inference 1', files))
    selection_right = str(right.selectbox('Select document to run inference 2', files))
    
    # select corresponding document image from media/documents
    selection_image_left = MEDIA_ROOT+selection_left+'.png'
    left.image(selection_image_left, use_column_width='always')

    selection_image_right = MEDIA_ROOT+selection_right+'.png'
    right.image(selection_image_right, use_column_width='always')
    return selection_image_left, selection_image_right

def upload_document():
    img1 = st.file_uploader("left_img")
    img2 = st.file_uploader("right_img")


p1, p2 = select_document()
img1 = Image.open(p1).convert('RGB')
img2 = Image.open(p2).convert('RGB')
#model = load_model()
transform_valid = transforms.Compose([
    transforms.ToTensor(),
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize((124, 124)),
])

# call YOLOv5 detection fn on all images in the document folder.
#detect.detect(MEDIA_ROOT)

st.session_state.predict = st.button('Detect Signature')
if st.session_state.predict:
    device = torch.device('cpu')
    backbone = models.efficientnet_v2_s(pretrained=True)
    model = SiameseNetwork(backbone)
    weights = torch.load("pages/siamese.pth", map_location=torch.device('cpu'))["model_state_dict"]
    model.load_state_dict(weights)
    img1, img2 = transform_valid(img1).unsqueeze(0), transform_valid(img2).unsqueeze(0)
    prediction = model(img1,img2)
    st.write(f"Score: {str(prediction)}")



# get the path where last detected results are stored.
#latest_detection = max(glob.glob(os.path.join(YOLO_RESULT, '*/')), key=os.path.getmtime)
# resize and add top and bottom padding to detected sigantures. 
# gan model expects ips in that particular format.

# selects and display the detections of the document which the user selected.
#selection_detection =latest_detection + YOLO_OP + selection + '.jpg'
#st.image(selection_detection)