import streamlit as st
import numpy as np
import tensorflow as tf
from tf_coder.value_search import colab_interface
from tf_coder.value_search import value_search_settings as settings_module

def get_problem():
  """Specifies a problem to run TF-Coder on. Edit this function!"""
  # A dict mapping input variable names to input tensors.
  inputs = {
      'rows': [10, 20, 30],
      'cols': [1, 2, 3, 4],
  }

  # The single desired output tensor.
  output = [[11, 12, 13, 14],
            [21, 22, 23, 24],
            [31, 32, 33, 34]]

  # A list of relevant scalar constants (if any).
  constants = []

  # An English description of the tensor manipulation.
  description = 'add two vectors with broadcasting to get a matrix'

  return inputs, output, constants, description

settings = settings_module.from_dict({
      'timeout': 300,
      'only_minimal_solutions': False,
      'max_solutions': 1,
      'require_all_inputs_used': True,
      'require_one_input_used': False,
  })

i = st.text_area("input tensor","second test")

inputs, output, constants, description = get_problem()
t = colab_interface.run_value_search_from_colab(inputs, output, constants, description, settings)
st.write(t.solution)
st.write(t.total_time)
st.write(t.statistics)