#| label: load-libraries-py
#| message: false
#| warning: false

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.spatial import Voronoi, voronoi_plot_2d
from scipy.spatial.distance import cdist
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
import warnings
warnings.filterwarnings('ignore')

#| label: load-tracking-py
#| eval: false
#| echo: true

# Load tracking data (example from Big Data Bowl)
tracking = pd.read_csv("data/tracking_week_1.csv")

# Examine structure
print(tracking.info())
print("\nFirst few rows:")
print(tracking.head())

# Sample columns:
# - gameId: Game identifier
# - playId: Play identifier
# - nflId: Player identifier (NA for football)
# - frameId: Frame number (1-N)
# - time: Timestamp
# - jerseyNumber: Player jersey number
# - team: Team abbreviation or 'football'
# - playDirection: Direction of play
# - x, y: Field coordinates
# - s: Speed (yards/second)
# - a: Acceleration (yards/second^2)
# - dis: Distance traveled (yards)
# - o: Orientation (degrees)
# - dir: Direction of movement (degrees)

#| label: synthetic-tracking-py
#| message: false
#| warning: false

# Create synthetic tracking data for one play
np.random.seed(123)

def generate_trajectory(start_x, start_y, end_x, end_y, frames=50):
    """Generate player trajectory with noise"""
    x = np.linspace(start_x, end_x, frames)
    y = np.linspace(start_y, end_y, frames)

    # Add noise for realism
    y = y + np.random.normal(0, 0.5, frames)

    return pd.DataFrame({
        'frameId': range(1, frames + 1),
        'x': x,
        'y': y
    })

# Generate receiver route (go route)
receiver = generate_trajectory(75, 20, 95, 22)
receiver['nflId'] = 1001
receiver['jerseyNumber'] = 88
receiver['team'] = 'OFF'
receiver['position'] = 'WR'

# Generate cornerback coverage
cornerback = generate_trajectory(75, 18, 93, 20)
cornerback['nflId'] = 2001
cornerback['jerseyNumber'] = 25
cornerback['team'] = 'DEF'
cornerback['position'] = 'CB'

# Generate safety help
safety = generate_trajectory(85, 26.65, 95, 24)
safety['nflId'] = 2002
safety['jerseyNumber'] = 43
safety['team'] = 'DEF'
safety['position'] = 'S'

# Combine all players
tracking_example = pd.concat([receiver, cornerback, safety], ignore_index=True)
tracking_example['gameId'] = 2023091000
tracking_example['playId'] = 1

# Calculate speed and acceleration
tracking_example = tracking_example.sort_values(['nflId', 'frameId'])

for nfl_id in tracking_example['nflId'].unique():
    mask = tracking_example['nflId'] == nfl_id

    # Calculate differences
    dx = tracking_example.loc[mask, 'x'].diff().fillna(0)
    dy = tracking_example.loc[mask, 'y'].diff().fillna(0)

    # Speed (yards/second)
    tracking_example.loc[mask, 's'] = np.sqrt(dx**2 + dy**2) * 10

    # Acceleration
    s_diff = tracking_example.loc[mask, 's'].diff().fillna(0)
    tracking_example.loc[mask, 'a'] = s_diff * 10

    # Direction
    tracking_example.loc[mask, 'dir'] = np.arctan2(dy, dx) * 180 / np.pi

# Display sample
print("Sample Tracking Data:")
print(tracking_example[tracking_example['frameId'] <= 3][
    ['frameId', 'nflId', 'jerseyNumber', 'position', 'x', 'y', 's']
].to_string(index=False))

#| label: speed-analysis-py
#| message: false
#| warning: false

# Calculate max speed for each player
max_speeds = (tracking_example
    .groupby(['nflId', 'jerseyNumber', 'position'])
    .agg(
        max_speed=('s', 'max'),
        avg_speed=('s', 'mean'),
        frames=('s', 'count')
    )
    .reset_index()
    .sort_values('max_speed', ascending=False)
)

print("\nPlayer Speed Summary:")
print(max_speeds.to_string(index=False))

#| label: fig-speed-time-py
#| fig-cap: "Player speed throughout the play - Python"
#| fig-width: 10
#| fig-height: 6
#| message: false
#| warning: false

# Filter out NA values
plot_data = tracking_example.dropna(subset=['s'])

# Create plot
plt.figure(figsize=(10, 6))

colors = {'WR': '#00BFC4', 'CB': '#F8766D', 'S': '#7CAE00'}
labels = {'WR': 'Wide Receiver', 'CB': 'Cornerback', 'S': 'Safety'}

