# covid_forecast_prophet.py
import pandas as pd
from prophet import Prophet
import matplotlib.pyplot as plt

# Load data
df = pd.read_csv("../data/united_kingdom_covid.csv")
df['Date'] = pd.date_range(start='2020-01-01', periods=len(df), freq='D')
df['Daily cases'] = df['Daily new confirmed cases of COVID-19'].fillna(0)

# Prepare data for Prophet (requires 'ds' and 'y' columns)
data = df[['Date', 'Daily cases']].rename(columns={'Date': 'ds', 'Daily cases': 'y'})

# Initialize and fit the model
model = Prophet()
model.fit(data)

# Create a dataframe of future dates (60 days ahead)
future = model.make_future_dataframe(periods=60)
forecast = model.predict(future)

# Plot forecast
fig1 = model.plot(forecast)
plt.title("Forecast of Daily COVID-19 Cases in the UK (Prophet)")
plt.xlabel("Date")
plt.ylabel("Number of Cases")
plt.savefig("../plots/forecast_prophet_cases1.png")
plt.close()
