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

# In[1]:


import sys
sys.path.append('../')
import pandas as pd
import numpy as np
import forestplot as fp
pd.set_option('display.max_rows', 100)

df = (pd.read_csv('../examples/data/sleep-mmodel.csv')
      .query("model=='all' | model=='young kids'")
     )
df.head(3)


# In[2]:


std_opts = dict(
    dataframe = df,
    estimate = "coef",
    ll  ="ll", 
    hl = "hl",
    varlabel = "var",
    model_col = "model",
    return_df=True,
)


# ## Normie usage

# In[3]:


# Vanilla
_df, ax = fp.mforestplot(**std_opts)


# In[4]:


# Vanilla
# Color rows and add borders
_df, ax = fp.mforestplot(**std_opts, color_alt_rows=True, despine=False)


# ## Make table

# In[5]:


# Vanilla
# Color rows
# Add annotation and make it a table
# Add right annotations
_df, ax = fp.mforestplot(**std_opts, color_alt_rows=True, 
                         annote=["label"], annoteheaders=["Description"], 
                         rightannote=["group"],
                         right_annoteheaders=["Group"],
                         table=True)


# ## Variable Groups

# In[6]:


std_opts = dict(
    dataframe = df,
    estimate = "coef",
    ll  ="ll", 
    hl = "hl",
    varlabel = "label",
    model_col = "model",
    return_df=True,
)


# In[7]:


# With groups
_df, ax = fp.mforestplot(**std_opts, color_alt_rows=True, 
                         groupvar="group",
                         rightannote=["group"],
                         right_annoteheaders=["Group"],
                         variable_header="Covariate",
                         table=True)


# In[8]:


# With groups and group ordering
group_order = ["family factors", "health factors", "labor factors", "age", "occupation", "area of residence", "other factors"]
_df, ax = fp.mforestplot(**std_opts, color_alt_rows=True, 
                         groupvar="group",
                         group_order=group_order,
                         rightannote=["group"],
                         right_annoteheaders=["Group"],
                         variable_header="Covariate",
                         table=True)


# In[9]:


# With groups and group ordering
group_order = ["family factors", "health factors", "labor factors", "age", "occupation", "area of residence", "other factors"]
_df, ax = fp.mforestplot(**std_opts, color_alt_rows=True, 
                         groupvar="group",
                         group_order=group_order,
                         rightannote=["group"],
                         right_annoteheaders=["Group"],
                         variable_header="Covariate",
                         capitalize="capitalize",                         
                         table=True)


# In[10]:


# Add groups
_df, ax = fp.mforestplot(**std_opts,
                      groupvar='group',
                     )


# In[11]:


# Add groups
# With border + color alternate rows
_df, ax = fp.mforestplot(**std_opts,
                         groupvar='group',
                         despine=False,
                         color_alt_rows=True,
                        )


# In[12]:


# Add groups
# Add annotation on right
_df, ax = fp.mforestplot(**std_opts,
                      groupvar='group',
                      rightannote=['group'],
                      right_annoteheaders=['Variable group'],
                     )


# ## More annotations on right

# In[13]:


std_opts = dict(
    dataframe = df,
    estimate = "coef",
    ll  ="ll", 
    hl = "hl",
    varlabel = "var",
    model_col = "model",
    rightannote=['label','group'],
    right_annoteheaders=['Description', 'Group'],
    return_df=True,
)


# In[14]:


# No groups
# Add multiple annotation on right
_df, ax = fp.mforestplot(**std_opts,
                     )


# In[15]:


# No groups
# Add annotation on right
# No flush
_df, ax = fp.mforestplot(**std_opts,
                      flush=False,
                     )


# In[16]:


# No groups
# Add annotation on right
# Add row colors
# No flush
_df, ax = fp.mforestplot(**std_opts,
                      flush=False,
                      color_alt_rows=True,
                     )


# In[17]:


# W/ groups
# Add annotation on right
# Add row colors
# No flush
std_opts = dict(
    dataframe = df,
    estimate = "coef",
    ll  ="ll", 
    hl = "hl",
    varlabel = "label",
    model_col = "model",
    rightannote=['var','group'],
    right_annoteheaders=['Field', 'Group'],
    return_df=True,
)
_df, ax = fp.mforestplot(**std_opts,
                      groupvar='group',
                      flush=False,
                      color_alt_rows=True,
                     )


# ## More annotations on left

# In[18]:


std_opts = dict(
    dataframe = df,
    estimate = "coef",
    ll  ="ll", 
    hl = "hl",
    varlabel = "label",
    model_col = "model",
    rightannote=['var','group'],
    right_annoteheaders=['Field', 'Group'],
    annote=['var','group'],
    annoteheaders=['Field', 'Group'],      
    return_df=True,
)


# In[19]:


# No groups
_df, ax = fp.mforestplot(**std_opts,
                     )


# In[20]:


# Add group
_df, ax = fp.mforestplot(**std_opts,
                      groupvar='group',
                     )


