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

# # Custom Legend with `manual_key` Option
# 
# In Lets-Plot, as in ggplot2, legends are automatically generated based on the aesthetic mappings in the plot. 
# 
# Sometimes, however, this automatic generation doesn't provide the precise control needed for complex visualizations. 
# 
# The new `manual_key` option in plot layer addresses this limitation by allowing you to create custom legend entries specific to individual plot layers.
# 
# The `manual_key` option can be used in two ways:
# 
# - Simple case: pass a string to serve as the label for the legend entry.
# - Advanced case: use the `layer_key()` function for more detailed customization:
# 
#     - `label` - text for the legend element
#     - `group` - key used to group elements in the legend
#     - `index` - position of the element within its legend group
#     - `kwargs` - dictionary of aesthetic parameters to be applied in the legend
# 

# In[1]:


import pandas as pd

from lets_plot import *


# In[2]:


LetsPlot.setup_html()


# In[3]:


mpg = pd.read_csv("https://raw.githubusercontent.com/JetBrains/lets-plot-docs/master/data/mpg.csv")
mpg.head(3)


# In[4]:


ggplot(mpg, aes('displ', 'hwy')) + \
    geom_point(color='#878787') + \
    geom_smooth(method="loess", se=False, color='#2166ac') + \
    geom_smooth(method="lm", se=False, color='#b2182b')


# #### 1. Simple Custom Legend

# In[5]:


ggplot(mpg, aes('displ', 'hwy')) + \
    geom_point(color='#878787', 
               manual_key="Observed Data"                   # <-- label for the legend entry
              ) + \
    geom_smooth(method="loess", se=False, color='#2166ac',
                manual_key="LOESS Trend Line") + \
    geom_smooth(method="lm", se=False, color='#b2182b',
                manual_key="Linear Model Trend Line") + \
    ggsize(800, 400)


# #### 2.  The `layer_key()` Function

# All constants for the visual representation of the legend key are inherited from the geometry. Use the `layer_key()` function to override these values. You can also specify the legend group in which the element is displayed and the position of the element within the legend group.

# In[6]:


p = ggplot(mpg, aes('displ', 'hwy')) + \
    geom_point(color='#878787',
               manual_key=layer_key("Observed Data",          # <-- label for the legend entry
                                    color='pen', size=8)      # <-- override aesthetics
              ) + \
    geom_smooth(method="loess", se=False, color='#2166ac',
                manual_key=layer_key("LOESS", 
                                     group="Methods",         # <-- assign a group (note the same group in the layer below)
                                     index=1, size=4)) + \
    geom_smooth(method="lm", se=False, color='#b2182b',
                manual_key=layer_key("LM", group="Methods", index=0, size=4)) + \
    ggsize(800, 400)
p


# #### 3. Change Title of Custom Legends

# In[7]:


p + labs(
        manual="Scatter Plot",  # <-- "manual" is a default 'key' referring to a custom legend
        Methods="Trend Lines"   # <-- "Methods" is the group name referring to the second custom legend
)


# #### 4. Compound Legend Entries in Custom Legend

# In[8]:


(
    ggplot(mapping=aes(as_discrete('class', order_by='..y..'), 'hwy')) +
        # 1999    
        stat_summary(data=mpg[mpg.year == 1999], fun='mean', 
                     geom='line', size=2, color='#D87093', 
                     manual_key="1999") +            # <-- label for the legend entry
                                          
        stat_summary(data=mpg[mpg.year == 1999], fun='mean', 
                     geom='point', size=5, shape=22, 
                     fill='#D87093', color='paper', 
                     manual_key=layer_key("1999",    # <-- Using the same label as above results in a composite entry in the legend 
                                          size=3)) +        
        # 2008    
        stat_summary(data=mpg[mpg.year == 2008], fun='mean', geom='line', size=2, color='#708090', 
                     manual_key="2008") +
        stat_summary(data=mpg[mpg.year == 2008], fun='mean', geom='point', size=7, shape=23, fill='#708090', color='paper', 
                     manual_key=layer_key("2008", size=5)) +

        guides(manual=guide_legend(
            override_aes=dict(stroke=0,       # <-- remove `stroke` in the "manual" legend
                              size=1)         # <-- update `size` if not specified otherwise in the layer itself
        )) +  
    
        ggsize(800, 400) +
        theme(legend_direction='horizontal',
              legend_position=[0.9, 0.95],
              legend_justification=[1, 1],
              legend_background=element_rect(size=1, linetype='dotted')
             )
)