import json
import sys
from pathlib import Path
from typing import List
from urllib.request import pathname2url
from xml.dom import minidom

from folium.plugins import BeautifyIcon
from folium.features import DivIcon

# import folium.plugins as plugins

import numpy as np
import pandas as pd
from scipy.signal import find_peaks
import streamlit as st

import folium
from streamlit_folium import st_folium
import altair as alt
from io import StringIO
import branca


def get_gpx(uploaded_file):
    data = StringIO(uploaded_file.getvalue().decode("utf-8"))

    xmldoc = minidom.parse(data)
    track = xmldoc.getElementsByTagName("trkpt")
    elevation = xmldoc.getElementsByTagName("ele")
    n_track = len(track)

    # Parsing GPX elements
    lon_list = []
    lat_list = []
    h_list = []
    for s in range(n_track):
        lon, lat = (
            track[s].attributes["lon"].value,
            track[s].attributes["lat"].value,
        )
        elev = elevation[s].firstChild.nodeValue
        lon_list.append(float(lon))
        lat_list.append(float(lat))
        h_list.append(float(elev))

    # Calculate average latitude and longitude
    ave_lat = sum(lat_list) / len(lat_list)
    ave_lon = sum(lon_list) / len(lon_list)

    return ave_lat, ave_lon, lon_list, lat_list, h_list


# From https://tomaugspurger.net/posts/modern-4-performance/
def gcd_vec(lat1, lng1, lat2, lng2):
    """
    Calculate great circle distance.
    http://www.johndcook.com/blog/python_longitude_latitude/

    Parameters
    ----------
    lat1, lng1, lat2, lng2: float or array of float

    Returns
    -------
    distance:
      distance from ``(lat1, lng1)`` to ``(lat2, lng2)`` in kilometers.
    """
    # python2 users will have to use ascii identifiers
    ϕ1 = np.deg2rad(90 - lat1)
    ϕ2 = np.deg2rad(90 - lat2)

    θ1 = np.deg2rad(lng1)
    θ2 = np.deg2rad(lng2)

    cos = np.sin(ϕ1) * np.sin(ϕ2) * np.cos(θ1 - θ2) + np.cos(ϕ1) * np.cos(ϕ2)
    arc = np.arccos(cos)
    return arc * 6373


CATEGORY_TO_COLOR = {
    5: "#68bd44",
    4: "#68bd44",
    3: "#fbaa1c",
    2: "#f15822",
    1: "#ed2125",
    0: "#800000",
}


def climb_category(climb_score):
    """Determine category of the climb based on the climb score as defined by Garmin"""
    if climb_score < 1_500:
        return 5  # Not categorised
    elif climb_score < 8_000:
        return 4
    elif climb_score < 16_000:
        return 3
    elif climb_score < 32000:
        return 2
    elif climb_score < 64000:
        return 1
    else:
        return 0  # Hors categorie


def grade_to_color(grade):
    """Determine the color of the climb based on its grade according to Garmin"""
    if grade < 3:
        return "lightgrey"
    elif grade < 6:
        return CATEGORY_TO_COLOR[3]
    elif grade < 9:
        return CATEGORY_TO_COLOR[2]
    elif grade < 12:
        return CATEGORY_TO_COLOR[1]
    else:
        return CATEGORY_TO_COLOR[0]


def find_climbs(df: pd.DataFrame) -> pd.DataFrame:
    """Detect all valleys and peaks. Filter out climbs and
    add meta data (lenght, meters climbed, average grade, climb_score, ...)
    """
    peaks, _ = find_peaks(df["smoothed_elevation"])
    df_peaks = df.iloc[peaks, :].assign(base=0).assign(kind="peak")
    valleys, _ = find_peaks(df["smoothed_elevation"].max() - df["smoothed_elevation"])
    df_valleys = df.iloc[valleys, :].assign(base=0).assign(kind="valley")
    df_elevation = pd.concat([df_valleys, df_peaks], axis=0).sort_values(
        by="distance_from_start"
    )

    # Climbscore acoording to Garmin:
    # https://s3.eu-central-1.amazonaws.com/download.navigation-professionell.de/
    # Garmin/Manuals/Understanding+ClimbPro+on+the+Edge.pdf

    df_peaks_filtered = (
        pd.concat(
            [df_elevation, df_elevation.shift(1).bfill().add_prefix("prev_")],
            axis=1,
        )
        .query("(kind=='peak') & (prev_kind=='valley')")
        .assign(
            length=lambda df_: df_["distance_from_start"]
            - df_["prev_distance_from_start"]
        )
        .assign(total_ascent=lambda df_: df_["elev"] - df_["prev_elev"])
        .assign(grade=lambda df_: (df_["total_ascent"] / df_["length"]) * 100)
        .assign(climb_score=lambda df_: df_["length"] * df_["grade"])
        .assign(hill_category=lambda df_: df_["climb_score"].map(climb_category))
        .query("climb_score >= 1_500")
        .assign(max_elevation=df["elev"].max().round(-1) + 10)
    )
    # Garmin rules
    # df_peaks_filtered = df_peaks_meta.query(
    #   "(climb_score >= 1_500) & (length >= 0.5) & (grade >= 3_000)"
    # )
    return df_peaks_filtered


