import geopandas
import matplotlib.pyplot as plt
import psycopg2


db = psycopg2.connect(user="joaorodrigues", password="", host="localhost", database="joaorodrigues", port="5440")
cursor = db.cursor()


def db_query(db, query):
    cursor.execute(query)
    result = cursor.fetchall()
    return result


def main():
    query = "select distinct(Distrito) from cont_aad_caop2017 order by 1;"
    results = db_query(db, query)
    execute(results)


def execute(results):
    districts = []
    for row in results:
        districts.append(row[0])
    ccount = []
    for i in districts:
        cursor.execute(
            "select count(final_point) from cont_aad_caop2017,taxi_services where st_contains(geom,final_point) and Distrito ilike '" + i + "';")
        results_1 = cursor.fetchone()
        ccount.append(results_1[0])
    db.close()

    url = '/home/tito/TABD/Cont_AAD_CAOP2017.shp'
    portugal = geopandas.read_file(url, encoding='utf-8')
    total_count = []

    for row in portugal.Distrito:
        i = 0
        while i < len(districts):
            if (row.encode('utf-8') == districts[i]):
                total_count.append(ccount[i])
            i = i + 1

    portugal['total'] = total_count
    portugal.plot(column='total', cmap='OrRd', scheme='quantiles')
    plt.show()


if __name__ == '__main__':
    main()