# ingest_timescaledb.py — Load AIS data from JSON into TimescaleDB
# Each script section is explained line-by-line below

import psycopg2         # PostgreSQL adapter for Python
from psycopg2.extras import execute_values  # Batch insert helper (faster than row-by-row)
import json             # Parse AISHub JSON responses
from datetime import datetime  # Convert timestamps

# Connect to TimescaleDB (which is just PostgreSQL under the hood)
# Parameters: host, port, database name, user, password
conn = psycopg2.connect(
    host="localhost",
    port=5432,
    dbname="aisdb",
    user="postgres",
    password="postgres"
)
# Create a cursor object to send SQL commands through the connection
cur = conn.cursor()

# Read AIS data from a JSON file (one array of vessel report objects)
# This is the format AISHub returns from their REST API
with open("ais_data.json", "r") as f:
    data = json.load(f)

# Prepare a list of tuples for batch insertion
# Each tuple matches the column order of the table definition above
rows = []
for report in data:
    # Convert the AISHub timestamp string (like "2026-07-22 12:09:25 GMT")
    # to a Python datetime object for PostgreSQL TIMESTAMPTZ type
    ts = datetime.strptime(report["TIME"], "%Y-%m-%d %H:%M:%S GMT")
    rows.append((
        report["MMSI"],         # INTEGER — vessel identifier
        ts,                     # TIMESTAMPTZ — report time with timezone
        report["LATITUDE"],     # DOUBLE PRECISION — decimal degrees
        report["LONGITUDE"],    # DOUBLE PRECISION — decimal degrees
        report.get("SOG"),      # Speed Over Ground (nullable — some reports omit it)
        report.get("COG"),      # Course Over Ground (nullable)
        report.get("HEADING"),  # Heading (nullable, 511 = unavailable)
        report.get("NAVSTAT"),  # Navigational status (nullable)
        report.get("TYPE"),     # Ship type code (nullable)
        report.get("STATION_ID"), # AISHub station (nullable)
        report.get("LENGTH"),   # Vessel length in meters (nullable)
    ))

# Batch insert using execute_values — much faster than individual INSERT statements
# The SQL template uses %s placeholders for safety (prevents SQL injection)
# execute_values splits the rows into batches of 100 internally (default)
insert_sql = """
    INSERT INTO ais_positions
    (mmsi, time, latitude, longitude, sog, cog, heading, navstat, ship_type, station_id, length)
    VALUES %s
"""
execute_values(cur, insert_sql, rows, page_size=1000)
# Commit the transaction so data is written to disk and visible to other connections
conn.commit()

# Verify the row count to make sure everything was inserted
cur.execute("SELECT count(*) FROM ais_positions")
print(f"Inserted {cur.fetchone()[0]} rows")

# Clean up: close cursor and connection
cur.close()
conn.close()

# setup_influxdb.py — Configure InfluxDB bucket and retention for AIS data
# Each line is explained below

from influxdb_client import InfluxDBClient  # Official InfluxDB 2.x Python client
from influxdb_client.client.write_api import SYNCHRONOUS  # Write mode: wait for ack
import datetime

# Connect to InfluxDB running on the default port 8086
# InfluxDB 2.x uses an org/bucket/token model instead of databases
# org = organization name (like a tenant/workspace)
# token = API token with read/write permissions to specific buckets
client = InfluxDBClient(
    url="http://localhost:8086",
    token="my-super-secret-token",
    org="ais-org"
)

# Get the buckets API for managing data storage buckets
buckets_api = client.buckets_api()

# List existing buckets (InfluxDB equivalent of databases)
# A bucket holds time-series data with a retention policy
for bucket in buckets_api.find_buckets().buckets:
    print(f"Bucket: {bucket.name}, Retention: {bucket.retention_rules}")

