from kats.models.prophet import ProphetModel, ProphetParams
from kats.consts import TimeSeriesData
import pandas as pd
import matplotlib.pyplot as plt
import gradio as gr

def build_model(time_series_csv, changepoint_range, changepoint_prior_scale, seasonality_prior_scale, seasonality_mode, test_size):
  
  # time_series_csv = pd.DataFrame(time_series_csv)
  time_series_csv = pd.read_csv(time_series_csv.name, delimiter=",")

  # Divisón entre entrenamiento y test (sintaxis de Kats)
  train_size = 1 - test_size
  split_idx = int(len(time_series_csv) * train_size)
  
  ts_train = TimeSeriesData(
            pd.DataFrame(
                {
                    "time": time_series_csv['time'][:split_idx],
                    "target": time_series_csv['value'][:split_idx]
                }
            )
        )
  
  ts_test = TimeSeriesData(
            pd.DataFrame(
                {
                    "time": time_series_csv['time'][split_idx:],
                    "target": time_series_csv['value'][split_idx:]
                }
            )
        )


  # Creación del modelo
  prophet_params = ProphetParams(
        changepoint_range = changepoint_range,
        changepoint_prior_scale = changepoint_prior_scale,
        seasonality_prior_scale = seasonality_prior_scale,
        seasonality_mode = seasonality_mode
    )
  
  model = ProphetModel(ts_train, prophet_params)
    
  # Entrenamiento del modelo
  model.fit()

  # Predicción
  forecasting = model.predict(
        steps=len(ts_test)
    )
  
  result = pd.concat([
            ts_test.to_dataframe().set_index('time'),
            forecasting.set_index('time')
          ], axis=1)

  fig = plt.figure()
  plt.plot(result)
  plt.title("Prediction")

  return fig


# Interfaz
iface = gr.Interface(
    fn=build_model, 
    inputs=[
            gr.File(),
            gr.Slider(0.8, 0.99),
            gr.Number(precision=2),
            gr.Number(precision=2),
            gr.Radio(choices=['additive','multiplicative']),
            gr.Slider(0.1,0.35,step=0.1,value=0.2)
    ], 
    outputs=gr.Plot(),
    title="Prophet model generator",
    description="Upload your .csv with a time series and set up your own Prophet model to make predictions. 🔮",
    allow_flagging='never',
    theme="peach"
)

iface.launch(debug=True)