import gradio as gr
from gradio_folium import Folium
import folium
import pandas as pd
import numpy as np
from geopy.distance import geodesic

def calculate_distances(coords):
    distances = []
    num_coords = len(coords)
    shortest_distances = []
    
    for i in range(num_coords):
        shortest_distance = float('inf')
        closest_pair = None
        for j in range(num_coords):
            if i != j:
                distance = geodesic(coords[i], coords[j]).kilometers
                distances.append((distance, coords[i], coords[j]))
                if distance < shortest_distance:
                    shortest_distance = distance
                    closest_pair = (coords[i], coords[j])
        if closest_pair:
            shortest_distances.append((shortest_distance, closest_pair[0], closest_pair[1]))
    
    return distances, shortest_distances

def plot_on_map(coords_list, shortest_distance, longest_shortest_distance, avg_distance):
    all_coords = [coord for coords in coords_list for coord in coords]
    avg_lat = np.mean([coord[0] for coord in all_coords])
    avg_lon = np.mean([coord[1] for coord in all_coords])
    m = folium.Map(location=[avg_lat, avg_lon], zoom_start=10)

    colors = ['blue', 'red', 'green', 'purple', 'orange', 'darkblue', 'lightgray', 'darkgreen', 'cadetblue', 'pink']
   
    for i, coords in enumerate(coords_list):
        group_marker_color = colors[i % len(colors)] 
        for coord in coords:
            store_name = f"Store: {coord}"
            if len(coords) == 1:
                store_name = "Only 1 store"
            folium.Marker(location=coord, popup=store_name, icon=folium.Icon(color=group_marker_color)).add_to(m)
    
    if shortest_distance:
        folium.PolyLine(locations=[shortest_distance[1], shortest_distance[2]], color='green', weight=5, tooltip=f'Shortest Distance: {shortest_distance[0]:.2f} km').add_to(m)
    
    if longest_shortest_distance:
        folium.PolyLine(locations=[longest_shortest_distance[1], longest_shortest_distance[2]], color='red', weight=5, tooltip=f'Longest Among the Shortest Distances : {longest_shortest_distance[0]:.2f} km').add_to(m)
    
    return m

def generate_map_with_stats(postcode_groups_input):
    df = pd.read_csv("speedmart99.csv")
    df['postcode'] = df['postcode'].astype(str).str.zfill(5)  # Ensure postcodes are 5 digits long

    postcode_groups = postcode_groups_input.split(',')

    coords_list = []
    for postcode_group_input in postcode_groups:
        key_length = len(postcode_group_input)
        # Filter rows where the postcode starts with the user input (considering the key length)
        group = df[df['postcode'].str[:key_length] == postcode_group_input]
        if not group.empty:
            coords = group['coords'].apply(lambda x: tuple(map(float, x.strip('()').split(',')))).tolist()
            coords_list.append(coords)

    # Combine all coordinates into a single list
    all_coords = [coord for coords in coords_list for coord in coords]

    if all_coords:
        distances, shortest_distances = calculate_distances(all_coords)

        if shortest_distances:
            shortest_distance = min(shortest_distances, key=lambda x: x[0])
            longest_shortest_distance = max(shortest_distances, key=lambda x: x[0])
            avg_distance = np.mean([dist[0] for dist in shortest_distances])

            map_object = plot_on_map([all_coords], shortest_distance, longest_shortest_distance, avg_distance)

            statistics = []
            for i, postcode_group in enumerate(postcode_groups):
                statistics.append(f"Postcode Group {i+1}: {postcode_group}")
            statistics.append(f"Shortest Distance (green): {shortest_distance[0]:.2f} km")
            statistics.append(f"Longest Among the Shortest Distances (red): {longest_shortest_distance[0]:.2f} km")
            statistics.append(f"Average of the Shortest Distances: {avg_distance:.2f} km")

            return map_object, "\n".join(statistics)
    else:
        return None, "No data found for the specified postcode groups."


iface = gr.Interface(
    fn=generate_map_with_stats,
    inputs="text",
    outputs=[Folium(), "text"],
    title="99 Speedmart Stores Analyzer",
    description="This app allows you to <b>input <a href='https://en.wikipedia.org/wiki/Postal_codes_in_Malaysia#:~:text=Postal%20codes%20in%20Malaysia%20%2D%20Wikipedia,History'>postcode areas</a></b> and <b>visualize</b> the locations of <a href='https://www.99speedmart.com.my/About'>99 Speedmart</a> stores within those areas.\n \nYou can <b>analyze the <span style='color: red;'>longest</span>, <span style='color: green;'>shortest</span> and average distances</b> between stores. \nEnter the first few digits of the postcode (up to 5 digits) and you may combine multiple postcode areas by separating them by commas.\n\n\nYou may try: \n\n 811,813 for Tebrau(811xx) and Kempas(813xx) \n\n 4 for Klang(4xxxx) \n\n 83 for Batu Pahat(83xxx)",
    article="\n\n\n*Store location data was last updated in February 2024.",
)

iface.launch(share=True)