# Create a new bucket for AIS data if it doesn't exist
# Retention rule: keep data for 90 days, then auto-delete
# InfluxDB is optimized for recent-data access and automatic expiry
bucket_name = "ais_positions"
# Find the organization ID (required to create buckets under the right org)
org_info = client.organizations_api().find_organizations(org="ais-org")[0]
# Check if the bucket already exists to avoid duplicate creation
existing = buckets_api.find_bucket_by_name(bucket_name)
if not existing:
    # Create the bucket with a 90-day retention period
    # Old data is automatically deleted by InfluxDB's retention enforcement
    retention = client.retention_rules(
        every_seconds=90 * 24 * 3600,  # 90 days in seconds
        type="expire"
    )
    bucket = buckets_api.create_bucket(
        bucket_name=bucket_name,
        org_id=org_info.id,
        retention_rules=[retention]
    )
    print(f"Created bucket: {bucket.name}")

client.close()

# AIS data modeled as InfluxDB points
# InfluxDB uses a measurement → tags → fields → timestamp structure:
#   - Measurement: the table name (ais_position)
#   - Tags: indexed metadata you filter/group by (mmsi, navstat, ship_type)
#   - Fields: numeric values you query (latitude, longitude, sog, cog)
#   - Timestamp: when the measurement was taken

from influxdb_client.client.write.point import Point
from influxdb_client import WritePrecision
from datetime import datetime

# Create an InfluxDB data point from a single AIS report
# Tags go in the "tag" set, fields go in the "field" set
# Tags are indexed (fast for GROUP BY), fields are not indexed
point = Point("ais_position") \                   # Measurement name
    .tag("mmsi", str(374651000)) \                 # Tag: vessel ID (converted to string)
    .tag("navstat", "0") \                         # Tag: navigational status (tagged for filtering)
    .tag("ship_type", "80") \                      # Tag: ship type code (tagged for filtering)
    .tag("station_id", "42") \                     # Tag: receiving station (tagged for geo-filtering)
    .field("latitude", 33.89075) \                 # Field: lat as float64 (queryable value)
    .field("longitude", -128.26693) \              # Field: lon as float64
    .field("sog", 14.7) \                          # Field: speed over ground in knots
    .field("cog", 122.6) \                         # Field: course over ground in degrees
    .field("heading", 119) \                       # Field: heading in degrees
    .field("length", 188) \                        # Field: vessel length in meters
    .time(datetime.utcnow(), WritePrecision.S)     # Timestamp (seconds precision)

# ingest_influxdb.py — Batch load AIS data into InfluxDB
# Compared to TimescaleDB, InfluxDB uses a different batch protocol

from influxdb_client import InfluxDBClient
from influxdb_client.client.write_api import SYNCHRONOUS  # Block until written
import json
from datetime import datetime

# Connect to InfluxDB on the default port
client = InfluxDBClient(
    url="http://localhost:8086",
    token="my-super-secret-token",
    org="ais-org"
)

# Get the write API — this sends data points to InfluxDB
# SYNCHRONOUS mode means write() blocks until the server acknowledges receipt
# Alternative: ASYNCHRONOUS (fire-and-forget) or BATCHING (buffer and batch)
write_api = client.write_api(write_options=SYNCHRONOUS)

# Read AIS data from JSON file
with open("ais_data.json", "r") as f:
    data = json.load(f)

# Build InfluxDB Point objects for each AIS report
points = []
for report in data:
    # Parse the AISHub timestamp into a Python datetime
    ts = datetime.strptime(report["TIME"], "%Y-%m-%d %H:%M:%S GMT")

    # Create a Point with tags on frequently-filtered columns
    # and fields for numeric measurements
    point = Point("ais_position") \
        .tag("mmsi", str(report["MMSI"])) \       # Tag: vessel ID
        .tag("navstat", str(report.get("NAVSTAT", "0"))) \  # Tag: status
        .tag("ship_type", str(report.get("TYPE", "0"))) \   # Tag: category
        .field("latitude", float(report["LATITUDE"])) \     # Field: lat
        .field("longitude", float(report["LONGITUDE"])) \   # Field: lon
        .field("sog", float(report.get("SOG", 0))) \        # Field: speed
        .field("cog", float(report.get("COG", 0))) \        # Field: course
        .field("heading", int(report.get("HEADING", 0))) \  # Field: heading
        .field("length", int(report.get("LENGTH", 0))) \    # Field: length
        .time(ts, WritePrecision.S)                         # Timestamp: seconds

    points.append(point)

# Batch write all points to the "ais_positions" bucket
# InfluxDB handles batching internally — it groups points by measurement and timestamp
# The bucket specifies retention policy (90 days auto-delete earlier if configured)
write_api.write(bucket="ais_positions", record=points)

