import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import matplotlib.pyplot as plt

class TCNBlock(layers.Layer):
    """
    A single TCN block consisting of two dilated causal convolutional layers
    with weight normalization, activation, dropout, and a residual connection.
    """
    def __init__(self, filters, kernel_size, dilation_rate, dropout_rate, return_sequences=False, **kwargs):
        super(TCNBlock, self).__init__(**kwargs)
        self.filters = filters
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.dropout_rate = dropout_rate
        self.return_sequences = return_sequences

        # First convolutional block
        self.conv1 = layers.Conv1D(
            filters=filters,
            kernel_size=kernel_size,
            dilation_rate=dilation_rate,
            padding='causal' # Ensures causal convolutions
        )
        self.batchnorm1 = layers.BatchNormalization()
        self.activation1 = layers.ReLU()
        self.dropout1 = layers.Dropout(dropout_rate)

        # Second convolutional block
        self.conv2 = layers.Conv1D(
            filters=filters,
            kernel_size=kernel_size,
            dilation_rate=dilation_rate,
            padding='causal' # Ensures causal convolutions
        )
        self.batchnorm2 = layers.BatchNormalization()
        self.activation2 = layers.ReLU()
        self.dropout2 = layers.Dropout(dropout_rate)

        # Shortcut connection (for residual block)
        self.downsample = None

    def build(self, input_shape):
        # If input and output filters don't match, we need a 1x1 convolution for the shortcut
        if input_shape[-1] != self.filters:
            self.downsample = layers.Conv1D(self.filters, kernel_size=1, padding='same')
        super(TCNBlock, self).build(input_shape)

    def call(self, inputs):
        # Store the original input for the residual connection
        x = inputs

        # Apply first convolutional block
        x = self.conv1(x)
        x = self.batchnorm1(x)
        x = self.activation1(x)
        x = self.dropout1(x)

        # Apply second convolutional block
        x = self.conv2(x)
        x = self.batchnorm2(x)
        x = self.activation2(x)
        x = self.dropout2(x)

        # Add residual connection
        if self.downsample is not None:
            res = self.downsample(inputs)
        else:
            res = inputs

        # Add the shortcut to the output of the convolutional blocks
        x = layers.add([x, res])
        x = layers.ReLU()(x) # Apply ReLU after addition for consistency with ResNet
        return x

    def get_config(self):
        config = super(TCNBlock, self).get_config()
        config.update({
            'filters': self.filters,
            'kernel_size': self.kernel_size,
            'dilation_rate': self.dilation_rate,
            'dropout_rate': self.dropout_rate,
            'return_sequences': self.return_sequences,
        })
        return config

def build_tcn_model(input_shape, output_units, num_blocks, filters, kernel_size, dropout_rate):
    """
    Builds a TCN model using the TCNBlock custom layer.

    Args:
        input_shape (tuple): Shape of the input sequence (e.g., (sequence_length, num_features)).
        output_units (int): Number of output units for the final prediction layer.
        num_blocks (int): Number of TCN blocks to stack.
        filters (int): Number of filters in each convolutional layer within a TCN block.
        kernel_size (int): Size of the convolutional kernel.
        dropout_rate (float): Dropout rate for regularization.

    Returns:
        keras.Model: A compiled Keras TCN model.
    """
    inputs = keras.Input(shape=input_shape)
    x = inputs

    for i in range(num_blocks):
        # Dilation rates typically grow exponentially
        dilation_rate = 2 ** i
        x = TCNBlock(
            filters=filters,
            kernel_size=kernel_size,
            dilation_rate=dilation_rate,
            dropout_rate=dropout_rate,
            name=f'tcn_block_{i+1}'
        )(x)

    # The output of the last TCNBlock will be a sequence. We need to decide how to process it.
    # For a many-to-one prediction (e.g., predicting next value), we might take the last timestep.
    # For many-to-many (e.g., sequence labeling), we keep the full sequence.
    # For this example, we'll assume many-to-one for simplicity, taking the last timestep.
    # If return_sequences=True in the last block, x would still be 3D. We take the last timestep.
    x = layers.Lambda(lambda x: x[:, -1, :])(x) # Take the output from the last timestep

    outputs = layers.Dense(output_units)(x)

    model = keras.Model(inputs, outputs)
    return model

# Example Usage:
# Define parameters
sequence_length = 50
num_features = 1
output_units = 1 # For predicting the next single value
num_blocks = 4
filters = 32
kernel_size = 3
dropout_rate = 0.2

