# Install dependencies
pip install pandas duckdb pyarrow psutil

# Download NYC taxi data in Parquet format
# Source: https://www.nyc.gov/site/tlc/about/tlc-trip-record-data.page

import pandas as pd
import duckdb
import time
import psutil
import os

def get_memory_usage():
    """Returns current process RSS memory in MB"""
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / 1024 / 1024

DATA_PATH = "nyc_taxi_2024.parquet"  # ~10GB

start_time = time.time()
mem_before = get_memory_usage()

df = pd.read_parquet(DATA_PATH)

mem_after = get_memory_usage()
load_time = time.time() - start_time

print(f"Pandas load time: {load_time:.2f}s")
print(f"Pandas memory: {mem_after - mem_before:.0f} MB")
print(f"DataFrame shape: {df.shape}")

start_time = time.time()
mem_before = get_memory_usage()

con = duckdb.connect()
con.execute(f"CREATE VIEW taxi AS SELECT * FROM '{DATA_PATH}'")

mem_after = get_memory_usage()
load_time = time.time() - start_time

print(f"DuckDB load time: {load_time:.2f}s")
print(f"DuckDB memory: {mem_after - mem_before:.0f} MB")

start_time = time.time()
mem_before = get_memory_usage()

result = (df.groupby(df['tpep_pickup_datetime'].dt.month)
            .agg({'total_amount': 'mean',
                  'trip_distance': 'mean',
                  'passenger_count': 'mean'})
            .reset_index())

mem_after = get_memory_usage()
query_time = time.time() - start_time

print(f"Pandas aggregation: {query_time:.2f}s")
print(f"Pandas peak memory: {mem_after - mem_before:.0f} MB")

start_time = time.time()
mem_before = get_memory_usage()

result = con.execute("""
    SELECT 
        month(tpep_pickup_datetime) AS month,
        AVG(total_amount) AS avg_fare,
        AVG(trip_distance) AS avg_distance,
        AVG(passenger_count) AS avg_passengers
    FROM taxi
    GROUP BY month
    ORDER BY month
""").fetchdf()

mem_after = get_memory_usage()
query_time = time.time() - start_time

print(f"DuckDB aggregation: {query_time:.2f}s")
print(f"DuckDB peak memory: {mem_after - mem_before:.0f} MB")

start_time = time.time()
mem_before = get_memory_usage()

df['pickup_hour'] = df['tpep_pickup_datetime'].dt.hour
df['is_rush'] = df['pickup_hour'].apply(
    lambda h: (7 <= h <= 9) or (17 <= h <= 19)
)

rush_data = df[df['is_rush']]
result = (rush_data.groupby(['PULocationID', 'pickup_hour'])
            .size()
            .reset_index(name='trip_count')
            .sort_values('trip_count', ascending=False)
            .head(20))

mem_after = get_memory_usage()
query_time = time.time() - start_time

print(f"Pandas complex query: {query_time:.2f}s")
print(f"Pandas peak memory: {mem_after - mem_before:.0f} MB")

start_time = time.time()
mem_before = get_memory_usage()

result = con.execute("""
    SELECT 
        PULocationID,
        EXTRACT(hour FROM tpep_pickup_datetime) AS pickup_hour,
        COUNT(*) AS trip_count
    FROM taxi
    WHERE EXTRACT(hour FROM tpep_pickup_datetime) BETWEEN 7 AND 9
       OR EXTRACT(hour FROM tpep_pickup_datetime) BETWEEN 17 AND 19
    GROUP BY PULocationID, pickup_hour
    ORDER BY trip_count DESC
    LIMIT 20
""").fetchdf()

mem_after = get_memory_usage()
query_time = time.time() - start_time

print(f"DuckDB complex query: {query_time:.2f}s")
print(f"DuckDB peak memory: {mem_after - mem_before:.0f} MB")

# Create zone dimension table
zones_df = pd.DataFrame({
    'LocationID': range(1, 266),
    'Borough': ['Manhattan', 'Brooklyn', 'Queens', 'Bronx', 'Staten Island'] * 53,
    'Zone': [f'Zone_{i}' for i in range(1, 266)]
})

