import streamlit as st
import cv2
import os
import io
import numpy as np    
from PIL import Image
from inferance import pipline
import pandas as pd


code = """
<style>
  .block-container{
    max-width: 100%;
    padding: 50px;
  }
  # [data-testid="stImage"], .e115fcil2, [data-testid="StyledFullScreenButton"], [data-testid="stFullScreenFrame"].e1vs0wn30, [data-testid="element-container"].e1f1d6gn4.element-container{
  #   width: fit-content !important;
  # }
  # [data-testid="stVerticalBlock"].e1f1d6gn2{
  #   flex-direction: row;
  #   flex-wrap: wrap;
  # }
  [data-testid="StyledFullScreenButton"]{
    display: none;
  }
  [data-testid="stVerticalBlockBorderWrapper"], [data-testid="stVerticalBlock"]{
    width: 100%;
  }
  .e115fcil2{
    justify-content: center;
    margin-top: 20px;
  }
</style>
"""
st.html(code)

st.title("Automated Surveillance System")

col1, col2 = st.columns([5, 5])
container = col2.container(height=800)
col3, col4= container.columns([1,1])

with col1:
  
  image = st.file_uploader("File upload", label_visibility="hidden")
  
  if image is not None:
    image = Image.open(io.BytesIO(image.getvalue()))
    image = np.asarray(image)
    
    cv2.imwrite("image.jpg", image)
    image = cv2.imread("image.jpg")
    
    results = pipline(image)
    
    for result in results:
      image = cv2.rectangle(image, result['updated_boxes']['top_left'], result['updated_boxes']['bottom_right'], (255, 0, 0), 1)

    st.image(image)
  
  else:
    image = cv2.imread("default_img.jpg")

    results = pipline(image)
    
    for result in results:
      image = cv2.rectangle(image, result['updated_boxes']['top_left'], result['updated_boxes']['bottom_right'], (255, 0, 0), 1)

    st.image(image)

if image is not None:
  with col2:
      results_1 = results[:len(results)//2]
      results_2 = results[len(results)//2:]
      with col4:
        for result in results_1:
          img = result['zoomed_img']
          
          df = pd.DataFrame(columns=['Object Type', 'Distance', 'Activity'])
          actual_width, actual_height = result['updated_boxes']['bottom_right'][0] - result['updated_boxes']['top_left'][0], result['updated_boxes']['bottom_right'][1] - result['updated_boxes']['top_left'][1]

          for box in result['actual_boxes']:
            top_left = (box['top_left'][0] - result['updated_boxes']['top_left'][0], (box['top_left'][1] - result['updated_boxes']['top_left'][1]))
            bottom_right = (box['bottom_right'][0] - result['updated_boxes']['top_left'][0], (box['bottom_right'][1] - result['updated_boxes']['top_left'][1]))

            print(img.shape, actual_height, actual_width)
            bottom_right = (bottom_right[0]*img.shape[0]//(actual_height), bottom_right[1]*img.shape[1]//(actual_width))
            top_left = (top_left[0]*img.shape[0]//(actual_height), top_left[1]*img.shape[1]//(actual_width))
            
            print(box['top_left'], result['updated_boxes']['top_left'], box['bottom_right'], result['updated_boxes']['bottom_right'], top_left, bottom_right)
            img = cv2.rectangle(img, top_left, bottom_right, (255, 0, 0), 1)
            img = cv2.putText(img, "ID: "+str(len(df)), top_left, 1, 1, (255, 255, 255))
            df.loc[len(df)] = [box['class'], box['distance'], box['activity']]
          
          st.image(img)
          st.table(df)

      with col3:
        for result in results_2:
          img = result['zoomed_img']
          
          df = pd.DataFrame(columns=['Object Type', 'Distance', 'Activity'])
          actual_width, actual_height = result['updated_boxes']['bottom_right'][0] - result['updated_boxes']['top_left'][0], result['updated_boxes']['bottom_right'][1] - result['updated_boxes']['top_left'][1]

          for box in result['actual_boxes']:
            top_left = (box['top_left'][0] - result['updated_boxes']['top_left'][0], (box['top_left'][1] - result['updated_boxes']['top_left'][1]))
            bottom_right = (box['bottom_right'][0] - result['updated_boxes']['top_left'][0], (box['bottom_right'][1] - result['updated_boxes']['top_left'][1]))

            print(img.shape, actual_height, actual_width)
            bottom_right = (bottom_right[0]*img.shape[0]//(actual_height), bottom_right[1]*img.shape[1]//(actual_width))
            top_left = (top_left[0]*img.shape[0]//(actual_height), top_left[1]*img.shape[1]//(actual_width))
            
            print(box['top_left'], result['updated_boxes']['top_left'], box['bottom_right'], result['updated_boxes']['bottom_right'], top_left, bottom_right)
            img = cv2.rectangle(img, top_left, bottom_right, (255, 0, 0), 1)
            img = cv2.putText(img, "ID: "+str(len(df)), top_left, 1, 1, (255, 255, 255))
            df.loc[len(df)] = [box['class'], box['distance'], box['activity']]
          
          st.image(img)
          st.table(df)