# 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, value_search_settings
import io
from contextlib import redirect_stdout
from streamlit_ace import st_ace
st.set_page_config(page_icon='👩‍💻', layout="wide", initial_sidebar_state="collapsed")

col1, col2, col3 = st.columns([5, 5, 3])
with col1:
    st.write('### Inputs')
    inputs = st_ace(placeholder="The input tensor(s) specified as a dictionary", value="{'rows': [10, 20, 30],\n'cols': [1,2,3,4]}", language="python", theme="solarized_dark")
with col2:
    st.write('### Output')
    output = st_ace(placeholder="The output tensor", value="[[11, 12, 13, 14],\n[21, 22, 23, 24],\n[31, 32, 33, 34]]", language="python", theme="solarized_dark")
with col3:
    st.write('### Constants')
    constants = st_ace(placeholder="Optional list of scalar constants", value="[]", language="python", theme="solarized_dark")
if any([inputs, output, constants]):
    i = eval(inputs)
    o = eval(output)
    c = eval(constants)
description = st.text_input(label="Description", placeholder="An optional natural language description of the operation", value="add two vectors with broadcasting to get a matrix")
st.sidebar.header("Settings:")
settings_kwargs = dict()
settings_kwargs["require_all_inputs_used"] = st.sidebar.checkbox("Require All Inputs", value=True)
settings_kwargs["only_minimal_solutions"] = st.sidebar.checkbox("Only Minimal Solutions", value=False)
settings_kwargs["max_solutions"] = st.sidebar.slider("Maximum number of solutions", value=1, min_value=1, step=1, max_value=256)
settings_kwargs["timeout"] = st.sidebar.slider("Timeout in seconds", value=300, min_value=1, step=10, max_value=300)
    
settings = value_search_settings.from_dict({
      'timeout': settings_kwargs["timeout"],
      'only_minimal_solutions': settings_kwargs["only_minimal_solutions"],
      'max_solutions': settings_kwargs["max_solutions"],
      'require_all_inputs_used': settings_kwargs["require_all_inputs_used"],
      'require_one_input_used': not settings_kwargs["require_all_inputs_used"],
  })

with io.StringIO() as buf, redirect_stdout(buf):
    constants = []
    results = colab_interface.run_value_search_from_colab(i, o, c, description, settings)
    stdout = buf.getvalue()
st.code(stdout, language='bash')