import gradio as gr
import folium
from folium.plugins import HeatMap
from ingest import load_data
from process import get_lat_lon
from gradio_folium import Folium
from folium import Map


def update_header(file_info):
    if file_info is not None:
        filename = file_info.split('/')[-1]  # Access the filename from the file_info dictionary
        header = f"<h1>UK Postcodes Map: {filename}</h1>"  # Update the Markdown content
    return header  # Continue to pass the file_info to the next function if necessary

def generate_map(file_path):
    # Load the postcodes
    postcode_mapping = load_data('data/ukpostcodes.csv')
    # Load the data (this needs to be adapted to work outside Flask)
    postcodes = load_data(file_path)

    # Get latitude, longitude, and count data for the specified postcodes
    lat_lon_data = get_lat_lon(postcodes, postcode_mapping)

    # Prepare data for different frequency bands
    low_freq_data = [
        [data['latitude'], data['longitude']] 
        for data in lat_lon_data 
        if data['count'] == 1 and data['latitude'] and data['longitude']
    ]
    med_freq_data = [
        [data['latitude'], data['longitude']] 
        for data in lat_lon_data 
        if 2 <= data['count'] <= 5 and data['latitude'] and data['longitude']
    ]
    high_freq_data = [
        [data['latitude'], data['longitude']] 
        for data in lat_lon_data 
        if data['count'] > 5 and data['latitude'] and data['longitude']
    ]

    # Create your map here using Folium
    map = Map(location=[51.505303, -0.13902], zoom_start=10)

    # Adding different heatmaps for different frequencies
    if low_freq_data:
        HeatMap(low_freq_data, radius=10, blur=10, gradient={0.8: 'blue', 1: 'lime'}).add_to(map)
    if med_freq_data:
        HeatMap(med_freq_data, radius=15, blur=10, gradient={0.8: 'orange', 1: 'lime'}).add_to(map)
    if high_freq_data:
        HeatMap(high_freq_data, radius=20, blur=10, gradient={0.8: 'red', 1: 'lime'}).add_to(map)
    
    return map 

# Define a Gradio interface
with gr.Blocks() as demo:

    with gr.Row(): 
        header = gr.Markdown(("<h1>UK Postcodes Map</h1>"))
    
    with gr.Row():
        map = Folium(value = Map(location=[51.505303, -0.13902], zoom_start=10), height=700)


    with gr.Row():
        file_uploader = gr.UploadButton(
            label=("Upload"),
            file_count="single",
            file_types=[".csv", ".xlsx", '.xls'],
            interactive=True,
            scale=1,
        )

    file_uploader.upload(fn = generate_map, 
                        inputs= file_uploader, 
                        outputs=map)
    
    file_uploader.upload(fn=update_header, 
                         inputs=file_uploader, 
                         outputs=header)

                    
if __name__ == "__main__":
    demo.launch()