#include <boost/python.hpp>
#include <boost/python/wrapper.hpp>
#include <iostream>

#include <lanelet2_core/primitives/Lanelet.h>
#include <lanelet2_traffic_rules/TrafficRules.h>
#include <lanelet2_routing/RoutingCost.h>

namespace bp = boost::python;

class PyRoutingCostDistance : public lanelet::routing::RoutingCostDistance, public bp::wrapper<lanelet::routing::RoutingCostDistance> {
public:
    PyRoutingCostDistance(double laneChangeCost, double minLaneChangeLength=0.0)
        : lanelet::routing::RoutingCostDistance(laneChangeCost, minLaneChangeLength) {}

    virtual double getCostSucceeding(const lanelet::traffic_rules::TrafficRules& tr,
                                     const lanelet::ConstLaneletOrArea& src,
                                     const lanelet::ConstLaneletOrArea& dst) const override {
        if (bp::override f = this->get_override("getCostSucceeding")) {
            //std::cout << "Calling custom getCostSucceeding" << std::endl;
            return f(boost::ref(tr), src, dst);
        }
        //std::cout << "Returning parent getCostSucceeding" << std::endl;
        return lanelet::routing::RoutingCostDistance::getCostSucceeding(tr, src, dst);
    }

    double default_getCostSucceeding(const lanelet::traffic_rules::TrafficRules& tr,
                                     const lanelet::ConstLaneletOrArea& src,
                                     const lanelet::ConstLaneletOrArea& dst) const {
        //std::cout << "Calling default getCostSucceeding" << std::endl;
        return lanelet::routing::RoutingCostDistance::getCostSucceeding(tr, src, dst);
    }

    virtual double getCostLaneChange(const lanelet::traffic_rules::TrafficRules& tr,
                                     const lanelet::ConstLanelets& src,
                                     const lanelet::ConstLanelets& dst) const noexcept override {
        if (bp::override f = this->get_override("getCostLaneChange")) {
            //std::cout << "Calling custom getCostLaneChange" << std::endl;
            return f(boost::ref(tr), src, dst);
        }
        //std::cout << "Returning parent getCostLaneChange" << std::endl;
        return lanelet::routing::RoutingCostDistance::getCostLaneChange(tr, src, dst);
    }

    double default_getCostLaneChange(const lanelet::traffic_rules::TrafficRules& tr,
                                     const lanelet::ConstLanelets& src,
                                     const lanelet::ConstLanelets& dst) const {
        //std::cout << "Calling default getCostLaneChange" << std::endl;
        return lanelet::routing::RoutingCostDistance::getCostLaneChange(tr, src, dst);
    }
};

BOOST_PYTHON_MODULE(custom_routing_cost) {
    bp::class_<PyRoutingCostDistance, std::shared_ptr<PyRoutingCostDistance>, boost::noncopyable>(
            "RoutingCostDistance",
            bp::init<double, double>((bp::arg("laneChangeCost"), bp::arg("minLaneChangeLength")=0.0))
        )
        .def("getCostSucceeding", &PyRoutingCostDistance::default_getCostSucceeding)
        .def("getCostLaneChange", &PyRoutingCostDistance::default_getCostLaneChange)
    ;
}

from .custom_routing_cost import RoutingCostDistance

class BuslaneRoutingCost(RoutingCostDistance):
    def __init__(self, laneChangeCost, minLaneChangeLength=0.0):
        super(BuslaneRoutingCost, self).__init__(laneChangeCost, minLaneChangeLength)
    
    def getCostSucceeding(self, trafficRules, src, dst):
        cost = super().getCostSucceeding(trafficRules, src, dst)
        if ('subtype' in src.attributes and src.attributes['subtype'] == 'bus_lane') or \
            ('subtype' in dst.attributes and dst.attributes['subtype'] == 'bus_lane'):
            # double the cost for bus lanes to not drive into bus stops
            return cost * 2;
        else:
            return cost