import numpy as np
import pandas as pd
import keras_tuner as kt

# Sample usage model for LSTM parameter optimization with Keras Tuner module

# Define the model building function for Keras Tuner
def build_model(hp):
    seq_length = hp.Int('seq_length', min_value=10, max_value=100, step=1)  # Adjust the minimum value
    steps_ahead = hp.Int('steps_ahead', min_value=1, max_value=seq_length, step=1)  # Ensure steps_ahead <= seq_length
    steps_ahead = round(min(steps_ahead, seq_length)/1.8);  # Ensure steps_ahead is not greater than seq_length
    feature_index = 0  # Index for 'log_returns'

    print(f"Hyperparameters - seq_length: {seq_length}, steps_ahead: {steps_ahead}")

    # Create sequences dynamically within the build_model function
    X, y = create_sequences(scaled_data, seq_length, steps_ahead, feature_index)
    print(f"Shapes after creating sequences - X: {X.shape}, y: {y.shape}")

    split = int(0.8 * len(X))
    X_train, X_val = X[:split], X[split:]
    y_train, y_val = y[:split], y[split:]

    input_layer = Input(shape=(seq_length, X.shape[2]))
    lstm_units = hp.Int('lstm_units', min_value=32, max_value=128, step=16)
    lstm_out = LSTM(lstm_units, return_sequences=True)(input_layer)
    lstm_out = Dropout(hp.Float('dropout_rate', min_value=0.1, max_value=0.5, step=0.1))(lstm_out)
    lstm_out = LSTM(lstm_units, return_sequences=True)(lstm_out)
    lstm_out = Dropout(hp.Float('dropout_rate', min_value=0.1, max_value=0.5, step=0.1))(lstm_out)

    # Attention mechanism
    attention = Attention()([lstm_out, lstm_out])
    attention_out = Concatenate()([lstm_out, attention])

    # Flatten the output before dense layers
    flat = Flatten()(attention_out)

    # Ensure the dense layer output size matches steps_ahead
    dense_units = hp.Int('dense_units', min_value=32, max_value=128, step=16)  # Dynamically set dense units
    dense_out = Dense(dense_units, activation='relu')(flat)

    # Reshape the output to match steps_ahead and features
    output = Dense(steps_ahead)(dense_out)
    output = Reshape((steps_ahead, 1))(output)  # Predicting 1 feature (log_returns) for steps_ahead

    model = Model(inputs=input_layer, outputs=output)

    learning_rate = hp.Float('learning_rate', min_value=1e-4, max_value=1e-2, sampling='log')
    model.compile(optimizer=Adam(learning_rate=learning_rate), loss='mean_squared_error')

    return model

# Custom Tuner class to override the run_trial method
class MyTuner(kt.tuners.BayesianOptimization):
    def run_trial(self, trial, *args, **kwargs):
        # Dynamically create the sequences for the current hyperparameters
        hp = trial.hyperparameters
        seq_length = hp.get('seq_length')
        steps_ahead = hp.get('steps_ahead')
        steps_ahead = min(steps_ahead, seq_length)  # Ensure steps_ahead is not greater than seq_length
        feature_index = 0  # Index for 'log_returns'

        X, y = create_sequences(scaled_data, seq_length, steps_ahead, feature_index)
        print(f"Shapes after creating sequences within run_trial - X: {X.shape}, y: {y.shape}")

        split = int(0.8 * len(X))
        X_train, X_val = X[:split], X[split:]
        y_train, y_val = y[:split], y[split:]

        # Print shapes to debug
        print(f"Inside run_trial - seq_length: {seq_length}, steps_ahead: {steps_ahead}")
        print(f"X_train shape: {X_train.shape}, y_train shape: {y_train.shape}")
        print(f"X_val shape: {X_val.shape}, y_val shape: {y_val.shape}")

        # Fit the model with the dynamically created data
        model = self.hypermodel.build(hp)
        kwargs['x'] = X_train
        kwargs['y'] = y_train
        kwargs['validation_data'] = (X_val, y_val)
        return super(MyTuner, self).run_trial(trial, *args, **kwargs)

# Initialize the tuner
tuning_dir = 'kt_dir'
if os.path.exists(tuning_dir):
    tuner = MyTuner(
        objective='val_loss',
        max_trials=20,
        directory=tuning_dir,
        project_name='lstm_attention_tuning',
        hypermodel=build_model
    )
    tuner.reload()
else:
    tuner = MyTuner(
        objective='val_loss',
        max_trials=20,
        directory=tuning_dir,
        project_name='lstm_attention_tuning',
        hypermodel=build_model
    )

# Perform the search
tuner.search(epochs=10, batch_size=16)

# Get the optimal hyperparameters
best_hps = tuner.get_best_hyperparameters(num_trials=1)[0]

# Retreive the optimal parameters
seq_length = best_hps.get('seq_length')
steps_ahead = np.min([best_hps.get('steps_ahead'), seq_length])