start_time = time.time()
mem_before = get_memory_usage()

result = (df.merge(zones_df, left_on='PULocationID', right_on='LocationID')
            .groupby('Borough')
            .agg({'total_amount': 'sum', 'trip_distance': 'sum'})
            .reset_index())

mem_after = get_memory_usage()
query_time = time.time() - start_time

print(f"Pandas JOIN: {query_time:.2f}s")
print(f"Pandas peak memory: {mem_after - mem_before:.0f} MB")

start_time = time.time()
mem_before = get_memory_usage()

con.register('zones', zones_df)

result = con.execute("""
    SELECT 
        z.Borough,
        SUM(t.total_amount) AS total_revenue,
        SUM(t.trip_distance) AS total_distance
    FROM taxi t
    JOIN zones z ON t.PULocationID = z.LocationID
    GROUP BY z.Borough
    ORDER BY total_revenue DESC
""").fetchdf()

mem_after = get_memory_usage()
query_time = time.time() - start_time

print(f"DuckDB JOIN: {query_time:.2f}s")
print(f"DuckDB peak memory: {mem_after - mem_before:.0f} MB")

import duckdb
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler

# 1. DuckDB handles heavy lifting (loading & aggregation)
con = duckdb.connect()
con.execute("CREATE VIEW taxi AS SELECT * FROM 'nyc_taxi_2024.parquet'")

# 2. DuckDB runs complex query, returns small result as DataFrame
df_result = con.execute("""
    SELECT 
        PULocationID,
        COUNT(*) AS trip_count,
        AVG(total_amount) AS avg_fare,
        SUM(total_amount) AS total_revenue
    FROM taxi
    WHERE total_amount > 0
    GROUP BY PULocationID
    HAVING COUNT(*) > 1000
    ORDER BY total_revenue DESC
    LIMIT 50
""").fetchdf()

# 3. Pandas handles visualization
plt.figure(figsize=(12, 6))
sns.barplot(data=df_result, x='PULocationID', y='total_revenue')
plt.title('Top 50 Pickup Locations by Revenue')
plt.tight_layout()
plt.show()

# 4. Pandas for ML preprocessing
features = df_result[['trip_count', 'avg_fare']]
scaled = StandardScaler().fit_transform(features)

# benchmark.py - DuckDB vs Pandas Full Benchmark
import pandas as pd
import duckdb
import time
import psutil
import os

DATA_PATH = "nyc_taxi_2024.parquet"

def get_memory():
    return psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024

def benchmark_pandas():
    mem_before = get_memory()
    t0 = time.time()
    df = pd.read_parquet(DATA_PATH)
    t1 = time.time()
    mem_after = get_memory()
    print(f"Pandas load: {t1-t0:.2f}s, memory: {mem_after-mem_before:.0f}MB")
    
    t2 = time.time()
    result = df.groupby(df['tpep_pickup_datetime'].dt.month)['total_amount'].mean()
    t3 = time.time()
    print(f"Pandas agg: {t3-t2:.2f}s")
    
    return df

def benchmark_duckdb():
    mem_before = get_memory()
    t0 = time.time()
    con = duckdb.connect()
    con.execute(f"CREATE VIEW taxi AS SELECT * FROM '{DATA_PATH}'")
    t1 = time.time()
    mem_after = get_memory()
    print(f"DuckDB load: {t1-t0:.2f}s, memory: {mem_after-mem_before:.0f}MB")
    
    t2 = time.time()
    result = con.execute("""
        SELECT month(tpep_pickup_datetime) AS m, AVG(total_amount)
        FROM taxi GROUP BY m ORDER BY m
    """).fetchdf()
    t3 = time.time()
    print(f"DuckDB agg: {t3-t2:.2f}s")
    
    return con

if __name__ == "__main__":
    print("=== Pandas Benchmark ===")
    df = benchmark_pandas()
    print("\n=== DuckDB Benchmark ===")
    con = benchmark_duckdb()