import argparse
import json

import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np


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

    # TODO: This should instead be a proper geo spatial projection
    coordinate_transform = np.array([[1.0, 1.6]]).T

    ax = plt.axes(aspect="equal")
    for bundesland in data["bundesland"].unique():
        bundesland_df = data[data["bundesland"] == bundesland]
        union = bundesland_df.unary_union
        for geom in getattr(union, "geoms", [union]):
            ax.fill(
                *(geom.exterior.xy * coordinate_transform),
                fc=map_background_color,
                ec=image_background_color,
                linewidth=0.1,
            )

        if not any(bundesland_df["highlighted"]):
            continue

        union = bundesland_df[bundesland_df["highlighted"]].unary_union
        for geom in getattr(union, "geoms", [union]):
            ax.fill(
                *(geom.exterior.xy * coordinate_transform),
                fc=map_accent_color,
                ec=None,
            )

    ax.set_facecolor(image_background_color)
    plt.axis("off")
    plt.tight_layout(pad=0.0, h_pad=0.0, w_pad=0.0)
    plt.savefig(
        output_file_path,
        dpi=300,
        facecolor=image_background_color,
        edgecolor="none",
        bbox_inches="tight",
    )


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,
    )