import gradio as gr
import subprocess
import os
import numpy as np
from PIL import Image

os.environ['data_raw'] = 'data_raw/'
os.environ['nnUNet_raw_data_base'] = 'nnUNet_raw_data_base/'
os.environ['nnUNet_preprocessed'] = 'nnUNet_preprocessed/'
os.environ['RESULTS_FOLDER'] = 'calvingfronts/'

model_path = 'Task500_Glacier_zonefronts'


def run_front_detection(input_img):
    image_gray = input_img.convert("L")
    image_gray.save('data_raw/test.png')

    subprocess.run(
        ['python3', 'nnunet/dataset_conversion/Task500_Glacier_inference.py', '-data_percentage', '100', '-base',
         os.environ['data_raw']])
    cmd = [
        'python3', 'nnunet/inference/predict_simple.py',
        '-i', os.path.join(os.environ['nnUNet_raw_data_base'], 'nnUNet_raw_data/Task500_Glacier_zonefronts/imagesTs/'),
        '-o', os.path.join(os.environ['RESULTS_FOLDER'], 'fold_0'),
        '-t', '500','-m','2d','-f','0','-p', 'nnUNetPlansv2.1', '-tr','nnUNetTrainerV2', '-model_folder_name',
        model_path
    ]
    subprocess.run(cmd)
    subprocess.run(['python3', 'nnunet/dataset_conversion/Task500_Glacier_reverse.py', '-i',
                    os.path.join(os.environ['RESULTS_FOLDER'], 'fold_0')])
    
    front = Image.open(os.path.join(os.environ['RESULTS_FOLDER'], 'tifs/test_front.png'))
    zone = Image.open(os.path.join(os.environ['RESULTS_FOLDER'], 'tifs/test_zone.png'))
    front = np.array(front)
    zone = np.array(zone)

    image_rgb = image_gray.convert('RGB')
    image_rgb = np.asarray(image_rgb).copy()

    image_rgb = np.array(image_rgb * 0.5, dtype=np.uint8)
    image_rgb[zone[:,:,0] == 0] += np.array(np.array([0, 0, 0]) / 2, dtype=np.uint8)
    image_rgb[zone[:,:,0] == 64] += np.array(np.array([52, 46, 55]) / 2, dtype=np.uint8)
    image_rgb[zone[:,:,0] == 127] += np.array(np.array([254, 254, 254]) / 2, dtype=np.uint8)
    image_rgb[zone[:,:,0] == 254] += np.array(np.array([60, 145, 230]) / 2, dtype=np.uint8)
    image_rgb[front[:, :, 0] == 255] = [255, 0, 0]
    # Convert NumPy array back to PIL image
    pil_image_modified = Image.fromarray(image_rgb)
    os.remove(os.path.join(os.environ['RESULTS_FOLDER'],'fold_0/test.nii.gz'))
    os.remove(os.path.join(os.environ['RESULTS_FOLDER'], 'tifs/test_front.png'))
    os.remove(os.path.join(os.environ['RESULTS_FOLDER'], 'tifs/test_zone.png'))
    os.remove('data_raw/test.png')
    os.remove('nnUNet_raw_data_base/nnUNet_raw_data/Task500_Glacier_zonefronts/imagesTs/test_0000.nii.gz')


    return pil_image_modified

demo = gr.Interface(run_front_detection, gr.Image(type='pil'), "image", title='Calving front detection with nnU-Net',
                    description="This app can delineate the glacier calving front in radar images from satellites (Sentinal-1, Tan-DEMX, ...). The method is described in \
                                'Out-of-the-box calving front detection method using deep-learning' by Herrmann et al.\
                                 https://tc.copernicus.org/preprints/tc-2023-34/. The project was build up on the nnU-Net project \
                                 by Isensee, F., Jaeger, P. F. (2020) https://github.com/MIC-DKFZ/nnUNet and  \
                                and was trained and tested with the dataset of Gourmelon et al. \
                                https://doi.pangaea.de/10.1594/PANGAEA.940950. <p> <p>\
                                Drag & Drop a PNG image and click 'Submit' to apply the calving front detection or clock on an example image below. <p> \
                                For images larger than 200KB we suggest to run on your own PC, this platform providese limited computational resources. <p>\
                                The output shows an overlay of the segmentation mask generated by the network over the input image. Area of ocean is colored in blue, glacier in white, rock in gray and radar shadow in black.\
                                The claving front is a thin red line.",
                    min_height=5000,
                    examples=[os.path.join(os.path.dirname(__file__), "test_sar_images/Mapple_2018-06-11_S1_20_2_009.png"),
                              os.path.join(os.path.dirname(__file__), "test_sar_images/COL_2011-06-18_TDX_7_1_024_small.png"),])





demo.launch()