#| label: setup-py
#| message: false
#| warning: false

# Import core data manipulation packages
# pandas is the primary library for working with tabular data (DataFrames)
# By convention, we import pandas as 'pd' - this is universal in Python data science
import pandas as pd

# numpy provides numerical computing capabilities (arrays, math functions, etc.)
# By convention, we import numpy as 'np'
import numpy as np

# Import NFL data package
# nfl_data_py is the Python equivalent of nflfastR
# By convention, we import it as 'nfl'
import nfl_data_py as nfl

# Import visualization packages
# matplotlib.pyplot is the core plotting library (like base R graphics)
import matplotlib.pyplot as plt

# seaborn provides statistical visualization built on matplotlib
# It creates more attractive plots with less code
import seaborn as sns

# Import datetime for temporal operations
from datetime import datetime, timedelta

# Configure pandas display options for better readability
# These settings control how DataFrames are displayed in notebooks/console
pd.set_option('display.max_columns', 20)     # Show up to 20 columns before truncating
pd.set_option('display.width', 120)          # Wider display width
pd.set_option('display.precision', 3)        # Show 3 decimal places for floats

# Display confirmation
print("✓ All Python packages loaded successfully")
print("✓ Ready to wrangle football data!")

#| label: load-sample-data-py
#| message: false
#| warning: false
#| cache: true

# Load 2023 play-by-play data
# import_pbp_data() is the primary nfl_data_py function for play-by-play data
# It accepts a list of season years, e.g., [2020, 2021, 2022, 2023]
# Like nflfastR, it caches data locally for fast subsequent loads
pbp = nfl.import_pbp_data([2023])

# Load roster data for 2023 season
# Rosters include player demographics, positions, and background information
rosters = nfl.import_rosters([2023])

# Load schedule data
# Schedules provide game-level information for joining and aggregation
schedules = nfl.import_schedules([2023])

# Display summary information about loaded datasets
print("=" * 70)
print("LOADED DATA SUMMARY")
print("=" * 70)

print("\nPlay-by-Play Data:")
print(f"  - {len(pbp):,} plays (rows)")
print(f"  - {len(pbp.columns)} variables (columns)")
print(f"  - Memory size: {pbp.memory_usage(deep=True).sum() / 1024**2:.1f} MB")
print(f"  - Weeks: {pbp['week'].min()} through {pbp['week'].max()}")

print("\nRoster Data:")
print(f"  - {len(rosters):,} player-season records (rows)")
print(f"  - {len(rosters.columns)} variables (columns)")

print("\nSchedule Data:")
print(f"  - {len(schedules)} games (rows)")
print(f"  - {len(schedules.columns)} variables (columns)")

# Display sample of key variables from play-by-play data
print("\n" + "=" * 70)
print("SAMPLE PLAY-BY-PLAY DATA (First 3 Plays)")
print("=" * 70)
sample_cols = ['game_id', 'posteam', 'defteam', 'down', 'ydstogo',
               'yardline_100', 'desc', 'yards_gained', 'epa']
print(pbp[sample_cols].head(3).to_string(index=False))

#| label: select-columns-py

# =====================================
# Basic Column Selection
# =====================================

# Select specific columns by name using list indexing
# This is the most explicit approach in pandas
# Pass a list of column names inside square brackets
pbp_basic = pbp[[
    # Identifiers
    'game_id', 'play_id', 'game_date',

    # Teams
    'posteam', 'defteam',

    # Situation
    'down', 'ydstogo', 'yardline_100',

    # Outcome
    'play_type', 'yards_gained',

    # Advanced metrics
    'epa', 'wpa', 'success'
]]

print(f"Basic selection: {len(pbp_basic.columns)} columns")

# =====================================
# Pattern-Based Selection
# =====================================

# Select columns matching patterns using list comprehension
# Python doesn't have built-in pattern matching like R's dplyr,
# but list comprehensions provide flexible pattern matching

# Columns starting with 'pass'
pass_cols = [col for col in pbp.columns if col.startswith('pass')]

# Columns ending with '_epa'
epa_cols = [col for col in pbp.columns if col.endswith('_epa')]

# Columns containing 'air_yards'
air_cols = [col for col in pbp.columns if 'air_yards' in col]

# Combine all selected columns with identifiers
pbp_pass = pbp[['game_id', 'play_id'] + pass_cols + epa_cols + air_cols]

print(f"Pass-focused selection: {len(pbp_pass.columns)} columns")

# =====================================
# Selection by Position (less common)
# =====================================

# Select columns by position using iloc
# Rarely used because position-based selection is brittle
# (column order might change, breaking your code)
pbp_range = pbp.iloc[:, 0:10]  # First 10 columns

