# ============================================
# OpenSky EDA Script (Polars Version)
# ============================================

from pathlib import Path

import polars as pl
import matplotlib.pyplot as plt

pl.Config.set_tbl_cols(-1)          # show all columns
pl.Config.set_tbl_rows(1000)         # number of visible rows
pl.Config.set_tbl_width_chars(200)  # wider table

BASE_DIR = Path.cwd()

data_dir = BASE_DIR / ".." / ".." / ".." / "data_pipeline" / "raw_data"
# extract_dir = BASE_DIR / "raw_data" / "extracted"

# INPUT_FILE = "bwi_states_smoothed_resampled.csv"

INPUT_FILE = "/home/jon/Documents/grad_school/OR568/project/OR568_ML_Project/data_pipeline/raw_data/bwi_states_2022_06_27_raw.csv"

# --------------------------------------------
# LOAD DATA
# --------------------------------------------
df = pl.read_csv(INPUT_FILE, try_parse_dates=True)

print("\n===== BASIC INFO =====")
print("Rows:", df.height)
print("Columns:", len(df.columns))

# --------------------------------------------
# TIME RANGE
# --------------------------------------------
print("\n===== TIME RANGE =====")
print(df.select([
    pl.col("time").min().alias("start_time"),
    pl.col("time").max().alias("end_time")
]))

# --------------------------------------------
# UNIQUE AIRCRAFT
# --------------------------------------------
print("\nUnique aircraft:", df.select(pl.col("icao24").n_unique()).item())

# --------------------------------------------
# NULL COUNTS
# --------------------------------------------
print("\n===== MISSING VALUES =====")
nulls = df.select([
    pl.all().null_count()
])
print(nulls)

# ============================================
# SQUAWK FREQUENCY DISTRIBUTION
# ============================================

print("\n===== SQUAWK FREQUENCY =====")

# -----------------------------------
# Frequency table
# -----------------------------------
squawk_freq = (
    df.filter(pl.col("squawk").is_not_null())
      .group_by("squawk")
      .len()
      .sort("len", descending=True)
)

print(squawk_freq.head(20))

top_sq = squawk_freq.head(20)

# -----------------------------------
# Highlight logic
# -----------------------------------

EMERGENCY_CODES = {"7500", "7600", "7700"}
IMPORTANT_CODES = {"1200", "7000"}

colors = []

for sq in top_sq["squawk"].cast(str):
    if sq in EMERGENCY_CODES:
        colors.append("red")          # emergencies
    elif sq in IMPORTANT_CODES:
        colors.append("orange")       # operational baseline
    else:
        colors.append("steelblue")    # normal

# -----------------------------------
# Plot
# -----------------------------------

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

plt.bar(
    top_sq["squawk"].cast(str),
    top_sq["len"],
    color=colors
)

plt.title("Top Squawk Codes Frequency (Emergencies Highlighted)")
plt.xlabel("Squawk Code")
plt.ylabel("Count")

plt.xticks(rotation=45)

# legend
import matplotlib.patches as mpatches

legend_handles = [
    mpatches.Patch(color="red", label="Emergency (7500/7600/7700)"),
    mpatches.Patch(color="orange", label="Operational Baseline (1200/7000)"),
    mpatches.Patch(color="steelblue", label="Other")
]

plt.legend(handles=legend_handles)

plt.tight_layout()
plt.show()

# # ============================================
# # MOVEMENT STATISTICS
# # ============================================

# print("\n===== NUMERIC SUMMARY =====")

# numeric_cols = [
#     "velocity",
#     "geoaltitude",
#     "baroaltitude",
#     "vertrate",
#     "heading"
# ]

# summary = df.select([
#     pl.col(c).mean().alias(f"{c}_mean") for c in numeric_cols
# ] + [
#     pl.col(c).std().alias(f"{c}_std") for c in numeric_cols
# ] + [
#     pl.col(c).min().alias(f"{c}_min") for c in numeric_cols
# ] + [
#     pl.col(c).max().alias(f"{c}_max") for c in numeric_cols
# ])

# print(summary)

# # ============================================
# # OPERATIONAL FLAGS
# # ============================================

# print("\n===== OPERATIONAL FLAGS =====")

# flags = df.select([
#     (pl.col("onground") == True).mean().alias("pct_onground"),
#     (pl.col("alert") == True).mean().alias("pct_alert"),
#     (pl.col("spi") == True).mean().alias("pct_spi")
# ])

# print(flags)

# # ============================================
# # PER AIRCRAFT SUMMARY (ML READY)
# # ============================================

# print("\n===== AIRCRAFT LEVEL SUMMARY =====")

