Update app.py

app.py

@@ -1,3 +1,292 @@
+# '''Copyright 2024 Ashok Kumar
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.'''
+
+# import os
+# import requests
+# import json
+# import pandas as pd
+# import numpy as np
+# import requests
+# import geopandas as gpd
+# import contextily as ctx
+# import tzlocal
+# import pytz 
+# from PIL import Image
+# from datetime import datetime
+# import matplotlib.pyplot as plt
+# from geopy.exc import GeocoderTimedOut
+# from geopy.geocoders import Nominatim
+# import warnings
+# warnings.filterwarnings('ignore')
+# from plotly.graph_objs import Marker
+# import plotly.express as px
+# import streamlit as st
+# from data import flight_data
+# from huggingface_hub import InferenceApi, login, InferenceClient
+
+
+# hf_token = os.getenv("HF_TOKEN")
+# if hf_token is None:
+#     raise ValueError("Hugging Face token not found. Please set the HF_TOKEN environment variable.")
+# login(hf_token)
+
+
+# API_URL = "https://api-inference.huggingface.co/models/google/tapas-base-finetuned-wtq"
+# headers = {"Authorization": f"Bearer {hf_token}"}
+
+# def query(payload):
+#     response = requests.post(API_URL, headers=headers, json=payload)
+#     return response.json()
+
+# def query_flight_data(geo_df, question):
+
+
+#     table_data = {
+#         "icao24": geo_df["icao24"].astype(str).iloc[:100].tolist(),
+#         "callsign": geo_df["callsign"].astype(str).replace({np.nan: None, np.inf: '0'}).iloc[:100].tolist(),
+#         "origin_country": geo_df["origin_country"].astype(str).replace({np.nan: None, np.inf: '0'}).iloc[:100].tolist(),
+#         "time_position": geo_df["time_position"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
+#         "last_contact": geo_df["last_contact"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
+#         "longitude": geo_df["longitude"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
+#         "latitude": geo_df["latitude"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
+#         "baro_altitude": geo_df["baro_altitude"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
+#         "on_ground": geo_df["on_ground"].astype(str).iloc[:100].tolist(),  # Assuming on_ground is boolean or categorical
+#         "velocity": geo_df["velocity"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
+#         "true_track": geo_df["true_track"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
+#         "vertical_rate": geo_df["vertical_rate"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
+#         "sensors": geo_df["sensors"].astype(str).replace({np.nan: None, np.inf: '0'}).iloc[:100].tolist(), # Assuming sensors can be None
+#         "geo_altitude": geo_df["geo_altitude"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
+#         "squawk": geo_df["squawk"].astype(str).replace({np.nan: None, np.inf: '0'}).iloc[:100].tolist(), # Assuming squawk can be None
+#         "spi": geo_df["spi"].astype(str).iloc[:100].tolist(),  # Assuming spi is boolean or categorical
+#         "position_source": geo_df["position_source"].astype(str).iloc[:100].tolist(),  # Assuming position_source is categorical
+#         "time": geo_df["time"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
+#         "geometry": geo_df["geometry"].astype(str).replace({np.nan: None, np.inf: '0'}).iloc[:100].tolist() # Assuming geometry can be None
+#     }
+
+
+#     # Construct the payload
+#     payload = {
+#         "inputs": {
+#             "query": question,
+#             "table": table_data,
+#         }
+#     }
+
+#     # Get the model response
+#     response = query(payload)
+
+#     # Check if 'answer' is in response and return it as a sentence
+#     if 'answer' in response:
+#         answer = response['answer']
+#         return f"The answer to your question '{question}': :orange[{answer}]"
+#     else:
+#         return "The model could not find an answer to your question."
+
+
+# def flight_tracking(flight_view_level, country, local_time_zone, flight_info, airport, color):
+#     geolocator = Nominatim(user_agent="flight_tracker")
+#     loc = geolocator.geocode(country)
+#     loc_box = loc[1]
+#     extend_left =+12*flight_view_level
+#     extend_right =+10*flight_view_level
+#     extend_top =+10*flight_view_level
+#     extend_bottom =+ 18*flight_view_level
+#     lat_min, lat_max = (loc_box[0] - extend_left), loc_box[0]+extend_right
+#     lon_min, lon_max = (loc_box[1] - extend_bottom), loc_box[1]+extend_top
+    
+#     tile_zoom = 8 # zoom of the map loaded by contextily
+#     figsize = (15, 15)
+#     columns = ["icao24","callsign","origin_country","time_position","last_contact","longitude","latitude",
+#             "baro_altitude","on_ground","velocity","true_track","vertical_rate","sensors","geo_altitude",
+#             "squawk","spi","position_source",]
+#     data_url = "https://raw.githubusercontent.com/ashok2216-A/ashok_airport-data/main/data/airports.dat"
+#     column_names = ["Airport ID", "Name", "City", "Country", "IATA/FAA", "ICAO", "Latitude", "Longitude",
+#                     "Altitude", "Timezone", "DST", "Tz database time zone", "Type", "Source"]
+#     airport_df = pd.read_csv(data_url, header=None, names=column_names)
+#     airport_locations = airport_df[["Name", "City", "Country", "IATA/FAA", "Latitude", "Longitude"]]
+#     airport_country_loc = airport_locations[airport_locations['Country'] == str(loc)]
+#     airport_country_loc = airport_country_loc[(airport_country_loc['Country'] == str(loc)) & (airport_country_loc['Latitude'] >= lat_min) &
+#                             (airport_country_loc['Latitude'] <= lat_max) & (airport_country_loc['Longitude'] >= lon_min) &
+#                             (airport_country_loc['Longitude'] <= lon_max)]
+#     def get_traffic_gdf():
+#         url_data = (
+#                 f"https://@opensky-network.org/api/states/all?"
+#                 f"lamin={str(lat_min)}"
+#                 f"&lomin={str(lon_min)}"
+#                 f"&lamax={str(lat_max)}"
+#                 f"&lomax={str(lon_max)}")
+#         json_dict = requests.get(url_data).json()
+
+#         unix_timestamp = int(json_dict["time"])
+#         local_timezone = pytz.timezone(local_time_zone) # get pytz timezone
+#         local_time = datetime.fromtimestamp(unix_timestamp, local_timezone).strftime('%Y-%m-%d %H:%M:%S')
+#         time = []
+#         for i in range(len(json_dict['states'])):
+#             time.append(local_time)
+#         df_time = pd.DataFrame(time,columns=['time'])
+#         state_df = pd.DataFrame(json_dict["states"],columns=columns)
+#         state_df['time'] = df_time
+#         gdf = gpd.GeoDataFrame(
+#                 state_df,
+#                 geometry=gpd.points_from_xy(state_df.longitude, state_df.latitude),
+#                 crs={"init": "epsg:4326"},  # WGS84
+#             )
+#         # banner_image = Image.open('banner.png')
+#         # st.image(banner_image, width=300)
+#         st.title("Live Flight Tracker")
+#         st.subheader('Flight Details', divider='rainbow')
+#         st.write('Location: {0}'.format(loc))
+#         st.write('Current Local Time: {0}-{1}:'.format(local_time, local_time_zone))
+#         st.write("Minimum_latitude is {0} and Maximum_latitude is {1}".format(lat_min, lat_max))
+#         st.write("Minimum_longitude is {0} and Maximum_longitude is {1}".format(lon_min, lon_max))
+#         st.write('Number of Visible Flights: {}'.format(len(json_dict['states'])))
+#         st.write('Plotting the flight: {}'.format(flight_info))
+#         st.subheader('Map Visualization', divider='rainbow')
+#         st.write('****Click ":orange[Update Map]" Button to Refresh the Map****')
+#         return gdf
+
+#     geo_df = get_traffic_gdf()
+#     if airport == 0:
+#         fig = px.scatter_mapbox(geo_df, lat="latitude", lon="longitude",color=flight_info,
+#                             color_continuous_scale=color, zoom=4,width=1200, height=600,opacity=1,
+#                             hover_name ='origin_country',hover_data=['callsign', 'baro_altitude',
+#         'on_ground', 'velocity', 'true_track', 'vertical_rate', 'geo_altitude'], template='plotly_dark')
+#     elif airport == 1:
+#         fig = px.scatter_mapbox(geo_df, lat="latitude", lon="longitude",color=flight_info,
+#                             color_continuous_scale=color, zoom=4,width=1200, height=600,opacity=1,
+#                             hover_name ='origin_country',hover_data=['callsign', 'baro_altitude',
+#         'on_ground', 'velocity', 'true_track', 'vertical_rate', 'geo_altitude'], template='plotly_dark')
+#         fig.add_trace(px.scatter_mapbox(airport_country_loc, lat="Latitude", lon="Longitude",
+#                                         hover_name ='Name', hover_data=["City", "Country", "IATA/FAA"]).data[0])
+#     else: None
+#     fig.update_layout(mapbox_style="carto-darkmatter")
+#     fig.update_layout(margin={"r": 0, "t": 0, "l": 0, "b": 0})
+#     # out = fig.show())
+#     out = st.plotly_chart(fig, theme=None)
+#     return out
+# st.set_page_config(
+#     layout="wide"
+# )
+# image = Image.open('logo.png')
+# add_selectbox = st.sidebar.image(
+#     image, width=150
+# )
+# add_selectbox = st.sidebar.subheader(
+#     "Configure Map",divider='rainbow'
+# )
+# with st.sidebar:
+#     Refresh = st.button('Update Map', key=1)
+#     on = st.toggle('View Airports')
+#     if on:
+#         air_port = 1
+#         st.write(':rainbow[Nice Work Buddy!]')
+#         st.write('Now Airports are Visible')
+#     else:
+#         air_port=0
+#     view = st.slider('Increase Flight Visibility',1,6,2)
+#     st.write("You Selected:", view)
+#     cou = st.text_input('Type Country Name', 'north america')
+#     st.write('The current Country name is', cou)
+#     time = st.text_input('Type Time Zone Name (Ex: America/Toronto, Europe/Berlin)', 'Asia/Kolkata')
+#     st.write('The current Time Zone is', time)
+#     info = st.selectbox(
+#     'Select Flight Information',
+#     ('baro_altitude',
+#         'on_ground', 'velocity',
+#         'geo_altitude'))
+#     st.write('Plotting the data of Flight:', info)
+#     clr = st.radio('Pick A Color for Scatter Plot',["rainbow","ice","hot"])
+#     if clr == "rainbow":
+#         st.write('The current color is', "****:rainbow[Rainbow]****")
+#     elif clr == 'ice':
+#         st.write('The current color is', "****:blue[Ice]****")
+#     elif clr == 'hot':
+#         st.write('The current color is', "****:red[Hot]****")
+#     else: None
+# # with st.spinner('Wait!, We Requesting API Data...'):
+# # try:
+# flight_tracking(flight_view_level=view, country=cou,flight_info=info,
+#             local_time_zone=time, airport=air_port, color=clr)
+# st.subheader('Ask your Questions!', divider='rainbow')
+# st.write("Google's TAPAS base LLM model 🤖")
+# geo_df = flight_data(flight_view_level = view, country= cou, flight_info=info, local_time_zone=time, airport=1)
+# question = st.text_input('Type your questions here', "What is the squawk code for SWR9XD?")
+# result = query_flight_data(geo_df, question)
+# st.markdown(result)
+# # except TypeError:
+# #     st.error(':red[Error: ] Please Re-run this page.', icon="🚨")
+# #     st.button('Re-run', type="primary")
+# # st.snow()
+
+
+
+# # import streamlit as st
+# # from huggingface_hub import InferenceClient
+# # import os
+
+# # hf_token = os.getenv("HF_TOKEN")
+# # # Set up the Hugging Face Inference Client
+# # client = InferenceClient(
+# #     provider="together",  # Replace with the correct provider if needed
+# #     api_key= hf_token  # Replace with your Hugging Face API key
+# # )
+
+# # # Streamlit app title
+# # st.title("🤖 Deepseek R1 Chatbot")
+# # st.write("Chat with the Deepseek R1 model powered by Hugging Face Inference API.")
+
+# # # Initialize session state to store chat history
+# # if "messages" not in st.session_state:
+# #     st.session_state.messages = []
+
+# # # Display chat history
+# # for message in st.session_state.messages:
+# #     with st.chat_message(message["role"]):
+# #         st.markdown(message["content"])
+
+# # # User input
+# # if prompt := st.chat_input("What would you like to ask?"):
+# #     # Add user message to chat history
+# #     st.session_state.messages.append({"role": "user", "content": prompt})
+# #     with st.chat_message("user"):
+# #         st.markdown(prompt)
+
+# #     # Generate response from Deepseek R1 model
+# #     with st.spinner("Thinking..."):
+# #         try:
+# #             # Prepare the messages for the model
+# #             messages = [{"role": m["role"], "content": m["content"]} for m in st.session_state.messages]
+            
+# #             # Call the Hugging Face Inference API
+# #             completion = client.chat.completions.create(
+# #                 model="deepseek-ai/DeepSeek-R1",  # Replace with the correct model name
+# #                 messages=messages,
+# #                 max_tokens=500
+# #             )
+            
+# #             # Extract the model's response
+# #             response = completion.choices[0].message.content
+
+# #             # Add model's response to chat history
+# #             st.session_state.messages.append({"role": "assistant", "content": response})
+# #             with st.chat_message("assistant"):
+# #                 st.markdown(response)
+        
+# #         except Exception as e:
+# #             st.error(f"An error occurred: {e}")
+
+
 '''Copyright 2024 Ashok Kumar
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -29,11 +318,70 @@ from geopy.exc import GeocoderTimedOut
 from geopy.geocoders import Nominatim
 import warnings
 warnings.filterwarnings('ignore')
