import numpy as np
import sys,os
import gradio as gr

# image parameters
height = 300
width  = 300

# this function reads and normalizes the images
def load_normalize_images(f_images):

    # load the images
    images = np.load(f_images)

    # normalize the values
    images = np.log10(images)
    minimum, maximum = np.min(images), np.max(images)
    images = (images-minimum)/(maximum-minimum)

    return images

f_images = 'Maps_T_IllustrisTNG_CV_z=0.00.npy'
images1 = load_normalize_images(f_images)

f_images = 'Maps_B_IllustrisTNG_CV_z=0.00.npy'
images2 = load_normalize_images(f_images)

f_images = 'Maps_Mgas_IllustrisTNG_CV_z=0.00.npy'
images3 = load_normalize_images(f_images)

f_images = 'Maps_MgFe_IllustrisTNG_CV_z=0.00.npy'
images4 = load_normalize_images(f_images)

f_images = 'Maps_Vcdm_IllustrisTNG_CV_z=0.00.npy'
images5 = load_normalize_images(f_images)

f_images = 'Maps_Z_IllustrisTNG_CV_z=0.00.npy'
images6 = load_normalize_images(f_images)

f_images = 'Maps_HI_IllustrisTNG_CV_z=0.00.npy'
images7 = load_normalize_images(f_images)

f_images = 'Maps_Mcdm_IllustrisTNG_CV_z=0.00.npy'
images8 = load_normalize_images(f_images)

f_images = 'Maps_ne_IllustrisTNG_CV_z=0.00.npy'
images9 = load_normalize_images(f_images)

f_images = 'Maps_P_IllustrisTNG_CV_z=0.00.npy'
images10 = load_normalize_images(f_images)

f_images = 'Maps_Vgas_IllustrisTNG_CV_z=0.00.npy'
images11 = load_normalize_images(f_images)


images = np.vstack([images1, images2, images3, images4, images5, images6,
                    images7, images8, images9, images10, images11])


# function that applies when the submit button is clicked
def guess_field(field, correct, wrong, state):

    labels = {0:'Temperature', 1:'Magnetic fields', 2:'Gas mass',
              3:'Mg/Fe', 4:'Dark matter velocity', 5:'Metallicity',
              6:'Neutral hydrogen', 7:'Dark matter mass', 8:'Electron density',
              9:'Gas pressure', 10:'Gas velocity'}
    
    block = state//num_images_field
    if field==labels[block]:
        text = "Correct! This is from a %s map"%labels[block]
        correct += 1
    else:
        text = "Wrong! This is from a %s map"%labels[block]
        wrong += 1
    return text, correct, wrong, gr.Button(interactive=False)

# function that applies when the next image button is clicked
def reset():
    index = np.random.choice(np.arange(num_images), size=1)[0]
    return gr.Image(value=images[index],type="numpy"), None, None, gr.Button(interactive=True), index
    #return gr.Image(value=images[index],label="Image id: %d"%index,type="numpy"), None, None, gr.Button(interactive=True), index

num_images = images.shape[0]
num_images_field = images1.shape[0]

index = np.random.choice(np.arange(num_images), size=1)[0]
image = images[index]
correct, wrong = 0, 0 

with gr.Blocks() as demo:

    state = gr.State(value=index)
    
    with gr.Row():
        with gr.Column():
            image = gr.Image(value=image, #label="Image id: %d"%index, 
                             type="numpy", height=height, width=width)
            results = gr.Textbox(label="Results", interactive=False)
        with gr.Column():
            options = gr.Radio(choices=["Temperature", "Magnetic fields", "Gas mass", "Mg/Fe", "Dark matter velocity", "Metallicity", "Neutral hydrogen", "Dark matter mass", "Electron density", "Gas pressure", "Gas velocity"],
                               label="Which field shows the image?")
            submit_button = gr.Button("Submit")
            next_image = gr.Button("Next image")
            with gr.Row():
                correct = gr.Number(label="Correct", value=correct, interactive=False)
                wrong   = gr.Number(label="Wrong",   value=wrong,   interactive=False)

    submit_button.click(guess_field, inputs=[options, correct, wrong, state],
                        outputs=[results, correct, wrong, submit_button])
    next_image.click(reset, inputs=[], outputs=[image, results, options, submit_button, state])

demo.launch()