# CNN
import os
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv1D, Input, Dense, Dropout, Flatten
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.optimizers import Adam
from pathlib import Path
import tensorflow.keras.backend as K

# Set global font size for plots and float precision for compatibility
plt.rcParams.update({'font.size': 20})
tf.keras.backend.set_floatx('float64')

# Define directories
input_dir = "./"  # Define input directory
output_dir = "./output"  # Define output directory
Path(output_dir).mkdir(parents=True, exist_ok=True)
input_path = f"{input_dir}/"
combined_dataset_path = os.path.join(input_path, "CombinedDatasetLongUnsmooth.csv")

# Window size and forecast horizon parameters
window_size = 50
forecast_horizon = 20

# Moving average kernel and function
def get_moving_average_kernel(window_size):
    kernel = tf.ones((window_size, 1, 1), dtype=tf.float64) / window_size
    return kernel

def moving_average(x, kernel):
    smoothed_x = tf.nn.conv1d(tf.expand_dims(x, axis=0), kernel, stride=1, padding='SAME')
    return tf.squeeze(smoothed_x, axis=0)

# Define custom loss function with smoothing
def custom_loss_with_smoothing_components(y_true, y_pred, gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10):
    y_true = K.cast(y_true, dtype=tf.float64)
    y_pred = K.cast(y_pred, dtype=tf.float64)
    y_true = K.reshape(y_true, (-1, 1))
    y_pred = K.reshape(y_pred, (-1, 1))
    kernel = get_moving_average_kernel(window_size)
    y_true_smooth = moving_average(y_true, kernel)
    y_pred_smooth = moving_average(y_pred, kernel)
    rmse = K.sqrt(K.mean(K.square(y_pred_smooth - y_true_smooth)))
    gradient_true = y_true_smooth[1:] - y_true_smooth[:-1]
    gradient_pred = y_pred_smooth[1:] - y_pred_smooth[:-1]
    gradient_true_smooth = moving_average(gradient_true, kernel)
    gradient_pred_smooth = moving_average(gradient_pred, kernel)
    threshold_5sigma = K.std(gradient_true_smooth) * 5
    threshold_2sigma = K.std(gradient_true_smooth) * 2
    abrupt_change_mask_5sigma = K.cast(K.greater(K.abs(gradient_true_smooth), threshold_5sigma), dtype=tf.float64)
    abrupt_change_mask_2sigma = K.cast(K.greater(K.abs(gradient_true_smooth), threshold_2sigma), dtype=tf.float64)
    gradient_diff_5sigma = K.mean(K.square(gradient_pred_smooth - gradient_true_smooth) * (1 + heavy_penalty * abrupt_change_mask_5sigma))
    gradient_diff_2sigma = K.mean(K.square(gradient_pred_smooth - gradient_true_smooth) * (1 + two_sigma_penalty * abrupt_change_mask_2sigma))
    total_loss = rmse + gradient_diff_5sigma + gradient_diff_2sigma
    return total_loss, rmse, gradient_diff_5sigma, gradient_diff_2sigma

# Define custom EDR function
def custom_edr(ts1, ts2, epsilon=0.1):
    len_ts1, len_ts2 = len(ts1), len(ts2)
    D = np.zeros((len_ts1 + 1, len_ts2 + 1))
    for i in range(1, len_ts1 + 1):
        D[i, 0] = i
    for j in range(1, len_ts2 + 1):
        D[0, j] = j
    for i in range(1, len_ts1 + 1):
        for j in range(1, len_ts2 + 1):
            cost = 0 if abs(ts1[i - 1] - ts2[j - 1]) <= epsilon else 1
            D[i, j] = min(D[i - 1, j] + 1, D[i, j - 1] + 1, D[i - 1, j - 1] + cost)
    return D[len_ts1, len_ts2]

# Load dataset
dataset = pd.read_csv(combined_dataset_path).drop(columns=['AMOC_20S', 'SALT_500m', 'ICEFRAC', 'AMOC'])

