import numpy as np
import tensorflow as tf
import pandas as pd
import matplotlib.pyplot as plt
import sklearn


class TimeSeries:

    # RNN setup
    rnn_units = 200
    epoch = 20
    dropout = 0.2

    def __init__(self):
        self.batch_size = 1000
        self.n_steps = 5


    def generate_time_series(self):
        freq1, freq2, offset1, offset2 = np.random.rand(4, self.batch_size, 1)
        time = np.linspace(0, 1,self.n_steps)
        series = 0.5 * np.sin((time - offset1)) * (freq1 * 10 + 10)  # Welle 1
        series += 0.2 * np.sin((time - offset2)) * (freq2 * 20 + 20)  # Welle 2
        series += 0.1 * (np.random.rand(self.batch_size, self.n_steps) - 0.5)  # Rauschen
        return series[..., np.newaxis].astype(np.float32)

    def training_data(self):

        frac_train = int(0.7*self.batch_size)
        frac_validation = int(0.2*self.batch_size)
        frac_test = int(0.1*self.batch_size)

        print("frac_train,frac_validation, frac_test",frac_train,frac_validation, frac_test)

        # The sequence of the data should be in the correct time order. No randomization!
        # Data sequences for training, validation, testing
        seq_train = frac_train
        seq_validation = seq_train+frac_validation
        seq_test = seq_validation + frac_test
        print("seq_train,seq_validation, seq_test",seq_train,seq_validation, seq_test)

        X_train, y_train = series[:seq_train, :self.n_steps], series[:seq_train, -1]
        X_valid, y_valid = series[seq_train:seq_validation, :self.n_steps], series[seq_train:seq_validation, -1]
        X_test, y_test = series[seq_validation:, :self.n_steps], series[seq_validation:, -1]
        print("X_train.shape ",X_train.shape)
        print("y_train.shape ",y_train.shape)
        print("X_valid.shape ",X_valid.shape)
        print("y_valid.shape ",y_valid.shape)
        print("X_test.shape ",X_test.shape)
        print("y_test.shape ",y_test.shape)
        return X_train, y_train, X_valid, y_valid, X_test, y_test


    def plot_series(self,series,xrange=10):
        x = np.linspace(-xrange, xrange, self.batch_size)
        fig = plt.figure()
        plt.plot(x, series[:, 0, 0], '-')

        plt.show()





    def MyRNN(self,X_train, y_train, X_valid, y_valid, X_test,y_test,dropout_rate=0.2,rnn_units=20, epochs=20):
        # Define RNN model
        model = tf.keras.models.Sequential([
            tf.keras.layers.SimpleRNN(rnn_units, return_sequences=True, input_shape=[None, 1]),
            tf.keras.layers.SimpleRNN(rnn_units),
            tf.keras.layers.Dropout(dropout_rate),
            tf.keras.layers.Dense(1)
        ])

        # Compile the model
        model.compile(loss="MSE", optimizer="Adam", metrics=["mse"])
        model.summary()
        model.layers
        # Train the model
        history = model.fit(X_train, y_train, epochs=epochs, validation_data=(X_valid, y_valid), verbose=1)

        # Plot training history
        pd.DataFrame(history.history).plot(figsize=(8, 5))
        plt.grid(True)
        # plt.gca().set_ylim(0,1) # Optionally set vertical range from 0 to 1
        plt.show()

        # Generate predictions on the test data
        from sklearn.metrics import mean_squared_error, mean_absolute_error
        y_pred = model.predict(X_test)
        rmse = np.sqrt(mean_squared_error(y_test, y_pred))/np.mean(y_test)*100
        mae = mean_absolute_error(y_test, y_pred)

        print("RMSE:", rmse)
        print("MAE:", mae)

        # Plot the actual values and the predictions
        plt.plot(y_test, label='Test data set', linestyle='dotted')
        plt.plot(y_pred, label=f'Predicted by RNN model RMSE/mean= {rmse: .1f} %')
        plt.xlabel('Time')
        plt.ylabel('Value')
        plt.legend()
        plt.show()



        return model, history

# Example usage:
# Assuming X_train, y_train, X_valid, y_valid are already defined
# Replace them with your actual data
# Also, you can adjust other parameters like rnn_units and epochs as needed






timeseries = TimeSeries()

series = timeseries.generate_time_series()
timeseries.plot_series(series)



class_ts = TimeSeries()
series = class_ts.generate_time_series()

X_train, y_train, X_valid, y_valid, X_test, y_test = class_ts.training_data()



self=""

trained_model, training_history = TimeSeries.MyRNN(self=self,X_train=X_train, y_train=y_train,
                                                   X_valid=X_valid, y_valid=y_valid,X_test=X_test,y_test=y_test,
                                                   dropout_rate=class_ts.dropout,rnn_units=class_ts.rnn_units, epochs=class_ts.epoch)