def generate_height_profile_json(df: pd.DataFrame) -> str:
    """Generate a height profile of the ride in Altair.
    Returns a string with json.
    """
    df_distance = (
        df.assign(lon_1=lambda df_: df["lon"].shift(1))
        .assign(lat_1=lambda df_: df["lat"].shift(1))
        .drop(columns=["elev"])
    )[["lat", "lon", "lat_1", "lon_1"]]
    df["distance"] = pd.Series(
        [gcd_vec(*x) for x in df_distance.itertuples(index=False)],
        index=df_distance.index,
    ).fillna(0)
    total_distance = df["distance"].sum()
    total_distance_round = np.round(total_distance)
    df["distance_from_start"] = df["distance"].cumsum()
    df["smoothed_elevation"] = df["elev"].rolling(10).mean().bfill()
    df["grade"] = (
        0.1
        * (df["elev"] - df["elev"].shift(1).bfill())
        / (df["distance_from_start"] - df["distance_from_start"].shift(1).bfill())
    )
    df["smoothed_grade"] = df["grade"].rolling(10).mean()
    df["smoothed_grade"] = df["smoothed_grade"].bfill()
    df["smoothed_grade_color"] = df["smoothed_grade"].map(grade_to_color)
    # df["grade_color"] = df["grade"].map(grade_to_color)

    elevation = (
        alt.Chart(
            df[
                [
                    "distance_from_start",
                    "smoothed_elevation",
                    "smoothed_grade_color",
                    "grade",
                ]
            ]
        )
        .mark_bar()
        .encode(
            x=alt.X("distance_from_start")
            .axis(
                grid=False,
                tickCount=10,
                labelExpr="datum.label + ' km'",
                title=None,
            )
            .scale(domain=(0, total_distance_round)),
            y=alt.Y("smoothed_elevation").axis(
                domain=False,
                ticks=False,
                tickCount=5,
                labelExpr="datum.label + ' m'",
                title=None,
            ),
            color=alt.Color("smoothed_grade_color").scale(None),
            tooltip=[
                alt.Tooltip(
                    "distance_from_start:Q", title="Distance (km)", format=".2f"
                ),
                alt.Tooltip("smoothed_elevation:Q", title="Elevation (m)", format="d"),
                alt.Tooltip("grade_percent:Q", title="Grade (%)", format=".0%"),
            ],
        )
        .transform_calculate(
            grade_percent="datum.grade/100",
        )
    )
    max_elevation = df["elev"].max().round(-1)
    # elevation = (
    #     alt.Chart(df)
    #     .mark_area(
    #         color=alt.Gradient(
    #             gradient="linear",
    #             stops=[
    #                 alt.GradientStop(color="lightgrey", offset=0),
    #                 alt.GradientStop(color="darkgrey", offset=1),
    #             ],
    #             x1=1,
    #             x2=1,
    #             y1=1,
    #             y2=0,
    #         ),
    #         line={"color": "darkgreen"},
    #     )
    #     .encode(
    #         x=alt.X(
    #             "distance_from_start",
    #             axis=alt.Axis(
    #                 domain=False,
    #                 ticks=False,
    #                 tickCount=10,
    #                 labelExpr="datum.label + ' km'",
    #             ),
    #             scale=alt.Scale(domain=(0, total_distance_round)),
    #         ),
    #         y=alt.Y(
    #             "elev",
    #             axis=alt.Axis(
    #                 domain=False,
    #                 ticks=False,
    #                 tickCount=5,
    #                 labelExpr="datum.label + ' m'",
    #             ),
    #             scale=alt.Scale(domain=(0, max_elevation)),
    #         ),
    #     )
    # )
    df_peaks_filtered = find_climbs(df)
    line_peaks = (
        alt.Chart(df_peaks_filtered[["distance_from_start", "elev", "max_elevation"]])
        .mark_rule(color="red")
        .encode(
            x=alt.X("distance_from_start:Q").scale(domain=(0, total_distance_round)),
            y="elev",
            y2="max_elevation",
        )
    )
    # line_peaks = (
    #     alt.Chart(df_peaks_filtered[["distance_from_start", "elev", "max_elevation"]])
    #     .mark_rule(color="red")
    #     .encode(
    #         x=alt.X(
    #             "distance_from_start:Q",
    #             scale=alt.Scale(domain=(0, total_distance_round)),
    #         ),
    #         y="elev",
    #         y2="max_elevation",
    #     )
    # )
    df_annot = (
        df_peaks_filtered.reset_index(drop=True)
        .assign(number=lambda df_: df_.index + 1)
        .assign(circle_pos=lambda df_: df_["max_elevation"] + 20)[
            [
                "distance_from_start",
                "max_elevation",
                "circle_pos",
                "number",
                "length",
                "total_ascent",
                "grade",
                "climb_score",
                "prev_distance_from_start",
            ]
        ]
    )
    # annotation = (
    #     alt.Chart(df_annot)
    #     .mark_text(align="center", baseline="bottom", fontSize=16, dy=-10)
    #     .encode(
    #         x=alt.X("distance_from_start:Q").scale(domain=(0, total_distance_round)),
    #         y="max_elevation",
    #         text="number",
    #     )
    # )
    annotation = (
        alt.Chart(df_annot)
        .mark_text(align="center", baseline="bottom", fontSize=16, dy=-10)
        .encode(
            x=alt.X(
                "distance_from_start:Q",
                scale=alt.Scale(domain=(0, total_distance_round)),
            ),
            y="max_elevation",
            text="number",
            tooltip=[
                alt.Tooltip(
                    "prev_distance_from_start:Q", title="Starts at (km)", format=".2f"
                ),
                alt.Tooltip("total_ascent:Q", title="Total ascent (m)", format="d"),
                alt.Tooltip("length:Q", title="Length (km)", format=".2f"),
                alt.Tooltip("grade_percent:Q", title="Average Grade", format=".0%"),
                alt.Tooltip("climb_score:Q", title="Climb score", format="d"),
            ],
        )
        .transform_calculate(
            grade_percent="datum.grade/(100*1000)",
            # total_ascent_int="Math.round(datum.total_ascent)",
        )
    )
    chart = (
        (elevation + line_peaks + annotation)
        .properties(width="container")
        .configure_view(
            strokeWidth=0,
        )
    )
    return chart, df_peaks_filtered


