#!/usr/bin/env python3
"""
Build GTFS database for aNavigator bus mode.
Reads TPER Bologna GTFS files, produces compact JSON for iOS app.
Now correctly groups trips into route variants by matching service patterns.
"""
import csv
import json
import sys
import os
from collections import OrderedDict, defaultdict


def _simplify_shape(coords, max_points):
    """Uniform sampling to reduce shape to max_points."""
    if len(coords) <= max_points:
        return coords
    step = len(coords) / max_points
    result = []
    for i in range(max_points):
        idx = min(int(i * step), len(coords) - 1)
        result.append(coords[idx])
    if result[-1] != coords[-1]:
        result[-1] = coords[-1]
    return result


GTFS_DIR = sys.argv[1] if len(sys.argv) > 1 else "/tmp/gtfs_tper"
OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/home/alina/aNavigator/src/gtfs_db.json"

print(f"Reading GTFS from {GTFS_DIR}...")

# ─── 1. Load routes ───
routes = {}
with open(os.path.join(GTFS_DIR, "routes.txt"), encoding="utf-8") as f:
    for row in csv.DictReader(f):
        routes[row["route_id"]] = row

# ─── 2. Load shapes ───
print("Loading shapes (this may take a moment)...")
shapes = defaultdict(list)
with open(os.path.join(GTFS_DIR, "shapes.txt"), encoding="utf-8") as f:
    for row in csv.DictReader(f):
        sid = row["shape_id"]
        shapes[sid].append([
            round(float(row["shape_pt_lat"]), 6),
            round(float(row["shape_pt_lon"]), 6)
        ])

print(f"  Loaded {len(shapes)} shapes")

# ─── 3. Load stops ───
stops = {}
with open(os.path.join(GTFS_DIR, "stops.txt"), encoding="utf-8") as f:
    for row in csv.DictReader(f):
        stops[row["stop_id"]] = {
            "name": row["stop_name"],
            "lat": round(float(row["stop_lat"]), 6),
            "lon": round(float(row["stop_lon"]), 6)
        }

print(f"  Loaded {len(stops)} stops")

# ─── 4. Index stop_times by trip_id ───
print("Indexing stop_times...")
stop_times_by_trip = defaultdict(list)
with open(os.path.join(GTFS_DIR, "stop_times.txt"), encoding="utf-8") as f:
    for row in csv.DictReader(f):
        # Arrival time in compact format (HH:MM)
        at = row.get("arrival_time", "")
        dt = row.get("departure_time", "")
        stop_times_by_trip[row["trip_id"]].append({
            "stop_id": row["stop_id"],
            "seq": int(row["stop_sequence"]),
            "arrival": at[:5] if len(at) >= 5 else at,
            "departure": dt[:5] if len(dt) >= 5 else dt
        })

print(f"  Indexed {len(stop_times_by_trip)} trips")

# ─── 5. Load calendar for service patterns ───
print("Loading calendar...")
calendar = {}
with open(os.path.join(GTFS_DIR, "calendar.txt"), encoding="utf-8") as f:
    for row in csv.DictReader(f):
        days = []
        day_names = ['monday','tuesday','wednesday','thursday','friday','saturday','sunday']
        for d in day_names:
            if row.get(d) == '1':
                days.append(d[:3])
        calendar[row["service_id"]] = ",".join(days) if days else ""

print(f"  Loaded {len(calendar)} calendar entries")

# ─── 6. Load trips ───
print("Loading trips...")
trips_by_route = defaultdict(list)
with open(os.path.join(GTFS_DIR, "trips.txt"), encoding="utf-8") as f:
    for row in csv.DictReader(f):
        trips_by_route[row["route_id"]].append(row)

# ─── 7. Build database with variant grouping ───
print("Building database (grouping by service pattern)...")
db = {"lines": OrderedDict()}

# Sort by route_short_name numeric then alpha
sorted_routes = sorted(routes.values(), key=lambda r: (
    int(r["route_short_name"]) if r["route_short_name"].isdigit() else 999,
    r["route_short_name"]
))