# Build the model
model = build_tcn_model(
    input_shape=(sequence_length, num_features),
    output_units=output_units,
    num_blocks=num_blocks,
    filters=filters,
    kernel_size=kernel_size,
    dropout_rate=dropout_rate
)

# Compile the model
model.compile(optimizer='adam', loss='mse')

# Print model summary to see the architecture
model.summary()

# Function to generate a simple time series
def generate_time_series(num_points, freq=0.05, noise_level=0.1):
    """
    Generates a synthetic sine wave time series with added noise.

    Args:
        num_points (int): The total number of data points to generate.
        freq (float): Frequency of the sine wave.
        noise_level (float): Standard deviation of the Gaussian noise.

    Returns:
        np.array: A 1D NumPy array representing the time series.
    """
    time = np.linspace(0, 100, num_points)
    series = np.sin(time * freq) + np.random.normal(scale=noise_level, size=num_points)
    return series

# Generate 1000 data points
num_data_points = 1000
synthetic_series = generate_time_series(num_data_points)

# Visualize the generated series
plt.figure(figsize=(12, 6))
plt.plot(synthetic_series)
plt.title('Synthetic Time Series (Sine Wave with Noise)')
plt.xlabel('Time Step')
plt.ylabel('Value')
plt.grid(True)
plt.show()

print(f"Generated time series with {len(synthetic_series)} points.")
print(f"First 10 points: {synthetic_series[:10]}")

# Function to create sequences from a time series
def create_sequences(data, sequence_length):
    """
    Transforms a 1D time series into sequences for supervised learning.

    Args:
        data (np.array): The 1D input time series.
        sequence_length (int): The length of the input sequence (number of past time steps).

    Returns:
        tuple: A tuple (X, y) where X is the input sequences and y is the target values.
               X will have shape (num_samples, sequence_length, 1) and y will have shape (num_samples, 1).
    """
    X, y = [], []
    for i in range(len(data) - sequence_length):
        X.append(data[i:(i + sequence_length)])
        y.append(data[i + sequence_length])
    return np.array(X), np.array(y)

# Prepare the data
X, y = create_sequences(synthetic_series, sequence_length)

# Reshape X for the TCN model (num_samples, sequence_length, num_features)
X = X.reshape(-1, sequence_length, num_features)
# Reshape y for the output layer (num_samples, output_units)
y = y.reshape(-1, output_units)

print(f"Original series length: {len(synthetic_series)}")
print(f"Input sequences shape (X): {X.shape}")
print(f"Target values shape (y): {y.shape}")
print(f"First input sequence:\n{X[0].flatten()}")
print(f"First target value: {y[0].flatten()}")

# Split data into training and testing sets
split_ratio = 0.8
split_index = int(len(X) * split_ratio)

X_train, X_test = X[:split_index], X[split_index:]
y_train, y_test = y[:split_index], y[split_index:]

print(f"\nTraining data shape (X_train): {X_train.shape}")
print(f"Training target shape (y_train): {y_train.shape}")
print(f"Test data shape (X_test): {X_test.shape}")
print(f"Test target shape (y_test): {y_test.shape}")

# Train the model
history = model.fit(
    X_train,
    y_train,
    epochs=20, # Number of training epochs
    batch_size=32, # Batch size for training
    validation_split=0.1, # Use 10% of training data for validation
    verbose=1
)

# Plot training & validation loss values
plt.figure(figsize=(12, 6))
plt.plot(history.history['loss'], label='Train Loss')
plt.plot(history.history['val_loss'], label='Validation Loss')
plt.title('Model Loss')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(loc='upper right')
plt.grid(True)
plt.show()

print("Model training complete.")

# Evaluate the model on the test data
loss = model.evaluate(X_test, y_test, verbose=0)
print(f'Test Loss (MSE): {loss:.4f}')

# Make predictions on the test set
y_pred = model.predict(X_test)

# Visualize actual vs. predicted values
plt.figure(figsize=(15, 7))
plt.plot(y_test, label='Actual Values')
plt.plot(y_pred, label='Predicted Values')
plt.title('TCN Model: Actual vs. Predicted Time Series Values')
plt.xlabel('Time Step (in test set)')
plt.ylabel('Value')
plt.legend()
plt.grid(True)
plt.show()

# Optionally, visualize a smaller portion for better detail
plt.figure(figsize=(15, 7))
plt.plot(y_test[:100], label='Actual Values (First 100)')
plt.plot(y_pred[:100], label='Predicted Values (First 100)')
plt.title('TCN Model: Actual vs. Predicted (First 100 Test Samples)')
plt.xlabel('Time Step (in test set)')
plt.ylabel('Value')
plt.legend()
plt.grid(True)
plt.show()