def generate_climb_profile(df_hill: pd.DataFrame, title: str):
    climb_profile = (
        alt.Chart(
            df_hill,
            title=alt.Title(
                title,
                anchor="start",
            ),
        )
        .mark_area()
        .encode(
            x=alt.X("distance_from_start")
            .axis(grid=False, tickCount=10, labelExpr="datum.label + ' m'", title=None)
            .scale(domain=(0, df_hill["distance_from_start"].max())),
            y=alt.Y("elev").axis(
                domain=False,
                ticks=False,
                tickCount=5,
                labelExpr="datum.label + ' m'",
                title=None,
            ),
            color=alt.Color("color_grade").scale(None),
            tooltip=[
                alt.Tooltip("distance_from_start:Q", title="Distance (m)", format="d"),
                alt.Tooltip("elev:Q", title="Elevation (m)", format="d"),
                alt.Tooltip("grade_percent:Q", title="Grade (%)", format=".0%"),
            ],
        )
        .transform_calculate(
            grade_percent="datum.grade/100",
        )
    )
    return climb_profile


gpx_file = st.file_uploader("Upload gpx file", type=["gpx"])

if gpx_file is not None:
    ave_lat, ave_lon, lon_list, lat_list, h_list = get_gpx(gpx_file)
    df = pd.DataFrame({"lon": lon_list, "lat": lat_list, "elev": h_list})
    route_map = folium.Map(location=[ave_lat, ave_lon], zoom_start=12, height=400)
    folium.PolyLine(
        list(zip(lat_list, lon_list)), color="red", weight=2.5, opacity=1
    ).add_to(route_map)

    chart, df_peaks = generate_height_profile_json(df)
    for index, row in df_peaks.reset_index(drop=True).iterrows():
        icon = BeautifyIcon(
            icon="arrow-down",
            icon_shape="marker",
            number=str(index + 1),
            border_color="red",
            background_color="white",
        )
        icon_div = DivIcon(
            icon_size=(150, 36),
            icon_anchor=(7, 20),
            html=f"<div style='font-size: 18pt; color : black'>{index+1}</div>",
        )
        length = (
            f"{row['length']:.1f} km"
            if row["length"] >= 1
            else f"{row['length']*1000:.0f} m"
        )
        popup_text = f"""Climb {index+1}<br>
                Lenght: {length}<br>
                Avg. grade: {row['grade']/1000:.1f}%<br>
                Total ascend: {int(row['total_ascent'])}m
        """
        popup = folium.Popup(popup_text, min_width=100, max_width=200)
        folium.Marker(
            [row["lat"], row["lon"]],
            popup=popup,
            icon=icon_div,
        ).add_to(route_map)

        df_hill = (
            df[
                df["distance_from_start"].between(
                    row["prev_distance_from_start"],
                    row["distance_from_start"],
                )
            ]
            .assign(
                distance_from_start=lambda df_: (
                    df_["distance_from_start"] - row["prev_distance_from_start"]
                )
                * 1_000
            )
            .assign(color_grade=lambda df_: df_["grade"].map(grade_to_color))
        )
        # df_hill_resample = df_hill.groupby((df_hill["distance_from_start"]*1000).round(-2)).agg({"elev":"mean", "grade":"mean"}).reset_index()
        # df_hill_resample["color_grade"] = df_resampled["grade"].map(grade_to_color)
        title = f"Climb {index+1}: {row['length']:.2f}km {(row['grade']/100_000):.2%} {int(row['total_ascent']):d}hm"
        climb_profile = generate_climb_profile(df_hill, title)
        climb_profile_json = json.loads(climb_profile.to_json())

        vega = folium.features.VegaLite(
            climb_profile_json,
            width=200,
            height=200,
        )
        circle = folium.CircleMarker(
            radius=15,
            location=[row["lat"], row["lon"]],
            # tooltip = label,
            color="crimson",
            fill=True,
        )
        # popup = folium.Popup()
        # vega.add_to(popup)
        # popup.add_to(circle)
        circle.add_to(route_map)
        # circle_marker = folium.CircleMarker(
        #     [row["lat"], row["lon"]],
        #     radius=15,
        #     popup=folium.Popup(max_width=400).add_child(
        #         folium.VegaLite(climb_profile_json, width=400, height=400)
        #     ),
        # )

    st.table(
        df_peaks[
            ["length", "total_ascent", "grade", "climb_score", "hill_category"]
        ].reset_index(drop=True)
    )

    st_data = st_folium(route_map, height=600, width=850)

    st.altair_chart(chart, use_container_width=True)

    for index, row in df_peaks.reset_index(drop=True).iterrows():
        df_hill = (
            df[
                df["distance_from_start"].between(
                    row["prev_distance_from_start"],
                    row["distance_from_start"],
                )
            ]
            .assign(
                distance_from_start=lambda df_: (
                    df_["distance_from_start"] - row["prev_distance_from_start"]
                )
                * 1_000
            )
            .assign(color_grade=lambda df_: df_["grade"].map(grade_to_color))
        )
        df_new_index = pd.DataFrame(
            index=pd.Index(np.arange(0, df_hill["distance_from_start"].max(), 10))
        )
        df_hill_resample = pd.concat(
            [df_hill.set_index("distance_from_start"), df_new_index], axis=0
        ).sort_index()
        df_hill_resample = df_hill_resample[["elev", "grade"]].interpolate()
        df_hill_resample["color_grade"] = df_hill_resample["grade"].map(grade_to_color)
        df_hill_resample = (
            df_hill_resample.reset_index()
            .rename(columns={"index": "distance_from_start"})
            .sort_values(by="distance_from_start")
        )
        max_grade = df_hill_resample["grade"].rolling(10).mean().max() / 100
        title = f"""Climb {index+1}, length:{row['length']:.2f}km
        Avg. grade: {(row['grade']/100_000):.2%}
        Max. grade: {max_grade:.2%}
        Total ascent: {int(row['total_ascent']):d}hm"""
        climb_profile = generate_climb_profile(df_hill_resample, title)
        st.altair_chart(climb_profile, use_container_width=True)