import os
import numpy as np
import cv2
import rasterio
from rasterio import shutil as rio_shutil

def tiff_to_array_and_profile(tiff_path):
    with rasterio.open(tiff_path) as src:
        array = src.read(1)  # Read the first band
        profile = src.profile  # Metadata
    return array, profile

def keep_largest_region(array, threshold=0.5, connectivity=8):
    # Binary mask where array > threshold
    binary = (array > threshold).astype(np.uint8)

    # Connected components labeling
    num_labels, labels = cv2.connectedComponents(binary, connectivity=connectivity)

    if num_labels <= 1:
        return np.zeros_like(array)  # No valid regions, return all zeros

    # Count area of each label (ignore background label 0)
    label_areas = [(labels == i).sum() for i in range(1, num_labels)]
    largest_label = np.argmax(label_areas) + 1  # +1 because label 0 is background

    # Keep only the largest region
    output = np.where(labels == largest_label, array, 0)
    return output

def process_folder_keep_largest(input_folder, output_folder):
    os.makedirs(output_folder, exist_ok=True)

    for filename in os.listdir(input_folder):
        if filename.endswith(".tif"):
            input_path = os.path.join(input_folder, filename)
            output_path = os.path.join(output_folder, filename)

            array, profile = tiff_to_array_and_profile(input_path)
            processed_array = keep_largest_region(array, threshold=0.5, connectivity=8)

            # Save result as GeoTIFF
            with rasterio.open(output_path, 'w', **profile) as dst:
                dst.write(processed_array, 1)

    print(f"Processing complete. Output saved to: {output_folder}")

input_folder = r"...\Predicted_Burned area"
output_folder = r"...\Predicted_Burned area_after_processing"

process_folder_keep_largest(input_folder, output_folder)