for position in plot_data['position'].unique():
    data = plot_data[plot_data['position'] == position]
    plt.plot(data['frameId'], data['s'],
             color=colors[position],
             label=labels[position],
             marker='o', markersize=4, alpha=0.6, linewidth=2)

plt.xlabel('Frame Number', fontsize=12)
plt.ylabel('Speed (yards/second)', fontsize=12)
plt.title('Player Speed Throughout the Play\nSpeed measured in yards per second',
          fontsize=14, fontweight='bold')
plt.legend(title='Position', loc='upper right')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

#| label: acceleration-analysis-py
#| message: false
#| warning: false

# Calculate acceleration metrics
acceleration_data = tracking_example.dropna(subset=['a'])

acceleration_summary = []
for nfl_id in acceleration_data['nflId'].unique():
    player_data = acceleration_data[acceleration_data['nflId'] == nfl_id]

    acceleration_summary.append({
        'nflId': nfl_id,
        'jerseyNumber': player_data['jerseyNumber'].iloc[0],
        'position': player_data['position'].iloc[0],
        'max_accel': player_data['a'].max(),
        'max_decel': player_data['a'].min(),
        'avg_accel': player_data[player_data['a'] > 0]['a'].mean()
    })

accel_df = pd.DataFrame(acceleration_summary).sort_values('max_accel', ascending=False)

print("\nPlayer Acceleration Summary:")
print("Measured in yards/second²")
print(accel_df.to_string(index=False))

#| label: distance-analysis-py
#| message: false
#| warning: false

# Calculate distance metrics
distance_summary = []

for nfl_id in tracking_example['nflId'].unique():
    player_data = tracking_example[tracking_example['nflId'] == nfl_id].sort_values('frameId')

    # Calculate distances
    dx = player_data['x'].diff().fillna(0)
    dy = player_data['y'].diff().fillna(0)
    total_dist = np.sqrt(dx**2 + dy**2).sum()

    straight_dist = np.sqrt(
        (player_data['x'].iloc[-1] - player_data['x'].iloc[0])**2 +
        (player_data['y'].iloc[-1] - player_data['y'].iloc[0])**2
    )

    distance_summary.append({
        'nflId': nfl_id,
        'jerseyNumber': player_data['jerseyNumber'].iloc[0],
        'position': player_data['position'].iloc[0],
        'total_distance': total_dist,
        'straight_line_distance': straight_dist,
        'efficiency': straight_dist / total_dist if total_dist > 0 else 0
    })

dist_df = pd.DataFrame(distance_summary).sort_values('total_distance', ascending=False)

print("\nPlayer Distance Summary:")
print(dist_df.to_string(index=False))

#| label: separation-calc-py
#| message: false
#| warning: false

def calculate_separation(tracking_data):
    """Calculate separation between receiver and nearest defender"""
    separations = []

    for frame in tracking_data['frameId'].unique():
        frame_data = tracking_data[tracking_data['frameId'] == frame]

        # Get receiver position
        wr_pos = frame_data[frame_data['position'] == 'WR'][['x', 'y']].values

        # Get defender positions
        def_pos = frame_data[frame_data['team'] == 'DEF'][['x', 'y']].values

        if len(wr_pos) > 0 and len(def_pos) > 0:
            # Calculate distances to all defenders
            distances = cdist(wr_pos, def_pos)
            min_distance = distances.min()

            separations.append({
                'frameId': frame,
                'separation': min_distance
            })

    return pd.DataFrame(separations)

# Calculate separation
separation_data = calculate_separation(tracking_example)

# Summary statistics
print("\nReceiver Separation Summary:")
print(f"Min: {separation_data['separation'].min():.2f} yards")
print(f"Max: {separation_data['separation'].max():.2f} yards")
print(f"Mean: {separation_data['separation'].mean():.2f} yards")
print(f"Median: {separation_data['separation'].median():.2f} yards")

#| label: fig-separation-time-py
#| fig-cap: "Receiver separation from nearest defender - Python"
#| fig-width: 10
#| fig-height: 6
#| message: false
#| warning: false

plt.figure(figsize=(10, 6))

plt.plot(separation_data['frameId'], separation_data['separation'],
         color='#00BFC4', linewidth=2, marker='o', markersize=4, alpha=0.6)

plt.axhline(y=2, color='red', linestyle='--', alpha=0.7)
plt.text(separation_data['frameId'].max() * 0.8, 2.3,
         'Tight Coverage (2 yards)', color='red', fontsize=10)