# aircraft_summary = (
#     df.group_by("icao24")
#       .agg([
#           pl.col("velocity").mean().alias("avg_velocity"),
#           pl.col("geoaltitude").max().alias("max_altitude"),
#           pl.col("time").min().alias("start_time"),
#           pl.col("time").max().alias("end_time"),
#           pl.len().alias("samples")
#       ])
# )

# print(aircraft_summary.head())

# # Optional save for ML
# aircraft_summary.write_csv("aircraft_summary.csv")

# print("\nSaved aircraft summary → aircraft_summary.csv")

===== BASIC INFO =====
Rows: 67831
Columns: 16

===== TIME RANGE =====
shape: (1, 2)
┌────────────┬────────────┐
│ start_time ┆ end_time   │
│ ---        ┆ ---        │
│ i64        ┆ i64        │
╞════════════╪════════════╡
│ 1656288010 ┆ 1656374390 │
└────────────┴────────────┘

Unique aircraft: 1068

===== MISSING VALUES =====
shape: (1, 16)
┌──────┬────────┬─────┬─────┬──────────┬─────────┬──────────┬──────────┬──────────┬───────┬─────┬────────┬──────────────┬─────────────┬───────────────┬─────────────┐
│ time ┆ icao24 ┆ lat ┆ lon ┆ velocity ┆ heading ┆ vertrate ┆ callsign ┆ onground ┆ alert ┆ spi ┆ squawk ┆ baroaltitude ┆ geoaltitude ┆ lastposupdate ┆ lastcontact │
│ ---  ┆ ---    ┆ --- ┆ --- ┆ ---      ┆ ---     ┆ ---      ┆ ---      ┆ ---      ┆ ---   ┆ --- ┆ ---    ┆ ---          ┆ ---         ┆ ---           ┆ ---         │
│ u32  ┆ u32    ┆ u32 ┆ u32 ┆ u32      ┆ u32     ┆ u32      ┆ u32      ┆ u32      ┆ u32   ┆ u32 ┆ u32    ┆ u32          ┆ u32         ┆ u32           ┆ u32         │
╞══════╪════════╪═════╪═════╪══════════╪═════════╪══════════╪══════════╪══════════╪═══════╪═════╪════════╪══════════════╪═════════════╪═══════════════╪═════════════╡
│ 0    ┆ 0      ┆ 0   ┆ 0   ┆ 2451     ┆ 2451    ┆ 2451     ┆ 1395     ┆ 0        ┆ 0     ┆ 0   ┆ 5627   ┆ 3164         ┆ 3729        ┆ 0             ┆ 0           │
└──────┴────────┴─────┴─────┴──────────┴─────────┴──────────┴──────────┴──────────┴───────┴─────┴────────┴──────────────┴─────────────┴───────────────┴─────────────┘

===== SQUAWK FREQUENCY =====
shape: (20, 2)
┌────────┬──────┐
│ squawk ┆ len  │
│ ---    ┆ ---  │
│ i64    ┆ u32  │
╞════════╪══════╡
│ 1200   ┆ 2721 │
│ 5154   ┆ 1044 │
│ 5113   ┆ 652  │
│ 5163   ┆ 490  │
│ 5101   ┆ 437  │
│ 371    ┆ 244  │
│ 5102   ┆ 231  │
│ 5214   ┆ 230  │
│ 4643   ┆ 220  │
│ 2132   ┆ 208  │
│ 3634   ┆ 207  │
│ 5322   ┆ 202  │
│ 5166   ┆ 199  │
│ 1431   ┆ 193  │
│ 4630   ┆ 184  │
│ 5107   ┆ 184  │
│ 2122   ┆ 176  │
│ 6216   ┆ 168  │
│ 7037   ┆ 164  │
│ 1063   ┆ 161  │
└────────┴──────┘

EMERGENCY_CODES = {"7500", "7600", "7700"}

emergency_df = (
    df.filter(
        pl.col("squawk")
          .cast(pl.Utf8)
          .is_in(EMERGENCY_CODES)
    )
)