print(f"Written {len(points)} points to InfluxDB")

# Verify by querying the row count using Flux
# InfluxDB's cardinality: number of distinct tag value combinations
query = 'from(bucket: "ais_positions") |> range(start: 0) |> count()'
result = client.query_api().query(query)
for table in result:
    for record in table.records:
        print(f"Records count field: {record.get_value()}")

client.close()

# query_influxdb.py — Run common AIS queries against InfluxDB using Flux
# Flux is InfluxDB's functional query language (pipeline of transformations)

from influxdb_client import InfluxDBClient

client = InfluxDBClient(
    url="http://localhost:8086",
    token="my-super-secret-token",
    org="ais-org"
)

query_api = client.query_api()

# 1. Get a vessel's trajectory for the last 24 hours (like TimescaleDB query 1)
# Flux pipelines data through functions using |> operator (pipe-forward)
flux_trajectory = '''
from(bucket: "ais_positions")                      // Source: our AIS bucket
  |> range(start: -24h)                             // Time window: last 24 hours
  |> filter(fn: (r) => r.mmsi == "374651000")       // Filter to one vessel by tag
  |> filter(fn: (r) => r._field == "latitude"       // Keep only lat and lon fields
      or r._field == "longitude"
      or r._field == "sog")
  |> pivot(rowKey:["_time"], columnKey:["_field"], valueColumn:"_value")  // Pivot to wide format
  |> keep(columns: ["_time", "latitude", "longitude", "sog"])  // Drop extra columns
  |> sort(columns: ["_time"], desc: true)            // Order by time descending
'''

print("=== Vessel trajectory (last 24h) ===")
result = query_api.query(flux_trajectory)
for table in result:
    for record in table.records:
        print(f"{record['_time']}: lat={record['latitude']}, lon={record['longitude']}, "
              f"sog={record['sog']}")

# 2. Average speed per hour for a vessel (aggregation)
# InfluxDB requires more steps than SQL for time bucketing
flux_hourly = '''
from(bucket: "ais_positions")                      // Source bucket
  |> range(start: -7d)                              // Last 7 days
  |> filter(fn: (r) => r.mmsi == "374651000")       // One vessel
  |> filter(fn: (r) => r._field == "sog")           // Only the sog field
  |> aggregateWindow(                               // Time-bucket window aggregation
      every: 1h,                                     // Bucket size: 1 hour
      fn: mean,                                      // Aggregation function: average
      createEmpty: false                             // Skip hours with no data
  )
  |> yield(name: "mean_sog")                        // Name the result stream
'''

print("\n=== Hourly average speed ===")
result = query_api.query(flux_hourly)
for table in result:
    for record in table.records:
        print(f"{record['_time']}: {record['_value']:.1f} knots")

# compare_storage.py — Compare storage usage between TimescaleDB and InfluxDB

import psycopg2
from influxdb_client import InfluxDBClient

print("=== Storage Comparison ===")
print(f"Dataset: ~10 million AIS position reports")
print()

# TimescaleDB storage size (via PostgreSQL's table size functions)
conn = psycopg2.connect(host="localhost", dbname="aisdb", user="postgres")
cur = conn.cursor()
# pg_total_relation_size returns the total disk space used by the table
# including indexes, TOAST (large field storage), and free space
cur.execute("SELECT pg_size_pretty(pg_total_relation_size('ais_positions'))")
tsdb_size = cur.fetchone()[0]
print(f"TimescaleDB (with indexes): {tsdb_size}")

# InfluxDB storage size (via bucket API query)
client = InfluxDBClient(url="http://localhost:8086", token="token", org="ais-org")
buckets_api = client.buckets_api()
bucket = buckets_api.find_bucket_by_name("ais_positions")
# InfluxDB reports storage usage per shard (time-partitioned storage unit)
# Total is sum of all shard sizes for this bucket
if bucket:
    # Query InfluxDB's internal _monitoring bucket for storage metrics
    query = f'''
    from(bucket: "_monitoring")
      |> range(start: -30d)
      |> filter(fn: (r) => r.bucket_id == "{bucket.id}")
      |> filter(fn: (r) => r._field == "disk_bytes")
      |> last()
    '''
    result = client.query_api().query(query)
    for table in result:
        for record in table.records:
            size_mb = record.get_value() / 1_000_000
            print(f"InfluxDB (uncompressed): {size_mb:.1f} MB")

