import streamlit as st
import pickle
import numpy as np

st.title("COVID19 GLOBAL FORECAST")
coc=st.number_input("Confirmed Cases")
lat=st.number_input("Latitude")
lon=st.number_input("Longitude")
model_selection=st.radio("Select your model",("Regression","Classification"))
if model_selection is "Regression":
    file_path = 'covid1_model.pkl'
    with open(file_path, 'rb') as file:
        model = pickle.load(file)
    data=[coc,lat,lon]
    if st.button("Predict"):
       data=np.array(data)
       if len(data.shape) == 1:
          data = np.expand_dims(data, axis=0) 
       prediction=model.predict(data)
       st.write(prediction)
else:
    file_path = 'covid2_model2.pkl'
    with open(file_path, 'rb') as file:
        model = pickle.load(file)
    data=[lat,lon]
    if st.button("Predict"):
       data=np.array(data)
       if len(data.shape) == 1:
          data = np.expand_dims(data, axis=0)
    prediction=model.predict(data)
    pred=np.argmax(prediction)
    st.write(pred)