# In[21]:


# Add group
# Color row
_df, ax = fp.mforestplot(**std_opts,
                      groupvar='group',
                      color_alt_rows=True,
                     )


# In[22]:


# Add group
# Keep borders
_df, ax = fp.mforestplot(**std_opts,
                      groupvar='group',
                      color_alt_rows=True,
                      **{'despine':False, "figsize":(4,8)}
                     )


# In[23]:


# Add group
# Keep borders
# Add X-label
_df, ax = fp.mforestplot(**std_opts,
                         groupvar='group',
                         color_alt_rows=True,
                         **{'despine':False, "figsize":(4,8)}
                        )


# ## Legends

# In[24]:


# Add group
# Color row
# Change legend label
# Add x-label
_df, ax = fp.mforestplot(**std_opts,
                         groupvar='group',
                         modellabels=["Model 1", "Model 2"],
                         xlabel="X-label",
                         color_alt_rows=True,
                        )


# In[25]:


# Add group
# Color row
# Change legend label
# Add x-label
# Legend at top
_df, ax = fp.mforestplot(**std_opts,
                         groupvar='group',
                         modellabels=["Model 1", "Model 2"],
                         xlabel="X-label",
                         color_alt_rows=True,
                         bbox_to_anchor=(.5,.95)
                        )


# In[26]:


# Add group
# Color row
# Change legend label
# Add x-label
_df, ax = fp.mforestplot(**std_opts,
                         groupvar='group',
                         modellabels=["Model 1", "Model 2"],
                         xlabel="X-label",
                         color_alt_rows=True,
#                          bbox_to_anchor=None,
                        )


# ## More models (cluttered)

# In[27]:


df = (pd.read_csv('../examples/data/sleep-mmodel.csv')
#       .query("model=='all' | model=='young kids'")
     )

df


# In[28]:


std_opts = dict(
    dataframe = df,
    estimate = "coef",
    ll  ="ll", 
    hl = "hl",
    varlabel = "label",
    model_col = "model",
    rightannote=['var','group'],
    right_annoteheaders=['Field', 'Group'],
    annote=['var','group'],
    annoteheaders=['Field', 'Group'],      
    return_df=True,
)


# In[29]:


# Add group
# Color row
# Change legend label
# Add x-label
_df, ax = fp.mforestplot(**std_opts,
                         groupvar='group',
                         xlabel="X-label",
                         color_alt_rows=True,
#                          bbox_to_anchor=None,
                        )


# In[30]:


# Add group
# Color row
# Change legend label
# Add x-label
modellabels=["Have young kids", "Men", "Women", "Full sample"]
_df, ax = fp.mforestplot(**std_opts,
                         groupvar='group',
                         modellabels=modellabels,
                         xlabel="X-label",
                         color_alt_rows=True,
                        )


# In[31]:


# Add group
# Color row
# Change legend label
# Add x-label
# Legend at top
modellabels=["Have young kids", "Women", "Men", "Full sample"]
_df, ax = fp.mforestplot(**std_opts,
                         groupvar='group',
                         modellabels=modellabels,
                         xlabel="X-label",
                         color_alt_rows=True,
                         bbox_to_anchor=(.5,.95)
                        )


# In[32]:


# Add group
# Color row
# Change legend label
# Add x-label
# Legend at top
modellabels=["Have young kids", "Women", "Men", "Full sample"]
_df, ax = fp.mforestplot(**std_opts,
                         modellabels=modellabels,
                         xlabel="X-label",
                         color_alt_rows=True,
                         flush=False,
                        )


# ## Camera ready

# In[33]:


df = (pd.read_csv('../examples/data/sleep-mmodel.csv')
      .query("model=='all' | model=='young kids'")
     )
std_opts = dict(
    dataframe = df,
    estimate = "coef",
    ll  ="ll", 
    hl = "hl",
    varlabel="label",
    model_col="model",
)
df.head()


# In[34]:


_df, ax = fp.mforestplot(**std_opts,
                         color_alt_rows=True,
                         groupvar="group",
                         table=True,
                         rightannote=["var", "group"],
                         right_annoteheaders=["Variable", "Variable group"],
                         xlabel="Coefficient (95% CI)",
                         modellabels=["Have young kids", "Full sample"],
                         mcolor=["#CC6677", "#4477AA"],
                         xticks=[-1200,-600, 0, 600],
                         return_df=True,
                         # Additional kwargs for customizations
                         **{"marker": "D",  # set maker symbol as diamond
                            "markersize": 35,  # adjust marker size
                            "xlinestyle": (0, (10, 5)),  # long dash for x-reference line 
                            "xlinecolor": "#808080",  # gray color for x-reference line
                            "xtick_size": 12,  # adjust x-ticker fontsize
                            "xlinestyle": (0, (10, 5)),  # long dash for x-reference line
                            "xlinecolor": ".8",  # gray color for x-reference line                           
                           }                           
                        )