-from plotly.graph_objs import Marker
-import plotly.express as px
+import folium
+from folium import plugins
 import streamlit as st
+import streamlit_folium as st_folium
 from data import flight_data
 from huggingface_hub import InferenceApi, login, InferenceClient
+import branca.colormap as cm
+from functools import lru_cache
+import time
+
+# Cache the airport data to avoid reloading it every time
+@st.cache_data(ttl=3600)  # Cache for 1 hour
+def load_airport_data():
+    data_url = "https://raw.githubusercontent.com/ashok2216-A/ashok_airport-data/main/data/airports.dat"
+    column_names = ["Airport ID", "Name", "City", "Country", "IATA/FAA", "ICAO", "Latitude", "Longitude",
+                    "Altitude", "Timezone", "DST", "Tz database time zone", "Type", "Source"]
+    return pd.read_csv(data_url, header=None, names=column_names)
+
+# Cache geocoding results
+@st.cache_data(ttl=3600)
+def get_location(country):
+    geolocator = Nominatim(user_agent="flight_tracker")
+    return geolocator.geocode(country)
+
+# Cache flight data fetching
+@st.cache_data(ttl=60)  # Cache for 1 minute
+def fetch_flight_data(lat_min, lat_max, lon_min, lon_max):
+    try:
+        # OpenSky Network API endpoint
+        url = "https://opensky-network.org/api/states/all"
+        
+        # Parameters for the request
+        params = {
+            'lamin': lat_min,
+            'lamax': lat_max,
+            'lomin': lon_min,
+            'lomax': lon_max
+        }
+        
+        # Make the request with a timeout
+        response = requests.get(url, params=params, timeout=10)
+        
+        # Check if the request was successful
+        response.raise_for_status()
+        
+        # Parse the JSON response
+        data = response.json()
+        
+        # Check if we got valid data
+        if not data or 'states' not in data:
+            st.warning("No flight data available for the selected area.")
+            return {'states': [], 'time': 0}
+            
+        return data
+        
+    except requests.exceptions.RequestException as e:
+        st.error(f"Error fetching flight data: {str(e)}")
+        return {'states': [], 'time': 0}
+    except json.JSONDecodeError as e:
+        st.error(f"Error parsing flight data: {str(e)}")
+        return {'states': [], 'time': 0}
+    except Exception as e:
+        st.error(f"Unexpected error: {str(e)}")
+        return {'states': [], 'time': 0}
 
 
 hf_token = os.getenv("HF_TOKEN")