print(f"Position-based selection: {len(pbp_range.columns)} columns")

# =====================================
# Negative Selection (Dropping Columns)
# =====================================

# Drop columns using the drop() method
# axis=1 indicates we're dropping columns (not rows)
# Use this when it's easier to specify what to remove

# Drop specific columns
pbp_no_ids = pbp.drop(columns=['old_game_id'])

# Drop columns matching patterns
player_id_cols = [col for col in pbp.columns if col.endswith('_player_id')]
td_cols = [col for col in pbp.columns if col.startswith('td_')]
pbp_no_ids = pbp.drop(columns=player_id_cols + td_cols)

print(f"After dropping columns: {len(pbp_no_ids.columns)} columns")

# =====================================
# Selection with Renaming
# =====================================

# Select and rename using rename() method
# First select columns, then rename them
pbp_renamed = pbp[[
    'game_id', 'posteam', 'defteam', 'down', 'ydstogo',
    'yardline_100', 'desc', 'yards_gained', 'epa'
]].rename(columns={
    'game_id': 'game',
    'posteam': 'offense',
    'defteam': 'defense',
    'ydstogo': 'distance',
    'yardline_100': 'field_position',
    'desc': 'play_description',
    'yards_gained': 'yards',
    'epa': 'expected_points_added'
})

print(f"Selection with renaming: {len(pbp_renamed.columns)} columns")

# Display column names from each selection approach
print("\n" + "=" * 70)
print("COLUMN NAMES FROM BASIC SELECTION:")
print("=" * 70)
print(", ".join(pbp_basic.columns))

print("\n" + "=" * 70)
print("SAMPLE OF PASS-FOCUSED COLUMNS:")
print("=" * 70)
print(", ".join(list(pbp_pass.columns)[:15]))
print(f"... and {len(pbp_pass.columns) - 15} more columns")

select(game_id, posteam, epa)

pbp[['game_id', 'posteam', 'epa']]

select(pass_attempt, pass_touchdown, pass_length, pass_location, ...)

select(starts_with("pass"))

#| label: filter-basic-py

# =====================================
# Single-Condition Filtering
# =====================================

# Filter to only pass plays
# pandas provides multiple filtering syntaxes:
# 1. Boolean indexing: df[df['column'] == value]
# 2. query() method: df.query("column == 'value'")
# We'll demonstrate both

# Boolean indexing approach
pass_plays = pbp[pbp['play_type'] == 'pass']

# query() approach (often more readable)
# pass_plays = pbp.query("play_type == 'pass'")

print(f"Total plays: {len(pbp):,}")
print(f"Pass plays: {len(pass_plays):,} ({100*len(pass_plays)/len(pbp):.1f}%)")

# =====================================
# Multiple-Condition Filtering (AND logic)
# =====================================

# Filter to passes on 3rd down
# Multiple conditions with & (AND) operator
# Note: Must wrap each condition in parentheses
third_down_passes = pbp[
    (pbp['play_type'] == 'pass') &      # Must be a pass AND
    (pbp['down'] == 3)                  # Must be 3rd down
]

# Alternative using query() method (often more readable)
# third_down_passes = pbp.query("play_type == 'pass' and down == 3")

print(f"\n3rd down passes: {len(third_down_passes):,}")

# =====================================
# Complex Filtering with Multiple Operators
# =====================================

# Filter to critical situations
# Demonstrates multiple comparison operators
critical_plays = pbp[
    (pbp['down'].isin([3, 4])) &            # 3rd or 4th down
    (pbp['ydstogo'] >= 5) &                 # 5+ yards to go
    (pbp['qtr'] == 4) &                     # 4th quarter
    (pbp['score_differential'] >= -8) &     # Within 8 points (trailing)
    (pbp['score_differential'] <= 8)        # Within 8 points (leading)
]

# Alternative using query() method
# critical_plays = pbp.query(
#     "down in [3, 4] and ydstogo >= 5 and qtr == 4 and "
#     "score_differential >= -8 and score_differential <= 8"
# )

print(f"\nCritical situation plays: {len(critical_plays):,}")

# =====================================
# Filtering with OR Logic
# =====================================

# Filter to touchdowns OR explosive gains
# Use the | operator for OR logic
# Note: Must wrap each condition in parentheses
explosive_plays = pbp[
    ((pbp['touchdown'] == 1) |              # Touchdown, OR
     (pbp['yards_gained'] >= 20)) &         # 20+ yard gain
    (pbp['epa'].notna())                    # Must have valid EPA
]

# Alternative using query() method
# explosive_plays = pbp.query(
#     "(touchdown == 1 or yards_gained >= 20) and epa.notna()"
# )

print(f"\nExplosive plays: {len(explosive_plays):,}")

