# pandas
import pandas as pd
pd.set_option('display.max_columns', None)
pd.set_option('display.max_colwidth', -1)

# plotting
import matplotlib.pyplot as plt
plt.style.use('ggplot')
def plt_remove_offset():
    plt.gca().ticklabel_format(useOffset=False);

# custom libraries
import scripps
import dfx

# R
%load_ext rpy2.ipython
%R source("../r/plot_help.R");

# setup output directories
scripps.setup_output_folder()

require(ggplot2, quietly=TRUE)
require(scales, quietly=TRUE)

save_and_display_figure = function(fig_path, fig_object=p) {
    # fig_path should be a string, as passed to png(file)
    # p should be either:
    #   - a ggplot object, such as p = ggplot(data) + geom_bars() + ...
    #   - a function that produces a plot, such as p=function(){hist(c(1, 1, 2, 5, 5))}
    # usage in iPython:
    # %%R -w 7 -h 5 -u in
    # p = ggplot(data) + geom_bars() + ...
    # save_and_display_figure("figures/some_figure")
    
    display_fig = function() {
        if ('ggplot' %in% class(fig_object)) {
            #cat("Displaying ggplot figure\n")
            print(fig_object)
        } else {
            #cat("Calling figure code\n")
            fig_object()
        }
    }
    
    # use the current display size for the size of the image file
    windows_size_inches = dev.size('in')
    #cat("Current window size:", windows_size_inches[1], "wide inches by", windows_size_inches[2], "high")
    
    # print to file
    png(file=fig_path, width=windows_size_inches[1], height=windows_size_inches[2], units='in', res=150)
    display_fig()
    dev.off();
    
    # print to screen
    display_fig()
}


theme_mcSimple <- function(base_size=14,
                           base_family="",
                           legend_title=TRUE,
                           legend=TRUE,
                           axes_lines=TRUE,
                           horz_gridlines=TRUE,
                           x_axis_font_size=NA) {
    
    # start with basic black and white theme, which specified font
    mc_theme = theme_bw(base_size=base_size, base_family=base_family)
        
    # figure is white background with optional horizontal gridlines
    if (horz_gridlines) {
        mc_theme = mc_theme + theme(
            panel.grid.minor = element_blank(),
            panel.grid.major.x = element_blank(),
            panel.grid.major.y = element_line(colour = "grey")
        )
    } else {
        mc_theme = mc_theme + theme(panel.grid = element_blank())
    }
    
    # turn off border, and maybe turn on left/bottom axes
    mc_theme = mc_theme + theme(panel.border = element_blank())        
    if ( axes_lines ) {
        mc_theme = mc_theme + theme(axis.line = element_line(colour = "black"))        
    } 
    
    # space between axis title and ticks
    mc_theme = mc_theme + theme(        
        plot.title = element_text(margin = margin(b = base_size)),
        axis.title.y=element_text(margin=margin(0,20,0,0)),
        axis.title.x=element_text(margin=margin(20,0,0,0))
    )

    # visibility of legend and legend title
    if ( ! legend ) {
        mc_theme = mc_theme + theme(legend.position="none")
    }
    else if ( ! legend_title) {
        mc_theme = mc_theme + theme(legend.title=element_blank())
    }
    
    # x-axis tick font size
    if (! is.na(x_axis_font_size)) {
        mc_theme = mc_theme + theme(axis.text.x = element_text(size=x_axis_font_size))
    }
    
    # hide all ticks
    mc_theme = mc_theme + theme(axis.ticks = element_blank())    

    return(mc_theme)
}