import tensorflow as tf
import numpy as np
from typing import List, Tuple, Dict, Any
import matplotlib.pyplot as plt

class TensorFlowFundamentals:
    """Core TensorFlow concepts and operations"""
    
    def __init__(self):
        print(f"TensorFlow version: {tf.__version__}")
        print(f"GPU available: {tf.config.list_physical_devices('GPU')}")
        
        # Configure GPU if available
        self.setup_gpu()
    
    def setup_gpu(self):
        """Configure GPU memory growth"""
        gpus = tf.config.experimental.list_physical_devices('GPU')
        if gpus:
            try:
                for gpu in gpus:
                    tf.config.experimental.set_memory_growth(gpu, True)
                print(f"Configured {len(gpus)} GPUs with memory growth")
            except RuntimeError as e:
                print(f"GPU configuration error: {e}")
    
    def tensor_operations(self):
        """Demonstrate basic tensor operations"""
        
        # Create tensors
        scalar = tf.constant(42)
        vector = tf.constant([1, 2, 3, 4, 5])
        matrix = tf.constant([[1, 2], [3, 4], [5, 6]])
        tensor_3d = tf.random.normal([2, 3, 4])
        
        print(f"Scalar shape: {scalar.shape}, dtype: {scalar.dtype}")
        print(f"Vector shape: {vector.shape}")
        print(f"Matrix shape: {matrix.shape}")
        print(f"3D Tensor shape: {tensor_3d.shape}")
        
        # Basic operations
        x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
        y = tf.constant([[5.0, 6.0], [7.0, 8.0]])
        
        operations = {
            'addition': tf.add(x, y),
            'multiplication': tf.multiply(x, y),
            'matrix_multiply': tf.matmul(x, y),
            'element_wise_square': tf.square(x),
            'transpose': tf.transpose(x),
            'reduce_sum': tf.reduce_sum(x),
            'reduce_mean': tf.reduce_mean(x, axis=1)
        }
        
        return operations
    
    def automatic_differentiation(self):
        """Demonstrate automatic differentiation with GradientTape"""
        
        # Simple function: f(x) = x^2 + 2x + 1
        def simple_function(x):
            return x**2 + 2*x + 1
        
        x = tf.Variable(3.0)
        
        with tf.GradientTape() as tape:
            y = simple_function(x)
        
        # Compute gradient: df/dx = 2x + 2
        grad = tape.gradient(y, x)
        print(f"f({x.numpy()}) = {y.numpy()}")
        print(f"f'({x.numpy()}) = {grad.numpy()}")
        
        # Multiple variables
        x1 = tf.Variable(2.0)
        x2 = tf.Variable(3.0)
        
        with tf.GradientTape() as tape:
            z = x1**2 + x2**3 + x1*x2
        
        gradients = tape.gradient(z, [x1, x2])
        print(f"∂z/∂x1 = {gradients[0].numpy()}")
        print(f"∂z/∂x2 = {gradients[1].numpy()}")
        
        return gradients
    
    def data_pipeline_example(self):
        """Create efficient data pipeline with tf.data"""
        
        # Generate sample data
        features = np.random.randn(1000, 10).astype(np.float32)
        labels = np.random.randint(0, 2, (1000, 1)).astype(np.float32)
        
        # Create dataset
        dataset = tf.data.Dataset.from_tensor_slices((features, labels))
        
        # Apply transformations
        dataset = (dataset
                  .shuffle(buffer_size=1000)
                  .batch(32)
                  .map(lambda x, y: (tf.nn.l2_normalize(x), y))
                  .prefetch(tf.data.AUTOTUNE))
        
        # Inspect first batch
        for batch_features, batch_labels in dataset.take(1):
            print(f"Batch shape: {batch_features.shape}")
            print(f"Labels shape: {batch_labels.shape}")
            print(f"Feature range: [{tf.reduce_min(batch_features):.3f}, {tf.reduce_max(batch_features):.3f}]")
        
        return dataset
    
    def custom_layer_example(self):
        """Create custom TensorFlow layer"""
        
        class DenseWithDropout(tf.keras.layers.Layer):
            def __init__(self, units, dropout_rate=0.2, **kwargs):
                super().__init__(**kwargs)
                self.units = units
                self.dropout_rate = dropout_rate
                self.dense = tf.keras.layers.Dense(units, activation='relu')
                self.dropout = tf.keras.layers.Dropout(dropout_rate)
            
            def call(self, inputs, training=None):
                x = self.dense(inputs)
                return self.dropout(x, training=training)
            
            def get_config(self):
                config = super().get_config()
                config.update({
                    'units': self.units,
                    'dropout_rate': self.dropout_rate
                })
                return config
        
        # Use custom layer
        model = tf.keras.Sequential([
            DenseWithDropout(64, dropout_rate=0.3),
            DenseWithDropout(32, dropout_rate=0.2),
            tf.keras.layers.Dense(1, activation='sigmoid')
        ])
        
        return model

# Example usage
def explore_tensorflow_fundamentals():
    tf_fundamentals = TensorFlowFundamentals()
    
    # Tensor operations
    operations = tf_fundamentals.tensor_operations()
    print("Tensor Operations:")
    for name, result in operations.items():
        print(f"{name}: {result.numpy()}")
    
    # Automatic differentiation
    print("\nAutomatic Differentiation:")
    tf_fundamentals.automatic_differentiation()
    
    # Data pipeline
    print("\nData Pipeline:")
    dataset = tf_fundamentals.data_pipeline_example()
    
    # Custom layer
    print("\nCustom Layer Model:")
    model = tf_fundamentals.custom_layer_example()
    print(model.summary())
    
    return tf_fundamentals