import os
import zipfile
import gradio as gr
import pandas as pd
import plotly.graph_objects as go

from geo_tools import (
    shapefile_to_latlong,
    mask_shapefile_to_grid_indices,
    points_to_shapefile,
    get_cached_grid_indices,
)

DATA_DIR = "data/"
MASK_PATH = os.path.join(DATA_DIR, "serp.shp")
HOSP_PATH = os.path.join(DATA_DIR, "hospitals.shp")
POLICE_PATH = os.path.join(DATA_DIR, "police.shp")
OUT_DIR = "out/"


def gr_generate_map(
    side_len: str,
    show_hospitals: bool = True,
    show_police: bool = True,
    region: str = "None",
):
    token = "pk.eyJ1IjoiZGlsaXRoIiwiYSI6ImNsaTZ3b3I4MjF6MmczZG80cXBmeTgyaGsifQ.JmrU3qbp2jlK_9Yl2il8pw"
    side_len = float(side_len)
    show_mask = False

    scattermaps = []
    grid_path = MASK_PATH[: -len(".shp")] + f"-gap={side_len}.shp"
    prefix = ".".join(grid_path.split(".")[:-1])
    if not os.path.exists(grid_path):
        indices, labels = mask_shapefile_to_grid_indices(MASK_PATH, side_len)
        points_to_shapefile(indices, labels, grid_path)

        file_list = [prefix + ext for ext in (".shp", ".shx", ".prj", ".dbf", ".cpg")]
        with zipfile.ZipFile(prefix + ".zip", "w") as fp:
            for file in file_list:
                fp.write(file, compress_type=zipfile.ZIP_DEFLATED)

        grid_point_df = pd.DataFrame(
            data=[
                [labels[i], f"{indices[i][1]}, {indices[i][0]}"]
                for i in range(len(indices))
            ],
            columns=["Name", "Coordinates"],
        )
        grid_point_df.to_csv(prefix + ".csv", index=False)
    else:
        indices, labels = get_cached_grid_indices(grid_path)
        grid_point_df = pd.read_csv(prefix + ".csv")
    box = go.Scattermapbox(
        lat=indices[:, 1],
        lon=indices[:, 0],
        mode="markers+text",
        text=labels,
        marker=go.scattermapbox.Marker(size=6),
    )
    box.name = "Grids"
    box.textfont.update({"color": "White"})
    scattermaps.append(box)

    if show_mask:
        contours = shapefile_to_latlong(MASK_PATH)
        for contour in contours:
            lons = contour[:, 0]
            lats = contour[:, 1]
            scattermaps.append(
                go.Scattermapbox(
                    fill="toself",
                    lat=lats,
                    lon=lons,
                    mode="markers",
                    marker=go.scattermapbox.Marker(size=6),
                )
            )

    if show_hospitals:
        indices, labels = get_cached_grid_indices(HOSP_PATH)
        box = go.Scattermapbox(
            lat=indices[:, 1],
            lon=indices[:, 0],
            mode="markers+text",
            text=labels,
            marker=go.scattermapbox.Marker(size=10),
        )
        box.name = "Hospitals"
        box.textfont.update({"color": "White"})
        scattermaps.append(box)

    if show_police:
        indices, labels = get_cached_grid_indices(POLICE_PATH)
        box = go.Scattermapbox(
            lat=indices[:, 1],
            lon=indices[:, 0],
            mode="markers+text",
            text=labels,
            marker=go.scattermapbox.Marker(size=10),
        )
        box.name = "Police Stations"
        box.textfont.update({"color": "White"})
        scattermaps.append(box)

    fig = go.Figure(scattermaps)

    center = (7.753769, 80.691730)
    if region == "Ussangoda":
        center = (6.0994295, 80.9860763)
    elif region == "Indikolapelessa":
        center = (6.3602253, 80.9371957)
    elif region == "Ginigalpelessa":
        center = (6.3846744, 80.8868755)
    elif region == "Yudhaganawa":
        center = (7.665643, 80.9529867)
    elif region == "Seruwila":
        center = (8.335057, 81.320460)
    elif region == "Rupaha":
        center = (7.0401336625287, 80.89709495963253)
    modebar_icons = [
        "lasso",
        "select",
        "pan",
        "zoomin",
        "zoomout",
        "toImage",
        "resetview",
    ]
    if token:
        fig.update_layout(
            mapbox=dict(
                style="satellite-streets",
                accesstoken=token,
                center=go.layout.mapbox.Center(lat=center[0], lon=center[1]),
                pitch=0,
                zoom=6 if region == "None" else 13,
            ),
            mapbox_layers=[
                {
                    # "below": "traces",
                    "sourcetype": "raster",
                    "sourceattribution": "United States Geological Survey",
                    "source": [
                        "https://basemap.nationalmap.gov/arcgis/rest/services/USGSImageryOnly/MapServer/tile/{z}/{y}/{x}"
                    ],
                }
            ],
            modebar_remove=modebar_icons,
        )
    else:
        fig.update_layout(
            mapbox_style="open-street-map",
            hovermode="closest",
            mapbox=dict(
                bearing=0,
                center=go.layout.mapbox.Center(lat=center[0], lon=center[1]),
                pitch=0,
                zoom=6 if region == "None" else 13,
            ),
            modebar_remove=modebar_icons,
        )

    return fig, prefix + ".zip", prefix + ".csv", grid_point_df


with gr.Blocks() as demo:
    gr.Markdown("""# Serpentinite Sampling Grid Generator""")

    with gr.Tab("Sampling"):
        grid_side_len = gr.Textbox(value="100", label="Sampling Gap (m)")

        grid_show_hosp = gr.Checkbox(True, label="Show Hospitals")
        grid_show_police = gr.Checkbox(True, label="Show Police Stations")

        grid_button = gr.Button("Generate Grid")

        grid_map = gr.Plot(label="Plot")

        grid_region = gr.Radio(
            label="Zoom to Region",
            choices=[
                "None",
                "Ussangoda",
                "Indikolapelessa",
                "Ginigalpelessa",
                "Yudhaganawa",
                "Seruwila",
                "Rupaha",
            ],
        )
        grid_shapefile = gr.File(label="Grid Shapefile")
        grid_point_info = gr.File(label="Grid Point Info")

        grid_point_table = gr.Dataframe(label="Grid Point Info Table")

    grid_button.click(
        gr_generate_map,
        inputs=[grid_side_len],
        outputs=[grid_map, grid_shapefile, grid_point_info, grid_point_table],
    )

    grid_region.change(
        gr_generate_map,
        inputs=[
            grid_side_len,
            grid_show_hosp,
            grid_show_police,
            grid_region,
        ],
        outputs=[grid_map, grid_shapefile, grid_point_info, grid_point_table],
    )

demo.queue(concurrency_count=10).launch(debug=True)