import stardog
from multiprocessing import Pool
from tqdm import tqdm

cities = ['http://dbpedia.org/resource/Vancouver',
 'http://dbpedia.org/resource/Zurich' 'http://dbpedia.org/resource/Vienna',
 'http://dbpedia.org/resource/Bern' 'http://dbpedia.org/resource/Sydney',
 'http://dbpedia.org/resource/Geneva',
 'http://dbpedia.org/resource/Auckland',
 'http://dbpedia.org/resource/Copenhagen',
 'http://dbpedia.org/resource/Helsinki',
 'http://dbpedia.org/resource/Amsterdam',
 'http://dbpedia.org/resource/Frankfurt',
 'http://dbpedia.org/resource/Munich',
 'http://dbpedia.org/resource/Melbourne',
 'http://dbpedia.org/resource/Honolulu',
 'http://dbpedia.org/resource/Dusseldorf',
 'http://dbpedia.org/resource/San_Francisco',
 'http://dbpedia.org/resource/Oslo' 'http://dbpedia.org/resource/Perth',
 'http://dbpedia.org/resource/Brussels',
 'http://dbpedia.org/resource/Toronto',
 'http://dbpedia.org/resource/Adelaide',
 'http://dbpedia.org/resource/Luxembourg',
 'http://dbpedia.org/resource/Stockholm',
 'http://dbpedia.org/resource/Paris',
 'http://dbpedia.org/resource/Nurnberg',
 'http://dbpedia.org/resource/Brisbane',
 'http://dbpedia.org/resource/Montreal',
 'http://dbpedia.org/resource/Tokyo' 'http://dbpedia.org/resource/Berlin',
 'http://dbpedia.org/resource/Hamburg',
 'http://dbpedia.org/resource/Wellington',
 'http://dbpedia.org/resource/Ludwigshafen',
 'http://dbpedia.org/resource/Lyon' 'http://dbpedia.org/resource/Atlanta',
 'http://dbpedia.org/resource/London',
 'http://dbpedia.org/resource/Calgary' 'http://dbpedia.org/resource/Kobe',
 'http://dbpedia.org/resource/Lexington',
 'http://dbpedia.org/resource/Seattle',
 'http://dbpedia.org/resource/Dublin',
 'http://dbpedia.org/resource/Winston_Salem',
 'http://dbpedia.org/resource/Los_Angeles',
 'http://dbpedia.org/resource/Yokohama',
 'http://dbpedia.org/resource/Singapore',
 'http://dbpedia.org/resource/Washington_Dc',
 'http://dbpedia.org/resource/Pittsburgh',
 'http://dbpedia.org/resource/Chicago',
 'http://dbpedia.org/resource/Boston' 'http://dbpedia.org/resource/Madrid',
 'http://dbpedia.org/resource/New_York',
 'http://dbpedia.org/resource/Minneapolis',
 'http://dbpedia.org/resource/Portland',
 'http://dbpedia.org/resource/Osaka' 'http://dbpedia.org/resource/Nagoya',
 'http://dbpedia.org/resource/Glasgow',
 'http://dbpedia.org/resource/St._Louis',
 'http://dbpedia.org/resource/Houston',
 'http://dbpedia.org/resource/Tsukuba',
 'http://dbpedia.org/resource/Cleveland',
 'http://dbpedia.org/resource/Birmingham',
 'http://dbpedia.org/resource/Barcelona',
 'http://dbpedia.org/resource/Lisbon',
 'http://dbpedia.org/resource/St._Peter_Port',
 'http://dbpedia.org/resource/Miami' 'http://dbpedia.org/resource/Detroit',
 'http://dbpedia.org/resource/Yokkaichi',
 'http://dbpedia.org/resource/Milan' 'http://dbpedia.org/resource/Omuta',
 'http://dbpedia.org/resource/Katsuyama',
 'http://dbpedia.org/resource/Hong_Kong',
 'http://dbpedia.org/resource/Rome',
 'http://dbpedia.org/resource/Montevideo',
 'http://dbpedia.org/resource/San_Juan',
 'http://dbpedia.org/resource/Buenos_Aires',
 'http://dbpedia.org/resource/Budapest',
 'http://dbpedia.org/resource/Leipzig',
 'http://dbpedia.org/resource/Prague',
 'http://dbpedia.org/resource/Tel_Aviv',
 'http://dbpedia.org/resource/Jerusalem',
 'http://dbpedia.org/resource/Port_Louis',
 'http://dbpedia.org/resource/Taipei',
 'http://dbpedia.org/resource/Cape_Town',
 'http://dbpedia.org/resource/Kuala_Lumpur',
 'http://dbpedia.org/resource/Limassol',
 'http://dbpedia.org/resource/Athens',
 'http://dbpedia.org/resource/Ljubljana',
 'http://dbpedia.org/resource/Santiago',
 'http://dbpedia.org/resource/Warsaw',
 'http://dbpedia.org/resource/Johannesburg',
 'http://dbpedia.org/resource/Vilnius' 'http://dbpedia.org/resource/Riga',
 'http://dbpedia.org/resource/Panama_City',
 'http://dbpedia.org/resource/Seoul',
 'http://dbpedia.org/resource/Istanbul',
 'http://dbpedia.org/resource/Port_Elizabeth',
 'http://dbpedia.org/resource/San_Jose',
 'http://dbpedia.org/resource/Dubai' 'http://dbpedia.org/resource/Papeete',
 'http://dbpedia.org/resource/Tunis',
 'http://dbpedia.org/resource/Monterrey',
 'http://dbpedia.org/resource/Manila',
 'http://dbpedia.org/resource/Abu_Dhabi',
 'http://dbpedia.org/resource/Brasilia',
 'http://dbpedia.org/resource/Victoria',
 'http://dbpedia.org/resource/Tallinn',
 'http://dbpedia.org/resource/Noumea',
 'http://dbpedia.org/resource/Johor_Baharu',
 'http://dbpedia.org/resource/Sao_Paulo',
 'http://dbpedia.org/resource/Bucharest',
 'http://dbpedia.org/resource/Asuncion',
 'http://dbpedia.org/resource/Rio_De_Janeiro',
 'http://dbpedia.org/resource/Bangkok' 'http://dbpedia.org/resource/Lima',
 'http://dbpedia.org/resource/Santo_Domingo',
 'http://dbpedia.org/resource/Quito' 'http://dbpedia.org/resource/Harare',
 'http://dbpedia.org/resource/Bratislava',
 'http://dbpedia.org/resource/Colombo' 'http://dbpedia.org/resource/Ulsan',
 'http://dbpedia.org/resource/Sofia' 'http://dbpedia.org/resource/Cairo',
 'http://dbpedia.org/resource/Caracas',
 'http://dbpedia.org/resource/Guayaquil',
 'http://dbpedia.org/resource/Casablanca',
 'http://dbpedia.org/resource/Manama' 'http://dbpedia.org/resource/Zagreb',
 'http://dbpedia.org/resource/Bogota' 'http://dbpedia.org/resource/Nassau',
 'http://dbpedia.org/resource/Amman',
 'http://dbpedia.org/resource/Kuwait_City',
 'http://dbpedia.org/resource/Mexico_City',
 'http://dbpedia.org/resource/Kingston',
 'http://dbpedia.org/resource/Accra' 'http://dbpedia.org/resource/Manaus',
 'http://dbpedia.org/resource/Rayong',
 'http://dbpedia.org/resource/Shanghai',
 'http://dbpedia.org/resource/Guatemala_City',
 'http://dbpedia.org/resource/Blantyre',
 'http://dbpedia.org/resource/La_Paz',
 'http://dbpedia.org/resource/Beijing',
 'http://dbpedia.org/resource/Medellin',
 'http://dbpedia.org/resource/Nairobi',
 'http://dbpedia.org/resource/Jakarta',
 'http://dbpedia.org/resource/Guangzhou',
 'http://dbpedia.org/resource/Lusaka',
 'http://dbpedia.org/resource/Libreville',
 'http://dbpedia.org/resource/Abidjan',
 'http://dbpedia.org/resource/Santa_Cruz',
 'http://dbpedia.org/resource/Islamabad',
 'http://dbpedia.org/resource/Riyadh' 'http://dbpedia.org/resource/Jeddah',
 'http://dbpedia.org/resource/Dakar' 'http://dbpedia.org/resource/Nanjing',
 'http://dbpedia.org/resource/Ho_Chi_Minh_City',
 'http://dbpedia.org/resource/Vientiane',
 'http://dbpedia.org/resource/Beirut' 'http://dbpedia.org/resource/Kiev',
 'http://dbpedia.org/resource/Skopje',
 'http://dbpedia.org/resource/St._Petersburg',
 'http://dbpedia.org/resource/Moscow',
 'http://dbpedia.org/resource/Shenyang',
 'http://dbpedia.org/resource/Mumbai' 'http://dbpedia.org/resource/Lahore',
 'http://dbpedia.org/resource/Karachi',
 'http://dbpedia.org/resource/Damascus',
 'http://dbpedia.org/resource/Tehran',
 'http://dbpedia.org/resource/Managua',
 'http://dbpedia.org/resource/Bangalore',
 'http://dbpedia.org/resource/Djibouti',
 'http://dbpedia.org/resource/Hanoi',
 'http://dbpedia.org/resource/San_Salvador',
 'http://dbpedia.org/resource/Jilin' 'http://dbpedia.org/resource/Madras',
 'http://dbpedia.org/resource/Tirana' 'http://dbpedia.org/resource/Minsk',
 'http://dbpedia.org/resource/Yangon',
 'http://dbpedia.org/resource/Yaounde',
 'http://dbpedia.org/resource/Cotonou',
 'http://dbpedia.org/resource/Novosibirsk',
 'http://dbpedia.org/resource/Tripoli',
 'http://dbpedia.org/resource/New_Delhi',
 'http://dbpedia.org/resource/Kazan',
 'http://dbpedia.org/resource/Port_Au_Prince',
 'http://dbpedia.org/resource/Dacca' 'http://dbpedia.org/resource/Algiers',
 'http://dbpedia.org/resource/Douala' 'http://dbpedia.org/resource/Almaty',
 'http://dbpedia.org/resource/Havana',
 'http://dbpedia.org/resource/Sarajevo' 'http://dbpedia.org/resource/Baku',
 'http://dbpedia.org/resource/Tashkent',
 'http://dbpedia.org/resource/Dar_Es_Salaam',
 'http://dbpedia.org/resource/Lome' 'http://dbpedia.org/resource/Maputo',
 'http://dbpedia.org/resource/Antananarivo',
 'http://dbpedia.org/resource/Conakry' 'http://dbpedia.org/resource/Lagos',
 'http://dbpedia.org/resource/Ndjamena',
 'http://dbpedia.org/resource/Sanaa',
 'http://dbpedia.org/resource/Nouakchott',
 'http://dbpedia.org/resource/Port_Harcourt',
 'http://dbpedia.org/resource/Belgrade',
 'http://dbpedia.org/resource/Niamey' 'http://dbpedia.org/resource/Bamako',
 'http://dbpedia.org/resource/Kinshasa',
 'http://dbpedia.org/resource/Ouagadougou',
 'http://dbpedia.org/resource/Luanda' 'http://dbpedia.org/resource/Bangui',
 'http://dbpedia.org/resource/Baghdad',
 'http://dbpedia.org/resource/Khartoum',
 'http://dbpedia.org/resource/Pointe_Noire',
 'http://dbpedia.org/resource/Brazzaville']