client.close()
cur.close()
conn.close()

# benchmark_queries.py — Time common operations against both databases
# This script runs the same logical queries on both and prints timing

import psycopg2
from influxdb_client import InfluxDBClient
import time

# Number of iterations per query to get stable averages
ITERATIONS = 5

def benchmark_tsdb():
    """Run benchmarks on TimescaleDB and return average timing per query."""
    conn = psycopg2.connect(host="localhost", dbname="aisdb", user="postgres")
    cur = conn.cursor()
    queries = {}

    # Query 1: Single vessel trajectory (24h range)
    # Measure how fast it takes to get a vessel's full track for replay
    start = time.time()
    for _ in range(ITERATIONS):
        cur.execute("""
            SELECT time, latitude, longitude, sog
            FROM ais_positions
            WHERE mmsi = 374651000
              AND time > NOW() - INTERVAL '24 hours'
            ORDER BY time DESC
        """)
        rows = cur.fetchall()  # Fetch all results to ensure query completes
    queries["trajectory_24h"] = (time.time() - start) / ITERATIONS

    # Query 2: Hourly average speed (7 days), with aggregation
    start = time.time()
    for _ in range(ITERATIONS):
        cur.execute("""
            SELECT date_trunc('hour', time), AVG(sog)
            FROM ais_positions
            WHERE mmsi = 374651000 AND time > NOW() - INTERVAL '7 days'
            GROUP BY date_trunc('hour', time)
            ORDER BY 1
        """)
        cur.fetchall()
    queries["hourly_avg_7d"] = (time.time() - start) / ITERATIONS

    # Query 3: Vessel count in bounding box (last hour)
    start = time.time()
    for _ in range(ITERATIONS):
        cur.execute("""
            SELECT COUNT(DISTINCT mmsi)
            FROM ais_positions
            WHERE time > NOW() - INTERVAL '1 hour'
              AND latitude BETWEEN 51.0 AND 52.0
              AND longitude BETWEEN 1.0 AND 3.0
        """)
        cur.fetchall()
    queries["bbox_count_1h"] = (time.time() - start) / ITERATIONS

    cur.close()
    conn.close()
    return queries

def benchmark_influx():
    """Run benchmarks on InfluxDB and return average timing per query."""
    client = InfluxDBClient(url="http://localhost:8086", token="token", org="ais-org")
    query_api = client.query_api()
    queries = {}

    # Query 1: Single vessel trajectory (24h)
    q1 = '''
    from(bucket: "ais_positions")
      |> range(start: -24h)
      |> filter(fn: (r) => r.mmsi == "374651000")
      |> filter(fn: (r) => r._field == "latitude" or r._field == "longitude" or r._field == "sog")
      |> pivot(rowKey:["_time"], columnKey:["_field"], valueColumn:"_value")
    '''
    start = time.time()
    for _ in range(ITERATIONS):
        result = query_api.query(q1)
        for table in result:
            for record in table.records:
                pass  # Consume all records to ensure full query execution
    queries["trajectory_24h"] = (time.time() - start) / ITERATIONS

    # Query 2: Hourly average speed (7 days)
    q2 = '''
    from(bucket: "ais_positions")
      |> range(start: -7d)
      |> filter(fn: (r) => r.mmsi == "374651000")
      |> filter(fn: (r) => r._field == "sog")
      |> aggregateWindow(every: 1h, fn: mean, createEmpty: false)
    '''
    start = time.time()
    for _ in range(ITERATIONS):
        result = query_api.query(q2)
        for table in result:
            for record in table.records:
                pass
    queries["hourly_avg_7d"] = (time.time() - start) / ITERATIONS

    # Query 3: Vessel count — InfluxDB doesn't have DISTINCT easily
    # This becomes complex: you need to find distinct mmsi values over the time range
    q3 = '''
    from(bucket: "ais_positions")
      |> range(start: -1h)
      |> filter(fn: (r) => r._field == "latitude")
      |> filter(fn: (r) => r.latitude >= 51.0 and r.latitude <= 52.0)
      |> group(columns: ["mmsi"])
      |> distinct(column: "mmsi")
      |> count()
    '''
    start = time.time()
    for _ in range(ITERATIONS):
        result = query_api.query(q3)
        for table in result:
            for record in table.records:
                pass
    queries["bbox_count_1h"] = (time.time() - start) / ITERATIONS

    client.close()
    return queries

