"""
Kriging / Gaussian Process Regression for wind-turbine fatigue-life prediction.

Inputs:
    wind_speed_mps
    turbulent_intensity
    yaw_misalignment_deg
    design_load_kN

Target:
    fatigue_life_hours

The target is modeled as log10(fatigue life) because fatigue life spans
orders of magnitude when provided in hours. This conversion can be avoided if
fatigue life is in years.

GaussianProcessRegressor is a practical ML form of Kriging for multi-
dimensional inputs.
"""
import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import (
    ConstantKernel, Matern, WhiteKernel
)
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

hr_in_yr = 365*24  # Set this value to 1.0 to predict life in hours.
# -----------------------------------------------------------------------------
# Load data into Panda dataframe, define feature columns and target column
# -----------------------------------------------------------------------------
start_cpu = time.process_time()

DATA_FILE = "WT_Dummy_Fatigue.csv"
df = pd.read_csv(DATA_FILE)

features = [
    "wind_speed_mps",
    "turb_intensity",
    "yaw_misalignment_deg",
    "design_load_kN",
]
target = "fatigue_life_hours"

'''
X = df[features].value shall throw following warning:
X has feature names, but StandardScaler was fitted without feature names. 
Scikit-learn expects consistency: while fitting the scaler without feature 
names, one should transform without them too.
'''
X = df[features]   
y = df[target].values

# Kriging is easier to train when the target is approximately Gaussian.
# Fatigue life is highly skewed, so predict log10(life).
y_log = np.log10(y)


# ------------------------------------------------------------
# Train/test split
# ------------------------------------------------------------
X_train, X_test, y_train, y_test = train_test_split(
    X, y_log, test_size=0.30, random_state=42
)

# ------------------------------------------------------------
# Standardize input variables
# ------------------------------------------------------------
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled  = scaler.transform(X_test)


# ------------------------------------------------------------
# Define Kriging / Gaussian Process model
# ------------------------------------------------------------
'''
Matern is often more flexible than a simple RBF for engineering data.
WhiteKernel represents measurement/model noise (nugget effect).
'''
kernel = (
    ConstantKernel(1.0, (1e-3, 1e3))
    * Matern(
        length_scale = np.ones(X_train_scaled.shape[1]),
        length_scale_bounds = (1e-2, 1e2),
        nu=2.5,
    )
    + WhiteKernel(
        noise_level=0.05,
        noise_level_bounds=(1e-8, 1.0)
    )
)

gpr = GaussianProcessRegressor(
    kernel=kernel,
    normalize_y=True,
    n_restarts_optimizer=5,
    random_state=42,
)

gpr.fit(X_train_scaled, y_train)


# ------------------------------------------------------------
# Predict fatigue life
# ------------------------------------------------------------
y_pred_log, y_std_log = gpr.predict(
    X_test_scaled,
    return_std=True
)

# Convert back from log10(hours) to hours
y_test_hours = 10**y_test
y_pred_hours = 10**y_pred_log

# ------------------------------------------------------------
# Evaluation
# ------------------------------------------------------------
mae = mean_absolute_error(y_test_hours, y_pred_hours)
mae_yrs = mae / hr_in_yr
rmse = np.sqrt(mean_squared_error(y_test_hours, y_pred_hours))
rmse_yrs = rmse / hr_in_yr
r2 = r2_score(y_test_hours, y_pred_hours)

log_mae = mean_absolute_error(y_test, y_pred_log)
log_mae_yrs = log_mae / hr_in_yr

print("\n===== Kriging / GPR Performance")
print(f"MAE  : {mae_yrs:,.02f} years")
print(f"RMSE : {rmse_yrs:,.02f} years")
print(f"R²   : {r2:.4f}")

end_cpu = time.process_time()
cpu_duration = np.round(end_cpu - start_cpu, 1)
# ------------------------------------------------------------
# Plot measured vs predicted fatigue life
# ------------------------------------------------------------
plt.figure(figsize=(7, 6))
plt.scatter(y_test_hours, y_pred_hours, alpha=0.75, s=25)

lims = [
    min(y_test_hours.min(), y_pred_hours.min()),
    max(y_test_hours.max(), y_pred_hours.max()),
]
plt.plot(lims, lims, "--")

plt.xscale("log")
plt.yscale("log")
plt.xlabel("Actual fatigue life [hours]")
plt.ylabel("Predicted fatigue life [hours]")
plt.title("Kriging / Gaussian Process Fatigue-Life Prediction")
plt.grid(True, which="both", alpha=0.25)
plt.tight_layout()
plt.savefig("WT_Fatigue_Kriging.png")


# ------------------------------------------------------------
# Example prediction for a new design load case
# ------------------------------------------------------------
new_case = pd.DataFrame({
    "wind_speed_mps": [18.0],
    "turb_intensity": [0.15],
    "yaw_misalignment_deg": [10.0],
    "design_load_kN": [2100.0],
})

new_scaled = scaler.transform(new_case[features])

pred_log, pred_std_log = gpr.predict(
    new_scaled,
    return_std=True
)

pred_hours = 10**pred_log[0]
pred_low = 10**(pred_log[0] - 1.96 * pred_std_log[0])
pred_high = 10**(pred_log[0] + 1.96 * pred_std_log[0])

print("\n=====--- New Load Case ---=====")
print(new_case.to_string(index=False))
print("\n=====--- ---- --- ---- ---=====")
print(f"Predicted fatigue life: {pred_hours/hr_in_yr:,.2f}")
print(
    f"Approx. 95% Kriging interval: "
    f"{pred_low/hr_in_yr:,.2f} to {pred_high/hr_in_yr:,.2f}\n"
)
print(f"CPU time to execute ML algorithm: {cpu_duration} [s]\n")