#!/usr/bin/env python3
"""
AirTag Trajectory Tracking & Filtering Pipeline (Adaptive Rate-Normalized Model)
with GeoJSON Output Formatting

Designed to handle real-world crowdsourced BLE location tracking (Apple Find My / AirTag):
  1. Stage 1: Robust Outlier Filtering via Hampel Geomedian & Kinematic Gating.
  2. Stage 2: Adaptive Rate-Normalized Two-Tier Stay-Point Tracker:
     - Dynamic Window Sizing (N_trailing): Automatically scales trailing comparison
       depth between 3 and 25 observations based on real-time arrival rate (λ).
     - Sample-Normalized Shift Thresholds: Scales distance boundary inversely with
       sample size (D_shift(n) = D_base + α/sqrt(n)) to cover typical ~35m BLE jitter.
     - Rate-Conditioned Hysteresis: Restricts time-based overrides to dead zones (<6/h).
  3. Stage 3: GeoJSON FeatureCollection Generation:
     - STAY nodes: Emitted as Point geometries (geomedian centroid) with SimpleStyle stars.
     - TRANSIT nodes: Emitted as LineString geometries containing all observation coordinates
       recorded during that movement segment.
"""

import argparse
import datetime
import json
import math
import sys
from typing import Dict, List, Optional, Tuple


