import gradio as gr
import math
import numpy as np
import cv2

def calculate_flow_direction(dtm_file):
    # Step 1: Read the DTM data and convert it to a 2D array of elevation values
    elevation_array = cv2.imread(dtm_file.name, cv2.IMREAD_GRAYSCALE).astype(float)

    # Step 2: Compute the slope and aspect of each cell in the DTM
    cell_size = 0.1  # Set the cell size in meters
    dx = dy = cell_size

    # Compute the x and y gradient values
    dzdx, dzdy = np.gradient(elevation_array, dx, dy)

    # Compute the slope angle in radians
    slope = np.arctan(np.sqrt(dzdx**2 + dzdy**2))

    # Compute the aspect angle in radians
    aspect = np.arctan2(-dzdy, dzdx)

    # Step 3: Compute the flow direction for each cell based on the slope and aspect values
    # Compute the x and y components of the flow direction vector
    flow_dir_x = np.sin(aspect)
    flow_dir_y = np.cos(aspect)

    # Normalize the flow direction vector to unit length
    flow_dir_length = np.sqrt(flow_dir_x**2 + flow_dir_y**2)
    flow_dir_x /= flow_dir_length
    flow_dir_y /= flow_dir_length

    # Step 4: Compute the flow accumulation for each cell based on the flow direction of neighboring cells
    def compute_flow_accumulation(flow_dir_x, flow_dir_y):
        # Initialize the flow accumulation array to zero
        flow_accum = np.zeros_like(elevation_array)

        # Find the cells with no upstream flow (i.e., cells with zero slope)
        no_upstream_flow = (slope == 0)
        flow_accum[no_upstream_flow] = 1

        # Recursively accumulate flow downstream
        rows, cols = elevation_array.shape
        for i in range(1, rows-1):
            for j in range(1, cols-1):
                if no_upstream_flow[i,j]:
                    continue

                # Find the downstream cell
                di = int(round(flow_dir_y[i,j]))
                dj = int(round(flow_dir_x[i,j]))
                downstream = (i+di, j+dj)

                # Add the flow accumulation from the downstream cell
                flow_accum[i,j] = flow_accum[downstream] + 1

        return flow_accum

    # Step 5: Compute the flow accumulation and direction maps
    flow_accum = compute_flow_accumulation(flow_dir_x, flow_dir_y)
    flow_accum_norm = flow_accum / np.max(flow_accum) * 255
    flow_accum_rgb = cv2.cvtColor(flow_accum_norm.astype(np.uint8), cv2.COLOR_GRAY2RGB)

    flow_dir_x_norm = (flow_dir_x + 1) / 2 * 255
    flow_dir_y_norm = (flow_dir_y + 1) / 2 * 255
    flow_dir_rgb = np.zeros((elevation_array.shape[0], elevation_array.shape[1], 3), dtype=np.uint8)
    flow_dir_rgb[:,:,0] = flow_dir_x_norm.astype(np.uint8)
    flow_dir_rgb[:,:,1] = flow_dir_y_norm.astype(np.uint8)

    return flow_dir_rgb, flow_accum_rgb

# Define the Gradio interface
iface = gr.Interface(
    fn=calculate_flow_direction,
    inputs="file",
    outputs=["image", "image"],
    title="DTM Flow Direction and Accumulation Calculator",
    description="This app calculates the flow direction and accumulation maps for a digital terrain model (DTM).",
)

# Launch the interface
iface.launch()