print(emergency_df)
shape: (25, 16)
┌────────────┬────────┬───────────┬────────────┬────────────┬────────────┬───────────┬──────────┬──────────┬───────┬───────┬────────┬──────────────┬─────────────┬───────────────┬─────────────┐
│ time       ┆ icao24 ┆ lat       ┆ lon        ┆ velocity   ┆ heading    ┆ vertrate  ┆ callsign ┆ onground ┆ alert ┆ spi   ┆ squawk ┆ baroaltitude ┆ geoaltitude ┆ lastposupdate ┆ lastcontact │
│ ---        ┆ ---    ┆ ---       ┆ ---        ┆ ---        ┆ ---        ┆ ---       ┆ ---      ┆ ---      ┆ ---   ┆ ---   ┆ ---    ┆ ---          ┆ ---         ┆ ---           ┆ ---         │
│ i64        ┆ str    ┆ f64       ┆ f64        ┆ f64        ┆ f64        ┆ f64       ┆ str      ┆ bool     ┆ bool  ┆ bool  ┆ i64    ┆ f64          ┆ f64         ┆ f64           ┆ f64         │
╞════════════╪════════╪═══════════╪════════════╪════════════╪════════════╪═══════════╪══════════╪══════════╪═══════╪═══════╪════════╪══════════════╪═════════════╪═══════════════╪═════════════╡
│ 1656297100 ┆ a97d63 ┆ 38.965899 ┆ -76.712097 ┆ 182.439135 ┆ 42.3715    ┆ 9.42848   ┆ JIA5065  ┆ false    ┆ false ┆ false ┆ 7500   ┆ 3977.64      ┆ 4175.76     ┆ 1.6563e9      ┆ 1.6563e9    │
│ 1656297110 ┆ a97d63 ┆ 38.978375 ┆ -76.697571 ┆ 184.691128 ┆ 41.951625  ┆ 9.42848   ┆ JIA5065  ┆ false    ┆ false ┆ false ┆ 7500   ┆ 4061.46      ┆ 4267.2      ┆ 1.6563e9      ┆ 1.6563e9    │
│ 1656297120 ┆ a97d63 ┆ 38.990851 ┆ -76.683228 ┆ 186.870677 ┆ 41.987212  ┆ 9.10336   ┆ JIA5065  ┆ false    ┆ false ┆ false ┆ 7500   ┆ 4152.9       ┆ 4366.26     ┆ 1.6563e9      ┆ 1.6563e9    │
│ 1656297130 ┆ a97d63 ┆ 39.003373 ┆ -76.668762 ┆ 187.979837 ┆ 41.893973  ┆ 11.05408  ┆ JIA5065  ┆ false    ┆ false ┆ false ┆ 7500   ┆ 4251.96      ┆ 4472.94     ┆ 1.6563e9      ┆ 1.6563e9    │
│ 1656297140 ┆ a97d63 ┆ 39.015793 ┆ -76.654179 ┆ 190.663801 ┆ 43.688112  ┆ 7.47776   ┆ JIA5065  ┆ false    ┆ false ┆ false ┆ 7500   ┆ 4335.78      ┆ 4556.76     ┆ 1.6563e9      ┆ 1.6563e9    │
│ 1656350000 ┆ abbe71 ┆ 38.93955  ┆ -76.354492 ┆ 244.404759 ┆ 23.705339  ┆ -9.7536   ┆ WMN557   ┆ false    ┆ false ┆ false ┆ 7700   ┆ 8237.22      ┆ 8663.94     ┆ 1.6563e9      ┆ 1.6564e9    │
│ 1656350010 ┆ abbe71 ┆ 38.958984 ┆ -76.343516 ┆ 242.784868 ┆ 23.740605  ┆ -10.07872 ┆ WMN557   ┆ false    ┆ false ┆ false ┆ 7700   ┆ 8145.78      ┆ 8572.5      ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656350020 ┆ abbe71 ┆ 38.979721 ┆ -76.331873 ┆ 241.63594  ┆ 23.727164  ┆ -9.7536   ┆ WMN557   ┆ false    ┆ false ┆ false ┆ 7700   ┆ 8054.34      ┆ 8473.44     ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656350030 ┆ abbe71 ┆ 38.998672 ┆ -76.320947 ┆ 240.685032 ┆ 25.579179  ┆ -10.07872 ┆ WMN557   ┆ false    ┆ false ┆ false ┆ 7700   ┆ 7962.9       ┆ 8382.0      ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656352940 ┆ ac508b ┆ 39.318067 ┆ -76.64801  ┆ 251.643058 ┆ 22.475702  ┆ 0.0       ┆ DAL1274  ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10058.4      ┆ 10576.56    ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656352950 ┆ ac508b ┆ 39.338177 ┆ -76.637329 ┆ 251.44684  ┆ 22.367383  ┆ 0.0       ┆ DAL1274  ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10058.4      ┆ 10576.56    ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656352960 ┆ ac508b ┆ 39.358978 ┆ -76.626295 ┆ 252.118502 ┆ 22.431008  ┆ 0.0       ┆ DAL1274  ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10058.4      ┆ 10576.56    ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656352970 ┆ ac508b ┆ 39.38026  ┆ -76.61499  ┆ 251.922654 ┆ 22.322858  ┆ 0.0       ┆ DAL1274  ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10058.4      ┆ 10576.56    ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656352980 ┆ ac508b ┆ 39.401321 ┆ -76.603845 ┆ 251.922654 ┆ 22.322858  ┆ 0.0       ┆ DAL1274  ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10058.4      ┆ 10576.56    ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656364440 ┆ a4bb8a ┆ 39.068686 ┆ -76.880371 ┆ 219.387899 ┆ 223.954919 ┆ -0.32512  ┆ AAL845   ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10370.82     ┆ 10888.98    ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656364450 ┆ a4bb8a ┆ 39.054657 ┆ -76.89761  ┆ 219.387899 ┆ 223.954919 ┆ 0.0       ┆ AAL845   ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10363.2      ┆ 10888.98    ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656364460 ┆ a4bb8a ┆ 39.040662 ┆ -76.914856 ┆ 219.402375 ┆ 223.764959 ┆ 0.32512   ┆ AAL845   ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10370.82     ┆ 10881.36    ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656364470 ┆ a4bb8a ┆ 39.026367 ┆ -76.932479 ┆ 219.402375 ┆ 223.764959 ┆ 0.0       ┆ AAL845   ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10363.2      ┆ 10888.98    ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656364480 ┆ a4bb8a ┆ 39.012032 ┆ -76.950073 ┆ 219.402375 ┆ 223.764959 ┆ 0.0       ┆ AAL845   ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10370.82     ┆ 10888.98    ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656364490 ┆ a4bb8a ┆ 38.997648 ┆ -76.967773 ┆ 219.402375 ┆ 223.764959 ┆ 0.0       ┆ AAL845   ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10370.82     ┆ 10888.98    ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656364500 ┆ a4bb8a ┆ 38.983475 ┆ -76.985142 ┆ 219.758531 ┆ 223.861823 ┆ -0.32512  ┆ AAL845   ┆ false    ┆ false ┆ false ┆ 7500   ┆ 10370.82     ┆ 10896.6     ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656367660 ┆ abd19e ┆ 38.931683 ┆ -76.544434 ┆ 133.375994 ┆ 350.900994 ┆ -4.8768   ┆ SWA2748  ┆ false    ┆ false ┆ false ┆ 7700   ┆ 1280.16      ┆ 1318.26     ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656367670 ┆ abd19e ┆ 38.943146 ┆ -76.546823 ┆ 131.852308 ┆ 350.794937 ┆ -5.20192  ┆ SWA2748  ┆ false    ┆ false ┆ false ┆ 7700   ┆ 1226.82      ┆ 1264.92     ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656367680 ┆ abd19e ┆ 38.954865 ┆ -76.549331 ┆ 130.499548 ┆ 350.240619 ┆ -4.55168  ┆ SWA2748  ┆ false    ┆ false ┆ false ┆ 7700   ┆ 1181.1       ┆ 1219.2      ┆ 1.6564e9      ┆ 1.6564e9    │
│ 1656367690 ┆ abd19e ┆ 38.966309 ┆ -76.551779 ┆ 128.215298 ┆ 350.764441 ┆ -2.60096  ┆ SWA2748  ┆ false    ┆ false ┆ false ┆ 7700   ┆ 1143.0       ┆ 1181.1      ┆ 1.6564e9      ┆ 1.6564e9    │
└────────────┴────────┴───────────┴────────────┴────────────┴────────────┴───────────┴──────────┴──────────┴───────┴───────┴────────┴──────────────┴─────────────┴───────────────┴─────────────┘
# ----------------------------------------
# 1. Define emergency squawk codes
# ----------------------------------------
EMERGENCY_CODES = {"7500", "7600", "7700"}