def haversine(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
  """Calculates distance between two GPS coordinates in meters."""
  r_earth = 6371000.0  # Earth's radius in meters
  phi1 = math.radians(lat1)
  phi2 = math.radians(lat2)
  delta_phi = math.radians(lat2 - lat1)
  delta_lambda = math.radians(lon2 - lon1)

  a = (math.sin(delta_phi / 2.0) ** 2 +
       math.cos(phi1) * math.cos(phi2) * math.sin(delta_lambda / 2.0) ** 2)
  c = 2.0 * math.atan2(math.sqrt(a), math.sqrt(1.0 - a))
  return r_earth * c


def geomedian(points: List[Tuple[float, float]], max_iter: int = 15) -> Tuple[float, float]:
  """
  Computes the Geometric Median of a set of (latitude, longitude) coordinates
  using Weiszfeld's algorithm with Haversine distance weighting.
  """
  if not points:
    return (0.0, 0.0)
  if len(points) == 1:
    return points[0]

  lats = sorted(p[0] for p in points)
  lons = sorted(p[1] for p in points)
  mid = len(points) // 2
  cur_lat = lats[mid] if len(points) % 2 != 0 else (lats[mid - 1] + lats[mid]) / 2.0
  cur_lon = lons[mid] if len(points) % 2 != 0 else (lons[mid - 1] + lons[mid]) / 2.0

  for _ in range(max_iter):
    num_lat, num_lon, denom = 0.0, 0.0, 0.0
    for lat, lon in points:
      dist = haversine(cur_lat, cur_lon, lat, lon)
      if dist < 1e-5:
        dist = 1e-5
      weight = 1.0 / dist
      num_lat += lat * weight
      num_lon += lon * weight
      denom += weight

    next_lat = num_lat / denom
    next_lon = num_lon / denom

    if haversine(cur_lat, cur_lon, next_lat, next_lon) < 0.001:
      break
    cur_lat, cur_lon = next_lat, next_lon

  return (cur_lat, cur_lon)


def parse_timestamp(ts_str: str) -> datetime.datetime:
  """Parses ISO timestamp string into timezone-aware datetime."""
  if ts_str.endswith('Z'):
    ts_str = ts_str[:-1] + '+00:00'
  return datetime.datetime.fromisoformat(ts_str)


class Observation:
  """Represents a clean, timestamped GPS observation."""
  def __init__(self, raw: Dict):
    self.id = raw.get('id', 0)
    self.lat = float(raw['latitude'])
    self.lon = float(raw['longitude'])
    self.time = parse_timestamp(raw.get('fixTime') or raw.get('deviceTime') or raw.get('serverTime'))
    self.hdop = float(raw.get('attributes', {}).get('hdop', 0.0))

  @property
  def coords(self) -> Tuple[float, float]:
    return (self.lat, self.lon)


def stage1_filter_outliers(
    obs_list: List[Observation],
    window_radius: int = 4,
    max_outlier_dist_m: float = 85.0,
    max_speed_mps: float = 40.0) -> List[Observation]:
  """
  Stage 1: Outlier Rejection via Hampel Geomedian Filter & Velocity Gating.
  Removes extreme cellular tower fallbacks and multi-path reflections (>85m spikes).
  """
  if not obs_list:
    return []

  n = len(obs_list)
  spatial_inliers = []
  rejected_count = 0

  # 1. Hampel Geomedian Filter
  for i in range(n):
    w_start = max(0, i - window_radius)
    w_end = min(n, i + window_radius + 1)
    neighbors = [o.coords for o in obs_list[w_start:w_end]]
    
    local_centroid = geomedian(neighbors)
    dist_from_median = haversine(obs_list[i].lat, obs_list[i].lon,
                                 local_centroid[0], local_centroid[1])

    if dist_from_median <= max_outlier_dist_m:
      spatial_inliers.append(obs_list[i])
    else:
      rejected_count += 1

  # 2. Kinematic Velocity Gate
  clean_obs = []
  if spatial_inliers:
    clean_obs.append(spatial_inliers[0])
    for i in range(1, len(spatial_inliers)):
      curr = spatial_inliers[i]
      prev = clean_obs[-1]
      delta_s = max(1.0, (curr.time - prev.time).total_seconds())
      dist_m = haversine(prev.lat, prev.lon, curr.lat, curr.lon)
      speed = dist_m / delta_s

      if speed > max_speed_mps and delta_s < 600:
        rejected_count += 1
      else:
        clean_obs.append(curr)

  print(f"--- Stage 1 Filter: Kept {len(clean_obs)} / {n} observations "
        f"({rejected_count} severe noise outliers rejected) ---")
  return clean_obs


class ClusterSegment:
  """Represents a consolidated trajectory waypoint or movement segment."""
  def __init__(self, observations: List[Observation], min_stay_sec: float = 300.0, max_stay_spread_m: float = 70.0):
    self.observations = observations
    self.count = len(observations)
    self.start_time = observations[0].time
    self.end_time = observations[-1].time
    self.duration_sec = (self.end_time - self.start_time).total_seconds()
    self.centroid = geomedian([o.coords for o in observations])
    
    distances = [haversine(o.lat, o.lon, *self.centroid) for o in observations]
    self.max_spread_m = max(distances) if distances else 0.0
    self.avg_spread_m = sum(distances) / len(distances) if distances else 0.0

    dur_hours = max(0.016, self.duration_sec / 3600.0)
    self.rate_per_hr = round(self.count / dur_hours, 1)

    if self.duration_sec >= min_stay_sec and self.avg_spread_m <= max_stay_spread_m and self.max_spread_m <= 180.0:
      self.seg_type = 'STAY'
    else:
      self.seg_type = 'TRANSIT'

  def to_dict(self) -> Dict:
    return {
        'type': self.seg_type,
        'start_time': self.start_time.isoformat(),
        'end_time': self.end_time.isoformat(),
        'duration_min': round(self.duration_sec / 60.0, 1),
        'observation_count': self.count,
        'obs_rate_per_hr': self.rate_per_hr,
        'latitude': round(self.centroid[0], 7),
        'longitude': round(self.centroid[1], 7),
        'avg_spread_meters': round(self.avg_spread_m, 1),
    }

  def to_geojson_feature(self, index: int) -> Dict:
    """
    Converts the segment into an RFC 7946 GeoJSON Feature:
      - STAY nodes: Emits Point geometry located at the geomedian centroid.
      - TRANSIT nodes: Emits LineString geometry containing all observation coordinates in the node.
    """
    is_stay = (self.seg_type == 'STAY')

    if is_stay or len(self.observations) < 2:
      geometry = {
          'type': 'Point',
          'coordinates': [round(self.centroid[1], 7), round(self.centroid[0], 7)]
      }
    else:
      geometry = {
          'type': 'LineString',
          'coordinates': [[round(o.lon, 7), round(o.lat, 7)] for o in self.observations]
      }

    properties = {
        'segment_index': index,
        'segment_type': self.seg_type,
        'start_time': self.start_time.isoformat(),
        'end_time': self.end_time.isoformat(),
        'duration_min': round(self.duration_sec / 60.0, 1),
        'observation_count': self.count,
        'obs_rate_per_hr': self.rate_per_hr,
        'avg_spread_meters': round(self.avg_spread_m, 1),
    }

    if is_stay:
      properties.update({
          'marker-color': '#d9534f',
          'marker-size': 'medium',
          'marker-symbol': 'star'
      })
    else:
      properties.update({
          'stroke': '#007cbf',
          'stroke-width': 4,
          'stroke-opacity': 0.9
      })

    return {
        'type': 'Feature',
        'geometry': geometry,
        'properties': properties
    }


def to_geojson_collection(segments: List[ClusterSegment]) -> Dict:
  """Generates a comprehensive GeoJSON FeatureCollection including path geometry."""
  features = []

  # 1. Overall background trajectory path connecting all stay points and transit observations
  if len(segments) > 1:
    full_path_coords = []
    for s in segments:
      if s.seg_type == 'STAY':
        full_path_coords.append([round(s.centroid[1], 7), round(s.centroid[0], 7)])
      else:
        for o in s.observations:
          full_path_coords.append([round(o.lon, 7), round(o.lat, 7)])
    
    # De-duplicate identical consecutive timestamps/coords
    clean_path = []
    for coord in full_path_coords:
      if not clean_path or clean_path[-1] != coord:
        clean_path.append(coord)

    if len(clean_path) > 1:
      features.append({
          'type': 'Feature',
          'geometry': {
              'type': 'LineString',
              'coordinates': clean_path
          },
          'properties': {
              'title': 'Overall Tracked Trajectory Path',
              'stroke': '#888888',
              'stroke-width': 2,
              'stroke-opacity': 0.6,
              'total_segments': len(segments),
              'start_time': segments[0].start_time.isoformat(),
              'end_time': segments[-1].end_time.isoformat()
          }
      })

  # 2. Segment Features (Points for STAY, LineStrings for TRANSIT)
  for i, seg in enumerate(segments):
    features.append(seg.to_geojson_feature(i))

  return {
      'type': 'FeatureCollection',
      'features': features
  }


def calculate_dynamic_window_size(bucket: List[Observation], target_minutes: float = 30.0) -> int:
  if len(bucket) < 3:
    return 3
  
  dur_min = max(0.1, (bucket[-1].time - bucket[0].time).total_seconds() / 60.0)
  lam = len(bucket) / dur_min
  n_target = int(math.ceil(lam * target_minutes))
  return max(4, min(25, n_target))


def sample_normalized_shift_threshold(n_samples: int, d_base: float = 26.0, alpha: float = 24.0) -> float:
  """
  Adaptive Spatial Confidence Bounds.
  Scales drift threshold inversely with sqrt(N).
  """
  return d_base + (alpha / math.sqrt(max(1, n_samples)))


def stage2_rate_normalized_tracking(
    obs_list: List[Observation],
    macro_shift_m: float = 75.0,
    macro_hysteresis: int = 3,
    min_stay_sec: float = 300.0) -> List[ClusterSegment]:
  if not obs_list:
    return []

  segments = []
  current_bucket = [obs_list[0]]
  
  macro_count = 0
  macro_start_idx = None

  micro_count = 0
  micro_start_idx = None

  for i in range(1, len(obs_list)):
    curr = obs_list[i]
    current_bucket.append(curr)

    dur_hr = max(0.01, (current_bucket[-1].time - current_bucket[0].time).total_seconds() / 3600.0)
    obs_rate_per_hr = len(current_bucket) / dur_hr

    # 1. Macro Shift Check
    if len(current_bucket) >= macro_hysteresis + 3:
      base_limit_macro = macro_start_idx if macro_start_idx else -macro_hysteresis
      base_macro = [o.coords for o in current_bucket[:base_limit_macro]]
      if not base_macro:
        base_macro = [o.coords for o in current_bucket[:-macro_hysteresis]]
      
      c_base_macro = geomedian(base_macro[:15] if len(base_macro) > 15 else base_macro)
      trailing_macro = [o.coords for o in current_bucket[-macro_hysteresis:]]
      c_macro = geomedian(trailing_macro)
      
      if haversine(*c_base_macro, *c_macro) > macro_shift_m:
        if macro_count == 0:
          macro_start_idx = len(current_bucket) - macro_hysteresis
        macro_count += 1
        
        drift_sec = (curr.time - current_bucket[macro_start_idx].time).total_seconds()
        if macro_count >= macro_hysteresis or (obs_rate_per_hr < 10.0 and macro_count >= 2 and drift_sec >= 180.0):
          old_obs = current_bucket[:macro_start_idx]
          if old_obs:
            segments.append(ClusterSegment(old_obs, min_stay_sec))
          current_bucket = current_bucket[macro_start_idx:]
          macro_count, micro_count = 0, 0
          macro_start_idx, micro_start_idx = None, None
          continue
      else:
        macro_count = 0
        macro_start_idx = None

    # 2. Adaptive Micro Shift Check
    n_micro = calculate_dynamic_window_size(current_bucket, target_minutes=30.0)
    
    if len(current_bucket) >= n_micro + 4:
      base_limit_micro = micro_start_idx if micro_start_idx else -n_micro
      base_micro = [o.coords for o in current_bucket[:base_limit_micro]]
      if not base_micro:
        base_micro = [o.coords for o in current_bucket[:-n_micro]]
      
      c_base_micro = geomedian(base_micro[:25] if len(base_micro) > 25 else base_micro)
      trailing_micro = [o.coords for o in current_bucket[-n_micro:]]
      c_micro = geomedian(trailing_micro)
      
      thresh_micro = sample_normalized_shift_threshold(n_micro, d_base=26.0, alpha=24.0)
      
      if haversine(*c_base_micro, *c_micro) > thresh_micro:
        if micro_count == 0:
          micro_start_idx = len(current_bucket) - n_micro
        micro_count += 1
        
        drift_dur_sec = (curr.time - current_bucket[micro_start_idx].time).total_seconds()
        req_counts = max(10, min(15, int(n_micro * 0.75)))
        is_volume_confirmed = (micro_count >= req_counts)
        is_time_confirmed = (obs_rate_per_hr < 6.0 and micro_count >= 3 and drift_dur_sec >= 900.0)
        
        if is_volume_confirmed or is_time_confirmed:
          old_obs = current_bucket[:micro_start_idx]
          if old_obs:
            segments.append(ClusterSegment(old_obs, min_stay_sec))
          current_bucket = current_bucket[micro_start_idx:]
          macro_count, micro_count = 0, 0
          macro_start_idx, micro_start_idx = None, None
          continue
      else:
        micro_count = 0
        micro_start_idx = None

  if current_bucket:
    segments.append(ClusterSegment(current_bucket, min_stay_sec))

  return segments


def format_summary_table(segments: List[ClusterSegment]):
  """Prints a clean human-readable table of the tracked trajectory."""
  print("\n=== EXTRACTED AIRTAG TRAJECTORY TIMELINE (Adaptive Rate-Normalized Model) ===")
  print(f"{'TYPE':<8} | {'START TIME (UTC)':<19} -> {'END TIME (UTC)':<17} | {'DURATION':<10} | {'OBS':<5} | {'RATE (obs/h)':<12} | {'LATITUDE':<10} | {'LONGITUDE':<11} | {'AVG SPREAD':<10} | {'SHIFT / NOTES'}")
  print("-" * 144)

  prev_centroid = None
  for seg in segments:
    t_start = seg.start_time.strftime("%Y-%m-%d %H:%M:%S")
    t_end = seg.end_time.strftime("%H:%M:%S")
    
    if seg.duration_sec >= 3600:
      dur_str = f"{seg.duration_sec / 3600.0:.1f} hr"
    else:
      dur_str = f"{seg.duration_sec / 60.0:.1f} min"

    shift_str = "-"
    if prev_centroid is not None:
      dist_from_prev = haversine(*prev_centroid, *seg.centroid)
      shift_str = f"Moved {dist_from_prev:.1f}m"
    prev_centroid = seg.centroid

    rate_str = f"{seg.rate_per_hr:5.1f}/h"
    print(f"{seg.seg_type:<8} | {t_start:<19} -> {t_end:<17} | {dur_str:<10} | {seg.count:<5} | {rate_str:<12} | {seg.centroid[0]:<10.6f} | {seg.centroid[1]:<11.6f} | {seg.avg_spread_m:4.1f}m      | {shift_str}")
  print("-" * 144)


def main():
  parser = argparse.ArgumentParser(description="AirTag Rate-Normalized Trajectory Tracker (GeoJSON Output)")
  parser.add_argument("input_file", help="Path to positions.json")
  parser.add_argument("--output_file", help="Path to save trajectory output (defaults to GeoJSON if ending in .geojson or .json)", default="trajectory_output.geojson")
  parser.add_argument("--format", choices=["geojson", "json"], default="geojson", help="Output serialization structure")
  args = parser.parse_args()

  with open(args.input_file, 'r') as f:
    raw_data = json.load(f)

  print(f"Loaded {len(raw_data)} raw location records from {args.input_file}.")
  
  observations = [Observation(r) for r in raw_data if 'latitude' in r and 'longitude' in r]
  observations.sort(key=lambda x: x.time)
  print(f"Time span: {observations[0].time} to {observations[-1].time}")

  # Stage 1: Filter outliers
  clean_obs = stage1_filter_outliers(observations, window_radius=4, max_outlier_dist_m=85.0)

  # Stage 2: Segment via Rate-Normalized Two-Tier Tracking
  segments = stage2_rate_normalized_tracking(clean_obs, macro_shift_m=75.0, min_stay_sec=300.0)

  format_summary_table(segments)

  # Stage 3: Save output
  if args.format == "geojson" or args.output_file.endswith(".geojson"):
    out_data = to_geojson_collection(segments)
    format_name = "GeoJSON FeatureCollection"
  else:
    out_data = [s.to_dict() for s in segments]
    format_name = "Standard JSON Array"

  with open(args.output_file, 'w') as f:
    json.dump(out_data, f, indent=2)
  print(f"\nSaved {format_name} trajectory to: {args.output_file}\n")


if __name__ == '__main__':
  main()