plt.xlabel('Frame Number', fontsize=12)
plt.ylabel('Separation (yards)', fontsize=12)
plt.title('Receiver Separation Throughout the Play\nDistance to nearest defender measured in yards',
          fontsize=14, fontweight='bold')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

#| label: coverage-quality-py
#| message: false
#| warning: false

# Calculate coverage quality metrics
defenders = tracking_example[tracking_example['team'] == 'DEF'].copy()
defenders = defenders.merge(separation_data[['frameId', 'separation']], on='frameId')

coverage_metrics = (defenders
    .groupby(['nflId', 'jerseyNumber', 'position'])
    .agg(
        avg_separation_allowed=('separation', 'mean'),
        min_separation=('separation', 'min'),
        tight_coverage_pct=('separation', lambda x: (x < 2).mean())
    )
    .reset_index()
    .sort_values('avg_separation_allowed')
)

print("\nDefender Coverage Quality:")
print("Tight coverage defined as < 2 yards")
print(coverage_metrics.to_string(index=False))

#| label: route-features-py
#| message: false
#| warning: false

def calculate_route_features(player_data):
    """Calculate features from route path"""
    player_data = player_data.sort_values('frameId')

    # Calculate distance
    dx = player_data['x'].diff().fillna(0)
    dy = player_data['y'].diff().fillna(0)
    total_distance = np.sqrt(dx**2 + dy**2).sum()

    # Direction changes
    dir_diff = player_data['dir'].diff().fillna(0).abs()
    direction_changes = (dir_diff > 45).sum()

    features = {
        'total_distance': total_distance,
        'depth': player_data['x'].max() - player_data['x'].min(),
        'width': player_data['y'].max() - player_data['y'].min(),
        'max_speed': player_data['s'].max(),
        'avg_speed': player_data['s'].mean(),
        'direction_changes': direction_changes,
        'end_x': player_data['x'].iloc[-1],
        'end_y': player_data['y'].iloc[-1]
    }

    return features

# Calculate features for receiver
wr_data = tracking_example[tracking_example['position'] == 'WR']
route_features = calculate_route_features(wr_data)

print("\nRoute Path Features:")
for feature, value in route_features.items():
    print(f"{feature:20s}: {value:.2f}")

#| label: route-clustering-py
#| eval: false
#| echo: true

# Create dataset of multiple routes (synthetic example)
np.random.seed(456)

def generate_route_set(n_routes=100):
    """Generate multiple routes of different types"""
    routes = []

    for i in range(n_routes):
        route_type = np.random.choice(['go', 'slant', 'out', 'post'])

        if route_type == 'go':
            route = generate_trajectory(75, 20, 95, 22 + np.random.normal(0, 2))
        elif route_type == 'slant':
            route = generate_trajectory(75, 20, 85, 26 + np.random.normal(0, 1))
        elif route_type == 'out':
            route = generate_trajectory(75, 20, 85, 15 + np.random.normal(0, 1))
        else:  # post
            route = generate_trajectory(75, 20, 95, 26 + np.random.normal(0, 1))

        route['route_id'] = i
        route['true_type'] = route_type
        routes.append(route)

    return pd.concat(routes, ignore_index=True)

# Generate routes
all_routes = generate_route_set()

# Calculate features for each route
route_features = (all_routes
    .groupby('route_id')
    .agg(
        depth=('x', lambda x: x.max() - x.min()),
        width=('y', lambda x: x.max() - x.min()),
        end_y=('y', 'last'),
        true_type=('true_type', 'first')
    )
    .reset_index()
)

# Perform k-means clustering
scaler = StandardScaler()
features_scaled = scaler.fit_transform(route_features[['depth', 'width', 'end_y']])

kmeans = KMeans(n_clusters=4, random_state=789, n_init=25)
route_features['cluster'] = kmeans.fit_predict(features_scaled)

# Visualize clusters
plt.figure(figsize=(10, 6))
for cluster in range(4):
    cluster_data = route_features[route_features['cluster'] == cluster]
    plt.scatter(cluster_data['depth'], cluster_data['end_y'],
               label=f'Cluster {cluster}', alpha=0.6, s=50)

plt.xlabel('Route Depth (yards)', fontsize=12)
plt.ylabel('Ending Y Position (yards)', fontsize=12)
plt.title('Route Clustering\nRoutes grouped by depth and ending position',
          fontsize=14, fontweight='bold')