@@ -41,62 +389,123 @@ if hf_token is None:
     raise ValueError("Hugging Face token not found. Please set the HF_TOKEN environment variable.")
 login(hf_token)
 
-
-API_URL = "https://api-inference.huggingface.co/models/google/tapas-base-finetuned-wtq"
+# Hugging Face model configuration
+HF_API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
 headers = {"Authorization": f"Bearer {hf_token}"}
 
-def query(payload):
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return response.json()
+def query_llm(prompt):
+    try:
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "max_new_tokens": 250,
+                "temperature": 0.1,
+                "top_p": 0.95,
+                "return_full_text": False
+            }
+        }
+        
+        response = requests.post(HF_API_URL, headers=headers, json=payload)
+        response.raise_for_status()
+        return response.json()[0]['generated_text']
+    except Exception as e:
+        st.error(f"Error querying language model: {str(e)}")
+        return None
 
 def query_flight_data(geo_df, question):
-
-
-    table_data = {
-        "icao24": geo_df["icao24"].astype(str).iloc[:100].tolist(),
-        "callsign": geo_df["callsign"].astype(str).replace({np.nan: None, np.inf: '0'}).iloc[:100].tolist(),
-        "origin_country": geo_df["origin_country"].astype(str).replace({np.nan: None, np.inf: '0'}).iloc[:100].tolist(),
-        "time_position": geo_df["time_position"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
-        "last_contact": geo_df["last_contact"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
-        "longitude": geo_df["longitude"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
-        "latitude": geo_df["latitude"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
-        "baro_altitude": geo_df["baro_altitude"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
-        "on_ground": geo_df["on_ground"].astype(str).iloc[:100].tolist(),  # Assuming on_ground is boolean or categorical
-        "velocity": geo_df["velocity"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
-        "true_track": geo_df["true_track"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
-        "vertical_rate": geo_df["vertical_rate"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
-        "sensors": geo_df["sensors"].astype(str).replace({np.nan: None, np.inf: '0'}).iloc[:100].tolist(), # Assuming sensors can be None
-        "geo_altitude": geo_df["geo_altitude"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
-        "squawk": geo_df["squawk"].astype(str).replace({np.nan: None, np.inf: '0'}).iloc[:100].tolist(), # Assuming squawk can be None
-        "spi": geo_df["spi"].astype(str).iloc[:100].tolist(),  # Assuming spi is boolean or categorical
-        "position_source": geo_df["position_source"].astype(str).iloc[:100].tolist(),  # Assuming position_source is categorical
-        "time": geo_df["time"].astype(str).replace({np.nan: '0', np.inf: '0'}).iloc[:100].tolist(),
-        "geometry": geo_df["geometry"].astype(str).replace({np.nan: None, np.inf: '0'}).iloc[:100].tolist() # Assuming geometry can be None
+    # Preprocess the question to extract key information
+    question = question.lower().strip()
+    
+    # Common flight information queries and their corresponding columns
+    query_mappings = {
+        'callsign': ['callsign'],
+        'altitude': ['baro_altitude', 'geo_altitude'],
+        'speed': ['velocity'],
+        'direction': ['true_track'],
+        'country': ['origin_country'],
+        'squawk': ['squawk'],
+        'icao': ['icao24'],
+        'vertical': ['vertical_rate'],
+        'ground': ['on_ground'],
+        'position': ['latitude', 'longitude'],
+        'time': ['time_position', 'last_contact']
     }
+    
+    # Extract the identifier (usually callsign or icao) from the question
+    identifier = None
+    if 'for' in question:
+        identifier = question.split('for')[-1].strip()
+    elif 'of' in question:
+        identifier = question.split('of')[-1].strip()
+    elif 'about' in question:
+        identifier = question.split('about')[-1].strip()
+    
+    if not identifier:
+        return "Please specify a flight identifier (callsign or ICAO code) in your question."
+    
+    # Try to find the flight by callsign or icao
+    flight_data = None
+    if identifier in geo_df['callsign'].values:
+        flight_data = geo_df[geo_df['callsign'] == identifier]
+    elif identifier in geo_df['icao24'].values:
+        flight_data = geo_df[geo_df['icao24'] == identifier]
+    
+    if flight_data is None or flight_data.empty:
+        return f"Could not find flight information for {identifier}. Please check the flight identifier and try again."
+    
+    # Prepare flight data for the LLM
+    flight_info = {}
+    for col in flight_data.columns:
+        if col in flight_data.columns:
+            value = flight_data[col].iloc[0]
+            if pd.notna(value):
+                if col == 'baro_altitude' or col == 'geo_altitude':
+                    flight_info[col] = f"{value} meters"
+                elif col == 'velocity':
+                    flight_info[col] = f"{value} m/s"
+                elif col == 'true_track':
+                    flight_info[col] = f"{value} degrees"
+                elif col == 'vertical_rate':
+                    flight_info[col] = f"{value} m/s"
+                elif col == 'latitude':
+                    flight_info[col] = f"{value}° N"
+                elif col == 'longitude':
+                    flight_info[col] = f"{value}° E"
+                else:
+                    flight_info[col] = str(value)
+    
+    if not flight_info:
+        return f"No information available for flight {identifier}."
+    
+    # Create a prompt for the LLM
+    prompt = f"""You are a flight information assistant. Answer the following question about flight {identifier} using the provided flight data.
 
+Question: {question}
 
-    # Construct the payload
-    payload = {
-        "inputs": {
-            "query": question,
-            "table": table_data,
-        }
-    }
+Flight Data:
+{json.dumps(flight_info, indent=2)}
 
-    # Get the model response
-    response = query(payload)
+Please provide a clear and concise answer focusing on the specific information requested in the question. If the question asks for information not available in the data, say so clearly."""
 
-    # Check if 'answer' is in response and return it as a sentence
-    if 'answer' in response:
-        answer = response['answer']
-        return f"The answer to your question '{question}': :orange[{answer}]"
+    # Get response from LLM
+    llm_response = query_llm(prompt)
+    
+    if llm_response:
+        return llm_response
     else:
-        return "The model could not find an answer to your question."
-
+        # Fallback to direct data response if LLM fails
+        response = f"Flight Information for {identifier}:\n"
+        for key, value in flight_info.items():
+            response += f"- {key.replace('_', ' ').title()}: {value}\n"
+        return response
 
 def flight_tracking(flight_view_level, country, local_time_zone, flight_info, airport, color):
-    geolocator = Nominatim(user_agent="flight_tracker")
-    loc = geolocator.geocode(country)
+    # Get cached location data
+    loc = get_location(country)
+    if not loc:
+        st.error("Could not find location. Please try a different country name.")
+        return
+    
     loc_box = loc[1]
     extend_left =+12*flight_view_level
     extend_right =+10*flight_view_level
@@ -105,76 +514,215 @@ def flight_tracking(flight_view_level, country, local_time_zone, flight_info, ai
     lat_min, lat_max = (loc_box[0] - extend_left), loc_box[0]+extend_right
     lon_min, lon_max = (loc_box[1] - extend_bottom), loc_box[1]+extend_top
     
-    tile_zoom = 8 # zoom of the map loaded by contextily
-    figsize = (15, 15)
     columns = ["icao24","callsign","origin_country","time_position","last_contact","longitude","latitude",
             "baro_altitude","on_ground","velocity","true_track","vertical_rate","sensors","geo_altitude",
             "squawk","spi","position_source",]
-    data_url = "https://raw.githubusercontent.com/ashok2216-A/ashok_airport-data/main/data/airports.dat"
-    column_names = ["Airport ID", "Name", "City", "Country", "IATA/FAA", "ICAO", "Latitude", "Longitude",
-                    "Altitude", "Timezone", "DST", "Tz database time zone", "Type", "Source"]
-    airport_df = pd.read_csv(data_url, header=None, names=column_names)
+    
+    # Get cached airport data
+    airport_df = load_airport_data()
     airport_locations = airport_df[["Name", "City", "Country", "IATA/FAA", "Latitude", "Longitude"]]
     airport_country_loc = airport_locations[airport_locations['Country'] == str(loc)]
-    airport_country_loc = airport_country_loc[(airport_country_loc['Country'] == str(loc)) & (airport_country_loc['Latitude'] >= lat_min) &
-                            (airport_country_loc['Latitude'] <= lat_max) & (airport_country_loc['Longitude'] >= lon_min) &
-                            (airport_country_loc['Longitude'] <= lon_max)]
+    airport_country_loc = airport_country_loc[(airport_country_loc['Country'] == str(loc)) & 
+                                            (airport_country_loc['Latitude'] >= lat_min) &
+                                            (airport_country_loc['Latitude'] <= lat_max) & 
+                                            (airport_country_loc['Longitude'] >= lon_min) &
+                                            (airport_country_loc['Longitude'] <= lon_max)]
+    
     def get_traffic_gdf():
-        url_data = (
-                f"https://@opensky-network.org/api/states/all?"
-                f"lamin={str(lat_min)}"
-                f"&lomin={str(lon_min)}"
-                f"&lamax={str(lat_max)}"
-                f"&lomax={str(lon_max)}")
-        json_dict = requests.get(url_data).json()
-
-        unix_timestamp = int(json_dict["time"])
-        local_timezone = pytz.timezone(local_time_zone) # get pytz timezone
-        local_time = datetime.fromtimestamp(unix_timestamp, local_timezone).strftime('%Y-%m-%d %H:%M:%S')
-        time = []
-        for i in range(len(json_dict['states'])):
-            time.append(local_time)
-        df_time = pd.DataFrame(time,columns=['time'])
-        state_df = pd.DataFrame(json_dict["states"],columns=columns)
-        state_df['time'] = df_time
-        gdf = gpd.GeoDataFrame(
+        # Get cached flight data
+        json_dict = fetch_flight_data(lat_min, lat_max, lon_min, lon_max)
+        
+        if not json_dict or not json_dict.get('states'):
+            st.warning("No flight data available for the selected area.")
+            return None
+
+        try:
+            unix_timestamp = int(json_dict["time"])
+            local_timezone = pytz.timezone(local_time_zone)
+            local_time = datetime.fromtimestamp(unix_timestamp, local_timezone).strftime('%Y-%m-%d %H:%M:%S')
+            
+            # Optimize DataFrame creation
+            state_df = pd.DataFrame(json_dict["states"], columns=columns)
+            state_df['time'] = local_time
+            
+            # Create GeoDataFrame more efficiently
+            gdf = gpd.GeoDataFrame(
                 state_df,
                 geometry=gpd.points_from_xy(state_df.longitude, state_df.latitude),
-                crs={"init": "epsg:4326"},  # WGS84
+                crs="EPSG:4326"
             )
-        # banner_image = Image.open('banner.png')
-        # st.image(banner_image, width=300)
-        st.title("Live Flight Tracker")
-        st.subheader('Flight Details', divider='rainbow')
-        st.write('Location: {0}'.format(loc))
-        st.write('Current Local Time: {0}-{1}:'.format(local_time, local_time_zone))
-        st.write("Minimum_latitude is {0} and Maximum_latitude is {1}".format(lat_min, lat_max))
-        st.write("Minimum_longitude is {0} and Maximum_longitude is {1}".format(lon_min, lon_max))
-        st.write('Number of Visible Flights: {}'.format(len(json_dict['states'])))
-        st.write('Plotting the flight: {}'.format(flight_info))
-        st.subheader('Map Visualization', divider='rainbow')
-        st.write('****Click ":orange[Update Map]" Button to Refresh the Map****')
-        return gdf
+            
+            # Display information
+            st.title("Live Flight Tracker")
+            st.subheader('Flight Details', divider='rainbow')
+            st.write('Location: {0}'.format(loc))
+            st.write('Current Local Time: {0}-{1}:'.format(local_time, local_time_zone))
+            st.write("Minimum_latitude is {0} and Maximum_latitude is {1}".format(lat_min, lat_max))
+            st.write("Minimum_longitude is {0} and Maximum_longitude is {1}".format(lon_min, lon_max))
+            st.write('Number of Visible Flights: {}'.format(len(json_dict['states'])))
+            st.write('Plotting the flight: {}'.format(flight_info))
+            st.subheader('Map Visualization', divider='rainbow')
+            st.write('****Click ":orange[Update Map]" Button to Refresh the Map****')
+            return gdf
+            
+        except Exception as e:
+            st.error(f"Error processing flight data: {str(e)}")
+            return None
 
+    # Get traffic data
     geo_df = get_traffic_gdf()
-    if airport == 0:
-        fig = px.scatter_mapbox(geo_df, lat="latitude", lon="longitude",color=flight_info,
-                            color_continuous_scale=color, zoom=4,width=1200, height=600,opacity=1,
-                            hover_name ='origin_country',hover_data=['callsign', 'baro_altitude',
-        'on_ground', 'velocity', 'true_track', 'vertical_rate', 'geo_altitude'], template='plotly_dark')
-    elif airport == 1:
-        fig = px.scatter_mapbox(geo_df, lat="latitude", lon="longitude",color=flight_info,
-                            color_continuous_scale=color, zoom=4,width=1200, height=600,opacity=1,
-                            hover_name ='origin_country',hover_data=['callsign', 'baro_altitude',
-        'on_ground', 'velocity', 'true_track', 'vertical_rate', 'geo_altitude'], template='plotly_dark')
-        fig.add_trace(px.scatter_mapbox(airport_country_loc, lat="Latitude", lon="Longitude",
-                                        hover_name ='Name', hover_data=["City", "Country", "IATA/FAA"]).data[0])
-    else: None
-    fig.update_layout(mapbox_style="carto-darkmatter")
-    fig.update_layout(margin={"r": 0, "t": 0, "l": 0, "b": 0})
-    # out = fig.show())
-    out = st.plotly_chart(fig, theme=None)
-    return out
+    if geo_df is None:
+        return
+    
+    # Create a base map
+    m = folium.Map(
+        location=[loc_box[0], loc_box[1]],
+        zoom_start=6,
+        tiles='CartoDB dark_matter'
+    )
+    
+    # Create colormap
+    if color == "rainbow":
+        colormap = cm.LinearColormap(
+            colors=['red', 'yellow', 'green', 'blue', 'purple'],
+            vmin=geo_df[flight_info].min(),
+            vmax=geo_df[flight_info].max()
+        )
+    elif color == "ice":
+        colormap = cm.LinearColormap(
+            colors=['white', 'lightblue', 'blue'],
+            vmin=geo_df[flight_info].min(),
+            vmax=geo_df[flight_info].max()
+        )
+    else:  # hot
+        colormap = cm.LinearColormap(
+            colors=['yellow', 'orange', 'red'],
+            vmin=geo_df[flight_info].min(),
+            vmax=geo_df[flight_info].max()
+        )
+    
+    # Pre-compute icon HTML template
+    icon_template = """
+    <div style="transform: rotate({rotation_angle}deg);">
+        <svg width="24" height="24" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
+            <path d="M21 16v-2l-8-5V3.5c0-.83-.67-1.5-1.5-1.5S10 2.67 10 3.5V9l-8 5v2l8-2.5V19l-2 1.5V22l3.5-1 3.5 1v-1.5L13 19v-5.5l8 2.5z" fill="{color_hex}"/>
+        </svg>
+    </div>
+    """
+    
+    # Pre-compute tooltip template
+    tooltip_template = """
+    <div style="font-size: 12px; font-family: Arial, sans-serif; max-width: 300px;">
+        <div style="font-weight: bold; font-size: 14px; margin-bottom: 5px; color: #2c3e50;">
+            Flight: {callsign}
+        </div>
+        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 5px;">
+            {rows}
+        </div>
+    </div>
+    """
+    
+    # Add flight markers
+    for idx, row in geo_df.iterrows():
+        if pd.notna(row['latitude']) and pd.notna(row['longitude']):
+            # Get color based on flight_info value
+            value = row[flight_info] if pd.notna(row[flight_info]) else geo_df[flight_info].min()
+            color_hex = colormap(value)
+            
+            # Create custom flight icon with rotation
+            rotation_angle = row['true_track'] if pd.notna(row['true_track']) else 0
+            icon_html = icon_template.format(rotation_angle=rotation_angle, color_hex=color_hex)
+            
+            # Create tooltip rows
+            tooltip_rows = []
+            for col in columns:
+                val = row[col] if pd.notna(row[col]) else 'N/A'
+                if col in ['baro_altitude', 'geo_altitude']:
+                    val = f"{val} m"
+                elif col == 'velocity':
+                    val = f"{val} m/s"
+                elif col == 'true_track':
+                    val = f"{val}°"
+                tooltip_rows.append(f'<div style="font-weight: bold;">{col}:</div><div>{val}</div>')
+            
+            # Create tooltip
+            tooltip_html = tooltip_template.format(
+                callsign=row['callsign'] if pd.notna(row['callsign']) else 'Unknown',
+                rows='\n'.join(tooltip_rows)
+            )
+            
+            # Create popup content
+            popup_content = f"""
+            <div style="font-size: 14px; font-family: Arial, sans-serif; max-width: 300px;">
+                <div style="font-weight: bold; font-size: 16px; margin-bottom: 10px; color: #2c3e50;">
+                    Flight: {row['callsign'] if pd.notna(row['callsign']) else 'Unknown'}
+                </div>
+                <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 5px;">
+                    <div style="font-weight: bold;">ICAO24:</div>
+                    <div>{row['icao24'] if pd.notna(row['icao24']) else 'N/A'}</div>
+                    <div style="font-weight: bold;">Origin Country:</div>
+                    <div>{row['origin_country'] if pd.notna(row['origin_country']) else 'N/A'}</div>
+                    <div style="font-weight: bold;">Time Position:</div>
+                    <div>{row['time_position'] if pd.notna(row['time_position']) else 'N/A'}</div>
+                    <div style="font-weight: bold;">Last Contact:</div>
+                    <div>{row['last_contact'] if pd.notna(row['last_contact']) else 'N/A'}</div>
+                    <div style="font-weight: bold;">Baro Altitude:</div>
+                    <div>{row['baro_altitude'] if pd.notna(row['baro_altitude']) else 'N/A'} m</div>
+                    <div style="font-weight: bold;">Geo Altitude:</div>
+                    <div>{row['geo_altitude'] if pd.notna(row['geo_altitude']) else 'N/A'} m</div>
+                    <div style="font-weight: bold;">Velocity:</div>
+                    <div>{row['velocity'] if pd.notna(row['velocity']) else 'N/A'} m/s</div>
+                    <div style="font-weight: bold;">True Track:</div>
+                    <div>{row['true_track'] if pd.notna(row['true_track']) else 'N/A'}°</div>
+                    <div style="font-weight: bold;">Vertical Rate:</div>
+                    <div>{row['vertical_rate'] if pd.notna(row['vertical_rate']) else 'N/A'} m/s</div>
+                    <div style="font-weight: bold;">Squawk:</div>
+                    <div>{row['squawk'] if pd.notna(row['squawk']) else 'N/A'}</div>
+                    <div style="font-weight: bold;">On Ground:</div>
+                    <div>{row['on_ground'] if pd.notna(row['on_ground']) else 'N/A'}</div>
+                    <div style="font-weight: bold;">SPI:</div>
+                    <div>{row['spi'] if pd.notna(row['spi']) else 'N/A'}</div>
+                    <div style="font-weight: bold;">Position Source:</div>
+                    <div>{row['position_source'] if pd.notna(row['position_source']) else 'N/A'}</div>
+                </div>
+            </div>
+            """
+            
+            # Create custom icon
+            icon = folium.DivIcon(
+                html=icon_html,
+                icon_size=(24, 24),
+                icon_anchor=(12, 12)
+            )
+            
+            # Add marker to map
+            folium.Marker(
+                location=[row['latitude'], row['longitude']],
+                icon=icon,
+                popup=folium.Popup(popup_content, max_width=300),
+                tooltip=tooltip_html
+            ).add_to(m)
+    
+    # Add airports if selected
+    if airport == 1:
+        for idx, row in airport_country_loc.iterrows():
+            folium.Marker(
+                location=[row['Latitude'], row['Longitude']],
+                icon=folium.Icon(icon='plane', prefix='fa', color='blue'),
+                popup=f"<b>{row['Name']}</b><br>IATA: {row['IATA/FAA']}<br>City: {row['City']}",
+                tooltip=f"Airport: {row['Name']}"
+            ).add_to(m)
+    
+    # Add colormap to the map
+    colormap.add_to(m)
+    
+    # Add a layer control
+    folium.LayerControl().add_to(m)
+    
+    # Display the map in Streamlit
+    st_folium.folium_static(m, width=1200, height=600)
+    return None
+
 st.set_page_config(
     layout="wide"
 )