import matplotlib.pyplot as plt
import numpy as np
import osmnx as ox
import pandas as pd

ox.config(log_console=True, use_cache=True)
weight_by_length = False

# define the study sites as label : query
places = {'Bucuresti'             : {'q':'Bucharest', 'exclude_place_ids':'298450'},
          'Cluj-Napoca'           : {'q':'Cluj Napoca', 'exclude_place_ids':'195112'},
          'Oradea'                : {'q':'Oradea', 'exclude_place_ids':'29122273' },
          'Constanta'             : 'Constanta',
          'Suceava'               : 'Suceava',
          'Iasi'                  : {'q':'Iasi', 'exclude_place_ids':'324553' },
          'Timisoara'             : {'q':'Timisoara', 'exclude_place_ids':'246522' },
          'Sibiu'                 : {'q':'Sibiu', 'exclude_place_ids':'223295999' },
          'Craiova'               : {'q': 'Craiova', 'exclude_place_ids':'3810059'},
          'Brasov'                : {'q': 'Brasov', 'exclude_place_ids':'785207'},
          'Galati'                : {'q':'Galati', 'exclude_place_ids':'11640887,305993'},
          'Ploiesti'              : {'q':'Ploiesti', 'exclude_place_ids':'251193'},
          'Braila'                : {'q':'Braila', 'exclude_place_ids':'1306402,885269'},
          'Arad'                  : 'Arad'                  ,
          'Pitesti'               : {'q':'Pitesti', 'exclude_place_ids':'326999'},
          'Bacau'                 : {'q':'Bacau', 'exclude_place_ids':'14936163'},
          'Targu Mures'           : {'q':'Targu Mures', 'exclude_place_ids':'130794'},
          'Baia Mare'             : {'q':'Baia Mare','exclude_place_ids':'321848'},
          'Buzau'                 : 'Buzau'                 ,
          'Botosani'              : 'Botosani'              ,
          'Satu Mare'             : 'Satu Mare'             ,
          'Ramnicu Valcea'        :  {'q':'Ramnicu Valcea','exclude_place_ids':'335988'},
          'Drobeta-Turnu Severin' : 'Drobeta-Turnu Severin' ,
          'Piatra Neamt'          : {'q':'Piatra Neamt','exclude_place_ids':'326492'},
          'Targu Jiu'             : 'Targu Jiu'             ,
          'Targoviste'            : {'q':'Targoviste', 'exclude_place_ids':'315265' },
          'Focsani'               : {'q':'Focsani', 'exclude_place_ids':'146569'},
          'Bistrita'              : 'Bistrita'              ,
}

# verify OSMnx geocodes each query to what you expect
gdf = ox.gdf_from_places(places.values())
print(gdf)

for place in places.keys():
    G = ox.graph_from_place(places[place])
    ox.plot_graph(ox.project_graph(G), save=True, show=False, file_format='svg', filename=place)

exit(0)


bearings = {}
for place in sorted(places.keys()):
    
    print("Processing", place)

    # get the graph
    query = places[place]
    G = ox.graph_from_place(query, network_type='drive')
    
    # calculate edge bearings
    G = ox.add_edge_bearings(G)
    
    if weight_by_length:
        # weight bearings by length (meters)
        streets = [(d['bearing'], int(d['length'])) for u, v, k, d in G.edges(keys=True, data=True)]
        city_bearings = []
        for street in streets:
            city_bearings.extend([street[0]] * street[1])
        bearings[place] = pd.Series(city_bearings)
    else:
        # don't weight bearings, just take one value per street segment
        bearings[place] = pd.Series([data['bearing'] for u, v, k, data in G.edges(keys=True, data=True)])
        
        
# function to draw a polar histogram for a set of edge bearings
def polar_plot(ax, bearings, n=36, title=''):

    bins = [ang * 360 / n for ang in range(0, n + 1)]
    count = count_and_merge(n, bearings)
    _, division = np.histogram(bearings, bins=bins)
    frequency = count / count.sum()
    division = division[0:-1]
    width =  2 * np.pi / n

    ax.set_theta_zero_location('N')
    ax.set_theta_direction('clockwise')

    x = division * np.pi / 180
    bars = ax.bar(x, height=frequency, width=width, align='center', bottom=0, zorder=2,
                  color='#003366', edgecolor='k', linewidth=0.5, alpha=0.7)
    
    ax.set_ylim(top=frequency.max())
    
    title_font = {'family':'Roboto', 'size':24, 'weight':'bold'}
    xtick_font = {'family':'Roboto', 'size':10, 'weight':'bold', 'alpha':1.0, 'zorder':3}
    ytick_font = {'family':'Roboto', 'size': 9, 'weight':'bold', 'alpha':0.2, 'zorder':3}
    
    ax.set_title(title.upper(), y=1.05, fontdict=title_font)
    
    ax.set_yticks(np.linspace(0, max(ax.get_ylim()), 5))
    yticklabels = ['{:.2f}'.format(y) for y in ax.get_yticks()]
    yticklabels[0] = ''
    ax.set_yticklabels(labels=yticklabels, fontdict=ytick_font)
    
    xticklabels = ['N', '', 'E', '', 'S', '', 'W', '']
    ax.set_xticklabels(labels=xticklabels, fontdict=xtick_font)
    ax.tick_params(axis='x', which='major', pad=-2)
    
def count_and_merge(n, bearings):
    # make twice as many bins as desired, then merge them in pairs
    # this prevents bin-edge effects around common values like 0 deg and 90 deg
    n = n * 2
    bins = [ang * 360 / n for ang in range(0, n + 1)]
    count, _ = np.histogram(bearings, bins=bins)
    
    # move the last bin to the front, so eg 0.01 deg and 359.99 deg will be binned together
    count = count.tolist()
    count = [count[-1]] + count[:-1]

    count_merged = []
    count_iter = iter(count)
    for count in count_iter:
        merged_count = count + next(count_iter)
        count_merged.append(merged_count)

    return np.array(count_merged)

# create figure and axes
n = len(places)
ncols = int(np.ceil(np.sqrt(n)))
nrows = int(np.ceil(n / ncols))
figsize = (ncols * 5, nrows * 5)
fig, axes = plt.subplots(nrows, ncols, figsize=figsize, subplot_kw={'projection':'polar'})
axes = [item for sublist in axes for item in sublist]

# plot each city's polar histogram
for ax, place in zip(axes, sorted(places.keys())):
    polar_plot(ax, bearings[place], title=place)

# add super title and save full image
suptitle_font = {'family':'Roboto', 'fontsize':60, 'fontweight':'normal', 'y':1.07}
fig.suptitle('Romanian City Street Network Orientation', **suptitle_font)
fig.tight_layout()
fig.subplots_adjust(hspace=0.35)
plt.gcf().savefig('images/street-orientations.png', dpi=120, bbox_inches='tight')
plt.close()