try:
    import traci
except ImportError:
    traci = None


class SUMODemandCalculator:
    """
    Calculates real-time traffic demand for movements and phases
    directly from the running SUMO simulation.

    The calculator uses the SUMO vehicle's current lane and route
    to determine which junction movement the vehicle belongs to.

    Movement matching priority:

        1. Exact (from_lane, next_route_edge) match
        2. Unique (from_edge, to_edge) match
        3. Otherwise ignore the vehicle as ambiguous

    This avoids assigning vehicles to movements when their route
    information is insufficient to determine the correct movement.
    """

    # ============================================================
    # MOVEMENT DEMAND
    # ============================================================

    def calculate_movement_demand(
        self,
        topology,
        movement_metadata=None
    ):
        """
        Calculate current vehicle demand for every allowed movement.

        Parameters
        ----------
        topology:
            JunctionTopology containing movement definitions.

        movement_metadata:
            Optional metadata generated by SUMO topology adapter.

            Expected format:

            {
                "A0_M0": {
                    "junction_id": "A0",
                    "link_index": 0,
                    "from_lane": "...",
                    "to_lane": "...",
                    "via_lane": "..."
                },
                ...
            }

        Returns
        -------
        dict

            Example:

            {
                "A0_M0": 3,
                "A0_M1": 5,
                "A0_M2": 0,
                ...
            }
        """

        # --------------------------------------------------------
        # Initialize every allowed movement with zero demand.
        # --------------------------------------------------------

        demand = {}

        allowed_movements = (
            topology.get_allowed_movements()
        )

        for movement in allowed_movements:

            demand[movement.movement_id] = 0

        # --------------------------------------------------------
        # Build lookup structures.
        # --------------------------------------------------------

        lane_lookup = {}

        edge_lookup = {}

        for movement in allowed_movements:

            movement_id = (
                movement.movement_id
            )

            metadata = None

            if movement_metadata is not None:

                metadata = movement_metadata.get(
                    movement_id
                )

            # ----------------------------------------------------
            # Lane-based lookup.
            # ----------------------------------------------------

            if metadata is not None:

                from_lane = metadata.get(
                    "from_lane"
                )

                if from_lane:

                    key = (
                        from_lane,
                    )

                    if key not in lane_lookup:

                        lane_lookup[key] = []

                    lane_lookup[key].append(
                        movement_id
                    )

            # ----------------------------------------------------
            # Edge-based fallback lookup.
            # ----------------------------------------------------

            from_edge = (
                movement.from_road
            )

            to_edge = (
                movement.to_road
            )

            key = (
                from_edge,
                to_edge
            )

            if key not in edge_lookup:

                edge_lookup[key] = []

            edge_lookup[key].append(
                movement_id
            )

        # --------------------------------------------------------
        # Examine active vehicles.
        # --------------------------------------------------------

        vehicle_ids = (
            traci.vehicle.getIDList()
        )

        for vehicle_id in vehicle_ids:

            try:

                lane_id = (
                    traci.vehicle.getLaneID(
                        vehicle_id
                    )
                )

                if not lane_id:
                    continue

                # ------------------------------------------------
                # Ignore internal SUMO junction lanes.
                # ------------------------------------------------

                from_edge = (
                    traci.lane.getEdgeID(
                        lane_id
                    )
                )

                if not from_edge:
                    continue

                if from_edge.startswith(":"):
                    continue

                # ------------------------------------------------
                # Get the vehicle's route.
                # ------------------------------------------------

                route = (
                    traci.vehicle.getRoute(
                        vehicle_id
                    )
                )

                if not route:
                    continue

                route_index = (
                    traci.vehicle.getRouteIndex(
                        vehicle_id
                    )
                )

                # ------------------------------------------------
                # Determine the next edge.
                # ------------------------------------------------

                next_edge = None

                if (
                    route_index is not None
                    and
                    route_index + 1 < len(route)
                ):

                    next_edge = (
                        route[
                            route_index + 1
                        ]
                    )

                # ------------------------------------------------
                # Ignore internal next edges.
                # ------------------------------------------------

                if (
                    next_edge is not None
                    and
                    next_edge.startswith(":")
                ):

                    next_edge = None

                matched_movement = None

                # =================================================
                # METHOD 1
                # Exact lane + next-edge matching.
                # =================================================

                if next_edge is not None:

                    lane_candidates = []

                    for movement in allowed_movements:

                        movement_id = (
                            movement.movement_id
                        )

                        metadata = None

                        if movement_metadata is not None:

                            metadata = (
                                movement_metadata.get(
                                    movement_id
                                )
                            )

                        if metadata is None:
                            continue

                        movement_from_lane = (
                            metadata.get(
                                "from_lane"
                            )
                        )

                        movement_to_lane = (
                            metadata.get(
                                "to_lane"
                            )
                        )

                        if (
                            movement_from_lane
                            == lane_id
                        ):

                            # Determine the edge reached
                            # by this movement.
                            movement_to_edge = (
                                movement.to_road
                            )

                            if (
                                movement_to_edge
                                == next_edge
                            ):

                                lane_candidates.append(
                                    movement_id
                                )

                            # Some SUMO networks expose
                            # a via/internal lane. In that
                            # case the to-lane may provide
                            # additional information but
                            # the destination edge remains
                            # the primary movement match.

                            elif (
                                movement_to_lane
                                == next_edge
                            ):

                                lane_candidates.append(
                                    movement_id
                                )

                    # Only accept an unambiguous match.

                    if len(lane_candidates) == 1:

                        matched_movement = (
                            lane_candidates[0]
                        )

                # =================================================
                # METHOD 2
                # Unique from-edge + to-edge matching.
                # =================================================

                if matched_movement is None:

                    if next_edge is not None:

                        candidates = (
                            edge_lookup.get(
                                (
                                    from_edge,
                                    next_edge
                                ),
                                []
                            )
                        )

                        # Only accept unique matches.
                        if len(candidates) == 1:

                            matched_movement = (
                                candidates[0]
                            )

                # =================================================
                # METHOD 3
                # No safe movement assignment.
                # =================================================

                if matched_movement is None:

                    continue

                # ------------------------------------------------
                # Increment movement demand.
                # ------------------------------------------------

                if matched_movement in demand:

                    demand[
                        matched_movement
                    ] += 1

            except (
                traci.TraCIException,
                KeyError,
                IndexError,
                TypeError,
                ValueError
            ):

                # A vehicle may disappear between TraCI calls.
                # Ignore that vehicle rather than stopping the
                # complete controller.
                continue

            except Exception:

                # Keep the demand layer robust against malformed
                # vehicle state.
                continue

        return demand

    # ============================================================
    # PHASE DEMAND
    # ============================================================

    def calculate_phase_demand(
        self,
        phases,
        movement_demand
    ):
        """
        Calculate demand for each signal phase.

        Parameters
        ----------
        phases:
            Iterable of phase objects.

            A phase is expected to contain either:

                phase_id
                movement_ids

            or equivalent attributes.

        movement_demand:
            Dictionary returned by calculate_movement_demand().

        Returns
        -------
        dict

            {
                phase_id: total_demand
            }
        """

        phase_demand = {}

        for phase in phases:

            phase_id = (
                phase.phase_id
            )

            movement_ids = (
                getattr(
                    phase,
                    "movement_ids",
                    []
                )
            )

            total = 0

            for movement_id in movement_ids:

                total += movement_demand.get(
                    movement_id,
                    0
                )

            phase_demand[
                phase_id
            ] = total

        return phase_demand

    # ============================================================
    # NORMALIZE PHASE DEMAND
    # ============================================================

    def normalize_phase_demand(
        self,
        phase_demand
    ):
        """
        Normalize phase demand to [0, 1].

        The phase with the highest demand receives 1.0.

        If every phase has zero demand, every normalized phase
        receives 0.0.
        """

        if not phase_demand:

            return {}

        maximum = max(
            phase_demand.values()
        )

        if maximum <= 0:

            return {
                phase_id: 0.0
                for phase_id in phase_demand
            }

        normalized = {}

        for phase_id, demand in (
            phase_demand.items()
        ):

            value = (
                float(demand)
                / float(maximum)
            )

            # Safety clamp.

            value = max(
                0.0,
                min(
                    1.0,
                    value
                )
            )

            normalized[
                phase_id
            ] = value

        return normalized


# ================================================================
# SIMPLE TEST
# ================================================================

def main():

    print()
    print("=" * 70)
    print("SUMO DEMAND CALCULATOR TEST")
    print("=" * 70)

    print()
    print(
        "This module is intended to be imported by the"
    )
    print(
        "integration controller after SUMO is running."
    )
    print()

    print(
        "Available class:"
    )

    print(
        "  SUMODemandCalculator"
    )

    print()

    print(
        "Methods:"
    )

    print(
        "  calculate_movement_demand()"
    )

    print(
        "  calculate_phase_demand()"
    )

    print(
        "  normalize_phase_demand()"
    )

    print()


if __name__ == "__main__":

    main()