# For neural networks
import keras
# For train-test splits
import sklearn.model_selection
# For random calculations
import numpy
# For help with saving and opening things
import os

# Disable eager execution because its bad
from tensorflow.python.framework.ops import disable_eager_execution
disable_eager_execution()

# Start a session for checking calculations and stuff
import tensorflow as tf
sess = tf.compat.v1.Session()

from keras import backend as K
K.set_session(sess)


# Do you want it loud? 
VERBOSE = 1

# This function loads a fuckton of data
def load_data():
    # Open all the files we downloaded at the beginning and take out hte good bits
    curves = numpy.load('/content/data_curves.npz')['curves']
    geometry = numpy.load('/content/data_geometry.npz')['geometry']
    constants = numpy.load('/content/constants.npz')
    S = constants['S']
    N = constants['N']
    D = constants['D']
    F = constants['F']
    G = constants['G']

    # Some of the good bits need additional processining
    new_curves = numpy.zeros((S*N, D * F))
    for i, curveset in enumerate(curves):
        new_curves[i, :] = curveset.T.flatten() / 1000000

    new_geometry = numpy.zeros((S*N, G * G * G))
    for i, geometryset in enumerate(geometry):
        new_geometry[i, :] = geometryset.T.flatten()

    # Return good bits to user
    return curves, geometry, S, N, D, F, G, new_curves, new_geometry
    
import gradio
import pandas

class Network(object):

  def __init__(self, structure, weights):
      # Instantiate variables
      self.curves = 0
      self.new_curves = 0
      self.geometry = 0
      self.new_geometry = 0
      self.S = 0
      self.N = 0
      self.D = 0
      self.F = 0
      self.G = 0

      # Load network
      with open(structure, 'r') as file:
          self.network = keras.models.model_from_json(file.read())
          self.network.load_weights(weights)

      # Load data
      self._load_data()

  def _load_data(self):
      self.curves, self.geometry, self.S, self.N, self.D, self.F, self.G, self.new_curves, self.new_geometry = load_data()

  def analysis(self, idx=None):
      print(idx)

      if idx is None:
          idx = numpy.random.randint(1, self.S * self.N)
      else:
        idx = int(idx)

      # Get the input
      data_input = self.new_geometry[idx:(idx+1), :]
      other_data_input = data_input.reshape((self.G, self.G, self.G), order='F')

      # Get the outputs
      predicted_output = self.network.predict(data_input)
      true_output = self.new_curves[idx].reshape((3, self.F))
      predicted_output = predicted_output.reshape((3, self.F))

      f = numpy.linspace(0.05, 2.0, 64)
      fd = pandas.DataFrame(f).rename(columns={0: "Frequency"})
      df_pred = pandas.DataFrame(predicted_output.transpose()).rename(columns={0: "Surge", 1: "Heave", 2: "Pitch"})
      df_true = pandas.DataFrame(true_output.transpose()).rename(columns={0: "Surge", 1: "Heave", 2: "Pitch"})

      # return idx, other_data_input, true_output, predicted_output
      return pandas.concat([fd, df_pred], axis=1), pandas.concat([fd, df_true], axis=1)

  def synthesis(self, idx=None):
      print(idx)

      if idx is None:
          idx = numpy.random.randint(1, self.S * self.N)
      else:
        idx = int(idx)

      # Get the input
      data_input = self.new_curves[idx:(idx+1), :]
      other_data_input = data_input.reshape((3, self.F))

      # Get the outputs
      predicted_output = self.network.predict(data_input)
      true_output = self.new_geometry[idx].reshape((self.G, self.G, self.G), order='F')
      predicted_output = predicted_output.reshape((self.G, self.G, self.G), order='F')

      # return idx, other_data_input, true_output, predicted_output
      return predicted_output, true_output

  def get_geometry(self, idx=None):

      if idx is None:
          idx = numpy.random.randint(1, self.S * self.N)
      else:
        idx = int(idx)

      idx = int(idx)

      # Get the input
      data_input = self.new_geometry[idx:(idx+1), :]
      other_data_input = data_input.reshape((self.G, self.G, self.G), order='F')

      # return idx, other_data_input, true_output, predicted_output
      return other_data_input


  def get_performance(self, idx=None):

      if idx is None:
          idx = numpy.random.randint(1, self.S * self.N)
      else:
        idx = int(idx)

      idx = int(idx)

      # Get the input
      data_input = self.new_curves[idx:(idx+1), :]
      other_data_input = data_input.reshape((3, self.F))

      f = numpy.linspace(0.05, 2.0, 64)
      fd = pandas.DataFrame(f).rename(columns={0: "Frequency"})
      df_pred = pandas.DataFrame(other_data_input.transpose()).rename(columns={0: "Surge", 1: "Heave", 2: "Pitch"})
      table = pandas.concat([fd, df_pred], axis=1)

      # return idx, other_data_input, true_output, predicted_output
      return table

