import os
import numpy as np
import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt
import matplotlib as mpl
import contextily as cx
from shapely import wkt

from ref_systems import osm_crs, projected_crs


categories =['healthcare', 'services', 'transport', 'outdoor', 'supplies', 'restaurant', 'culture', 'education', 'physical']


if __name__ == "__main__":

    # full path to AGILE_2025_data_code folder
    fullpath = os.path.abspath(os.path.join(os.getcwd(), '..'))

    # input path
    path_roi = os.path.join(fullpath, 'data', 'BEZIRKSGRENZEOGD', 'BEZIRKSGRENZEOGDPolygon.shp')

    # output path
    path_plots = os.path.join(fullpath, 'results', 'plots')

    # create the output directory if it doesn't exist
    if not os.path.exists(path_plots):
        os.makedirs(path_plots)

    # read accessibility results per mode and per category
    all_data = {'walk': None, 'walk_PT': None}
    for mode in ['walk', 'walk_PT']:
        print('-----', mode, '-----')
        
        # input path (mode dependent)
        path_access = os.path.join(fullpath, 'results', 'accessibility', mode, 'terrain_yes')
        
        for category in categories:
            print(category)
            # read accessibility results
            data = pd.read_csv(os.path.join(path_access, category + '.csv'))
            data.replace([np.inf, -np.inf], np.nan, inplace=True)

            # add to combined results
            if all_data[mode] is None:
                all_data[mode] = data
                all_data[mode].rename(columns={'time': category}, inplace=True)
            else:
                all_data[mode][category] = data.time

    # compute city-level accessibility for each mode
    for mode in ['walk', 'walk_PT']:
        all_data[mode]['avg'] = all_data[mode][categories].mean(axis=1, skipna=True)

    # split into individual mode variables
    walk_data = all_data['walk']
    walk_data['geometry'] = walk_data['geometry'].apply(wkt.loads)
    walk_data = gpd.GeoDataFrame(walk_data, geometry='geometry', crs=projected_crs)
    pt_data = all_data['walk_PT']
    pt_data['geometry'] = pt_data['geometry'].apply(wkt.loads)
    pt_data = gpd.GeoDataFrame(pt_data, geometry='geometry', crs=projected_crs)

    # difference between walk and walk_PT
    diff = walk_data.avg - pt_data.avg

    # ------------------------------------ PLOTS ------------------------------------
    print('--------------')
    print('start plotting')
    
    # define colorscale for heatmaps
    low, high = 0, 30
    norm = mpl.colors.Normalize(vmin=low, vmax=high, clip=True)
    mapper = mpl.cm.ScalarMappable(norm=norm, cmap=mpl.cm.coolwarm)

    #--------------- heatmap for WALK ---------------
    fig, ax = plt.subplots()
    ax.set_axis_off()
    walk_data.plot(ax=ax, facecolor=[mapper.to_rgba(i) for i in walk_data.avg.values], edgecolor="None", alpha=0.6)
    cbar = fig.colorbar(mapper, ax=ax, fraction=0.04, pad=0.04)
    cbar.ax.set_yticks([low, high])
    cbar.set_label('[min]')
    cx.add_basemap(ax=ax, crs=projected_crs, source=cx.providers.OpenStreetMap.Mapnik, alpha=0.6)
    plt.savefig(os.path.join(path_plots, 'heatmap_walk.png'), dpi=600, bbox_inches='tight')

    # --------------- heatmap for WALK PLUS PUBLIC TRANSPORT ---------------
    fig, ax = plt.subplots()
    ax.set_axis_off()
    pt_data.plot(ax=ax, facecolor=[mapper.to_rgba(i) for i in pt_data.avg.values], edgecolor="None", alpha=0.6)
    cbar = fig.colorbar(mapper, ax=ax, fraction=0.04, pad=0.04)
    cbar.ax.set_yticks([low, high])
    cbar.set_label('[min]')
    cx.add_basemap(ax=ax, crs=projected_crs, source=cx.providers.OpenStreetMap.Mapnik, alpha=0.6)
    plt.savefig(os.path.join(path_plots, 'heatmap_walk_PT.png'), dpi=600, bbox_inches='tight')

    # --------------- heatmap of DIFFERENCES ---------------

    # define colorscale (bins)
    bins = [0, 0.5, 2, 5, 10, diff.max()]
    colors = ['#440154', '#31688E', '#35B779', '#FDE725', '#FF5733']
    cmap = mpl.colors.ListedColormap(colors)
    norm = mpl.colors.BoundaryNorm(bins, cmap.N)

    # assign colors based on bins
    diff_values_binned = np.clip(np.digitize(diff.values, bins) - 1, 0, len(colors) - 1)

    # heatmap of differences
    fig, ax = plt.subplots()
    ax.set_axis_off()
    walk_data.plot(ax=ax, facecolor=[colors[i] for i in diff_values_binned], edgecolor="None", alpha=0.6)
    cx.add_basemap(ax=ax, crs=projected_crs, source=cx.providers.OpenStreetMap.Mapnik, alpha=0.6)
    cbar = fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap), ax=ax, fraction=0.04, pad=0.04)
    cbar.set_label('[min]')
    cbar.set_ticks([(bins[i] + bins[i + 1]) / 2 for i in range(len(bins) - 1)])
    cbar.set_ticklabels(['0-0.5', '0.5-2.0', '2.0-5.0', '5.0-10.0', '>10.0'])
    plt.savefig(os.path.join(path_plots, 'heatmap_differences.png'), dpi=600, bbox_inches='tight')

    # --------------- line-plots DISTANCE-BANDS ---------------

    # one plot for each scenario
    for mode in ['walk', 'walk_PT']:
        f1, ax1 = plt.subplots()
        ax1.set_xlabel('radius [m] around city centre')
        ax1.set_ylabel('average time [min]')
        ax1.set_ylim([0, 25])

        for category in categories:
            # filter for correct subset of data
            data = all_data[mode][[category]].copy()
            data['centroids'] = all_data[mode]['centroids'].apply(wkt.loads).values
            data = gpd.GeoDataFrame(data, geometry='centroids', crs=projected_crs)

            # read ROI shapefile and get city centre (centroid of 1. district)
            borders = gpd.read_file(path_roi)
            borders.to_crs(projected_crs, inplace=True)
            city_centre = borders[borders.NAMEK == 'Innere Stadt'].geometry.iloc[0].centroid

            val_avg = []
            for radius in range(150, 18000, 100):
                # increase polygon around city centre
                polygon = city_centre.buffer(radius)

                # find hexagons that fall within this circle
                hex_within = data[data.geometry.within(polygon)]

                # compute average accessibility value
                val_avg.append(hex_within[category].mean())

            ax1.plot(range(150, 18000, 100), val_avg, '.-', label=category)
        ax1.legend()
        ax1.axhline(y=15, color='k')  # 15-minute cutoff value
        plt.savefig(os.path.join(path_plots, 'lineplot_'+ mode + '.png'), dpi=600, bbox_inches='tight')

    print('done')