# Written by Dr Daniel Buscombe, Marda Science LLC
#
# MIT License
#
# Copyright (c) 2022, Marda Science LLC
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import gradio as gr
import os
import pandas as pd
import matplotlib.pyplot as plt
# from skimage.transform import resize
plt.style.use('fivethirtyeight')

from funcs import *

#================================================================
def do_dgs(input_img):

   maxscale=5 #20
   verbose=True
   x=0 #-0.5
   resolution=1

   w=input_img.shape[0]
   h=input_img.shape[1]

   # input_img = resize(input_img, (int(w/2), int(h/2)), preserve_range=True, clip=True) 
   # resolution = 2

   data_out = dgs(input_img, resolution, maxscale, verbose, x)

   ## parse out dict into three separate dictionaries
   stats = dict(list(data_out.items())[:4])
   percentiles = dict(list(data_out.items())[4:6])
   freqs_bins = dict(list(data_out.items())[6:])

   if resolution!=1:
      freqs_bins['grain size bins']*=resolution
      percentiles['percentile_values']*=resolution

      for k in stats.keys():
            stats[k] = stats[k]*resolution

   tmp = list(stats.keys())
   d = {} 
   for k in range(len(tmp)):
       d.update( {tmp[k]: str(stats[tmp[k]])} )

   pd.DataFrame(data=d.values(), index=d.keys()).to_csv('stats.csv')

   tmp = list(percentiles.keys())
   d2 = {} 
   for k in range(len(tmp)):
       d2.update( {tmp[k]: str(percentiles[tmp[k]])} )

   pd.DataFrame(data=d2.values(), index=d2.keys()).to_csv('percentiles.csv')

   # write each to csv file
   pd.DataFrame.from_dict(freqs_bins).to_csv('freqs_bins.csv')

   # plt.clf()
   fig1=plt.figure(figsize=(4,4))
   plt.plot(freqs_bins['grain size bins'], freqs_bins['grain size frequencies'], lw=2) 
   plt.xlabel('Grain Size (pixels)')
   plt.ylabel('Frequency')
   plt.axvline(x=percentiles['percentile_values'][5], color="black", linestyle="--")
   plt.text(percentiles['percentile_values'][5], .1, 'd50')
   # plt.title('Grain Size Distribution')
   # plt.xlim(np.maximum(0,stats['mean grain size']-(2*stats['grain size sorting'])), stats['mean grain size']+(2*stats['grain size sorting']))
   # plt.savefig('psd.png', dpi=300, bbox_inches='tight')

   # plt.clf()
   fig2=plt.figure(figsize=(4,4))
   plt.imshow(input_img,cmap='gray')
   plt.plot([50, 50+stats['mean grain size']], [50, 50], 'r')
   try:
      plt.xlim(0,200)
      plt.ylim(0,200)
   except:
      pass
   plt.axis("off")

   out = {}
   out.update( {'mean grain size': float(str(stats['mean grain size'])[:4]) } )
   out.update( {'grain size sorting': float(str(stats['grain size sorting'])[:4]) } )
   out.update( {'d16': float(str(percentiles['percentile_values'][2])[:4]) } )
   out.update( {'d50': float(str(percentiles['percentile_values'][5])[:4]) } )
   out.update( {'d84': float(str(percentiles['percentile_values'][7])[:4]) } )
   
   return pd.DataFrame(data=out.values(), index=out.keys()).transpose(), fig1, fig2, 'stats.csv', 'percentiles.csv', 'freqs_bins.csv'

 
  
title = "Digital Grain Size" 
description = "Upload an image of sediment. Download grain size stats. Clear between images. See https://github.com/dbuscombe-usgs/pyDGS. May take a very long time for large images" 
    

examples= [[l] for l in glob('examples/*.jpg')]

inp = gr.Image(label='Upload images one-by-one')
out1 = gr.Dataframe(label='Summary statistics', headers=["mean (px)", "sorting (px)", "d16 (px)", "d50 (px)", "d84 (px)"], type='pandas')
#out1 = gr.outputs.Image(type='numpy')
out2a = gr.Plot(label='Grain Size Distribution', type='auto') #'matplotlib')
out2b = gr.Plot(label='Red Line = Mean Grain Length', type='auto') #'matplotlib')

out4 = gr.File(label='Click to download grain size summary stats (px)')
out5 = gr.File(label='Click to download grain size percentiles (px)')
out6 = gr.File(label='Click to download grain size distribution: bins (px) and frequencies')

Segapp = gr.Interface(do_dgs, inp, [out1, out2a, out2b, out4, out5, out6], title = title, description = description, examples=examples)

Segapp.launch(enable_queue=True)