plt.legend(title='Cluster')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

#| label: xcomp-features-py
#| eval: false
#| echo: true

def calculate_completion_features(tracking_data, target_frame):
    """Calculate features at target time for completion model"""
    target_data = tracking_data[tracking_data['frameId'] == target_frame]

    # Get receiver and defenders
    receiver = target_data[target_data['position'] == 'WR'].iloc[0]
    defenders = target_data[target_data['team'] == 'DEF']

    # Calculate distances
    distances = np.sqrt(
        (defenders['x'] - receiver['x'])**2 +
        (defenders['y'] - receiver['y'])**2
    )

    features = {
        'separation': distances.min(),
        'defenders_within_1yd': (distances < 1).sum(),
        'defenders_within_2yd': (distances < 2).sum(),
        'defenders_within_3yd': (distances < 3).sum(),
        'receiver_speed': receiver['s'],
        'closest_defender_speed': defenders.iloc[distances.argmin()]['s'],
        'target_x': receiver['x'],
        'target_y': receiver['y'],
        'depth_of_target': receiver['x'] - 75  # assuming LOS at 75
    }

    return features

# Example usage
target_frame = 30  # Frame when ball arrives
completion_features = calculate_completion_features(tracking_example, target_frame)

print("\nCompletion Probability Features:")
for feature, value in completion_features.items():
    print(f"{feature:25s}: {value:.2f}")

#| label: xcomp-model-py
#| eval: false
#| echo: true

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, roc_curve

# Simulate features and outcomes
np.random.seed(999)
n_plays = 1000

completion_data = pd.DataFrame({
    'separation': np.abs(np.random.normal(2.5, 1.5, n_plays)),
    'defenders_within_2yd': np.random.poisson(0.8, n_plays),
    'receiver_speed': np.abs(np.random.normal(10, 2, n_plays)),
    'depth_of_target': np.abs(np.random.normal(15, 8, n_plays)),
    'air_yards': np.abs(np.random.normal(12, 6, n_plays))
})

# Simulate completion probability
from scipy.special import expit

completion_data['xcomp'] = expit(
    0.5 +
    0.3 * completion_data['separation'] -
    0.4 * completion_data['defenders_within_2yd'] -
    0.02 * completion_data['depth_of_target'] +
    0.01 * completion_data['receiver_speed']
)

completion_data['completion'] = np.random.binomial(
    1, completion_data['xcomp']
)

# Prepare features and target
features = ['separation', 'defenders_within_2yd', 'receiver_speed',
            'depth_of_target', 'air_yards']
X = completion_data[features]
y = completion_data['completion']

# Split data
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Fit logistic regression
xcomp_model = LogisticRegression(random_state=42, max_iter=1000)
xcomp_model.fit(X_train, y_train)

# Predictions
y_pred_proba = xcomp_model.predict_proba(X_test)[:, 1]

# Evaluate
auc_score = roc_auc_score(y_test, y_pred_proba)
print(f"\nModel AUC: {auc_score:.3f}")

# Feature importance
print("\nFeature Coefficients:")
for feature, coef in zip(features, xcomp_model.coef_[0]):
    print(f"{feature:25s}: {coef:7.3f}")

#| label: formation-detection-py
#| eval: false
#| echo: true

def detect_defensive_formation(tracking_data, snap_frame=1):
    """Detect defensive formation from pre-snap positioning"""
    # Get defender positions at snap
    defenders = (tracking_data
        .query(f"frameId == {snap_frame} & team == 'DEF'")
        .sort_values('x')
    )

    # Count defenders by region
    box_count = ((defenders['y'] > 14) &
                 (defenders['y'] < 40) &
                 (defenders['x'] < 80)).sum()

    deep_safeties = (defenders['x'] > 85).sum()

    # Classify formation
    if deep_safeties >= 2 and box_count <= 6:
        formation = "Cover 2"
    elif deep_safeties == 1 and box_count <= 6:
        formation = "Cover 3"
    elif deep_safeties == 0:
        formation = "Cover 0/1"
    elif box_count >= 8:
        formation = "Heavy Box"
    else:
        formation = "Other"

    return {
        'formation': formation,
        'box_count': box_count,
        'deep_safeties': deep_safeties,
        'defender_positions': defenders[['x', 'y', 'position']]
    }

# Example usage
formation_info = detect_defensive_formation(tracking_example)
print(f"\nDetected Formation: {formation_info['formation']}")
print(f"Box Count: {formation_info['box_count']}")
print(f"Deep Safeties: {formation_info['deep_safeties']}")