print("=== TimescaleDB Benchmarks ===")
tsdb_results = benchmark_tsdb()
for query, timing in tsdb_results.items():
    print(f"  {query}: {timing*1000:.1f}ms")

print("\n=== InfluxDB Benchmarks ===")
influx_results = benchmark_influx()
for query, timing in influx_results.items():
    print(f"  {query}: {timing*1000:.1f}ms")

# ingest_benchmark.py — Measure how fast each database ingests 1M AIS rows

import time
import json
from datetime import datetime

# Load sample AIS data (1M rows)
with open("ais_data_1m.json", "r") as f:
    data = json.load(f)

print("=== Ingestion Speed: 1M AIS Records ===\n")

# --- TimescaleDB ---
import psycopg2
from psycopg2.extras import execute_values

conn = psycopg2.connect(host="localhost", dbname="aisdb", user="postgres")
cur = conn.cursor()

# Build rows as tuples in the same format as the table schema
rows = []
for report in data:
    ts = datetime.strptime(report["TIME"], "%Y-%m-%d %H:%M:%S GMT")
    rows.append((
        report["MMSI"], ts,
        report["LATITUDE"], report["LONGITUDE"],
        report.get("SOG"), report.get("COG"),
        report.get("HEADING"), report.get("NAVSTAT"),
        report.get("TYPE"), report.get("STATION_ID"),
        report.get("LENGTH"),
    ))

# Time the batch insert with page_size=1000 (good balance of speed and memory)
start = time.time()
execute_values(cur, """
    INSERT INTO ais_positions
    (mmsi, time, latitude, longitude, sog, cog, heading, navstat, ship_type, station_id, length)
    VALUES %s
""", rows, page_size=1000)
conn.commit()
tsdb_time = time.time() - start
print(f"TimescaleDB: {tsdb_time:.2f}s ({1_000_000/tsdb_time:.0f} rows/sec)")

cur.close()
conn.close()

# --- InfluxDB ---
from influxdb_client import InfluxDBClient, Point, WritePrecision
from influxdb_client.client.write_api import SYNCHRONOUS

client = InfluxDBClient(url="http://localhost:8086", token="token", org="ais-org")
# BATCHING mode for production throughput (SYNCHRONOUS would be too slow for 1M rows)
# BATCHING mode: accumulates points in a buffer, flushes every 5000 points or 1 second
# This gives much higher throughput by reducing HTTP round-trips
from influxdb_client.client.write_api import WriteApi
write_api = client.write_api(write_options=WriteApi.BATCHING)

points = []
for report in data:
    ts = datetime.strptime(report["TIME"], "%Y-%m-%d %H:%M:%S GMT")
    point = Point("ais_position") \
        .tag("mmsi", str(report["MMSI"])) \
        .tag("navstat", str(report.get("NAVSTAT", "0"))) \
        .tag("ship_type", str(report.get("TYPE", "0"))) \
        .field("latitude", float(report["LATITUDE"])) \
        .field("longitude", float(report["LONGITUDE"])) \
        .field("sog", float(report.get("SOG", 0))) \
        .field("cog", float(report.get("COG", 0))) \
        .field("heading", int(report.get("HEADING", 0))) \
        .field("length", int(report.get("LENGTH", 0))) \
        .time(ts, WritePrecision.S)
    points.append(point)

start = time.time()
write_api.write(bucket="ais_positions", record=points)
# close() will flush any remaining buffered points and close the connection
# Without close(), batched points might still be in the buffer
client.close()
influx_time = time.time() - start
print(f"InfluxDB (BATCHING): {influx_time:.2f}s ({1_000_000/influx_time:.0f} rows/sec)")

print()
print(f"InfluxDB is {tsdb_time/influx_time:.1f}x faster at ingestion")