import warnings
import numpy as np
import pandas as pd
from sklearn.metrics import r2_score
from sklearn.preprocessing import MinMaxScaler
from skopt import gp_minimize
from tensorflow.keras import Sequential, layers
from tensorflow.keras.callbacks import ModelCheckpoint
import matplotlib.pyplot as plt

warnings.filterwarnings('ignore')

# 1. Data preprocessing
# 1.1 Load data
dataset = pd.read_excel(r"E:\Dataset\InputData\75-19_monthly.xlsx", sheet_name="Sheet1", usecols='B:U', header=0)

# Handle NaN values by filling with 0
dataset = dataset.fillna(0)

# Define columns: 19 rainfall columns as inputs, and Q as the target variable
columns = ["rainoff1", "rainoff2", "rainoff3", "rainoff4", "rainoff5", "rainoff6", "rainoff7", "rainoff8", "rainoff9",
           "rainoff10", "rainoff11", "rainoff12", "rainoff13", "rainoff14", "rainoff15", "rainoff16", "rainoff17",
           "rainoff18", "rainoff19", "Q"]

# Normalize the columns using MinMaxScaler
for col in columns:
    scaler = MinMaxScaler()
    dataset[col] = scaler.fit_transform(dataset[col].values.reshape(-1, 1))

# 1.3 Split dataset into training, validation, and test sets
num_samples = dataset["Q"].values.shape[0]
num_train = round(num_samples * 0.8)  # 80% of data for training
num_test = round(num_samples * 0.1)   # 10% of data for testing
num_val = num_samples - num_test - num_train  # Remaining 10% for validation
dataset_train, dataset_val, dataset_test = (dataset[:num_train], dataset[num_train:num_train + num_val],
                                            dataset[num_train:])

# 1.4 Convert time series data to supervised learning format
# Function to create feature and target datasets for supervised learning
def create_dataset(x, y, seq_len_=3, pred_horizon_=1):
    # seq_len = lookback window, the length of the input sequence (number of past observations the model needs to
    # consider) pred_horizon = prediction horizon, the number of future time steps the model needs to predict
    features = []
    targets = []
    for i in range(0, len(x) - seq_len_ - (pred_horizon_ - 1)):
        data = x.iloc[i:i + seq_len_]  # Sequence data
        label = y.iloc[i + seq_len_ + (pred_horizon_ - 1)]  # Target data
        features.append(data)
        targets.append(label)
    return np.array(features), np.array(targets)

# Set sequence length and prediction horizon
seq_len = 3
pred_horizon = 1

# Prepare the datasets for training, validation, and testing
x_train, x_val, x_test = dataset_train, dataset_val, dataset_test
y_train, y_val, y_test = dataset_train["Q"], dataset_val["Q"], dataset_test["Q"]
train_dataset, train_labels = create_dataset(x_train, y_train, seq_len, pred_horizon)
val_dataset, val_labels = create_dataset(x_val, y_val, seq_len, pred_horizon)
test_dataset, test_labels = create_dataset(x_test, y_test, seq_len, pred_horizon)

# Define the objective function for Bayesian optimization
def objective1(params):
    # Extract batch size from the optimization parameters
    batch_size = int(params[0])

    def train_and_evaluate(train_dataset, train_labels, val_dataset, val_labels, test_dataset, test_labels):
        # Define the model, with units set to a fixed value of 70 for the LSTM layers
        model = Sequential([
            layers.LSTM(units=70, input_shape=train_dataset.shape[-2:], return_sequences=True),
            layers.Dropout(0.4),  # Apply dropout to prevent overfitting
            layers.LSTM(units=70),
            layers.Dropout(0.4),  # Apply dropout again
            layers.Dense(1)  # Output layer with a single unit
        ])

        # Compile the model with Adam optimizer and mean squared error loss
        model.compile(optimizer="adam", loss="mse")

        # Define a callback for saving the best model based on loss
        checkpoint_file = "best_model.hdf5"
        checkpoint_callback = ModelCheckpoint(filepath=checkpoint_file,
                                              monitor="loss",
                                              mode="min",
                                              save_best_only=True,
                                              save_weights_only=True)

        # Train the model, using batch_size as the hyperparameter to optimize
        history = model.fit(train_dataset, train_labels, epochs=80, batch_size=batch_size, verbose=1, shuffle=False,
                            validation_data=(val_dataset, val_labels), callbacks=[checkpoint_callback])

        # Evaluate the model on the test dataset
        test_preds = model.predict(test_dataset)
        test_labels, test_preds = test_labels.reshape(-1, 1), test_preds.reshape(-1, 1)

        # Calculate NSE (Nash-Sutcliffe Efficiency)
        fenzi = np.sum((test_labels - test_preds) ** 2)
        fenmu = np.sum((test_labels - np.mean(test_labels)) ** 2)
        nse = 1 - fenzi / fenmu

        # Return negative NSE for optimization, as Bayesian optimization minimizes the objective function
        return -nse

    # Call the training and evaluation function for the current batch size
    nse = train_and_evaluate(train_dataset, train_labels, val_dataset, val_labels, test_dataset, test_labels)
    return nse

# Define the search space for batch_size
param_space1 = [(16, 128)]

# Perform Bayesian optimization to find the best batch size
opt_result1 = gp_minimize(func=objective1, dimensions=param_space1, n_calls=20, random_state=42)

# Get the best parameters and corresponding result
best_params1 = {'batch_size': int(opt_result1.x[0])}
best_nse1 = -opt_result1.fun  # Negate the result back to the actual NSE
print("Best NSE: ", best_nse1)