#| label: fig-formation-viz-py
#| fig-cap: "Defensive alignment at snap - Python"
#| fig-width: 12
#| fig-height: 8
#| eval: false
#| echo: true

def plot_defensive_formation(tracking_data, snap_frame=1):
    """Visualize defensive formation"""
    snap_data = tracking_data[tracking_data['frameId'] == snap_frame]

    fig, ax = plt.subplots(figsize=(12, 8))

    # Field background
    ax.add_patch(plt.Rectangle((0, 0), 120, 53.3,
                               facecolor='#196f0c', alpha=0.3))

    # Yard lines
    for x in range(10, 111, 10):
        ax.axvline(x, color='white', alpha=0.3, linewidth=0.5)

    # Line of scrimmage
    ax.axvline(75, color='yellow', linewidth=2)

    # Plot defenders
    defenders = snap_data[snap_data['team'] == 'DEF']
    ax.scatter(defenders['x'], defenders['y'],
              c='red', s=200, alpha=0.8, zorder=3)
    for _, player in defenders.iterrows():
        ax.text(player['x'], player['y'], str(int(player['jerseyNumber'])),
               color='white', fontsize=10, fontweight='bold',
               ha='center', va='center', zorder=4)

    # Plot offense
    offense = snap_data[snap_data['team'] == 'OFF']
    ax.scatter(offense['x'], offense['y'],
              c='blue', s=200, alpha=0.8, zorder=3)
    for _, player in offense.iterrows():
        ax.text(player['x'], player['y'], str(int(player['jerseyNumber'])),
               color='white', fontsize=10, fontweight='bold',
               ha='center', va='center', zorder=4)

    ax.set_xlim(60, 95)
    ax.set_ylim(0, 53.3)
    ax.set_aspect('equal')
    ax.set_xlabel('Field Position (yards)', fontsize=12)
    ax.set_ylabel('Field Width (yards)', fontsize=12)
    ax.set_title(f'Defensive Alignment at Snap\nFormation: {formation_info["formation"]}',
                fontsize=14, fontweight='bold')
    ax.grid(False)

    plt.tight_layout()
    plt.show()

plot_defensive_formation(tracking_example)

#| label: voronoi-py
#| eval: false
#| echo: true

def create_voronoi_plot(tracking_data, frame_num):
    """Create Voronoi diagram for spatial control"""
    frame_data = tracking_data[tracking_data['frameId'] == frame_num]

    # Get player positions
    points = frame_data[['x', 'y']].values

    # Create Voronoi diagram
    vor = Voronoi(points)

    fig, ax = plt.subplots(figsize=(12, 8))

    # Field background
    ax.add_patch(plt.Rectangle((0, 0), 120, 53.3,
                               facecolor='#196f0c', alpha=0.3))

    # Plot Voronoi regions
    voronoi_plot_2d(vor, ax=ax, show_vertices=False,
                   line_colors='black', line_width=1.5)

    # Color regions by team
    for region_idx, point_idx in enumerate(vor.point_region):
        region = vor.regions[point_idx]
        if not -1 in region and len(region) > 0:
            polygon = [vor.vertices[i] for i in region]
            team = frame_data.iloc[region_idx]['team']
            color = 'blue' if team == 'OFF' else 'red'
            ax.fill(*zip(*polygon), alpha=0.3, color=color)

    # Plot players
    for team in ['OFF', 'DEF']:
        team_data = frame_data[frame_data['team'] == team]
        color = 'blue' if team == 'OFF' else 'red'
        ax.scatter(team_data['x'], team_data['y'],
                  c=color, s=200, edgecolors='white', linewidths=2, zorder=3)
        for _, player in team_data.iterrows():
            ax.text(player['x'], player['y'],
                   str(int(player['jerseyNumber'])),
                   color='white', fontsize=10, fontweight='bold',
                   ha='center', va='center', zorder=4)

    ax.set_xlim(60, 100)
    ax.set_ylim(0, 53.3)
    ax.set_aspect('equal')
    ax.set_xlabel('Field Position (yards)', fontsize=12)
    ax.set_ylabel('Field Width (yards)', fontsize=12)
    ax.set_title(f'Spatial Control via Voronoi Diagram\nFrame: {frame_num}',
                fontsize=14, fontweight='bold')

    plt.tight_layout()
    plt.show()

create_voronoi_plot(tracking_example, frame_num=25)