"""
.. _example_gtfs_import_and_asssignment:

Public transport assignment with Optimal Strategies
===========

In this example, we import a GTFS feed to our model, create a public transport network, create project match connectors, and perform a Spiess & Florian assignment.

We use data from Coquimbo, a city in La Serena Metropolitan Area in Chile.
"""
# %%
# Imports for example construction
from uuid import uuid4
from os import remove
from os.path import join
from tempfile import gettempdir

from aequilibrae.paths import TransitAssignment, TransitClass
from aequilibrae.utils.create_example import create_example
import numpy as np

# Imports for GTFS import
from aequilibrae.transit import Transit

# Imports for SF transit graph construction
from aequilibrae.project.database_connection import database_connection
from aequilibrae.transit.transit_graph_builder import TransitGraphBuilder


# %%
# Let's create an empty project on an arbitrary folder.
fldr = join(gettempdir(), uuid4().hex)

project = create_example(fldr, "coquimbo")

# %%
# As the Coquimbo example already has a complete GTFS model, we shall remove its public transport
# database for the sake of this example.
remove(join(fldr, "public_transport.sqlite"))

# %%
# Let's import the GTFS feed.
dest_path = join(fldr, "gtfs_coquimbo.zip")


# %%
# Now we create our Transit object and import the GTFS feed into our model.
# This will automatically create a new public transport database.
data = Transit(project)

transit = data.new_gtfs_builder(agency="LISANCO", file_path=dest_path)

# %%
# To load the data, we must choose one date. We're going to continue with 2016-04-13 but feel free
# to experiment with any other available dates. Transit class has a function allowing you to check
# dates for the GTFS feed. It should take approximately 2 minutes to load the data.
transit.load_date("2016-04-13")

# %%
# Let's save this model for later use.
transit.save_to_disk()

# %%
# Graph building
# --------------
# Let's build the transit network. We'll disable `outer_stop_transfers` and `walking_edges` because Coquimbo doesn't have any parent stations.
# For the OD connections we'll use the `overlapping_regions` method and create some accurate line geometry later.
# Creating the graph should only take a moment. By default zoning information is pulled from the project network. If you have your own zoning information add it using `graph.add_zones(zones)` then `graph.create_graph()`. We drop gemoetry here for the sake of display
graph = data.create_graph(with_outer_stop_transfers=False, with_walking_edges=False, blocking_centroid_flows=False, connector_method="overlapping_regions")

# %%
graph.vertices.drop(columns="geometry")

# %%
graph.edges

# %%
# The graphs also also stored in the `Transit.graphs` dictionary. They are keyed by the `period_id` they were created for.
# A graph for a different `period_id` can be created by providing `period_id=` in the `Transit.create_graph` call. You can view previously created periods with the `Periods` object.
periods = project.network.periods
periods.data

# %%
# Connector project matching
# ~~~~~~~~~~~~~~~~~~~~~~~~~~
project.network.build_graphs()

# %%
# Now we'll create the line strings for the access connectors, this step is optinal but provides more accurate distance estimations and better looking geometry. Because Coquimbo doesn't have many walking edges we'll match onto the `"c"` graph.
graph.create_line_geometry(method="connector project match", graph="c")

# %%
# Saving and reloading
# ~~~~~~~~~~~~~~~~~~~~~~~~~~
# Lets save all graphs to the `public_transport.sqlite` database.
data.save_graphs()

# %%
# We can reload the saved graphs with `data.load`. This will create new `TransitGraphBuilder`s based on the `period_id` of the saved graphs.
# The graph configuration is stored in the `transit_graph_config` table in `project_database.sqlite` as serialised JSON.
data.load()

# %%
# Links and nodes are stored in a similar manner to the `project_database.sqlite` database.

# %%
# Reading back into AequilibraE
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# You can create back in a particular graph via it's `period_id`.
pt_con = database_connection("transit")
graph_db = TransitGraphBuilder.from_db(pt_con, periods.default_period.period_id)
graph_db.vertices.drop(columns="geometry")

# %%
graph_db.edges

# %%
# Converting to a AequilibraE graph object
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# To perform an assignment we need to convert the graph builder into a graph.
transit_graph = graph.to_transit_graph()

# %%
# Spiess & Florian assignment
# ---------------------------

# %%
# Mock demand matrix
# ~~~~~~~~~~~~~~~~~~
# We'll create a mock demand matrix with demand `1` for every zone.
# We'll also need to convert from `zone_id`s to `node_id`s.
from aequilibrae.matrix import AequilibraeMatrix

zones_in_the_model = len(transit_graph.centroids)

names_list = ['pt']

mat = AequilibraeMatrix()
mat.create_empty(zones=zones_in_the_model,
                 matrix_names=names_list,
                 memory_only=True)
mat.index = transit_graph.centroids[:]
mat.matrices[:, :, 0] = np.full((zones_in_the_model, zones_in_the_model), 1.0)
mat.computational_view()

# %%
# Hyperpath generation/assignment
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# We'll create a `TransitAssignment` object as well as a `TransitClass`
assig = TransitAssignment()

# Create the assignment class
assigclass = TransitClass(name="pt", graph=transit_graph, matrix=mat)
assig.add_class(assigclass)

# %%
# We need to tell AequilbraE where to find the appropriate fields we want to use, as well as the assignment algorithm to use.
assig.set_time_field("trav_time")
assig.set_frequency_field("freq")

assig.set_algorithm("os")

# %%
# When there's multiple matrix cores we'll also need to set the core to use for the demand.
assigclass.set_demand_matrix_core("pt")

# %%
# Lets perform the assignment with the mock demand matrx for all `TransitClass`s added.
assig.execute()

# %%
# View the results
assig.results()

# %%
# We can also access the `TransitAssignmentResults` object from the `TransitClass`
assigclass.results

# %%
# Saving results
# ~~~~~~~~~~~~~~
# We'll be saving the results to another sqlite db called `results_database.sqlite`. The `results` table with `project_database.sqlite` contains some metadata about each table in `results_database.sqlite`.
assig.save_results(table_name='hyperpath example')
# %%
# Wrapping up
project.close()