from lets_plot import *
import lets_plot.mapping as pm

LetsPlot.setup_html()

df = {
    'x': [0, 5, 10, 15],
    'y': [0, 5, 10, 15],
    'a': [1, 2, 3, 4],
    'b': [4, 5, 6, 7]
}

# as_discrete, no scale
p = ggplot(df, aes(x='x', y='y')) + geom_point(aes(color=pm.as_discrete('a', label='custom name')), size=9)
p

# as_discrete, scale
p = ggplot(df, aes(x='x', y='y')) \
    + geom_point(aes(color='a', fill=pm.as_discrete('b')), shape=21, size=9, stroke=5) \
    + scale_color_discrete()
p

# data in mappings, scale_color_discrete
p = ggplot() + geom_point(aes(x=[0, 5, 10], y=[0, 5, 10], color=[1, 2, 4]), size=9) + scale_color_discrete()
p

p3 = ggplot(df, aes(x='x', y='y', color=pm.as_discrete('a'))) + geom_point(size=9)
p3

df = {
    'x': [0, 5, 10, 15],
    'y': [0, 5, 10, 15],
    'a': [0, 0, 1, 1],
    'c': ['a', 'a', 'b', 'b']
}

p4 = ggplot(df, aes(x='x', y='y')) + geom_line(aes(color=pm.as_discrete('a')), size=3)
p4

# expected: 2 lines ('c' is a discrete variable)
p5 = ggplot(df, aes(x='x', y='y')) + geom_line(aes(color='c'), size=3)
p5

# expected: 2 lines (`group` is defined manually)
p6 = ggplot(df, aes(x='x', y='y')) + geom_line(aes(color='a', group='a'), size=3)
p6

import pandas as pd
mpg_df = pd.read_csv('https://raw.githubusercontent.com/JetBrains/lets-plot-docs/master/data/mpg.csv')
mpg_plot = ggplot(mpg_df, aes(x='displ', y='hwy'))

cyl_factor = pm.as_discrete('cyl')
mpg_plot + geom_point(aes(color=cyl_factor)) \
+ geom_smooth(aes(color=cyl_factor), method='lm', size=1, se=False)

# expected: seperate regression line for each group
mpg_plot + geom_point(aes(color=cyl_factor)) \
+ geom_smooth(aes(color='cyl', group='cyl'), method='lm', size=1, se=False)

mpg_plot + geom_point(aes(color=cyl_factor)) \
+ geom_smooth(aes(color=cyl_factor, group='cyl'), method='lm', size=1, se=False)

mpg_plot + geom_point(aes(color=cyl_factor)) \
+ scale_color_discrete(breaks=[4, 5, 6, 8])

mpg_plot + geom_point(aes(color=pm.as_discrete('cyl', label='cyl')))