#!/usr/bin/env python
# coding: utf-8

# # Folium Map Plotting in MSTICPy
# 
# ## Introduction
# This module contains a class and functions that wraps the `folium` package to plot geo-location data.
# 
# Read the [MSTICPy Folium documentation](https://msticpy.readthedocs.io/visualization/FoliumMap.html)
# 
# Read the [Folium Python package documentation](https://python-visualization.github.io/folium/)
# 
# You must have msticpy installed to run this notebook:
# ```
# %pip install --upgrade msticpy
# ```

# ## Contents
# 
# - Introduction
# - Using Folium with MSTICPy
# - Plotting from a Pandas DataFrame
#   - Basic plot
#   - Plotting Coordinates
#   - Map Layers
#   - Adding custom tooltip and popup data
#   - Using custom icons
# - Using the plot_map function 
# - Folium map class
#   - Underlying folium map object
#   - Adding IP entities
#   - Adding IP addresses
#   - Adding locations
#   - Plotting IpAddress entities with location data
#   - Using different colors and icons
#   - Using custom icons
# - Utility functions
#   - Calculating the center point of locations
#   - Calculating the distance between locations  
# 

# In[2]:


# Imports
import msticpy
msticpy.init_notebook(verbosity=0)

from msticpy.vis.foliummap import FoliumMap


# ## Using Folium with MSTICPy
# 
# You can use Folium via a pandas accessor, the plot_map function
# or directly interacting with our FoliumMap class.

# ## Plotting from a Pandas DataFrame
# 
# MSTICPy uses pandas accessors to expose a lot of its 
# visualization functions. 
# 
# Plotting with Folium can be done directly from a pandas
# DataFrame using the `mp_plot.folium_map` extension.
# 
# This function returns an instance of the `FoliumMap`
# class that you can further customize (see FoliumMap class later in the document).
# 
# ### Basic plot

# In[3]:


# read in a DataFrame from a csv file
geo_loc_df = (
    pd
    .read_csv("data/ip_locs.csv", index_col=0)
    .dropna(subset=["Latitude", "Longitude", "IpAddress"])  # We need to remove an NaN values
)
    
display(geo_loc_df.head(5))

geo_loc_df.mp_plot.folium_map(ip_column="IpAddress")


# ### Plotting from coordinates
# 
# If you already have coordinates in the data you can use 
# these (rather than looking up the IP location again) using
# the `lat_column` and `long_column` parameters.

# In[4]:


geo_loc_df.mp_plot.folium_map(
    lat_column="Latitude", long_column="Longitude", zoom_start=10
)


# ### Plotting Layers
# 
# You can use the Folium layers feature, specifying a column
# value on which to group each layer with the `layer_column` parameter.

# In[5]:


geo_loc_df.mp_plot.folium_map(
    ip_column="IpAddress", layer_column="CountryName", zoom_start=2
)


# ### Adding custom tooltip and popup data
# 
# You use DataFrame column values to populate the tooltip
# and popup elements for each marker with the
# `tooltip_columns` and `popup_columns` parameters.

# In[6]:


