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="Drag & Drop an PNG image to apply the calving fron detection <p> \
                                we sugget to use images smaller than 200KB <p> \
                                trained and tested with the dataset of Gourmelon et al. \
                                https://doi.pangaea.de/10.1594/PANGAEA.940950 ")





demo.launch()