#!/usr/local/bin/python3.12
# Eclipse SUMO, Simulation of Urban MObility; see https://eclipse.org/sumo
# Copyright (C) 2019-2019 German Aerospace Center (DLR) and others.
# This program and the accompanying materials
# are made available under the terms of the Eclipse Public License v2.0
# which accompanies this distribution, and is available at
# http://www.eclipse.org/legal/epl-v20.html
# SPDX-License-Identifier: EPL-2.0

# @file    personGenerator.py
# @author  tarek chouaki
# @date    2019-03-22
# @version $Id$

"""
This tool allows to generate flows of persons for a SUMO simulation which is currently not possible in SUMO route files.
It does so by converting an xml file (usually having the ``.pflow.xml`` extension) to a sumo route file
containing the generated <peron> elements.

Here is an example ``.pflow.xml`` :

.. code-block:: xml

    <routes>
        <personRoute id="route-1">
            <walk from="e1" busStop="1" />
            <probability>
                <probabilityItem probability="0.5">
                    <ride busStop="2" modes="public" />
                    <probability>
                        <probabilityItem probability="0.5">
                            <stop busStop="2" duration="10" />
                        </probabilityItem>
                        <probabilityItem probability="0.5" />
                    </probability>
                </probabilityItem>
                <probabilityItem probability="0.5">
                    <ride busStop="3" modes="public">
                </probabilityItem>
            </probability>
        </personRoute>
        <personFlow id="forward" begin="0" end="3600" number="7" perWave="10" departPos="0" route="forward" />

        <personFlow id="backward" begin="0" end="3600" period="600" perWave="10" departPos="0">
            <walk from="e3" busStop="3" />
            <ride busStop="1" modes="public"/>
            <stop busStop="1" duration="50"/>
        </personFlow>
    </routes>

The example above allows to generate two flows of persons :

-The first flow consists of persons taking a bus from stop 1 to either stop 2 or stop 3 (with a 50% chance for each).
The persons of this flow are spawned in 7 waves (equally separated in time) and each wave consists of 10 persons.
For the persons going to bus stop 2, there's a 50% chance they'll stay there during 10 ticks.
The route followed by the persons of this flow is defined separately in a ``<personRoute>`` element and referenced by its ID.
-The second flow consists of persons taking a bus from stop 3 to stop 1 and then stopping there for 50 ticks.
The persons of this flow are spawned in periodic waves with 10 persons pere wave.
The route followed by the persons is defined directly under the ``<personFlow>``

How to Use
----------

Via Command Line
~~~~~~~~~~~~~~~~

This script can be accessed directly by command line passing an input `.pflow.xml`` file's path
and an output ``.rou.xml`` file's path.

.. code-block:: bash

    python personGenerator.py pedestrians.pflow.xml pedestrians.rou.xml

Note that the output file is overwritten without asking for permission.

In your script
~~~~~~~~~~~~~~
You can import the classes and methods in this module and use them in your own python script.

See the documentation below for more details.
"""

from lxml import etree
import argparse
import random


class PersonGenerationElement(object):
    """
    This class serves as a base for person generation elements
    """

    def __init__(self, xml_element):
        self.xml_element = xml_element
        if self.xml_element.tag != self.get_xml_tag():
            raise Exception("Bad tag")

    @classmethod
    def get_xml_tag(cls):
        """
        This class method is meant to be implemented by subclasses
        It should return the xml tag for elements described by the current class
        """
        raise NotImplementedError

    def generate(self):
        """
        This method is meant to be implemented by subclasses
        It should return a list of elements generated by the element
        """
        raise NotImplementedError

    @classmethod
    def wrap_elements(cls, elements, *args, **kwargs):
        """
        Replaces xml elements with the appropriate tag (the one defined in get_xml_tag)
        with an object of the current class.
        The given list is modified, be careful
        :param elements: a list of xml elements
        :type elements: list
        """
        for i in range(len(elements)):
            if not isinstance(elements[i], PersonGenerationElement) and elements[i].tag == cls.get_xml_tag():
                elements[i] = cls(elements[i], *args, **kwargs)

    @staticmethod
    def generate_multiple(elements):
        """
        Loops over a list containing xml elements and PersonGenerationElement objects.
        The PersonGenerationElement objects are replaced by elements generated from them
        The given list is not modified
        :param elements: A list containing xml elements and PersonGenerationElement objects.
        :type elements: list
        :return: a list of resulting xml elements
        :rtype list
        """
        result = list()
        for element in elements:
            if isinstance(element, PersonGenerationElement):
                result.extend(element.generate())
            else:
                result.append(element.__copy__())
        return result


