import argparse
import json

import cartopy.crs as ccrs
import geopandas as gpd
import geoplot
import geoplot.crs as gcrs
import matplotlib.pyplot as plt


def main(
    plz_data_path: str,
    plz_prefixes_path: str,
    output_file_path: str,
    map_background_color: str,
    map_accent_color: str,
    image_background_color: str,
):
    if not output_file_path.endswith(".png"):
        raise ValueError(f"Output file needs to be a .png but is {output_file_path}")

    with open(plz_prefixes_path, "r") as fh:
        plz_prefixes = json.load(fh)
    with open(plz_data_path, "r") as fh:
        plz_geojson_data = json.load(fh)

    data = gpd.GeoDataFrame.from_features(plz_geojson_data["features"])
    data["highlighted"] = False
    for plz_prefix in plz_prefixes:
        data.loc[data["plz"].str.startswith(plz_prefix), "highlighted"] = True

    ax = plt.axes(projection=ccrs.TransverseMercator())
    geoplot.polyplot(
        data,
        projection=gcrs.TransverseMercator(),
        edgecolor=map_background_color,
        facecolor=map_background_color,
        linewidth=0.3,
        ax=ax,
    )
    geoplot.polyplot(
        data[data["highlighted"]],
        projection=gcrs.TransverseMercator(),
        edgecolor=map_accent_color,
        facecolor=map_accent_color,
        linewidth=0.3,
        ax=ax,
    )

    ax.set_facecolor(image_background_color)
    plt.tight_layout()
    plt.savefig(
        output_file_path, dpi=300, facecolor=image_background_color, edgecolor="none"
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--plz-data",
        required=True,
        help="A GEOJSON file with a 'plz' feature, like the one for Germany that can be"
        + " downloaded from "
        + "https://opendata-esri-de.opendata.arcgis.com/datasets/5b203df4357844c8a6715d7d411a8341_0",
    )
    parser.add_argument(
        "--plz-highlight",
        required=True,
        help="A JSON file with a list of PLZ (prefixes) that should be highlighted.",
    )
    parser.add_argument(
        "--out",
        required=True,
        help="A .png file path where the generated image should be stored.",
    )
    parser.add_argument(
        "--img-bg-color",
        type=str,
        default="#EDEFEB",
        help="Color outside of the map, the image file background.",
    )
    parser.add_argument(
        "--map-bg-color",
        type=str,
        default="#202020",
        help="Color code for the background map, non highlighted areas.",
    )
    parser.add_argument(
        "--map-accent-color",
        type=str,
        default="#EF7C21",
        help="Color code for the highlighted areas on the map.",
    )
    args = parser.parse_args()

    main(
        plz_data_path=args.plz_data,
        plz_prefixes_path=args.plz_highlight,
        output_file_path=args.out,
        map_background_color=args.map_bg_color,
        map_accent_color=args.map_accent_color,
        image_background_color=args.img_bg_color,
    )