"""
Neural network modeling of real air heater.
Air heater home page: http://techteach.no/lab/air_heater/
Code generated by ChatGPT, but with a few modifications
Finn Aakre Haugen, finn@techteach.no
2025 07 09
"""

import numpy as np
import matplotlib.pyplot as plt
from sklearn.neural_network import MLPRegressor
from sklearn.preprocessing import StandardScaler

# --- Load data ---
datafile_url = 'https://techteach.no/control/python/air_heater_v01.csv'
data = np.loadtxt(datafile_url, delimiter=",", comments="#")
t = data[:, 0]
u = data[:, 1]
T = data[:, 2]

# --- Select time interval ---
mask = (t >= 70) & (t <= 220)
t = t[mask]
u = u[mask]
T = T[mask]

# --- Create lagged dataset ---
lags = 30  # 2 sec delay, 0.1 s timestep

def create_lagged_dataset(T, u, t, lags):
    X, y, t_out, u_out = [], [], [], []
    for i in range(lags, len(T)):
        # Features: [T[i-1], ..., T[i-lags], u[i-1], ..., u[i-lags]]
        features = []
        for j in range(1, lags+1):
            features.append(T[i-j])
            features.append(u[i-j])
        X.append(features)
        y.append(T[i])
        t_out.append(t[i])
        u_out.append(u[i])
    return np.array(X), np.array(y), np.array(t_out), np.array(u_out)

X, y, t_out, u_out = create_lagged_dataset(T, u, t, lags)

# --- Split data ---
n = len(X)
train_end = int(0.6 * n)
test_end = int(0.8 * n)

X_train, X_test, X_sim = X[:train_end], X[train_end:test_end], X[test_end:]
y_train, y_test, y_sim = y[:train_end], y[train_end:test_end], y[test_end:]
t_train, t_test, t_sim = t_out[:train_end], t_out[train_end:test_end], t_out[test_end:]
u_train, u_test, u_sim = u_out[:train_end], u_out[train_end:test_end], u_out[test_end:]

# --- Scale ---
scaler_X = StandardScaler().fit(X_train)
scaler_y = StandardScaler().fit(y_train.reshape(-1,1))
X_train_scaled = scaler_X.transform(X_train)
X_test_scaled = scaler_X.transform(X_test)
X_sim_scaled = scaler_X.transform(X_sim)
y_train_scaled = scaler_y.transform(y_train.reshape(-1,1)).ravel()
y_test_scaled = scaler_y.transform(y_test.reshape(-1,1)).ravel()

# --- Train MLP ---
model = MLPRegressor(hidden_layer_sizes=(96), activation='relu', solver='adam',
                     max_iter=2000, early_stopping=True, random_state=1)
model.fit(X_train_scaled, y_train_scaled)

# --- Predict train/test ---
y_train_pred = scaler_y.inverse_transform(model.predict(X_train_scaled).reshape(-1,1)).ravel()
y_test_pred = scaler_y.inverse_transform(model.predict(X_test_scaled).reshape(-1,1)).ravel()

# --- Recursive simulation ---
def recursive_predict(model, T_seed, u_seq, lags, scaler_X, scaler_y):
    y_pred = []
    T_hist = list(T_seed)
    for i in range(lags, len(u_seq)):
        features = []
        for j in range(1, lags+1):
            features.append(T_hist[-j])
            features.append(u_seq[i-j])
        X_input = np.array(features).reshape(1, -1)
        X_scaled = scaler_X.transform(X_input)
        y_scaled = model.predict(X_scaled)
        y_val = scaler_y.inverse_transform(y_scaled.reshape(-1,1)).ravel()[0]
        y_pred.append(y_val)
        T_hist.append(y_val)
    return np.array(y_pred)

T_seed = list(T[test_end-lags:test_end])
u_sim_full = u[test_end-lags:]
y_sim_pred = recursive_predict(model, T_seed, u_sim_full, lags, scaler_X, scaler_y)

# --- Plotting ---
sim_len = len(y_sim_pred)
t_recursive = t[test_end:test_end+sim_len]
T_true_recursive = T[test_end:test_end+lags+sim_len]
u_recursive = u[test_end:test_end+lags+sim_len]

fig, axs = plt.subplots(2, 3, figsize=(16, 8))

axs[0, 0].plot(t_train, y_train, label='True T (train)')
axs[0, 0].plot(t_train, y_train_pred, label='Predicted T (train)', color='red')
axs[0, 0].set_title('Training: True vs Predicted T', color='blue')
axs[0, 0].set_xlabel('Time [s]')
axs[0, 0].set_ylabel('Temperature [°C]')
axs[0, 0].legend()
axs[0, 0].grid(True)

axs[0, 1].plot(t_test, y_test, label='True T (test)', color='blue')
axs[0, 1].plot(t_test, y_test_pred, label='Predicted T (test)', color='red')
axs[0, 1].set_title('Test: True vs Predicted T')
axs[0, 1].set_xlabel('Time [s]')
axs[0, 1].set_ylabel('Temperature [°C]')
axs[0, 1].legend()
axs[0, 1].grid(True)

axs[0, 2].plot(t_recursive, T_true_recursive, label='True T (sim)', color='blue')
axs[0, 2].plot(t_recursive, y_sim_pred, label='Recursive Predicted T', color='red')
axs[0, 2].set_title('Simulation: Recursive Prediction')
axs[0, 2].set_xlabel('Time [s]')
axs[0, 2].set_ylabel('Temperature [°C]')
axs[0, 2].legend()
axs[0, 2].grid(True)

axs[1, 0].plot(t_train, u_train, label='u (train)', color='green')
axs[1, 0].set_title('Training Input u')
axs[1, 0].set_xlabel('Time [s]')
axs[1, 0].set_ylabel('u [V]')
axs[1, 0].legend()
axs[1, 0].grid(True)

axs[1, 1].plot(t_test, u_test, label='u (test)', color='green')
axs[1, 1].set_title('Test Input u')
axs[1, 1].set_xlabel('Time [s]')
axs[1, 1].set_ylabel('u [V]')
axs[1, 1].legend()
axs[1, 1].grid(True)

axs[1, 2].plot(t_recursive, u_recursive, label='u (sim)', color='green')
axs[1, 2].set_title('Simulation Input u')
axs[1, 2].set_xlabel('Time [s]')
axs[1, 2].set_ylabel('u [V]')
axs[1, 2].legend()
axs[1, 2].grid(True)

plt.tight_layout()

plt.savefig('plot_dnn_real_airheater.pdf')
plt.show()