variant_counter = 0

for route in sorted_routes:
    short_name = route["route_short_name"].strip()
    route_id = route["route_id"]
    long_name = route["route_long_name"].strip()

    trips = trips_by_route.get(route_id, [])
    if not trips:
        continue

    # Group trips by shape_id + service pattern into distinct variants
    # A "variant" is a pair: dir=0 trip + dir=1 trip that share a service pattern
    
    # First, group all trips by their shape_id
    shape_groups = defaultdict(lambda: {"dir0": [], "dir1": []})
    for t in trips:
        d = t["direction_id"]
        sid = t["shape_id"]
        if d == "0":
            shape_groups[sid]["dir0"].append(t)
        else:
            shape_groups[sid]["dir1"].append(t)

    # For each shape, determine which day pattern it serves
    shape_service_map = {}
    for shape_id, grp in shape_groups.items():
        all_trips = grp["dir0"] + grp["dir1"]
        if all_trips:
            sample_svc = all_trips[0]["service_id"]
            day_pattern = calendar.get(sample_svc, "unknown")
            shape_service_map[shape_id] = day_pattern

    # Now intelligently pair dir=0 and dir=1 shapes that share day patterns
    # First pass: pair matching service patterns
    used_shapes = set()
    variants = []
    
    # Group shapes by their day patterns
    weekday_shapes = {s for s, p in shape_service_map.items() if 'mon' in p or 'tue' in p or 'wed' in p or 'thu' in p or 'fri' in p}
    weekend_shapes = {s for s, p in shape_service_map.items() if 'sat' in p or 'sun' in p}
    all_day_shapes = {s for s, p in shape_service_map.items() if 'mon' in p and 'sat' in p and 'sun' in p}
    
    # Sort shapes to ensure deterministic results
    for dir0_shape in sorted(shape_groups.keys()):
        if not shape_groups[dir0_shape]["dir0"]:
            continue
        if dir0_shape in used_shapes:
            continue
            
        dir0_trip = shape_groups[dir0_shape]["dir0"][0]
        dir0_pattern = shape_service_map.get(dir0_shape, "unknown")
        
        # Find matching dir=1 shape with same service pattern
        best_dir1_shape = None
        for dir1_shape in sorted(shape_groups.keys()):
            if not shape_groups[dir1_shape]["dir1"]:
                continue
            if dir1_shape in used_shapes:
                continue
            dir1_pattern = shape_service_map.get(dir1_shape, "unknown")
            # Match if same pattern or both weekday/weekend
            if dir1_pattern == dir0_pattern:
                best_dir1_shape = dir1_shape
                break
        
        # If no exact match, try weekday/weekend grouping
        if not best_dir1_shape:
            is_weekend = 'sat' in dir0_pattern or 'sun' in dir0_pattern
            for dir1_shape in sorted(shape_groups.keys()):
                if not shape_groups[dir1_shape]["dir1"]:
                    continue
                if dir1_shape in used_shapes:
                    continue
                dir1_pattern = shape_service_map.get(dir1_shape, "unknown")
                dir1_is_weekend = 'sat' in dir1_pattern or 'sun' in dir1_pattern
                if is_weekend == dir1_is_weekend:
                    best_dir1_shape = dir1_shape
                    break
        
        if not best_dir1_shape:
            continue
        
        dir1_trip = shape_groups[best_dir1_shape]["dir1"][0]
        
        used_shapes.add(dir0_shape)
        used_shapes.add(best_dir1_shape)
        variants.append((dir0_trip, dir1_trip))

    # For shapes without a pair, add them as solo trips
    for shape_id in sorted(shape_groups.keys()):
        if shape_id in used_shapes:
            continue
        grp = shape_groups[shape_id]
        if grp["dir0"] and not grp["dir1"]:
            # Single direction trip
            dir0_trip = grp["dir0"][0]
            # If there's an unpaired dir=1 shape, pair them
            for other_shape in sorted(shape_groups.keys()):
                if other_shape in used_shapes:
                    continue
                if shape_groups[other_shape]["dir1"] and not shape_groups[other_shape]["dir0"]:
                    used_shapes.add(other_shape)
                    used_shapes.add(shape_id)
                    variants.append((dir0_trip, shape_groups[other_shape]["dir1"][0]))
                    break
            else:
                used_shapes.add(shape_id)
                variants.append((dir0_trip, None))
        elif grp["dir1"] and not grp["dir0"]:
            dir1_trip = grp["dir1"][0]
            used_shapes.add(shape_id)
            variants.append((None, dir1_trip))

    # Create line entries for each variant
    for idx, (dir0_trip, dir1_trip) in enumerate(variants):
        line_data = {"name": long_name, "trips": []}
        
        for t, d in [(dir0_trip, 0), (dir1_trip, 1)]:
            if t is None:
                continue
            
            sid = t["shape_id"]
            
            # Get shape
            shape_coords = shapes.get(sid, [])
            if len(shape_coords) < 2:
                continue
            
            # Simplify shape
            shape_simple = _simplify_shape(shape_coords, 150)
            
            # Get stops for this trip
            trip_stops = stop_times_by_trip.get(t["trip_id"], [])
            trip_stops.sort(key=lambda x: x["seq"])
            
            stop_info = []
            for s in trip_stops:
                sid_stop = s["stop_id"]
                if sid_stop in stops:
                    entry = dict(stops[sid_stop])
                    if s.get("arrival"):
                        entry["arrival"] = s["arrival"]
                    if s.get("departure"):
                        entry["departure"] = s["departure"]
                    stop_info.append(entry)
            
            if not stop_info:
                continue
            
            # Headsign
            headsign = t.get("trip_headsign", "").strip()
            if not headsign and stop_info:
                headsign = stop_info[0]["name"] + " → " + stop_info[-1]["name"]
            
            # Determina il giorno di servizio
            trip_shape_id = t["shape_id"]
            day_pattern = shape_service_map.get(trip_shape_id, "")
            if day_pattern == "mon,tue,wed,thu,fri":
                service_label = "Feriale"
            elif day_pattern == "sat":
                service_label = "Sabato"
            elif day_pattern == "sun":
                service_label = "Domenica"
            elif "sat" in day_pattern and "sun" in day_pattern:
                service_label = "Sabato e Domenica"
            elif "sat" in day_pattern:
                service_label = "Sabato"
            elif "sun" in day_pattern:
                service_label = "Domenica"
            elif "mon" in day_pattern:
                service_label = "Feriale"
            else:
                service_label = ""
            
            line_data["trips"].append({
                "dir": d,
                "headsign": headsign,
                "service": service_label,
                "shape": shape_simple,
                "stops": stop_info
            })
        
        if not line_data["trips"]:
            continue
        
        # Determine variant suffix based on service pattern
        pattern = shape_service_map.get(dir0_trip["shape_id"] if dir0_trip else (dir1_trip["shape_id"] if dir1_trip else ""), "")
        
        if len(variants) == 1:
            # Single variant - use short_name as-is
            key = short_name
        else:
            # Multiple variants - add suffix
            if idx == 0:
                # First variant gets no suffix if it's the default
                key = short_name
            else:
                # Subsequent variants get _variant suffix
                key = f"{short_name}_v{idx}"
        
        # Avoid duplicate keys
        if key in db["lines"]:
            # Append to existing (some routes might have same key)
            # Actually, merge trips
            for tt in line_data["trips"]:
                if tt["dir"] not in [x["dir"] for x in db["lines"][key]["trips"]]:
                    db["lines"][key]["trips"].append(tt)
                else:
                    key = f"{short_name}_v{idx}"
                    db["lines"][key] = line_data
        else:
            db["lines"][key] = line_data

print(f"  Built {len(db['lines'])} line entries")

# ─── 8. Write output ───
with open(OUTPUT, "w", encoding="utf-8") as f:
    json.dump(db, f, ensure_ascii=False)

size_kb = os.path.getsize(OUTPUT) / 1024
print(f"Written to {OUTPUT} ({size_kb:.0f} KB)")
print("Done!")