#!/usr/bin/python2.5
 
# Copyright (C) 2007 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
 
import bisect
import cStringIO as StringIO
import datetime
import itertools
import os
try:
  import sqlite3 as sqlite
except ImportError:
  from pysqlite2 import dbapi2 as sqlite
import tempfile
import time
import warnings
# Objects in a schedule (Route, Trip, etc) should not keep a strong reference
# to the Schedule object to avoid a reference cycle. Schedule needs to use
# __del__ to cleanup its temporary file. The garbage collector can't handle
# reference cycles containing objects with custom cleanup code.
import weakref
import zipfile
 
import gtfsfactory
import problems as problems_module
from transitfeed.util import defaultdict
import util
 
class Schedule:
  """Represents a Schedule, a collection of stops, routes, trips and
  an agency.  This is the main class for this module."""
 
  def __init__(self, problem_reporter=None,
               memory_db=True, check_duplicate_trips=False,
               gtfs_factory=None):
    if gtfs_factory is None:
      gtfs_factory = gtfsfactory.GetGtfsFactory()
    self._gtfs_factory = gtfs_factory
 
    # Map from table name to list of columns present in this schedule
    self._table_columns = {}
 
    self._agencies = {}
    self.stops = {}
    self.routes = {}
    self.trips = {}
    self.service_periods = {}
    self.fares = {}
    self.fare_zones = {}  # represents the set of all known fare zones
    self._shapes = {}  # shape_id to Shape
    # A map from transfer._ID() to a list of transfers. A list is used so
    # there can be more than one transfer with each ID. Once GTFS explicitly
    # prohibits duplicate IDs this might be changed to a simple dict of
    # Transfers.
    self._transfers = defaultdict(lambda: [])
    self._default_service_period = None
    self._default_agency = None
    if problem_reporter is None:
      self.problem_reporter = problems_module.default_problem_reporter
    else:
      self.problem_reporter = problem_reporter
    self._check_duplicate_trips = check_duplicate_trips
    self.ConnectDb(memory_db)
 
  def AddTableColumn(self, table, column):
    """Add column to table if it is not already there."""
    if column not in self._table_columns[table]:
      self._table_columns[table].append(column)
 
  def AddTableColumns(self, table, columns):
    """Add columns to table if they are not already there.
 
    Args:
      table: table name as a string
      columns: an iterable of column names"""
    table_columns = self._table_columns.setdefault(table, [])
    for attr in columns:
      if attr not in table_columns:
        table_columns.append(attr)
 
  def GetTableColumns(self, table):
    """Return list of columns in a table."""
    return self._table_columns[table]
 
  def __del__(self):
    self._connection.cursor().close()
    self._connection.close()
    if hasattr(self, '_temp_db_filename'):
      os.remove(self._temp_db_filename)
 
  def ConnectDb(self, memory_db):
    if memory_db:
      self._connection = sqlite.connect(":memory:")
    else:
      try:
        self._temp_db_file = tempfile.NamedTemporaryFile()
        self._connection = sqlite.connect(self._temp_db_file.name)
      except sqlite.OperationalError:
        # Windows won't let a file be opened twice. mkstemp does not remove the
        # file when all handles to it are closed.
        self._temp_db_file = None
        (fd, self._temp_db_filename) = tempfile.mkstemp(".db")
        os.close(fd)
        self._connection = sqlite.connect(self._temp_db_filename)
 
    cursor = self._connection.cursor()
    cursor.execute("""CREATE TABLE stop_times (
                                           trip_id CHAR(50),
                                           arrival_secs INTEGER,
                                           departure_secs INTEGER,
                                           stop_id CHAR(50),
                                           stop_sequence INTEGER,
                                           stop_headsign VAR CHAR(100),
                                           pickup_type INTEGER,
                                           drop_off_type INTEGER,
                                           shape_dist_traveled FLOAT);""")
    cursor.execute("""CREATE INDEX trip_index ON stop_times (trip_id);""")
    cursor.execute("""CREATE INDEX stop_index ON stop_times (stop_id);""")
 
  def GetStopBoundingBox(self):
    return (min(s.stop_lat for s in self.stops.values()),
            min(s.stop_lon for s in self.stops.values()),
            max(s.stop_lat for s in self.stops.values()),
            max(s.stop_lon for s in self.stops.values()),
           )
 
  def AddAgency(self, name, url, timezone, agency_id=None):
    """Adds an agency to this schedule."""
    agency = self._gtfs_factory.Agency(name, url, timezone, agency_id)
    self.AddAgencyObject(agency)
    return agency
 
  def AddAgencyObject(self, agency, problem_reporter=None, validate=False):
    assert agency._schedule is None
 
    if not problem_reporter:
      problem_reporter = self.problem_reporter
 
    if agency.agency_id in self._agencies:
      problem_reporter.DuplicateID('agency_id', agency.agency_id)
      return
 
    self.AddTableColumns('agency', agency._ColumnNames())
    agency._schedule = weakref.proxy(self)
 
    if validate:
      agency.Validate(problem_reporter)
    self._agencies[agency.agency_id] = agency
 
  def GetAgency(self, agency_id):
    """Return Agency with agency_id or throw a KeyError"""
    return self._agencies[agency_id]
 
  def GetDefaultAgency(self):
    """Return the default Agency. If no default Agency has been set select the
    default depending on how many Agency objects are in the Schedule. If there
    are 0 make a new Agency the default, if there is 1 it becomes the default,
    if there is more than 1 then return None.
    """
    if not self._default_agency:
      if len(self._agencies) == 0:
        self.NewDefaultAgency()
      elif len(self._agencies) == 1:
        self._default_agency = self._agencies.values()[0]
    return self._default_agency
 
  def NewDefaultAgency(self, **kwargs):
    """Create a new Agency object and make it the default agency for this Schedule"""
    agency = self._gtfs_factory.Agency(**kwargs)
    if not agency.agency_id:
      agency.agency_id = util.FindUniqueId(self._agencies)
    self._default_agency = agency
    self.SetDefaultAgency(agency, validate=False)  # Blank agency won't validate
    return agency
 
  def SetDefaultAgency(self, agency, validate=True):
    """Make agency the default and add it to the schedule if not already added"""
    assert isinstance(agency, self._gtfs_factory.Agency)
    self._default_agency = agency
    if agency.agency_id not in self._agencies:
      self.AddAgencyObject(agency, validate=validate)
 
  def GetAgencyList(self):
    """Returns the list of Agency objects known to this Schedule."""
    return self._agencies.values()
 
  def GetServicePeriod(self, service_id):
    """Returns the ServicePeriod object with the given ID."""
    return self.service_periods[service_id]
 
  def GetDefaultServicePeriod(self):
    """Return the default ServicePeriod. If no default ServicePeriod has been
    set select the default depending on how many ServicePeriod objects are in
    the Schedule. If there are 0 make a new ServicePeriod the default, if there
    is 1 it becomes the default, if there is more than 1 then return None.
    """
    if not self._default_service_period:
      if len(self.service_periods) == 0:
        self.NewDefaultServicePeriod()
      elif len(self.service_periods) == 1:
        self._default_service_period = self.service_periods.values()[0]
    return self._default_service_period
 
  def NewDefaultServicePeriod(self):
    """Create a new ServicePeriod object, make it the default service period and
    return it. The default service period is used when you create a trip without
    providing an explict service period. """
    service_period = self._gtfs_factory.ServicePeriod()
    service_period.service_id = util.FindUniqueId(self.service_periods)
    # blank service won't validate in AddServicePeriodObject
    self.SetDefaultServicePeriod(service_period, validate=False)
    return service_period
 
  def SetDefaultServicePeriod(self, service_period, validate=True):
    assert isinstance(service_period, self._gtfs_factory.ServicePeriod)
    self._default_service_period = service_period
    if service_period.service_id not in self.service_periods:
      self.AddServicePeriodObject(service_period, validate=validate)
 
  def AddServicePeriodObject(self, service_period, problem_reporter=None,
                             validate=True):
    if not problem_reporter:
      problem_reporter = self.problem_reporter
 
    if service_period.service_id in self.service_periods:
      problem_reporter.DuplicateID('service_id', service_period.service_id)
      return
 
    if validate:
      service_period.Validate(problem_reporter)
    self.service_periods[service_period.service_id] = service_period
 
  def GetServicePeriodList(self):
    return self.service_periods.values()
 
  def GetDateRange(self):
    """Returns a tuple of (earliest, latest) dates on which the service
    periods in the schedule define service, in YYYYMMDD form."""
 
    ranges = [period.GetDateRange() for period in self.GetServicePeriodList()]
    starts = filter(lambda x: x, [item[0] for item in ranges])
    ends = filter(lambda x: x, [item[1] for item in ranges])
 
    if not starts or not ends:
      return (None, None)
 
    return (min(starts), max(ends))
 
  def GetServicePeriodsActiveEachDate(self, date_start, date_end):
    """Return a list of tuples (date, [period1, period2, ...]).
 
    For each date in the range [date_start, date_end) make list of each
    ServicePeriod object which is active.
 
    Args:
      date_start: The first date in the list, a date object
      date_end: The first date after the list, a date object
 
    Returns:
      A list of tuples. Each tuple contains a date object and a list of zero or
      more ServicePeriod objects.
    """
    date_it = date_start
    one_day = datetime.timedelta(days=1)
    date_service_period_list = []
    while date_it < date_end:
      periods_today = []
      date_it_string = date_it.strftime("%Y%m%d")
      for service in self.GetServicePeriodList():
        if service.IsActiveOn(date_it_string, date_it):
          periods_today.append(service)
      date_service_period_list.append((date_it, periods_today))
      date_it += one_day
    return date_service_period_list
 
 
  def AddStop(self, lat, lng, name, stop_id=None):
    """Add a stop to this schedule.
 
    Args:
      lat: Latitude of the stop as a float or string
      lng: Longitude of the stop as a float or string
      name: Name of the stop, which will appear in the feed
      stop_id: stop_id of the stop or None, in which case a unique id is picked
 
    Returns:
      A new Stop object
    """
    if stop_id is None:
      stop_id = util.FindUniqueId(self.stops)
    stop = self._gtfs_factory.Stop(stop_id=stop_id, lat=lat, lng=lng, name=name)
    self.AddStopObject(stop)
    return stop
 
  def AddStopObject(self, stop, problem_reporter=None):
    """Add Stop object to this schedule if stop_id is non-blank."""
    assert stop._schedule is None
    if not problem_reporter:
      problem_reporter = self.problem_reporter
 
    if not stop.stop_id:
      return
 
    if stop.stop_id in self.stops:
      problem_reporter.DuplicateID('stop_id', stop.stop_id)
      return
 
    stop._schedule = weakref.proxy(self)
    self.AddTableColumns('stops', stop._ColumnNames())
    self.stops[stop.stop_id] = stop
    if hasattr(stop, 'zone_id') and stop.zone_id:
      self.fare_zones[stop.zone_id] = True
 
  def GetStopList(self):
    return self.stops.values()
 
  def AddRoute(self, short_name, long_name, route_type, route_id=None):
    """Add a route to this schedule.
 
    Args:
      short_name: Short name of the route, such as "71L"
      long_name: Full name of the route, such as "NW 21st Ave/St Helens Rd"
      route_type: A type such as "Tram", "Subway" or "Bus"
      route_id: id of the route or None, in which case a unique id is picked
    Returns:
      A new Route object
    """
    if route_id is None:
      route_id = util.FindUniqueId(self.routes)
    route = self._gtfs_factory.Route(short_name=short_name, long_name=long_name,
                        route_type=route_type, route_id=route_id)
    route.agency_id = self.GetDefaultAgency().agency_id
    self.AddRouteObject(route)
    return route
 
  def AddRouteObject(self, route, problem_reporter=None):
    if not problem_reporter:
      problem_reporter = self.problem_reporter
 
    if route.route_id in self.routes:
      problem_reporter.DuplicateID('route_id', route.route_id)
      return
 
    if route.agency_id not in self._agencies:
      if not route.agency_id and len(self._agencies) == 1:
        # we'll just assume that the route applies to the only agency
        pass
      else:
        problem_reporter.InvalidValue('agency_id', route.agency_id,
                                      'Route uses an unknown agency_id.')
        return
 
    self.AddTableColumns('routes', route._ColumnNames())
    route._schedule = weakref.proxy(self)
    self.routes[route.route_id] = route
 
  def GetRouteList(self):
    return self.routes.values()
 
  def GetRoute(self, route_id):
    return self.routes[route_id]
 
  def AddShapeObject(self, shape, problem_reporter=None):
    if not problem_reporter:
      problem_reporter = self.problem_reporter
 
    shape.Validate(problem_reporter)
 
    if shape.shape_id in self._shapes:
      problem_reporter.DuplicateID('shape_id', shape.shape_id)
      return
 
    self._shapes[shape.shape_id] = shape
 
  def GetShapeList(self):
    return self._shapes.values()
 
  def GetShape(self, shape_id):
    return self._shapes[shape_id]
 
  def AddTripObject(self, trip, problem_reporter=None, validate=False):
    if not problem_reporter:
      problem_reporter = self.problem_reporter
 
    if trip.trip_id in self.trips:
      problem_reporter.DuplicateID('trip_id', trip.trip_id)
      return
 
    self.AddTableColumns('trips', trip._ColumnNames())
    trip._schedule = weakref.proxy(self)
    self.trips[trip.trip_id] = trip
 
    # Call Trip.Validate after setting trip._schedule so that references
    # are checked. trip.ValidateChildren will be called directly