"""
Neural network modeling of real DC motor.
DC motor home page: http://techteach.no/lab/dcmotor/
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

# --- Download and load data ---
import urllib.request
url = "https://techteach.no/control/python/dcmotor_v01.csv"
local_filename = "dcmotor_v01.csv"
urllib.request.urlretrieve(url, local_filename)
data = np.loadtxt(local_filename, delimiter=",", comments="#")
t = data[:, 0]
S = data[:, 2]
u = data[:, 3]

# --- Select time interval ---
mask = (t >= 30) & (t <= 48)
t = t[mask]
S = S[mask]
u = u[mask]

# --- Create lagged dataset ---
lags = 10  # number of lag steps (0.05s * 10 = 0.5s window)
def create_lagged_dataset(S, u, t, lags):
    X, y, t_out, u_out = [], [], [], []
    for i in range(lags, len(S)):
        features = []
        for j in range(1, lags+1):
            features.append(S[i-j])
            features.append(u[i-j])
        X.append(features)
        y.append(S[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(S, u, t, lags)

# --- Split data (60% train, 20% test, 20% sim) ---
n = len(X)
train_end = int(0.6 * n)
test_end = int(0.8 * n)
sim_start = test_end

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:]

# --- Standardize (fit only on train) ---
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 the deep NN ---
model = MLPRegressor(
    hidden_layer_sizes=(100, 80, 50), activation='relu',
    solver='adam', max_iter=2000, early_stopping=True, random_state=42
)
model.fit(X_train_scaled, y_train_scaled)

# --- Predict on 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 (closed-loop) ---
def recursive_predict(model, S_seed, u_seq, lags, scaler_X, scaler_y, n_pred):
    y_pred = []
    S_hist = list(S_seed)
    for i in range(n_pred):
        features = []
        for j in range(1, lags+1):
            features.append(S_hist[-j])
            features.append(u_seq[i + lags - 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)
        S_hist.append(y_val)
    return np.array(y_pred)

# Seed with last lags S from end of test set
S_seed = list(S[test_end-lags:test_end])
# u_sim_full = u[test_end-lags:]  # Original code
u_sim_full = u[test_end:]  #  Modified by FAH
n_sim = len(X_sim)
y_sim_pred = recursive_predict(model, S_seed, u_sim_full, lags, scaler_X, scaler_y, n_sim)

# --- For plotting: ensure matching lengths
t_recursive = t_out[sim_start:sim_start + n_sim]
S_true_recursive = y[sim_start:sim_start + n_sim]
u_recursive = u_out[sim_start:sim_start + n_sim]

# --- Plotting ---
fig, axs = plt.subplots(2, 3, figsize=(12, 8))

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

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

axs[0, 2].plot(t_recursive, S_true_recursive, label='True S (sim)',
               color='blue')
axs[0, 2].plot(t_recursive, y_sim_pred, label='Recursive Predicted S',
               color='red')
axs[0, 2].set_title('Simulation: Recursive Prediction')
axs[0, 2].set_xlabel('Time [s]')
axs[0, 2].set_ylabel('S [krpm]')
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_nn_dcmotor.pdf')
plt.show()