class Connector():
    def __init__(self, conn_details, database):
        self.details = conn_details
        self.host = conn_details['endpoint']
        self.db = database
        self.connection = None


    def query(self, q_str):
        """
        Execute a query on the initialized Stardog database
        :param q_str: Query string.
        :type q_str: str
        :return: Dictionary generated from the ['results']['bindings'] json.
        :rtype: dict
        """

        if self.connection is None:
            print("new connection")
            self.connection = stardog.Connection(self.db, **self.details)

        r = self.connection.select(q_str)
        return r['results']['bindings']

    def close(self):
        print("connection closed")
        self.connection.close()

class Capture():
    def __init__(self, connector):
        self.connector=connector

    def neighborhood_request(self, noi):
        try:
            if noi[0] == '<':
                noi = noi[1:]
            if noi[-1] == '>':
                noi = noi[:-1]
            print(noi)
            q = 'SELECT ?p ?o ?dt WHERE { BIND( IRI("' + noi + '") AS ?s ) ?s ?p ?o. BIND (datatype(?o) AS ?dt) }'

            return self.connector.query(q)
        except Exception as e:
            print(e)
            return []

    def _create_neighbour_paths(self, t):
        noi, depth = t
        res = self._define_neighborhood(noi, depth)

        return noi, res

    def _define_neighborhood(self, value, depth,):
        n_e = value
        res = self.neighborhood_request(n_e)
        next_noi = []
        for row in res:
            o = row['o']['value']
            if not 'dt' in row:
                next_noi.append('<' + o + '>')

        if depth - 1 > 0:
            print(n_e, len((next_noi)))
            [self._define_neighborhood(value, depth - 1,) for value in next_noi]
            self.connector.close()


    def extract_neighborhoods(self, data, depth, jobs=1):
        seq =[(r, depth) for r in data]
        if jobs > 1:
            with Pool(jobs, maxtasksperchild=1) as pool:
                list(tqdm(pool.imap_unordered(self._create_neighbour_paths, seq, chunksize=1),
                                total=len(data)))
                pool.close()
                pool.join()
        else:
            res = []
            for s in tqdm(seq, total=len(data)):
                self._create_neighbour_paths(s)


if __name__ == '__main__':
    details = {'endpoint': 'PLACE_URL_HERE'}
    connector = Connector(details, "dbpedia")

    cap = Capture(connector)

    cap.extract_neighborhoods(cities,depth=3,jobs=30)