import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.arima.model import ARIMA

# Load dataset
data = pd.read_csv("Electric_Production.csv")

# Auto-detect columns
date_col = data.columns[0]
value_col = data.columns[1]

# Convert date
data[date_col] = pd.to_datetime(data[date_col], errors='coerce')

# Clean data
data = data.dropna(subset=[date_col, value_col])
data = data.sort_values(date_col)
data.set_index(date_col, inplace=True)

# Create time series
ts = data[value_col].asfreq('MS')   # monthly (change if needed)
ts = ts.ffill()                     # handle missing values

# AR model (ARIMA(p,0,0))
model = ARIMA(ts, order=(2,0,0))
result = model.fit()

# Forecast
fc = result.get_forecast(steps=20).summary_frame()

# Print output
print(fc)

# Plot
plt.figure()
plt.plot(ts, label="Historical")
plt.plot(fc['mean'], label="Forecast")
plt.fill_between(fc.index, fc['mean_ci_lower'], fc['mean_ci_upper'], alpha=0.3)
plt.legend()
plt.show()
#pip install pandas matplotlib statsmodels