# Create some data to display
data_df = pd.DataFrame({
    "Status": ["Home", "Office", "Vacation"] * (len(geo_loc_df) // 3),
    "Friendliness": ["Warm", "Cold", "Medium"] * (len(geo_loc_df) // 3),
    "Flavor": ["Chocolate", "Cinnamon", "Mango"] * (len(geo_loc_df) // 3),
    "SpiceLevel": [1, 2, 3] * (len(geo_loc_df) // 3)
})
geo_loc_data_df = pd.concat([geo_loc_df, data_df], axis=1).dropna(subset=["IpAddress"])
geo_loc_data_df.head(3)


# In[7]:


geo_loc_data_df.mp_plot.folium_map(
    ip_column="IpAddress",
    layer_column="CountryName",
    tooltip_columns=["Status", "Flavor"],
    popup_columns=["Friendliness", "SpiceLevel", "Status", "Flavor"],
    zoom_start=2,
)


# ### Using custom icons
# 
# You can also control the icons used for each marker with the
# `icon_column` parameters. If you happen to have a column in your
# data that contains names of FontAwesome or GlyphIcons icons.
# More typically you would combine the `icon_column` with the
# `icon_map` parameter. You can specify either a dictionary or a
# function. For a dictionary, the value of the row in `icon_column`
# is used as a key - the value is a dictionary of icon parameters
# passed to the Folium.Icon class. For a method, the `icon_column`
# value is passed as a single parameter and the return value
# should be a dictionary of valid parameters for the `Icon` class.
# You can read the documentation for this function in the doc.
# 
# If `icon_map` is a dict it should contain keys that map to the
# value of `icon_col` and values that a dicts of valid
# folium Icon properties ("color", "icon_color", "icon", "angle", "prefix").
# The dict should include a "default" entry that will be used if the
# value in the DataFrame[icon_col] doesn't match any key.
# For example:
# 
# ```python
# 
#     icon_map = {
#         "high": {
#             "color": "red",
#             "icon": "warning",
#         },
#         "medium": {
#             "color": "orange",
#             "icon": "triangle-exclamation",
#             "prefix": "fa",
#         },
#         "default": {
#             "color": "blue",
#             "icon": "info-sign",
#         },
#     }
# ```
# 
# If icon_map is a function it should take a single str parameter
# (the item key) and return a dict of icon properties. It should
# return a default set of values if the key does not match a known
# key. The `icon_col` value for each row will be passed to this
# function and the return value used to populate the Icon arguments.
# 
# For example:
# 
# ```python
# 
#     def icon_mapper(icon_key):
#         if icon_key.startswith("bad"):
#             return {
#                 "color": "red",
#                 "icon": "triangle-alert",
#             }
#         ...
#         else:
#             return {
#                 "color": "blue",
#                 "icon": "info-sign",
#             }
# ```
# 
# Check out the possible names for icons:
# 
# - FontAwesome icon (prefix "fa") names are available at https://fontawesome.com/
# - GlyphIcons icons (prefix "glyphicon") are available at https://www.glyphicons.com/
# 

# In[8]:


icon_map = {
    "US": {
        "color": "green",
        "icon": "flash",
    },
    "GB": {
        "color": "purple",
        "icon": "flash",
    },
    "default": {
        "color": "blue",
        "icon": "info-sign",
    },
}

geo_loc_df.mp_plot.folium_map(
    ip_column="AllExtIPs",
    icon_column="CountryCode",
    icon_map=icon_map,
    zoom_start=2,
)


# ## Using the `plot_map` function
# 
# The `plot_map` function is identical to the mp_plot.folium
# map accessor. You can import this directly and use in place
# of the pandas accessor.

# In[9]:


from msticpy.vis.foliummap import plot_map

plot_map(
    data=geo_loc_df,
    ip_column="AllExtIPs",
    icon_column="CountryCode",
    icon_map=icon_map,
    zoom_start=2,
)


# ## FoliumMap class
# 
# Use the Folium Map class when you want to build up data
# clusters and layers incrementally.
# 
# It now supports multiple data types for entry:
# - IpAddress (`map.add_ip_cluster`) and Geolocation (`map.add_geoloc_cluster`) entities
# - IP addresses (`map.add_ips`)
# - Locations (`map.add_locations`)
# - GeoHashes (`map.add_geo_hashes`)
# 
# You can also use other member functions to add layers and cluster groups.
# 
# ```
# FoliumMap(
#     title: str = 'layer1',
#     zoom_start: float = 2.5,
#     tiles=None,
#     width: str = '100%',
#     height: str = '100%',
#     location: list = None,
# )
# Wrapper class for Folium/Leaflet mapping.
# 
# Parameters
# ----------
# title : str, optional
#     Name of the layer (the default is 'layer1')
# zoom_start : int, optional
#     The zoom level of the map (the default is 7)
# tiles : [type], optional
#     Custom set of tiles or tile URL (the default is None)
# width : str, optional
#     Map display width (the default is '100%')
# height : str, optional
#     Map display height (the default is '100%')
# location : list, optional
#     Location to center map on
# 
# Attributes
# ----------
# folium_map : folium.Map
# ```

# In[10]:


folium_map = FoliumMap(location=(47.5982328,-122.331), zoom_start=14)
folium_map


# The underlying folium map object is accessible as the `folium_map` attribute

# In[11]:


type(folium_map.folium_map)


# ### Adding IP Entities to the map
# 
# ```
# fol_map.add_ip_cluster(
#     ip_entities: Iterable[msticpy.datamodel.entities.IpAddress],
#     **kwargs,
# )
# 
# ```

# In[12]:


import pickle
with open(b"data/ip_entities.pkl", "rb") as fh:
    ip_entities = pickle.load(fh)
ip_entities = [ip for ip in ip_entities if ip.Location and ip.Location.Latitude]

folium_map = FoliumMap(zoom_start=9)
folium_map.add_ip_cluster(ip_entities=ip_entities, color='orange')
folium_map.center_map()
folium_map


# ### Adding IP addresses
# 
# ```
# ip_map.add_ips(ip_addresses: Iterable[str], **kwargs)
# ```

# In[13]:


ips = geo_loc_df.query("State == 'California'").AllExtIPs.values

print("IP dataset", ips[:3], "...")
ip_map = FoliumMap(zoom_start="3")
ip_map.add_ips(ips)
ip_map.center_map()
ip_map


# ### Adding locations
# 
# ```
# ip_map.add_locations(locations: Iterable[Tuple[float, float]], **kwargs)
# ```

# In[14]:


locations = geo_loc_df.query("CountryCode != 'US'").apply(lambda x: (x.Latitude, x.Longitude), axis=1).values

print("Location dataset", locations[:3], "...")
ip_map.add_locations(locations)
ip_map.center_map()
ip_map


# ### Plotting IPAddress entities with location data

# In[15]:


# Create IP and GeoLocation Entities from the dataframe
def create_ip_entity(row):
    ip_ent = IpAddress(Address=row["AllExtIPs"])
    geo_loc = create_geo_entity(row)
    ip_ent.Location = geo_loc
    return ip_ent

def create_geo_entity(row):
    # get subset of fields for GeoLocation
    loc_props = row[["CountryCode", "CountryName","State", "City", "Longitude", "Latitude"]]
    return GeoLocation(**loc_props.to_dict())


geo_locs = list(geo_loc_df.apply(create_geo_entity, axis=1).values)
ip_ents = list(geo_loc_df.apply(create_ip_entity, axis=1).values)

fmap_ips = FoliumMap()
fmap_ips.add_ip_cluster(ip_entities=ip_ents[:20], color='blue')
fmap_ips.center_map()
fmap_ips


# ### Using different colors and icons

# In[16]:


fmap_ips.add_ip_cluster(ip_entities=ip_ents[30:40], color='red', icon="flash")
fmap_ips.center_map()
fmap_ips


# ### Using custom Icons
# 
# By default folium uses the information icon (i).
# Icons can be taken from the default Bootstrap set. See the default list here [glyphicons](https://www.w3schools.com/icons/bootstrap_icons_glyphicons.asp)
# 
# Alternatively you can use icons from the [Font Awesome collection](https://fontawesome.com/icons?d=gallery)
# by adding prefx="fa" and icon="icon_name" to the call to add_ip_cluster or add_geo_cluster.

# In[17]:


fmap_ips.add_geoloc_cluster(
    geo_locations=geo_locs[40:50],
    color='darkblue',
    icon="desktop",
    prefix="fa"
)
fmap_ips.center_map()
fmap_ips


# ## Utility Functions

# ### Calculate center point of entity locations

# In[18]:


from msticpy.vis.foliummap import get_map_center, get_center_ip_entities, get_center_geo_locs

print(get_center_geo_locs(geo_locs))
print(get_center_geo_locs(geo_locs, mode="mean"))

# get_map_center Will accept iterable of any entity type that is either
# an IpAddress entity or an entity that has properties of type IpAddress
print(get_map_center(ip_ents[30:40]))
print(get_map_center(ip_ents[:20]))

print(get_center_ip_entities(ip_ents[:20])) 


# ### Calculate distance between entity locations

# In[19]:


from msticpy.context.geoip import entity_distance
print("Distance between")
print(f"{ip_ents[0].Address} ({ip_ents[0].Location.City})")
print(f"{ip_ents[1].Address} ({ip_ents[1].Location.City})")
print(entity_distance(ip_ents[0], ip_ents[1]), "km", "\n")

print("Distance between")
print(f"{ip_ents[0].Address} ({ip_ents[0].Location.City})")
print(f"{ip_ents[13].Address} ({ip_ents[13].Location.City})")
print(entity_distance(ip_ents[0], ip_ents[13]), "km", "\n")