# ----------------------------------------
# 2. Filter emergency rows
# ----------------------------------------
emergency_df = (
    df.filter(
        pl.col("squawk")
          .cast(pl.Utf8)
          .is_in(EMERGENCY_CODES)
    )
)

# ----------------------------------------
# 3. Frequency count per code
# ----------------------------------------
freq_df = (
    emergency_df
    .with_columns(pl.col("squawk").cast(pl.Utf8))
    .group_by("squawk")
    .count()
    .sort("count", descending=True)
)

print(freq_df)

# ----------------------------------------
# 4. Plot frequency graph
# ----------------------------------------
plot_df = freq_df.to_pandas()

plt.figure(figsize=(6,4))
plt.bar(plot_df["squawk"], plot_df["count"])
plt.title("Frequency of Emergency Squawk Codes")
plt.xlabel("Squawk Code")
plt.ylabel("Count")
plt.grid(axis="y", alpha=0.3)

plt.show()
/tmp/ipykernel_56177/402016319.py:24: DeprecationWarning: `GroupBy.count` was renamed; use `GroupBy.len` instead
  .count()
shape: (2, 2)
┌────────┬───────┐
│ squawk ┆ count │
│ ---    ┆ ---   │
│ str    ┆ u32   │
╞════════╪═══════╡
│ 7500   ┆ 17    │
│ 7700   ┆ 8     │
└────────┴───────┘