import os
import rasterio
import geopandas as gpd
from shapely.geometry import box
from rasterio.mask import mask
from PIL import Image
import numpy as np
import warnings
from rasterio.errors import NodataShadowWarning
import sys

warnings.filterwarnings("ignore", category=NodataShadowWarning)

def cut_trees(output_dir, geojson_path, tif_path):
    # create output directory if it doesnt exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Load the GeoDataFrame
    gdf = gpd.read_file(geojson_path)

    # Clear the terminal screen
    os.system('cls' if os.name == 'nt' else 'clear')

    # Open the .tif file
    with rasterio.open(tif_path) as src:
        # Get the bounds of the .tif image
        tif_bounds = box(*src.bounds)

        # Get the CRS (Coordinate Reference System) of the .tif image
        tif_crs = src.crs
        
        # Reproject the GeoDataFrame to the CRS of the .tif file
        gdf = gdf.to_crs(tif_crs)

        # Loop through each polygon in the GeoDataFrame
        N = len(gdf)
        n = int(N/10)
        image_counter = 0
        for idx, row in gdf.iterrows():
            if idx % n == 0:
                progress = f"{round(idx/N*100)} % complete --> {idx}/{N}"
                sys.stdout.write('\r' + progress)
                sys.stdout.flush()

            # Extract the geometry (polygon)
            geom = row['geometry']
            name = row['id']
            
            # Check if the polygon intersects the image bounds
            if geom.intersects(tif_bounds):
                # Create a mask for the current polygon
                out_image, out_transform = mask(src, [geom], crop=True)
                
                # Convert the masked image to a numpy array
                out_image = out_image.transpose(1, 2, 0)  # rearrange dimensions for PIL (H, W, C)

                # Ensure the array is not empty
                if out_image.size == 0:
                    message = f"{round(idx/N*100)} % complete --> {idx}/{N} | Polygon {idx} resulted in an empty image and will be skipped."
                    sys.stdout.write('\r' + message)
                    sys.stdout.flush()
                    continue
                
                # Remove the zero-padded areas (optional)
                mask_array = (out_image[:, :, 0] != src.nodata)
                non_zero_rows = np.any(mask_array, axis=1)
                non_zero_cols = np.any(mask_array, axis=0)

                # Ensure there are non-zero rows and columns
                if not np.any(non_zero_rows) or not np.any(non_zero_cols):
                    message = f"{round(idx/N*100)} % complete --> {idx}/{N} | Polygon {idx} resulted in an invalid image area and will be skipped."
                    sys.stdout.write('\r' + message)
                    sys.stdout.flush()
                    continue

                out_image = out_image[non_zero_rows][:, non_zero_cols]

                # Convert to a PIL Image and save as PNG
                out_image = Image.fromarray(out_image.astype(np.uint8))  # Ensure correct type for PIL
                output_path = os.path.join(output_dir, f'tree_{name}.png')
                out_image.save(output_path)
                image_counter += 1
            else:
                message = f"{round(idx/N*100)} % complete --> {idx}/{N} | Polygon {idx} is outside the image bounds and will be skipped."
                sys.stdout.write('\r' + message)
                sys.stdout.flush()
                
    print(f'\n {image_counter}/{N} Tree images have been successfully saved in the "detected_trees" folder.')


def resize_images(input_folder, output_folder, target_size):
    # Create the output folder if it doesn't exist
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    counter = 0
    # Loop through all files in the input folder
    for filename in os.listdir(input_folder):
        if filename.endswith('.png'):  # Check for PNG files
            # Open image
            with Image.open(os.path.join(input_folder, filename)) as img:
                # Resize image while preserving aspect ratio
                img.thumbnail(target_size, Image.LANCZOS)
                # Calculate paste position to center image in canvas
                paste_pos = ((target_size[0] - img.size[0]) // 2, (target_size[1] - img.size[1]) // 2)
                # Create a new blank canvas with the target size and black background
                new_img = Image.new("RGBA", target_size, (0, 0, 0, 255))
                # Paste resized image onto the canvas
                new_img.paste(img, paste_pos, img)
                # Convert to RGB to remove transparency by merging with black background
                new_img = new_img.convert("RGB")
                # Save resized image to output folder
                new_img.save(os.path.join(output_folder, filename))
                
            counter += 1
            # Display the counter
            if counter % 50 == 0:
                message = f"Processed {counter} images"
                print(message, end='\r')
    
    # Final message after processing all images
    print(f"Processed a total of {counter} images.")


# THIS IS THE FUNCTION TO IMPORT
def generate_tree_images(geojson_path, tif_path, target_size = (224, 224)):
    """
    INPUT: geojson path, tif_path that contain the trees, optional target_size of the resulting images
    
    RETURNS: nothing
    
    Action: It creates two folders: + "detected trees" --> the cut tree images
                                    + "tree_images" --> the processed cut tree images, ready to use for species recognition
    """


    # Set input and output folders
    folder_cut_trees = "detected_trees"
    folder_finished_images = "tree_images"
    # Set target size (width, height)    
    cut_trees(geojson_path = geojson_path, tif_path = tif_path, output_dir = folder_cut_trees)
    resize_images(input_folder = folder_cut_trees, output_folder = folder_finished_images, target_size = target_size)