class ProbabilityElement(PersonGenerationElement):
    """
    This class describes probability elements that are used to generate alternatives with given probabilities.
    In XML it looks like:

    .. code-block:: xml

        <probability>
            <probabilityItem probability="0.5">fist alternative</probabilityItem>
            <probabilityItem probability="0.5">second alternative</probabilityItem>
        </probability>


    Each time the element is asked to generate,
    it returns the children of one of its alternatives according to the probabilities.

    Probability elements can be nested, so you can have:

    .. code-block:: xml

        <probability>
            <probabilityItem probability="0.5">
                <probability>
                    ...
                </probability>
                ...Possibly other stuff
            </probabilityItem>
            <probabilityItem probability="0.5">
                second alternative
            </probabilityItem>
        </probability>

    This allows you to define conditional probabilities.

    Note that the nested <probability> element should be a direct child of <probabilityItem>
    """

    def __init__(self, xml_element):
        """
        :param xml_element: The source xml element
        """
        super().__init__(xml_element)
        self.possibilities = []
        for sub_element in list(self.xml_element):
            if sub_element.tag != "probabilityItem":
                raise Exception("Only probabilityItem elements are allowed inside probability")
            try:
                proba = float(sub_element.get("probability"))
                if proba < 0 or proba > 1:
                    raise ValueError("")
                possibility = (proba, list(sub_element))
                ProbabilityElement.wrap_elements(possibility[1])
                self.possibilities.append(possibility)
            except (KeyError, ValueError):
                raise ValueError("probabilityItem element requires attribute probability between 0 and 1")

        if sum([child[0] for child in self.possibilities]) != 1:
            raise ValueError("Probabilities not summing up to 1 at line : " + str(self.xml_element.sourceline))

    @classmethod
    def get_xml_tag(cls):
        """
        :return: The tag of xml element coresponding to this class (probability)
        """
        return "probability"

    def generate(self):
        """
        :return: One of the alternatives according to the given probabilities
        """
        result = []
        cumulated_probability = 0
        p = random.random()
        for possibility in self.possibilities:
            cumulated_probability += float(possibility[0])
            if p <= cumulated_probability:
                result.extend(self.generate_multiple(possibility[1]))
                break
        return result


class PersonRouteElement(PersonGenerationElement):
    """
    This class describes xml elements that are used to define person routes separately.

    .. code-block:: xml

        <personRoute id="route">
            <walk />
            <stop />
            <ride />
        </personRoute>

    The content of the route is then copied to each person using it.

    You can use probabilities inside the **personRoute** element to have different alternatives.
    Basically, you can have:

    .. code-block:: xml

        <personRoute id="route">
            <walk from="edge1" busStop="1">
            <probability>
                <probabilityItem probability="0.5">
                    <ride busStop="2" modes="public" />
                </probabilityItem>
                <probabilityItem probability="0.5">
                    <ride busStop="3" modes="public" />
                </probabilityItem>
            </probability>
        </personRoute>
    """

    def __init__(self, xml_element):
        super().__init__(xml_element)
        self.id = self.xml_element.get("id")
        self.children = list(self.xml_element)
        ProbabilityElement.wrap_elements(self.children)

    @classmethod
    def get_xml_tag(cls):
        """
        :return: The tag of the xml elements corresponding to this class (personRoute)
        """
        return "personRoute"

    @staticmethod
    def get_route_by_id(routes, route_id):
        """
        :param routes:
        :type routes: collections.Iterable
        :param route_id:
        :type route_id: str
        :return: The PersonRouteElement object having the given id from the given iterable. None if not found
        """
        for route in routes:
            if isinstance(route, PersonRouteElement) and route.id == route_id:
                return route
        return None

    def generate(self):
        """
        :return: A copy of the sub elements of the original personRoute element
                probability elements are taken into account & used to generate an alternative
        """
        return self.generate_multiple(self.children)


