import streamlit as st
import torch
import numpy as np
from PIL import Image
import rasterio
from rasterio.windows import Window
from tqdm.auto import tqdm
import io
import zipfile

# Assuming you have these functions defined elsewhere
from your_module import preprocess, best_model, DEVICE

def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4):
    tiles = []
    
    with rasterio.open(map_file) as src:
        height = src.height
        width = src.width
        
        effective_tile_size = tile_size - overlap
        
        for y in tqdm(range(0, height, effective_tile_size)):
            for x in range(0, width, effective_tile_size):
                batch_images = []
                batch_metas = []
                
                for i in range(batch_size):
                    curr_y = y + (i * effective_tile_size)
                    if curr_y >= height:
                        break
                    
                    window = Window(x, curr_y, tile_size, tile_size)
                    out_image = src.read(window=window)
                    
                    if out_image.shape[0] == 1:
                        out_image = np.repeat(out_image, 3, axis=0)
                    elif out_image.shape[0] != 3:
                        raise ValueError("The number of channels in the image is not supported")
                    
                    out_image = np.transpose(out_image, (1, 2, 0))
                    tile_image = Image.fromarray(out_image.astype(np.uint8))
                    
                    out_meta = src.meta.copy()
                    out_meta.update({
                        "driver": "GTiff",
                        "height": tile_size,
                        "width": tile_size,
                        "transform": rasterio.windows.transform(window, src.transform)
                    })
                    tile_image = np.array(tile_image)
                    
                    preprocessed_tile = preprocess(image=tile_image)['image']
                    batch_images.append(preprocessed_tile)
                    batch_metas.append(out_meta)
                
                if not batch_images:
                    break
                
                # Concatenate batch images
                batch_tensor = torch.cat([img.unsqueeze(0).to(DEVICE) for img in batch_images], dim=0)
                # Perform inference on the batch
                with torch.no_grad():
                    batch_masks = model(batch_tensor.to(DEVICE))
                
                batch_masks = torch.sigmoid(batch_masks)
                batch_masks = (batch_masks > 0.6).float()
                
                # Process each mask in the batch
                for j, mask_tensor in enumerate(batch_masks):
                    mask_resized = torch.nn.functional.interpolate(mask_tensor.unsqueeze(0), size=(tile_size, tile_size), mode='bilinear', align_corners=False).squeeze(0)
                    
                    mask_array = mask_resized.squeeze().cpu().numpy()
                    
                    if mask_array.any() == 1:
                        tiles.append([mask_array, batch_metas[j]])
    
    return tiles

def main():
    st.title("TIF File Processor")

    uploaded_file = st.file_uploader("Choose a TIF file", type="tif")

    if uploaded_file is not None:
        st.write("File uploaded successfully!")

        # Process button
        if st.button("Process File"):
            st.write("Processing...")
            
            # Save the uploaded file temporarily
            with open("temp.tif", "wb") as f:
                f.write(uploaded_file.getbuffer())
            
            # Process the file
            best_model.float()
            tiles = extract_tiles("temp.tif", best_model, tile_size=512, overlap=15, batch_size=4)
            
            st.write("Processing complete!")

            # Prepare zip file for download
            zip_buffer = io.BytesIO()
            with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED, False) as zip_file:
                for i, (mask_array, meta) in enumerate(tiles):
                    # Save each tile as a separate TIF file
                    with rasterio.open(f"tile_{i}.tif", 'w', **meta) as dst:
                        dst.write(mask_array, 1)
                    
                    # Add the tile to the zip file
                    zip_file.write(f"tile_{i}.tif")

            # Offer the zip file for download
            st.download_button(
                label="Download processed tiles",
                data=zip_buffer.getvalue(),
                file_name="processed_tiles.zip",
                mime="application/zip"
            )

if __name__ == "__main__":
    main()