import gradio as gr
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
import os
import leafmap
from samgeo import SamGeo, show_image, download_file, overlay_images, tms_to_geotiff
from zipfile import ZipFile


#initialize the model
sam = SamGeo(
    model_type="vit_h",
    checkpoint='sam_vit_h_4b8939.pth',
    sam_kwargs=None,
)

def get_all_file_paths(directory):

    # initializing empty file paths list
    file_paths = []

    # crawling through directory and subdirectories
    for root, directories, files in os.walk(directory):
        for filename in files:
            # join the two strings in order to form the full filepath.
            filepath = os.path.join(root, filename)
            file_paths.append(filepath)

    # returning all file paths
    return file_paths

def get_shape_files(image):
  filename, file_extension = os.path.splitext(image.name)
  # print(filename,file_extension)
  if file_extension in ['.tif','tiff']:
    #mkdir for app
    storage='app_store'
    if os.path.exists(storage):
        pass
    else:
      os.mkdir(storage)

    zip_folder=f'{storage}/zip_folder'
    if os.path.exists(zip_folder):
        pass
    else:
      os.mkdir(zip_folder)

    out_file=f'{storage}/out_files'
    if os.path.exists(out_file):
        pass
    else:
      os.mkdir(out_file)

    # load_model()
    sam.generate(image.name, output=f"{out_file}/masks.tif", foreground=True, unique=True)
    sam.tiff_to_vector(f"{out_file}/masks.tif", out_file)
    sam.show_anns(axis="off", alpha=1, output=f"{out_file}/annotations.tif")
    
    #converting tif to jpg
    image = cv2.imread(f"{out_file}/annotations.tif")
    cv2.imwrite("{}/{}.jpg".format(out_file,'converted'), image)

    # writing files to a zipfile
    file_paths=get_all_file_paths(out_file)

    # for file_name in file_paths:
    #       print(file_name)

    with ZipFile(f'{zip_folder}/shapefile.zip','w') as zip:
        # writing each file one by one
        for file in file_paths:
            zip.write(file)

    print('All files zipped successfully!')

    return f"{out_file}/converted.jpg",f'{zip_folder}/shapefile.zip'
  # get_shape_files(image)
  else:
    return "Try uploading .tif file for processing!."

my_app = gr.Blocks()
with my_app:
  gr.Markdown("<center><h1>Segmenting Satellite Image</h1></center>")
  gr.Markdown("<center><h3>Processing time depends on the file size and since the instance is running on CPU it takes longer time.</h></center>")
  with gr.TabItem("Get Shapefiles"):
    with gr.Row():
      with gr.Column():
          img_source = gr.File(label="Please select source tif")
          source_image_loader = gr.Button("Get Shape File")
      with gr.Column():
          image=gr.Image()
          output = gr.outputs.File("zip")
  source_image_loader.click(get_shape_files,img_source,[image,output])
my_app.launch(debug = True)