import streamlit as st
from matplotlib import pyplot as plt
from utils import * 
import warnings

st.write(" ")
raq_ra = st.number_input("# raq/ra ", value=None, placeholder="betweel 0 and 10")
fkt = st.number_input("# FKT ", value=None, placeholder="betweel 1e+6 and 1e+10")
fkv = st.number_input("# FKV ", value=None, placeholder="betweel 1 and 100")
num_points =st.number_input("# number of profile points ", value=128, placeholder=" e.g. 128")
num_points = int(num_points)
st.write("raq/ra = ", raq_ra, "FKT = ", fkt, "FKV = ", fkv, "profile points = ", num_points)

if raq_ra is not None and fkt is not None and fkv is not None:
    with open('numpy_networks/mlp_[128, 128, 128, 128, 128].pkl', 'rb') as file: 
        mlp = pickle.load(file) 
    
    r_list = [raq_ra]
    t_list = [fkt]
    v_list = [fkv]
    
    for i in range(len(r_list)):
        if r_list[i]<0 or  r_list[i]>9.5:
            st.warning('RaQ/Ra is outside the range of the training dataset', icon="⚠️") 
        if t_list[i]<1e+6 or  t_list[i]>5e+9:
            st.warning('FKT is outside the range of the training dataset', icon="⚠️")
        if v_list[i]<1 or  v_list[i]>95:
            st.warning('FKV is outside the range of the training dataset', icon="⚠️")
    
    y_prof = np.concatenate(([1], np.linspace(0+1/(num_points*2),1-1/(num_points*2),num_points-2)[::-1], [0]))
    
    ### calculates temperature profile ###
    x_in = get_input(r_list, t_list, v_list, y_prof)
    y_pred_nn_pointwise = get_profile(x_in, mlp, num_sims=len(r_list), prof_points=num_points) 
    ### calculates temperature profile ###
    
    ### writes out temperature profile ###
    profs = np.concatenate((y_prof.reshape(num_points,1),y_pred_nn_pointwise.reshape(num_points,1)), axis=1)
    table = '''
    Depth       Temperature
    '''
    for p in profs:
        table += str(p[0]) + "   " + str(p[1]) + "\n"
    
    fname = "profile_raq_ra" + str(r_list[0]) + "_fkt" + str(t_list[0]) + "_fkv" + str(v_list[0]) + ".txt"
    st.download_button('Download temperature profile', table, fname)
    ### writes out temperature profile ###
    
    ### plots temperature profile ###
    for i in range(len(r_list)):
        fig = plt.figure()
        plt.plot(y_pred_nn_pointwise[i,:], y_prof, 'k-', linewidth=3.0, label="neural network")
        plt.ylim([1,0])
        plt.xlabel("Temperature")
        plt.ylabel("Depth")
        plt.legend()
        plt.grid()
        st.pyplot(fig)
    ### plots temperature profile ###