"""
Lazy Evaluation and Query Optimization in Polars
Demonstrates lazy evaluation, query planning, and optimization techniques

Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
"""

import polars as pl
import numpy as np
import os
from datetime import datetime, timedelta

def demonstrate_lazy_evaluation():
    """Demonstrate lazy evaluation and query optimization"""
    
    print("=" * 60)
    print("LAZY EVALUATION AND QUERY OPTIMIZATION")
    print("=" * 60)
    
    # Create a sample DataFrame
    print("\n1. Creating a sample DataFrame:")
    df = pl.DataFrame({
        'id': range(1, 10001),
        'category': np.random.choice(['A', 'B', 'C', 'D', 'E'], 10000),
        'value1': np.random.randn(10000) * 100,
        'value2': np.random.randn(10000) * 50,
        'value3': np.random.randint(1, 1000, 10000)
    })
    print(f"DataFrame shape: {df.shape}")
    
    # Convert to LazyFrame
    print("\n2. Converting to LazyFrame:")
    lazy_df = df.lazy()
    print(f"Type: {type(lazy_df)}")
    print("Operations are not executed until .collect() is called")
    
    # Build a complex query
    print("\n3. Building a complex query (not executed yet):")
    query = (lazy_df
        .filter(pl.col('value1') > 50)
        .filter(pl.col('value2') < 20)
        .select(['id', 'category', 'value1', 'value2'])
        .group_by('category')
        .agg([
            pl.col('value1').mean().alias('avg_value1'),
            pl.col('value2').mean().alias('avg_value2'),
            pl.count().alias('count')
        ])
        .sort('avg_value1', descending=True)
    )
    print("Query built successfully!")
    
    # Show query plan
    print("\n4. Query Plan (optimized):")
    print(query.explain())
    
    # Execute the query
    print("\n5. Executing the query:")
    result = query.collect()
    print("Query executed!")
    print(result)
    
    # Demonstrate query optimization benefits
    print("\n6. Query Optimization Benefits:")
    print("- Predicate pushdown: Filters are applied early")
    print("- Projection pushdown: Only needed columns are selected")
    print("- Predicate combination: Multiple filters are combined")
    print("- Join reordering: Joins are optimized for performance")
    
    # Lazy CSV reading
    print("\n7. Lazy CSV Reading (more efficient for large files):")
    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        project_root = os.path.dirname(script_dir)
        csv_path = os.path.join(project_root, 'data', 'sample_data.csv')
        lazy_from_csv = pl.scan_csv(csv_path)
        print("LazyFrame from CSV created")
        print("\nQuery plan for CSV operation:")
        csv_query = lazy_from_csv.filter(pl.col('price') > 100).select(['name', 'price'])
        print(csv_query.explain())
        print("\nExecuting CSV query:")
        csv_result = csv_query.collect()
        print(csv_result.head())
    except FileNotFoundError:
        print("Sample data file not found. Run data_generator.py first.")
    
    print("\n" + "=" * 60)
    print("Lazy evaluation demonstration complete!")
    print("=" * 60)

if __name__ == "__main__":
    demonstrate_lazy_evaluation()