# https://github.com/google-research/tensorflow-coder/blob/master/tf_coder/tf_coder_main.py
import streamlit as st
from tf_coder.value_search import colab_interface
from tf_coder.value_search import value_search_settings
import io
from contextlib import redirect_stdout

inputs = st.text_area('The input tensor(s) specified as key-value pairs', placeholder="{'rows': [10, 20, 30],'cols': [1,2,3,4]}")
  # The single desired output tensor.

st.sidebar.header("Generation settings:")
gen_kwargs["do_sample"] = st.sidebar.radio("Decoding strategy",  ["Require All", "Require One"]) == "Require All"
gen_kwargs["max_new_tokens"] = st.sidebar.slider("Number of tokens to generate", value=default_length, min_value=8, step=8, max_value=256)
if gen_kwargs["do_sample"]:
    gen_kwargs["temperature"] = st.sidebar.slider("Temperature", value = 0.2, min_value = 0.0, max_value=2.0, step=0.05)
    gen_kwargs["top_k"] = st.sidebar.slider("Top-k", min_value = 0, max_value=100, value = 0)
    gen_kwargs["top_p"] = st.sidebar.slider("Top-p", min_value = 0.0, max_value=1.0, step = 0.01, value = 0.95)
    
settings = value_search_settings.from_dict({
      'timeout': 300,
      'only_minimal_solutions': False,
      'max_solutions': 1,
      'require_all_inputs_used': True,
      'require_one_input_used': False,
  })

with io.StringIO() as buf, redirect_stdout(buf):
    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'
    results = colab_interface.run_value_search_from_colab(eval(inputs), output, constants, description, settings)
    stdout = buf.getvalue()
st.code(stdout, language='bash')