#cython: language_level=3
#cython: boundscheck=False
#cython: wraparound=False
#cython: c_string_type=unicode, c_string_encoding=utf8

cimport cython
from libcpp cimport bool
from libcpp.vector cimport vector
from libcpp.string cimport string
from libcpp.pair cimport pair
from libc.stdint cimport int64_t

import numpy as np
cimport numpy as np
np.import_array()


cdef extern from "accessibility.h" namespace "MTC::accessibility":
    cdef cppclass Accessibility:
        Accessibility(int64_t, vector[vector[int64_t]], vector[vector[double]], bool) except +
        vector[string] aggregations
        vector[string] decays
        void initializeCategory(double, int64_t, string, vector[int64_t])
        pair[vector[vector[double]], vector[vector[int64_t]]] findAllNearestPOIs(
            float, int64_t, string, int64_t)
        void initializeAccVar(string, vector[int64_t], vector[double])
        vector[double] getAllAggregateAccessibilityVariables(
            float, string, string, string, int64_t)
        vector[int64_t] Route(int64_t, int64_t, int64_t)
        vector[vector[int64_t]] Routes(vector[int64_t], vector[int64_t], int64_t)
        double Distance(int64_t, int64_t, int64_t)
        vector[double] Distances(vector[int64_t], vector[int64_t], int64_t)
        vector[vector[pair[int64_t, float]]] Range(vector[int64_t], float, int64_t, vector[int64_t])
        void precomputeRangeQueries(double)


cdef np.ndarray[double] convert_vector_to_array_dbl(vector[double] vec):
    cdef int n = vec.size()
    cdef np.ndarray[double] arr = np.zeros(n, dtype=np.float64)
    cdef int i
    for i in range(n):
        arr[i] = vec[i]
    return arr


cdef np.ndarray[double, ndim=2] convert_2D_vector_to_array_dbl(
        vector[vector[double]] vec):
    if vec.size() == 0:
        return np.empty((0, 0), dtype=np.float64)
    
    cdef int rows = vec.size()
    cdef int cols = vec[0].size() if rows > 0 else 0
    cdef np.ndarray[double, ndim=2] arr = np.empty((rows, cols), dtype=np.float64)
    
    cdef int i, j
    for i in range(rows):
        if vec[i].size() != cols:
            raise ValueError(f"Inconsistent row sizes: expected {cols}, got {vec[i].size()} at row {i}")
        for j in range(cols):
            arr[i, j] = vec[i][j]
    return arr


cdef np.ndarray[int64_t, ndim=2] convert_2D_vector_to_array_int(
        vector[vector[int64_t]] vec):
    if vec.size() == 0:
        return np.empty((0, 0), dtype=np.int64)
    
    cdef int rows = vec.size()
    cdef int cols = vec[0].size() if rows > 0 else 0
    cdef np.ndarray[int64_t, ndim=2] arr = np.empty((rows, cols), dtype=np.int64)
    
    cdef int i, j
    for i in range(rows):
        if vec[i].size() != cols:
            raise ValueError(f"Inconsistent row sizes: expected {cols}, got {vec[i].size()} at row {i}")
        for j in range(cols):
            arr[i, j] = <int64_t>vec[i][j]
    return arr


cdef class cyaccess:
    cdef Accessibility * access

    def __cinit__(
        self,
        np.ndarray[int64_t] node_ids,
        np.ndarray[double, ndim=2] node_xys,
        np.ndarray[int64_t, ndim=2] edges,
        np.ndarray[double, ndim=2] edge_weights,
        bool twoway=True
    ):
        """
        node_ids: vector of node identifiers
        node_xys: the spatial locations of the same nodes
        edges: a pair of node ids which comprise each edge
        edge_weights: the weights (impedances) that apply to each edge
        twoway: whether the edges should all be two-way or whether they
            are directed from the first to the second node
        """
        self.access = new Accessibility(len(node_ids), edges, edge_weights, twoway)
        if self.access == NULL:
            raise MemoryError("Failed to allocate Accessibility object")

    def __dealloc__(self):
        if self.access != NULL:
            del self.access

    def initialize_category(
        self,
        double maxdist,
        int64_t maxitems,
        string category,
        np.ndarray[int64_t] node_ids
    ):
        """
        maxdist - the maximum distance that will later be used in
            find_all_nearest_pois
        maxitems - the maximum number of items that will later be requested
            in find_all_nearest_pois
        category - the category name
        node_ids - an array of nodeids which are locations where this poi occurs
        """
        
        self.access.initializeCategory(maxdist, maxitems, category, node_ids)

    def find_all_nearest_pois(
        self,
        double radius,
        int64_t num_of_pois,
        string category,
        int64_t impno=0
    ):
        """
        radius - search radius
        num_of_pois - number of pois to search for
        category - the category name
        impno - the impedance id to use
        return_nodeids - whether to return the nodeid locations of the nearest
            not just the distances
        """
        ret = self.access.findAllNearestPOIs(radius, num_of_pois, category, impno)

        return convert_2D_vector_to_array_dbl(ret.first),\
            convert_2D_vector_to_array_int(ret.second)

    def initialize_access_var(
        self,
        string category,
        np.ndarray[int64_t] node_ids,
        np.ndarray[double] values
    ):
        """
        category - category name
        node_ids: vector of node identifiers
        values: vector of values that are location at the nodes
        """
        self.access.initializeAccVar(category, node_ids, values)

    def get_available_aggregations(self):
        return self.access.aggregations

    def get_available_decays(self):
        return self.access.decays

    def get_all_aggregate_accessibility_variables(
        self,
        double radius,
        category,
        aggtyp,
        decay,
        int64_t impno=0,
    ):
        """
        radius - search radius
        category - category name
        aggtyp - aggregation type, see docs
        decay - decay type, see docs
        impno - the impedance id to use
        """
        ret = self.access.getAllAggregateAccessibilityVariables(
            radius, category, aggtyp, decay, impno)

        return convert_vector_to_array_dbl(ret)

    def shortest_path(self, int64_t srcnode, int64_t destnode, int64_t impno=0):
        """
        srcnode - node id origin
        destnode - node id destination
        impno - the impedance id to use
        """
        return self.access.Route(srcnode, destnode, impno)

    def shortest_paths(self, np.ndarray[int64_t] srcnodes, 
            np.ndarray[int64_t] destnodes, int64_t impno=0):
        """
        srcnodes - node ids of origins
        destnodes - node ids of destinations
        impno - impedance id
        """
        return self.access.Routes(srcnodes, destnodes, impno)

    def shortest_path_distance(self, int64_t srcnode, int64_t destnode, int64_t impno=0):
        """
        srcnode - node id origin
        destnode - node id destination
        impno - the impedance id to use
        """
        return self.access.Distance(srcnode, destnode, impno)

    def shortest_path_distances(self, np.ndarray[int64_t] srcnodes, 
            np.ndarray[int64_t] destnodes, int64_t impno=0):
        """
        srcnodes - node ids of origins
        destnodes - node ids of destinations
        impno - impedance id
        """
        return self.access.Distances(srcnodes, destnodes, impno)
    
    def precompute_range(self, double radius):
        self.access.precomputeRangeQueries(radius)

    def nodes_in_range(self, vector[int64_t] srcnodes, float radius, int64_t impno, 
            np.ndarray[int64_t] ext_ids):
        """
        srcnodes - node ids of origins
        radius - maximum range in which to search for nearby nodes
        impno - the impedance id to use
        ext_ids - all node ids in the network
        """
        return self.access.Range(srcnodes, radius, impno, ext_ids)