import base64
import json
import os, shutil
import re
import time
import uuid

import cv2

import numpy as np
import streamlit as st
from PIL import Image
# from extract_video import extract_method_single_video

import shlex
import subprocess
from file_picker import st_file_selector

import os

from inference import classify_fake, heatmap_analysis

DEBUG = True
SAMPLE_FOLDER = 'examples'

def main():
    st.markdown("###")    
    uploaded_file = st.file_uploader('Upload a picture', type=['jpg', 'jpeg', 'png'], accept_multiple_files=False)

    with st.spinner(f'Loading samples...'):
        while not os.path.isdir(SAMPLE_FOLDER):
            time.sleep(1)
    st.markdown("### or")
    selected_file = st_file_selector(st, path=SAMPLE_FOLDER, key = 'selected', label = 'Choose a sample image')

    if uploaded_file: 
        img = Image.open(uploaded_file).convert('RGB')
        st.image(img)
    elif selected_file:
        img = Image.open(os.path.join(SAMPLE_FOLDER, selected_file)).convert('RGB')
        st.image(img)
    else:
        return

   

    

    
    
    with st.spinner(f'Analyzing image...'):
        try:
            modified_probability = classify_fake(img)

        except Exception as e:
            if DEBUG:
                st.write(e)
            else:
                st.text("Encountered a problem while analyzing image 🚨")
                return
            
    if modified_probability > 0.6:
        st.error('  MODIFIED IMAGE! ', icon="🚨")
    else:
        st.success("  REAL IMAGE! ", icon="✅")
            
    st.text("modified probability {:.2f}".format(modified_probability))


    if modified_probability > 0.6:
        with st.spinner(f'Analyzing heatmap...'):
            try:
                modified, reverse, heatmap = heatmap_analysis(img)

            except Exception as e:
                if DEBUG:
                    st.write(e)
                else:
                    st.text("Encountered a problem while analyzing image 🚨")
                    return
        
        st.write("### Heatmap")
        st.image(heatmap)

        st.write("### Reversed stretch imgae")
        st.image(reverse)


if __name__ == "__main__":
    st.set_page_config(
        page_title="Nodeflux Photosop Detection", page_icon=":pencil2:"
    )
    st.title("Photosop Detection")
    main()