import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error
from statsmodels.tsa.arima.model import ARIMA
import matplotlib.pyplot as plt

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

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

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

# Rename for simplicity
df = df.rename(columns={value_col: 'Value'})

# Create lag features
df['Lag1'] = df['Value'].shift(1)
df['Lag2'] = df['Value'].shift(2)
df = df.dropna()

# Train-test split (last 20%)
split = int(len(df)*0.8)
train, test = df.iloc[:split], df.iloc[split:]

X_train, y_train = train[['Lag1','Lag2']], train['Value']
X_test, y_test = test[['Lag1','Lag2']], test['Value']

# Linear Regression
lr = LinearRegression().fit(X_train, y_train)
lr_pred = lr.predict(X_test)

# Random Forest
rf = RandomForestRegressor(n_estimators=100, random_state=42).fit(X_train, y_train)
rf_pred = rf.predict(X_test)

# ARIMA (ARIMA needs only Value)
arima = ARIMA(train['Value'], order=(2,1,2)).fit()
arima_pred = arima.forecast(steps=len(test))

# Evaluation
def eval_model(y, pred, name):
    print(name, "MAE:", mean_absolute_error(y, pred),
          "RMSE:", np.sqrt(mean_squared_error(y, pred)))

eval_model(y_test, lr_pred, "Linear Regression")
eval_model(y_test, rf_pred, "Random Forest")
eval_model(y_test, arima_pred, "ARIMA")

# Plot
plt.figure()
plt.plot(test[date_col], y_test, label="Actual")
plt.plot(test[date_col], lr_pred, label="LR")
plt.plot(test[date_col], rf_pred, label="RF")
plt.plot(test[date_col], arima_pred, label="ARIMA")
plt.legend()
plt.show()
#pip install pandas numpy scikit-learn statsmodels matplotlib