# 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('data_curves.npz')['curves']
    geometry = numpy.load('data_geometry.npz')['geometry']
    constants = numpy.load('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 synthesis_from_spectrum(self, other_data_input):
      # Get the input
      data_input = other_data_input.reshape((1, 3*self.F))
 
      # Get the outputs
      predicted_output = self.network.predict(data_input)
      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

  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("16forward_structure.json", "16forward_weights.h5")
    return net.analysis(index)
    
def simple_synthesis(index):
    net = Network("16inverse_structure.json", "16inverse_weights.h5")
    pred, true = net.synthesis(index)
    return plotly_fig(pred), plotly_fig(true)
    
def synthesis_from_spectrum(df):
    net = Network("16inverse_structure.json", "16inverse_weights.h5")
    pred = net.synthesis_from_spectrum(df.to_numpy())
    return plotly_fig(pred)
    
import plotly.graph_objects as go
import numpy as np

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

def geometry(index):
    net = Network("16forward_structure.json", "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])
    
    
with gradio.Blocks() as synthesis_demo2:
    with gradio.Row():
        perf = gradio.Timeseries(x="Frequency", y=['Surge', 'Heave', 'Pitch'], label="Performance")

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

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

    btn2.click(fn=synthesis_from_spectrum, inputs=[perf], outputs=[pred])
    
    
with gradio.Blocks() as synthesis_demo3:
    with gradio.Row():
        perf = gradio.DataFrame(headers=['Surge', 'Heave', 'Pitch'], value=numpy.zeros((64, 3)).tolist())

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

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

    btn2.click(fn=synthesis_from_spectrum, inputs=[perf], outputs=[pred])
    
    
with gradio.Blocks() as intro:
    with gradio.Row():
        with gradio.Column():
            with gradio.Row():
                title = gradio.Markdown("# Toward the Rapid Design of Engineered Systems Through Deep Neural Networks")
            with gradio.Row():
                gradio.Markdown("The design of a system commits a significant portion of the final cost of that system. Many computational approaches have been developed to assist designers in the analysis (e.g., computational fluid dynamics) and synthesis (e.g., topology optimization) of engineered systems. However, many of these approaches are computationally intensive, taking significant time to complete an analysis and even longer to iteratively synthesize a solution. The current work proposes a methodology for rapidly evaluating and synthesizing engineered systems through the use of deep neural networks. The proposed methodology is applied to the analysis and synthesis of offshore structures such as oil platforms. These structures are constructed in a marine environment and are typically designed to achieve specific dynamics in response to a known spectrum of ocean waves. Results show that deep learning can be used to accurately and rapidly synthesize and analyze offshore structures.\n\nThe paper linked to the left provides details about the implementation. This site contains demos of the trained networks.")
        with gradio.Column():
            download = gradio.File(value="McComb2019_Chapter_TowardTheRapidDesignOfEngineer.pdf")
    
    
all_synthesis_demos = gradio.TabbedInterface([synthesis_demo, synthesis_demo2, synthesis_demo3], ["Spectrum from Dataset", "Spectrum from File", "Spectrum from DataFrame"])

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

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