# =====================================
# Negative Filtering (Exclusion)
# =====================================

# Filter to clean offensive plays
# Use ~ for negation (NOT) in boolean indexing
# .notna() checks for non-missing values (opposite of .isna())
clean_plays = pbp[
    (pbp['epa'].notna()) &                  # EPA must not be missing
    (pbp['play_type'].isin(['pass', 'run'])) # Must be pass or run
]

# Alternative using query() method
# clean_plays = pbp.query(
#     "epa.notna() and play_type in ['pass', 'run']"
# )

print(f"\nClean offensive plays: {len(clean_plays):,}")

# =====================================
# Filtering by Field Position Zones
# =====================================

# Filter to red zone plays
red_zone = pbp[
    (pbp['yardline_100'] <= 20) &           # 20 or fewer yards to goal
    (pbp['yardline_100'] > 0) &             # Not in end zone
    (pbp['play_type'].isin(['pass', 'run'])) &
    (pbp['epa'].notna())
]

print(f"\nRed zone plays: {len(red_zone):,}")

# =====================================
# Filtering by Team
# =====================================

# Filter to Kansas City Chiefs offensive plays
chiefs_offense = pbp[
    (pbp['posteam'] == 'KC') &
    (pbp['play_type'].isin(['pass', 'run'])) &
    (pbp['epa'].notna())
]

print(f"\nKC offensive plays: {len(chiefs_offense):,}")

# =====================================
# Demonstrate Filter Logic with Summary
# =====================================

print("\n" + "=" * 70)
print("FILTERING IMPACT SUMMARY")
print("=" * 70)

filter_summary = pd.DataFrame({
    'Filter': [
        'All plays',
        'Pass plays only',
        '3rd down passes',
        'Critical situations',
        'Explosive plays',
        'Clean offensive plays',
        'Red zone plays',
        'KC offensive plays'
    ],
    'Rows': [
        len(pbp),
        len(pass_plays),
        len(third_down_passes),
        len(critical_plays),
        len(explosive_plays),
        len(clean_plays),
        len(red_zone),
        len(chiefs_offense)
    ]
})

filter_summary['Percent of Total'] = (
    filter_summary['Rows'] / len(pbp) * 100
)
filter_summary['Rows Kept'] = filter_summary['Rows'].apply(
    lambda x: f"{x:,}"
)
filter_summary['% Kept'] = filter_summary['Percent of Total'].apply(
    lambda x: f"{x:.1f}%"
)

print(filter_summary[['Filter', 'Rows Kept', '% Kept']].to_string(index=False))

pbp %>% filter(down == 3)

# R: Comma means AND
filter(down == 3, ydstogo >= 5)

# Python: & operator means AND
pbp[(pbp['down'] == 3) & (pbp['ydstogo'] >= 5)]

# R: | operator means OR
filter(down == 3 | down == 4)

# Python: | operator means OR
pbp[(pbp['down'] == 3) | (pbp['down'] == 4)]

# WRONG: This creates a syntax error
filter(down == 3, down == 4)  # A play can't be both!

# CORRECT:
filter(down == 3 | down == 4)  # OR logic

# WRONG: Operator precedence causes unexpected behavior
pbp[pbp['down'] == 3 & pbp['ydstogo'] >= 5]  # Precedence issue!

# CORRECT: Parentheses ensure proper order
pbp[(pbp['down'] == 3) & (pbp['ydstogo'] >= 5)]

# This DOESN'T work as you might expect:
filter(epa == NA)  # Returns zero rows!

# This is correct:
filter(is.na(epa))  # Returns rows where epa is missing

# To keep non-missing:
filter(!is.na(epa))  # Returns rows where epa is NOT missing

# R: Efficient way to test multiple values
filter(down %in% c(3, 4))

# Python: .isin() method
pbp[pbp['down'].isin([3, 4])]

# This is equivalent to but cleaner than:
# R: filter(down == 3 | down == 4)
# Python: pbp[(pbp['down'] == 3) | (pbp['down'] == 4)]

# Option 1: Two conditions with AND
filter(yardline_100 >= 20, yardline_100 <= 50)

# Option 2: between() helper function (R only)
filter(between(yardline_100, 20, 50))

# Verify filter
filtered_data <- pbp %>% filter(down == 3)

# Check that all rows have down == 3
all(filtered_data$down == 3)  # Should be TRUE

# Check range of other variables makes sense
summary(filtered_data$ydstogo)

# Boolean indexing (traditional approach)
result = pbp[
    (pbp['down'] == 3) &
    (pbp['ydstogo'] >= 5) &
    (pbp['play_type'] == 'pass')
]

# query() method (often more readable)
result = pbp.query("down == 3 and ydstogo >= 5 and play_type == 'pass'")