class PersonFlowElement(PersonGenerationElement):
    """
    This class describes xml elements that are used to generate flows of persons as it is already possible for vehicles.
    For example, this xml code:

    .. code-block:: xml

        <personFlow id="flow" begin="0" end="3600" number="7" perWave="10">
            <walk />
            <ride />
            <stop />
        </personFlow>

    will generate person elements having the same children (walk, ride, stop).
    The generated persons will be in 7 waves each containing 10 persons.
    These waves will be equally separated in time between 0 and 3600
    The complete attributes list is:

    - id
    - begin : the time at which the flow starts
    - end : the time at which the flow ends. Not mandatory, default is 3600.
    - period : The time (in seconds) between two consecutive waves.
      Not mandatory, if not given, number will be used
    - number : the number of waves. Only meaningful when period is not specified
    - perWave : the number of persons in each wave. Not mandatory, default is 1
    - route : the id of the route that the persons will follow
      Not mandatory, if not given, uses the children of the <personFlow> element

    The id of generated persons will be `<id>_<person_index>` where `<person_index>` is the index
    of the person in the flow (starting from 0)
    """
    default_end = 3600
    id_attribute_key = "id"
    begin_attribute_key = "begin"
    end_attribute_key = "end"
    period_attribute_key = "period"
    number_attribute_key = "number"
    per_wave_attribute_key = "perWave"
    route_attribute_key = "route"

    def __init__(self, xml_element, routes):
        """
        :param xml_element: The xml element
        :param routes: An iterable where to look for routes
        :type routes: collections.Iterable
        """
        super().__init__(xml_element)
        self.routes = routes
        self.attributes = {item[0]: item[1] for item in self.xml_element.items()}
        self.children = list(self.xml_element)
        ProbabilityElement.wrap_elements(self.children)
        self.id = None
        self.begin = None
        self.period = None
        self.route = None
        # We check for the attributes that concern us & we leave the others
        try:
            self.id = self.attributes.pop(self.id_attribute_key)
        except KeyError:
            print("No id attribute in personFlow, quitting")
            exit(-1)
        try:
            self.begin = int(self.attributes.pop(self.begin_attribute_key))
        except KeyError:
            print("No begin in personFlow " + str(id) + ", quitting")
            exit(-1)
        try:
            self.end = int(self.attributes.pop(self.end_attribute_key))
        except KeyError:
            self.end = self.default_end

        try:
            self.period = int(self.attributes.pop(self.period_attribute_key))
        except KeyError:
            try:
                self.number = int(self.attributes.pop(self.number_attribute_key))
                if self.number == 1:
                    self.period = (self.end - self.begin) * 2 + 1
                else:
                    self.period = (self.end - self.begin) / (self.number - 1)
            except KeyError:
                print("Neither period nor number given for personFlow " + str(id) + ", quitting")
                exit(-1)

        try:
            self.per_wave = int(self.attributes.pop(self.per_wave_attribute_key))
        except KeyError:
            self.per_wave = 1
        try:
            route_id = self.attributes.pop(self.route_attribute_key)
            self.route = PersonRouteElement.get_route_by_id(routes, route_id)
            if self.route is None:
                raise Exception("Route with id " + route_id + " not found at line " + str(self.xml_element.sourceline))
        except KeyError:
            pass

    @classmethod
    def get_xml_tag(cls):
        """
        :return: The tag of the xml elements corresponding to the current class (personFlow)
        """
        return "personFlow"

    def generate(self):
        """
        :return: The persons of the flow
        """
        begin = self.begin
        p_id = 0
        elements = list()
        while begin <= self.end:
            for i in range(self.per_wave):
                element = etree.Element("person", self.attributes)
                element.set("depart", str(int(begin)))
                element.set("id", self.id + "_" + str(p_id))
                if self.route is not None:
                    element.extend(self.route.generate())
                else:
                    element.extend(self.generate_multiple(self.children))
                elements.append(element)
                p_id += 1
            begin += self.period
        return elements


def generate_persons(input_file, output_file):
    """
    Core method of the script, parses <personFlow> tags in an XML file and generates <person> elements.
    The generated <person> elements are sorted by their depart time.
    The original file is not modified and the result is written in another file.
    The resulting file will not contain the <personFlow> elements.
    Note that the output file is overwritten if it is already exist
    :param input_file: The path of the input file
    :param output_file: The path of the output file
    """
    # Parse the input file
    tree = etree.parse(input_file)

    routes = tree.getroot()
    children = list(routes)
    for child in children:
        routes.remove(child)
    PersonRouteElement.wrap_elements(children)
    person_routes = [child for child in children if isinstance(child, PersonRouteElement)]
    PersonFlowElement.wrap_elements(children, routes=person_routes)
    for person_route in person_routes:
        children.remove(person_route)
    person_elements = PersonGenerationElement.generate_multiple(children)
    person_elements.sort(key=lambda e: int(e.get('depart')))
    routes.extend(person_elements)

    with open(output_file, "w") as f:
        f.write(etree.tostring(routes).decode())
        f.close()


if __name__ == "__main__":
    # Parses the command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("source")
    parser.add_argument("destination")
    source, destination = parser.parse_args()
    generate_persons(source, destination)