def simple_analysis(index):
    net = Network("/content/16forward_structure.json", "/content/16forward_weights.h5")
    return net.analysis(index)
    
def simple_synthesis(index):
    net = Network("/content/16inverse_structure.json", "/content/16inverse_weights.h5")
    pred, true = net.synthesis(index)
    return plotly_fig(pred), plotly_fig(true)
    
import plotly.graph_objects as go
import numpy as np

def performance(index):
    net = Network("/content/16forward_structure.json", "/content/16forward_weights.h5")
    return net.get_performance(index)

def geometry(index):
    net = Network("/content/16forward_structure.json", "/content/16forward_weights.h5")
    values = net.get_geometry(index)
    return plotly_fig(values)


def plotly_fig(values):
    X, Y, Z = np.mgrid[0:1:32j, 0:1:32j, 0:1:32j]
    fig = go.Figure(data=go.Volume(
        x=X.flatten(),
        y=Y.flatten(),
        z=Z.flatten(),
        value=values.flatten(),
        isomin=-0.1,
        isomax=0.8,
        opacity=0.1, # needs to be small to see through all surfaces
        surface_count=21, # needs to be a large number for good volume rendering
        ))
    return fig 
    
with gradio.Blocks() as analysis_demo:
    with gradio.Row():
        with gradio.Column():
            num = gradio.Number(42, label="data index")
            btn1 = gradio.Button("Select")    
        with gradio.Column():
            geo = gradio.Plot(label="Geometry")

    with gradio.Row():
        btn2 = gradio.Button("Estimate Spectrum")

    with gradio.Row():
      with gradio.Column():
          pred = gradio.Timeseries(x="Frequency", y=['Surge', 'Heave', 'Pitch'], label="Predicted")

      with gradio.Column():
          true = gradio.Timeseries(x="Frequency", y=['Surge', 'Heave', 'Pitch'], label="True")

    btn1.click(fn=geometry, inputs=[num], outputs=[geo])
    btn2.click(fn=simple_analysis, inputs=[num], outputs=[pred, true])
    
with gradio.Blocks() as synthesis_demo:
    with gradio.Row():
        with gradio.Column():
            num = gradio.Number(42, label="data index")
            btn1 = gradio.Button("Select")    
        with gradio.Column():
            perf = gradio.Timeseries(x="Frequency", y=['Surge', 'Heave', 'Pitch'], label="Performance")

    with gradio.Row():
        btn2 = gradio.Button("Synthesize Geometry")

    with gradio.Row():
      with gradio.Column():
          pred = gradio.Plot(label="Predicted")

      with gradio.Column():
          true = gradio.Plot(label="True")

    btn1.click(fn=performance, inputs=[num], outputs=[perf])
    btn2.click(fn=simple_synthesis, inputs=[num], outputs=[pred, true])
    
all_synthesis_demos = gradio.TabbedInterface([synthesis_demo], ["Random Spectrum from Data"])

all_analysis_demos = gradio.TabbedInterface([analysis_demo], ["Random Geometry from Data"])

demo = gradio.TabbedInterface([all_analysis_demos, all_synthesis_demos], ["Analysis", "Synthesis"])
demo.launch()