## Daniel Buscombe, Marda Science LLC 2023
# This file contains many functions originally from Doodleverse https://github.com/Doodleverse programs

import gradio as gr
import numpy as np
import sys, json, os
import pandas as pd

sys.path.insert(1, 'app_files'+os.sep+'src')
from sedinet_eval import *

###===================================================
def estimate_siso_simo_1image(vars, im, greyscale,
                       dropout, weights_path): # numclass, name, mode, res_folder,
                    #    batch_size, ):#, scale): #
    """
    This function uses a sedinet model for continuous prediction on 1 image
    """
    SM = make_sedinet_siso_simo(vars, greyscale, dropout)
    SM.load_weights(weights_path)

    # im = Image.open(image).convert('LA')
    #im = im.resize((IM_HEIGHT, IM_HEIGHT))
    im = Image.fromarray(im)
    im = np.array(im)[:,:,0]
    nx,ny = np.shape(im)
    if (nx!=IM_HEIGHT) or (ny!=IM_HEIGHT):
        im = im[int(nx/2)-int(IM_HEIGHT/2):int(nx/2)+int(IM_HEIGHT/2), int(ny/2)-int(IM_HEIGHT/2):int(ny/2)+int(IM_HEIGHT/2)]

    if DO_STANDARDIZE==True:
       im = do_standardize(im)
    else:
       im = np.array(im) / 255.0

    result = SM.predict(np.expand_dims(np.expand_dims(im, axis=2), axis=0))
    result = [float(r[0]) for r in result]

    return result

###===================================================
def grainsize(input_img, dims=(1024, 1024)):

    configfile = 'weights/config_usace_combined2021_2022_v12.json'
    weights_path = 'weights/sandsnap_merged_1024_modelrevOct2022_v12_simo_batch10_im1024_9vars_mse_noaug.hdf5'

    # load the user configs
    with open(os.getcwd()+os.sep+configfile) as f:
      config = json.load(f)

    ###===================================================
    dropout = config["dropout"]
    greyscale = config['greyscale']

    try:
       greyscale = config['greyscale']
    except:
       greyscale = 'true'

    #output variables
    vars = [k for k in config.keys() if not np.any([k.startswith('base'), k.startswith('MAX_LR'),
            k.startswith('MIN_LR'), k.startswith('DO_AUG'), k.startswith('SHALLOW'),
            k.startswith('res_folder'), k.startswith('train_csvfile'), k.startswith('csvfile'),
            k.startswith('test_csvfile'), k.startswith('name'), k.startswith('val_csvfile'),
            k.startswith('greyscale'), k.startswith('aux_in'),
            k.startswith('dropout'), k.startswith('N'),k.startswith('scale'),
            k.startswith('numclass')])]
    vars = sorted(vars)

    #this relates to 'mimo' and 'miso' modes that are planned for the future but not currently implemented
    auxin = [k for k in config.keys() if k.startswith('aux_in')]

    if len(auxin) > 0:
       auxin = config[auxin[0]]   ##at least for now, just one 'auxilliary' (numerical/categorical) input in addition to imagery
       if len(vars) ==1:
          mode = 'miso'
       elif len(vars) >1:
          mode = 'mimo'
    else:
       if len(vars) ==1:
          mode = 'siso'
       elif len(vars) >1:
          mode = 'simo'

    print("Mode: %s" % (mode))

    result = estimate_siso_simo_1image(vars, input_img, greyscale,
                       dropout, weights_path)
    result_str = [str(i)[:5] for i in result]

    result = np.array(result)
    print(result)
    plt.clf()    
    plt.plot(np.hstack((result[:3], result[4:])),[10,16,25,50,65,75,84,90],  'k-o')
    plt.xlabel('Grain size (pixels)')
    plt.ylabel('Percent finer')
    plt.savefig("psd.png", dpi=300, bbox_inches="tight")


    return 'mean grain size = %f pixels' % (result[4]), '90th percentile grain size = %f pixels' % (result[-1]), plt, pd.DataFrame(data=result_str, index=['10','16','25','50','mean','65','75','84','90']).transpose(),

title = "SandSnap/SediNet Model Demo- Measure grain size from image of sand!" 
description = "Allows upload of imagery and download of grain size statistics. Statistics are unscaled (i.e. in pixels)" 

examples = [
['examples/IMG_20210922_170908944_cropped.jpg'],
['examples/20210208_172834_cropped.jpg'],
['examples/20220101_165359_cropped.jpg']
] 

inp = gr.Image()
out1 = gr.Plot(type='matplotlib')
out2 = gr.Dataframe(label='Summary statistics', headers=['10','16','25','50','mean','65','75','84','90'], type='pandas')

Segapp = gr.Interface(grainsize, inp, ["text", "text", out1, out2], title = title, description = description, examples=examples)
#, allow_flagging='manual', flagging_options=["bad", "ok", "good", "perfect"], flagging_dir="flagged")

Segapp.launch(enable_queue=True)