################## Run on LUMI 20241218 ##########################################
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, Bidirectional, 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 LUMI
plt.rcParams.update({'font.size': 20})
tf.keras.backend.set_floatx('float64')

# Define directories
# input_dir = "/groups/ocean/qifan/draft"
# output_dir = "/groups/ocean/qifan/draft"
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 = f"{input_path}/CombinedDatasetLongUnsmooth.csv"
combined_dataset_path = os.path.join(input_path, "CombinedDatasetLongUnsmooth.csv")

# Define window size and forecast horizon combinations
# window_sizes = [20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140, 150, 160, 170, 180, 190, 200]
# forecast_horizons = [0, 5, 10, 20, 50, 75, 100, 150, 200]

# Retrieve window size and forecast horizon from command-line arguments
# Command python try.py 50 20
window_size = int(sys.argv[1])
forecast_horizon = int(sys.argv[2])

# Set the number of runs for averaging
n_runs = 10

# Create a kernel for 5-year smoothing (5-point moving average)
def get_moving_average_kernel(window_size):
    kernel = tf.ones((window_size, 1, 1), dtype=tf.float64) / window_size
    return kernel

# Define the moving average function
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'])

# Define directories for saving results
results_path = os.path.join(output_dir, f"bilstm_w={window_size}_f={forecast_horizon}_results.csv")
model_save_path_template = os.path.join(output_dir, f"bilstm_model_run_{{}}_w{window_size}_f{forecast_horizon}.h5")
comparison_plot_path_template = os.path.join(output_dir, f"bilstm_comparison_zoomed_run_{{}}_w{window_size}_f{forecast_horizon}.png")
full_plot_path_template = os.path.join(output_dir, f"bilstm_full_test_comparison_run_{{}}_w{window_size}_f{forecast_horizon}.png")

# Initialize results DataFrame
results_df = pd.DataFrame(columns=['Run', 'Window_Size', 'Forecast_Horizon', 'RMSE', 'Custom_Performance_Metric', 'EDR', '5_Sigma_Gradient_Penalty', '2_Sigma_Gradient_Penalty'])

for run in range(1, n_runs + 1):
    print(f"Run {run}/{n_runs}...")
    train_data, test_data = train_test_split(dataset, test_size=0.3, shuffle=False)
    train_features, test_features = train_data.drop(columns=['AMOC_40N']), test_data.drop(columns=['AMOC_40N'])
    train_labels, test_labels = train_data['AMOC_40N'], test_data['AMOC_40N']

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

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

    window_slider = WindowSlider(window_size=window_size, forecast_horizon=forecast_horizon)
    train_windows, train_labels = window_slider.collect_windows(train_features, train_labels)
    test_windows, test_labels = window_slider.collect_windows(test_features, test_labels)

    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))

    def build_and_train_bilstm(input_shape, X_train, y_train, batch_size=16):
        model = Sequential([
            Input(shape=input_shape),
            Bidirectional(LSTM(input_shape[0], activation='relu', return_sequences=True)),
            Dropout(0.3),
            Bidirectional(LSTM(max(1, round(input_shape[0] / 2)), 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=batch_size, verbose=1, validation_split=0.2,
                  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

    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)

    rmse = np.sqrt(mean_squared_error(test_labels, predictions))
    total_loss, rmse_component, five_sigma_penalty, two_sigma_penalty = 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
    )
    edr_distance = custom_edr(test_labels.values.flatten(), predictions.flatten(), epsilon=0.1)

    results_df = pd.concat([results_df, pd.DataFrame({
        'Run': [run],
        'Window_Size': [window_size],
        'Forecast_Horizon': [forecast_horizon],
        'RMSE': [rmse],
        'Custom_Performance_Metric': [total_loss.numpy()],
        'EDR': [edr_distance],
        '5_Sigma_Gradient_Penalty': [five_sigma_penalty.numpy()],
        '2_Sigma_Gradient_Penalty': [two_sigma_penalty.numpy()]
    })], ignore_index=True)

    run_results_path = os.path.join(output_dir, f"bilstm_results_run_{run}_w{window_size}_f{forecast_horizon}.csv")
    results_df.to_csv(run_results_path, index=False)
    print(f"Results for Run {run} saved to {run_results_path}")
    print(f"CSV content for Run {run}:")
    print(results_df)

    model_save_path = model_save_path_template.format(run)
    bilstm_model.save(model_save_path)
    print(f"Model for run {run} saved to: {model_save_path}")

    full_plot_path = full_plot_path_template.format(run)
    comparison_plot_path = comparison_plot_path_template.format(run)

    plt.figure(figsize=(12, 8), dpi=600)
    plt.plot(range(len(train_labels)), train_labels, label='Training Data', color="#8FD89F")
    plt.plot(range(len(train_labels), len(train_labels) + len(test_labels)), test_labels, label='Test Data', color="#8FD89F")
    plt.plot(range(len(train_labels), len(train_labels) + len(predictions)), predictions, label='BiLSTM Predictions', color="#191970")
    plt.xlabel('Time Index')
    plt.ylabel('AMOC')
    plt.title('BiLSTM Full Results')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(full_plot_path)
    plt.close()
    print(f"Full plot for run {run} saved to: {full_plot_path}")

    zoom_start, zoom_end = 23000, 25000
    plt.figure(figsize=(12, 8), dpi=600)
    plt.plot(range(len(train_labels)), train_labels, label='Training Data', color="#8FD89F")
    plt.plot(range(len(train_labels), len(train_labels) + len(test_labels)), test_labels, label='Test Data', color="#8FD89F")
    plt.plot(range(len(train_labels), len(train_labels) + len(predictions)), predictions, label='BiLSTM Predictions', color="#191970")
    plt.xlim(zoom_start, zoom_end)
    plt.xlabel('Time Index')
    plt.ylabel('AMOC')
    plt.title(f'BiLSTM Zoomed Results (Index {zoom_start} to {zoom_end})')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(comparison_plot_path)
    plt.close()
    print(f"Zoomed plot for run {run} saved to: {comparison_plot_path}")

results_df.to_csv(results_path, index=False)
print(f"Results for Window Size: {window_size}, Forecast Horizon: {forecast_horizon} saved to {results_path}")