import gradio as gr
import pandas as pd
from neuralprophet import NeuralProphet
import warnings
import torch.optim as optim
from torch.optim.lr_scheduler import LambdaLR

warnings.filterwarnings("ignore", category=UserWarning)

url = "VN Index Historical Data.csv"
df = pd.read_csv(url)
df = df[["Date", "Price"]]
df = df.rename(columns={"Date": "ds", "Price": "y"})
df.fillna(method='ffill', inplace=True)
df.dropna(inplace=True)


class CustomNeuralProphet(NeuralProphet):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.optimizer = None

    def lr_scheduler_step(self, epoch, batch_idx, optimizer):
        # Custom logic for LR scheduler step
        for lr_scheduler in optimizer.param_groups[0]['lr_scheduler']:
            lr_scheduler.step()

m = CustomNeuralProphet(
    n_forecasts=30,
    n_lags=12,
    changepoints_range=1,
    num_hidden_layers=3,
    yearly_seasonality=True,
    n_changepoints=150,
    trend_reg_threshold=False,
    d_hidden=3,
    global_normalization=True,
    seasonality_reg=1,
    unknown_data_normalization=True,
    seasonality_mode="multiplicative",
    drop_missing=True,
    learning_rate=0.03,
)

# Set the custom LR scheduler
m.fit(df, freq='D')  # Fit the model first before accessing the optimizer
m.optimizer = optim.Adam(m.model.parameters(), lr=0.03)  # Example optimizer, adjust as needed

lr_scheduler = LambdaLR(m.optimizer, lambda epoch: 0.95 ** epoch)  # Example LR scheduler, adjust as needed
m.optimizer.param_groups[0]['lr_scheduler'] = [lr_scheduler]

future = m.make_future_dataframe(df, periods=30, n_historic_predictions=True)
forecast = m.predict(future)


def predict_vn_index(option=None):
    fig = m.plot(forecast)
    path = "forecast_plot.png"
    fig.savefig(path)
    disclaimer = "Quý khách chỉ xem đây là tham khảo, công ty không chịu bất cứ trách nhiệm nào về tình trạng đầu tư của quý khách."
    return path, disclaimer


if __name__ == "__main__":
    dropdown = gr.inputs.Dropdown(["VNIndex"], label="Choose an option", default="VNIndex")
    image_output = gr.outputs.Image(type="file", label="Forecast Plot")
    disclaimer_output = gr.outputs.Textbox(label="Disclaimer")
    interface = gr.Interface(fn=predict_vn_index, inputs=dropdown, outputs=[image_output, disclaimer_output], title="Dự báo VN Index 30 ngày tới")
    interface.launch()