# Sliding window class
class WindowSlider:
    def __init__(self, window_size, forecast_horizon):
        self.window_size = window_size
        self.forecast_horizon = forecast_horizon

    def collect_windows(self, features, labels):
        X, y = [], []
        for i in range(self.window_size, len(features) - self.forecast_horizon + 1):
            X.append(features.iloc[i - self.window_size:i].values.flatten())
            y.append(labels.iloc[i + self.forecast_horizon - 1])
        return pd.DataFrame(X), pd.Series(y)

# Prepare data
slider = WindowSlider(window_size, forecast_horizon)
train_features, test_features = train_test_split(dataset.drop(columns=['AMOC_40N']), test_size=0.3, shuffle=False)
train_labels, test_labels = train_test_split(dataset['AMOC_40N'], test_size=0.3, shuffle=False)
train_windows, train_labels = slider.collect_windows(train_features, train_labels)
test_windows, test_labels = slider.collect_windows(test_features, test_labels)

# Normalize data
scaler_X, scaler_y = MinMaxScaler(), MinMaxScaler()
train_windows_scaled = scaler_X.fit_transform(train_windows)
test_windows_scaled = scaler_X.transform(test_windows)
train_labels_scaled = scaler_y.fit_transform(train_labels.values.reshape(-1, 1))
test_labels_scaled = scaler_y.transform(test_labels.values.reshape(-1, 1))
train_windows_3d = train_windows_scaled.reshape((-1, window_size, train_windows_scaled.shape[1] // window_size))
test_windows_3d = test_windows_scaled.reshape((-1, window_size, test_windows_scaled.shape[1] // window_size))

# Build and train CNN model
def build_and_train_cnn(input_shape, X_train, y_train):
    model = Sequential([
        Input(shape=input_shape),
        Conv1D(filters=64, kernel_size=2, activation='relu'),
        Dropout(0.3),
        Conv1D(filters=32, kernel_size=2, activation='relu'),
        Flatten(),
        Dense(50, activation='relu'),
        Dense(1)
    ])
    model.compile(optimizer=Adam(learning_rate=0.001), loss=lambda y_true, y_pred: custom_loss_with_smoothing_components(y_true, y_pred)[0])
    model.fit(X_train, y_train, epochs=100, batch_size=16, validation_split=0.2, verbose=1,
              callbacks=[EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True),
                         ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=3, min_lr=0.0001)])
    return model

# Train CNN
cnn_model = build_and_train_cnn((window_size, train_windows_3d.shape[2]), train_windows_3d, train_labels_scaled)
predictions_scaled = cnn_model.predict(test_windows_3d)
predictions = scaler_y.inverse_transform(predictions_scaled)

# Save predictions to CSV
predictions_df = pd.DataFrame({
    'Time_Index': range(len(predictions)),
    'CNN_Predictions': predictions.flatten()
})
predictions_csv_path = os.path.join(output_dir, 'cnn_predictions.csv')
predictions_df.to_csv(predictions_csv_path, index=False)

# Save test labels to CSV for alignment
original_test_labels_df = pd.DataFrame({
    'Time_Index': range(len(test_labels)),
    'AMOC_40N': test_labels.values
})
original_test_labels_path = os.path.join(output_dir, 'original_test_labels.csv')
original_test_labels_df.to_csv(original_test_labels_path, index=False)

# Save performance metrics to CSV
metrics_df = pd.DataFrame({
    'Metric': ['RMSE', 'Custom_Performance_Metric', '5_Sigma_Gradient_Penalty', '2_Sigma_Gradient_Penalty'],
    'Value': [
        np.sqrt(mean_squared_error(test_labels, predictions)),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[0].numpy(),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[2].numpy(),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[3].numpy()
    ]
})
metrics_csv_path = os.path.join(output_dir, 'cnn_metrics.csv')
metrics_df.to_csv(metrics_csv_path, index=False)

print(f"Predictions saved to: {predictions_csv_path}")
print(f"Original test labels saved to: {original_test_labels_path}")
print(f"Metrics saved to: {metrics_csv_path}")

# RNN
import os
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import SimpleRNN, Input, Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.optimizers import Adam
from pathlib import Path
import tensorflow.keras.backend as K

# Set global font size for plots and float precision for compatibility
plt.rcParams.update({'font.size': 20})
tf.keras.backend.set_floatx('float64')

# Define directories
input_dir = "./"
output_dir = "./output"
Path(output_dir).mkdir(parents=True, exist_ok=True)
input_path = f"{input_dir}/"
combined_dataset_path = os.path.join(input_path, "CombinedDatasetLongUnsmooth.csv")

# Window size and forecast horizon parameters
window_size = 50
forecast_horizon = 20

# Moving average kernel and function
def get_moving_average_kernel(window_size):
    kernel = tf.ones((window_size, 1, 1), dtype=tf.float64) / window_size
    return kernel

def moving_average(x, kernel):
    smoothed_x = tf.nn.conv1d(tf.expand_dims(x, axis=0), kernel, stride=1, padding='SAME')
    return tf.squeeze(smoothed_x, axis=0)

# Custom loss with smoothing
def custom_loss_with_smoothing_components(y_true, y_pred, gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10):
    y_true = K.cast(y_true, dtype=tf.float64)
    y_pred = K.cast(y_pred, dtype=tf.float64)
    kernel = get_moving_average_kernel(window_size)
    y_true_smooth = moving_average(K.reshape(y_true, (-1, 1)), kernel)
    y_pred_smooth = moving_average(K.reshape(y_pred, (-1, 1)), kernel)
    rmse = K.sqrt(K.mean(K.square(y_pred_smooth - y_true_smooth)))
    gradient_true = y_true_smooth[1:] - y_true_smooth[:-1]
    gradient_pred = y_pred_smooth[1:] - y_pred_smooth[:-1]
    gradient_true_smooth = moving_average(gradient_true, kernel)
    gradient_pred_smooth = moving_average(gradient_pred, kernel)
    threshold_5sigma = K.std(gradient_true_smooth) * 5
    threshold_2sigma = K.std(gradient_true_smooth) * 2
    abrupt_change_mask_5sigma = K.cast(K.greater(K.abs(gradient_true_smooth), threshold_5sigma), dtype=tf.float64)
    abrupt_change_mask_2sigma = K.cast(K.greater(K.abs(gradient_true_smooth), threshold_2sigma), dtype=tf.float64)
    gradient_diff_5sigma = K.mean(K.square(gradient_pred_smooth - gradient_true_smooth) * (1 + heavy_penalty * abrupt_change_mask_5sigma))
    gradient_diff_2sigma = K.mean(K.square(gradient_pred_smooth - gradient_true_smooth) * (1 + two_sigma_penalty * abrupt_change_mask_2sigma))
    total_loss = rmse + gradient_diff_5sigma + gradient_diff_2sigma
    return total_loss, rmse, gradient_diff_5sigma, gradient_diff_2sigma

# Load dataset
dataset = pd.read_csv(combined_dataset_path).drop(columns=['AMOC_20S', 'SALT_500m', 'ICEFRAC', 'AMOC'])

# Sliding window class
class WindowSlider:
    def __init__(self, window_size, forecast_horizon):
        self.window_size = window_size
        self.forecast_horizon = forecast_horizon

    def collect_windows(self, features, labels):
        X, y = [], []
        for i in range(self.window_size, len(features) - self.forecast_horizon + 1):
            X.append(features.iloc[i - self.window_size:i].values.flatten())
            y.append(labels.iloc[i + self.forecast_horizon - 1])
        return pd.DataFrame(X), pd.Series(y)

# Prepare data
slider = WindowSlider(window_size, forecast_horizon)
train_features, test_features = train_test_split(dataset.drop(columns=['AMOC_40N']), test_size=0.3, shuffle=False)
train_labels, test_labels = train_test_split(dataset['AMOC_40N'], test_size=0.3, shuffle=False)
train_windows, train_labels = slider.collect_windows(train_features, train_labels)
test_windows, test_labels = slider.collect_windows(test_features, test_labels)

# Normalize data
scaler_X, scaler_y = MinMaxScaler(), MinMaxScaler()
train_windows_scaled = scaler_X.fit_transform(train_windows)
test_windows_scaled = scaler_X.transform(test_windows)
train_labels_scaled = scaler_y.fit_transform(train_labels.values.reshape(-1, 1))
test_labels_scaled = scaler_y.transform(test_labels.values.reshape(-1, 1))
train_windows_3d = train_windows_scaled.reshape((-1, window_size, train_windows_scaled.shape[1] // window_size))
test_windows_3d = test_windows_scaled.reshape((-1, window_size, test_windows_scaled.shape[1] // window_size))

# Build and train RNN model
def build_and_train_rnn(input_shape, X_train, y_train):
    model = Sequential([
        Input(shape=input_shape),
        SimpleRNN(50, activation='relu', return_sequences=True),
        Dropout(0.3),
        SimpleRNN(25, activation='relu'),
        Dense(1)
    ])
    model.compile(optimizer=Adam(learning_rate=0.001), loss=lambda y_true, y_pred: custom_loss_with_smoothing_components(y_true, y_pred)[0])
    model.fit(X_train, y_train, epochs=100, batch_size=16, validation_split=0.2, verbose=1,
              callbacks=[EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True),
                         ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=3, min_lr=0.0001)])
    return model

# Train RNN
rnn_model = build_and_train_rnn((window_size, train_windows_3d.shape[2]), train_windows_3d, train_labels_scaled)
predictions_scaled = rnn_model.predict(test_windows_3d)
predictions = scaler_y.inverse_transform(predictions_scaled)

# Save predictions to CSV
predictions_df = pd.DataFrame({
    'Time_Index': range(len(predictions)),
    'RNN_Predictions': predictions.flatten()
})
predictions_csv_path = os.path.join(output_dir, 'rnn_predictions.csv')
predictions_df.to_csv(predictions_csv_path, index=False)

# Save test labels to CSV for alignment
original_test_labels_df = pd.DataFrame({
    'Time_Index': range(len(test_labels)),
    'AMOC_40N': test_labels.values
})
original_test_labels_path = os.path.join(output_dir, 'original_test_labels.csv')
original_test_labels_df.to_csv(original_test_labels_path, index=False)

# Save performance metrics to CSV
metrics_df = pd.DataFrame({
    'Metric': ['RMSE', 'Custom_Performance_Metric', '5_Sigma_Gradient_Penalty', '2_Sigma_Gradient_Penalty'],
    'Value': [
        np.sqrt(mean_squared_error(test_labels, predictions)),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[0].numpy(),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[2].numpy(),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[3].numpy()
    ]
})
metrics_csv_path = os.path.join(output_dir, 'rnn_metrics.csv')
metrics_df.to_csv(metrics_csv_path, index=False)

print(f"Predictions saved to: {predictions_csv_path}")
print(f"Original test labels saved to: {original_test_labels_path}")
print(f"Metrics saved to: {metrics_csv_path}")

# LSTM
import os
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Input, Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.optimizers import Adam
from pathlib import Path
import tensorflow.keras.backend as K

# Set global font size for plots and float precision for compatibility
plt.rcParams.update({'font.size': 20})
tf.keras.backend.set_floatx('float64')

# Define directories
input_dir = "./"  # Define input directory
output_dir = "./output"  # Define output directory
Path(output_dir).mkdir(parents=True, exist_ok=True)
input_path = f"{input_dir}/"
combined_dataset_path = os.path.join(input_path, "CombinedDatasetLongUnsmooth.csv")

# Window size and forecast horizon parameters
window_size = 50
forecast_horizon = 20

# Moving average kernel and function
def get_moving_average_kernel(window_size):
    kernel = tf.ones((window_size, 1, 1), dtype=tf.float64) / window_size
    return kernel

def moving_average(x, kernel):
    smoothed_x = tf.nn.conv1d(tf.expand_dims(x, axis=0), kernel, stride=1, padding='SAME')
    return tf.squeeze(smoothed_x, axis=0)

# Define custom loss function with smoothing
def custom_loss_with_smoothing_components(y_true, y_pred, gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10):
    y_true = K.cast(y_true, dtype=tf.float64)
    y_pred = K.cast(y_pred, dtype=tf.float64)
    y_true = K.reshape(y_true, (-1, 1))
    y_pred = K.reshape(y_pred, (-1, 1))
    kernel = get_moving_average_kernel(window_size)
    y_true_smooth = moving_average(y_true, kernel)
    y_pred_smooth = moving_average(y_pred, kernel)
    rmse = K.sqrt(K.mean(K.square(y_pred_smooth - y_true_smooth)))
    gradient_true = y_true_smooth[1:] - y_true_smooth[:-1]
    gradient_pred = y_pred_smooth[1:] - y_pred_smooth[:-1]
    gradient_true_smooth = moving_average(gradient_true, kernel)
    gradient_pred_smooth = moving_average(gradient_pred, kernel)
    threshold_5sigma = K.std(gradient_true_smooth) * 5
    threshold_2sigma = K.std(gradient_true_smooth) * 2
    abrupt_change_mask_5sigma = K.cast(K.greater(K.abs(gradient_true_smooth), threshold_5sigma), dtype=tf.float64)
    abrupt_change_mask_2sigma = K.cast(K.greater(K.abs(gradient_true_smooth), threshold_2sigma), dtype=tf.float64)
    gradient_diff_5sigma = K.mean(K.square(gradient_pred_smooth - gradient_true_smooth) * (1 + heavy_penalty * abrupt_change_mask_5sigma))
    gradient_diff_2sigma = K.mean(K.square(gradient_pred_smooth - gradient_true_smooth) * (1 + two_sigma_penalty * abrupt_change_mask_2sigma))
    total_loss = rmse + gradient_diff_5sigma + gradient_diff_2sigma
    return total_loss, rmse, gradient_diff_5sigma, gradient_diff_2sigma

# Define custom EDR function
def custom_edr(ts1, ts2, epsilon=0.1):
    len_ts1, len_ts2 = len(ts1), len(ts2)
    D = np.zeros((len_ts1 + 1, len_ts2 + 1))
    for i in range(1, len_ts1 + 1):
        D[i, 0] = i
    for j in range(1, len_ts2 + 1):
        D[0, j] = j
    for i in range(1, len_ts1 + 1):
        for j in range(1, len_ts2 + 1):
            cost = 0 if abs(ts1[i - 1] - ts2[j - 1]) <= epsilon else 1
            D[i, j] = min(D[i - 1, j] + 1, D[i, j - 1] + 1, D[i - 1, j - 1] + cost)
    return D[len_ts1, len_ts2]

# Load dataset
dataset = pd.read_csv(combined_dataset_path).drop(columns=['AMOC_20S', 'SALT_500m', 'ICEFRAC', 'AMOC'])

# Sliding window class
class WindowSlider:
    def __init__(self, window_size, forecast_horizon):
        self.window_size = window_size
        self.forecast_horizon = forecast_horizon

    def collect_windows(self, features, labels):
        X, y = [], []
        for i in range(self.window_size, len(features) - self.forecast_horizon + 1):
            X.append(features.iloc[i - self.window_size:i].values.flatten())
            y.append(labels.iloc[i + self.forecast_horizon - 1])
        return pd.DataFrame(X), pd.Series(y)

# Prepare data
slider = WindowSlider(window_size, forecast_horizon)
train_features, test_features = train_test_split(dataset.drop(columns=['AMOC_40N']), test_size=0.3, shuffle=False)
train_labels, test_labels = train_test_split(dataset['AMOC_40N'], test_size=0.3, shuffle=False)
train_windows, train_labels = slider.collect_windows(train_features, train_labels)
test_windows, test_labels = slider.collect_windows(test_features, test_labels)

# Normalize data
scaler_X, scaler_y = MinMaxScaler(), MinMaxScaler()
train_windows_scaled = scaler_X.fit_transform(train_windows)
test_windows_scaled = scaler_X.transform(test_windows)
train_labels_scaled = scaler_y.fit_transform(train_labels.values.reshape(-1, 1))
test_labels_scaled = scaler_y.transform(test_labels.values.reshape(-1, 1))
train_windows_3d = train_windows_scaled.reshape((-1, window_size, train_windows_scaled.shape[1] // window_size))
test_windows_3d = test_windows_scaled.reshape((-1, window_size, test_windows_scaled.shape[1] // window_size))

# Build and train LSTM model
def build_and_train_lstm(input_shape, X_train, y_train):
    model = Sequential([
        Input(shape=input_shape),
        LSTM(50, activation='relu', return_sequences=True),
        Dropout(0.3),
        LSTM(25, activation='relu'),
        Dense(1)
    ])
    model.compile(optimizer=Adam(learning_rate=0.001), loss=lambda y_true, y_pred: custom_loss_with_smoothing_components(y_true, y_pred)[0])
    model.fit(X_train, y_train, epochs=100, batch_size=16, validation_split=0.2, verbose=1,
              callbacks=[EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True),
                         ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=3, min_lr=0.0001)])
    return model

# Train LSTM
lstm_model = build_and_train_lstm((window_size, train_windows_3d.shape[2]), train_windows_3d, train_labels_scaled)
predictions_scaled = lstm_model.predict(test_windows_3d)
predictions = scaler_y.inverse_transform(predictions_scaled)

# Save predictions to CSV
predictions_df = pd.DataFrame({
    'Time_Index': range(len(predictions)),
    'LSTM_Predictions': predictions.flatten()
})
predictions_csv_path = os.path.join(output_dir, 'lstm_predictions.csv')
predictions_df.to_csv(predictions_csv_path, index=False)

# Save test labels to CSV for alignment
original_test_labels_df = pd.DataFrame({
    'Time_Index': range(len(test_labels)),
    'AMOC_40N': test_labels.values
})
original_test_labels_path = os.path.join(output_dir, 'original_test_labels.csv')
original_test_labels_df.to_csv(original_test_labels_path, index=False)

# Save performance metrics to CSV
metrics_df = pd.DataFrame({
    'Metric': ['RMSE', 'Custom_Performance_Metric', '5_Sigma_Gradient_Penalty', '2_Sigma_Gradient_Penalty'],
    'Value': [
        np.sqrt(mean_squared_error(test_labels, predictions)),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[0].numpy(),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[2].numpy(),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[3].numpy()
    ]
})
metrics_csv_path = os.path.join(output_dir, 'lstm_metrics.csv')
metrics_df.to_csv(metrics_csv_path, index=False)

print(f"Predictions saved to: {predictions_csv_path}")
print(f"Original test labels saved to: {original_test_labels_path}")
print(f"Metrics saved to: {metrics_csv_path}")

#BiLSTM
import os
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Bidirectional, LSTM, Input, Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.optimizers import Adam
from pathlib import Path
import tensorflow.keras.backend as K

# Set global font size for plots and float precision for compatibility
plt.rcParams.update({'font.size': 20})
tf.keras.backend.set_floatx('float64')

# Define directories
input_dir = "./"  # Define input directory
output_dir = "./output"  # Define output directory
Path(output_dir).mkdir(parents=True, exist_ok=True)
input_path = f"{input_dir}/"
combined_dataset_path = os.path.join(input_path, "CombinedDatasetLongUnsmooth.csv")

# Window size and forecast horizon parameters
window_size = 50
forecast_horizon = 20

# Moving average kernel and function
def get_moving_average_kernel(window_size):
    kernel = tf.ones((window_size, 1, 1), dtype=tf.float64) / window_size
    return kernel

def moving_average(x, kernel):
    smoothed_x = tf.nn.conv1d(tf.expand_dims(x, axis=0), kernel, stride=1, padding='SAME')
    return tf.squeeze(smoothed_x, axis=0)

# Define custom loss function with smoothing
def custom_loss_with_smoothing_components(y_true, y_pred, gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10):
    y_true = K.cast(y_true, dtype=tf.float64)
    y_pred = K.cast(y_pred, dtype=tf.float64)
    y_true = K.reshape(y_true, (-1, 1))
    y_pred = K.reshape(y_pred, (-1, 1))
    kernel = get_moving_average_kernel(window_size)
    y_true_smooth = moving_average(y_true, kernel)
    y_pred_smooth = moving_average(y_pred, kernel)
    rmse = K.sqrt(K.mean(K.square(y_pred_smooth - y_true_smooth)))
    gradient_true = y_true_smooth[1:] - y_true_smooth[:-1]
    gradient_pred = y_pred_smooth[1:] - y_pred_smooth[:-1]
    gradient_true_smooth = moving_average(gradient_true, kernel)
    gradient_pred_smooth = moving_average(gradient_pred, kernel)
    threshold_5sigma = K.std(gradient_true_smooth) * 5
    threshold_2sigma = K.std(gradient_true_smooth) * 2
    abrupt_change_mask_5sigma = K.cast(K.greater(K.abs(gradient_true_smooth), threshold_5sigma), dtype=tf.float64)
    abrupt_change_mask_2sigma = K.cast(K.greater(K.abs(gradient_true_smooth), threshold_2sigma), dtype=tf.float64)
    gradient_diff_5sigma = K.mean(K.square(gradient_pred_smooth - gradient_true_smooth) * (1 + heavy_penalty * abrupt_change_mask_5sigma))
    gradient_diff_2sigma = K.mean(K.square(gradient_pred_smooth - gradient_true_smooth) * (1 + two_sigma_penalty * abrupt_change_mask_2sigma))
    total_loss = rmse + gradient_diff_5sigma + gradient_diff_2sigma
    return total_loss, rmse, gradient_diff_5sigma, gradient_diff_2sigma

# Define custom EDR function
def custom_edr(ts1, ts2, epsilon=0.1):
    len_ts1, len_ts2 = len(ts1), len(ts2)
    D = np.zeros((len_ts1 + 1, len_ts2 + 1))
    for i in range(1, len_ts1 + 1):
        D[i, 0] = i
    for j in range(1, len_ts2 + 1):
        D[0, j] = j
    for i in range(1, len_ts1 + 1):
        for j in range(1, len_ts2 + 1):
            cost = 0 if abs(ts1[i - 1] - ts2[j - 1]) <= epsilon else 1
            D[i, j] = min(D[i - 1, j] + 1, D[i, j - 1] + 1, D[i - 1, j - 1] + cost)
    return D[len_ts1, len_ts2]

# Load dataset
dataset = pd.read_csv(combined_dataset_path).drop(columns=['AMOC_20S', 'SALT_500m', 'ICEFRAC', 'AMOC'])

# Sliding window class
class WindowSlider:
    def __init__(self, window_size, forecast_horizon):
        self.window_size = window_size
        self.forecast_horizon = forecast_horizon

    def collect_windows(self, features, labels):
        X, y = [], []
        for i in range(self.window_size, len(features) - self.forecast_horizon + 1):
            X.append(features.iloc[i - self.window_size:i].values.flatten())
            y.append(labels.iloc[i + self.forecast_horizon - 1])
        return pd.DataFrame(X), pd.Series(y)

# Prepare data
slider = WindowSlider(window_size, forecast_horizon)
train_features, test_features = train_test_split(dataset.drop(columns=['AMOC_40N']), test_size=0.3, shuffle=False)
train_labels, test_labels = train_test_split(dataset['AMOC_40N'], test_size=0.3, shuffle=False)
train_windows, train_labels = slider.collect_windows(train_features, train_labels)
test_windows, test_labels = slider.collect_windows(test_features, test_labels)

# Normalize data
scaler_X, scaler_y = MinMaxScaler(), MinMaxScaler()
train_windows_scaled = scaler_X.fit_transform(train_windows)
test_windows_scaled = scaler_X.transform(test_windows)
train_labels_scaled = scaler_y.fit_transform(train_labels.values.reshape(-1, 1))
test_labels_scaled = scaler_y.transform(test_labels.values.reshape(-1, 1))
train_windows_3d = train_windows_scaled.reshape((-1, window_size, train_windows_scaled.shape[1] // window_size))
test_windows_3d = test_windows_scaled.reshape((-1, window_size, test_windows_scaled.shape[1] // window_size))

# Build and train BiLSTM model
def build_and_train_bilstm(input_shape, X_train, y_train):
    model = Sequential([
        Input(shape=input_shape),
        Bidirectional(LSTM(50, activation='relu', return_sequences=True)),
        Dropout(0.3),
        Bidirectional(LSTM(25, activation='relu')),
        Dense(1)
    ])
    model.compile(optimizer=Adam(learning_rate=0.001), loss=lambda y_true, y_pred: custom_loss_with_smoothing_components(y_true, y_pred)[0])
    model.fit(X_train, y_train, epochs=100, batch_size=16, validation_split=0.2, verbose=1,
              callbacks=[EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True),
                         ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=3, min_lr=0.0001)])
    return model

# Train BiLSTM
bilstm_model = build_and_train_bilstm((window_size, train_windows_3d.shape[2]), train_windows_3d, train_labels_scaled)
predictions_scaled = bilstm_model.predict(test_windows_3d)
predictions = scaler_y.inverse_transform(predictions_scaled)

# Save predictions to CSV
predictions_df = pd.DataFrame({
    'Time_Index': range(len(predictions)),
    'BiLSTM_Predictions': predictions.flatten()
})
predictions_csv_path = os.path.join(output_dir, 'bilstm_predictions.csv')
predictions_df.to_csv(predictions_csv_path, index=False)

# Save test labels to CSV for alignment
original_test_labels_df = pd.DataFrame({
    'Time_Index': range(len(test_labels)),
    'AMOC_40N': test_labels.values
})
original_test_labels_path = os.path.join(output_dir, 'original_test_labels.csv')
original_test_labels_df.to_csv(original_test_labels_path, index=False)

# Save performance metrics to CSV
metrics_df = pd.DataFrame({
    'Metric': ['RMSE', 'Custom_Performance_Metric', '5_Sigma_Gradient_Penalty', '2_Sigma_Gradient_Penalty'],
    'Value': [
        np.sqrt(mean_squared_error(test_labels, predictions)),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[0].numpy(),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[2].numpy(),
        custom_loss_with_smoothing_components(
            tf.convert_to_tensor(test_labels.values.flatten(), dtype=tf.float64),
            tf.convert_to_tensor(predictions.flatten(), dtype=tf.float64),
            gradient_penalty=300.0, heavy_penalty=9999.0, two_sigma_penalty=5000.0, window_size=10
        )[3].numpy()
    ]
})
metrics_csv_path = os.path.join(output_dir, 'bilstm_metrics.csv')
metrics_df.to_csv(metrics_csv_path, index=False)

print(f"Predictions saved to: {predictions_csv_path}")
print(f"Original test labels saved to: {original_test_labels_path}")
print(f"Metrics saved to: {metrics_csv_path}")