import cPickle
import matplotlib.pyplot as plt
from numpy import array, mean, std, sqrt, linspace, shape
import os
import sys

from explaupoppydiva.utils.drawer import Drawer

logs = '../../logs/'
log_dirs = [
#             '2015-06-15_11-22-59-CompetentExplo',
#             '2015-06-14_11-37-19-CompetentExplo',
#            '2015-06-14_11-02-21-CompetentExplo',
#            '2015-06-14_21-34-23-CompetentExplo'

'2015-07-25_21-24-36-CompetentExplo'
            
            
            
#'2015-05-01_09-30-57-explaupoppydiva-discretized_progress-cube3',
#'2015-04-30_21-43-43-explaupoppydiva-random-cube3',
#'2015-04-30_12-54-30-explaupoppydiva-tree-cube3'
#'2015-05-13_20-46-00-Test-Diva-seq_10000'
            ]


conditions_to_plot = [
                     "TOP-DOWN-CMA-Tree-seq_250", 
                      "SEQ-seq_5000",
                     "MS2-seq_5000", 
                    "MB-seq_5000",   
                      ]

labels = {
          "SEQ-seq_5000":"SIMPLEST-FIRST",
          "MB-seq_5000":"CONTROL-MOTOR-BABBLING",
          "MS2-seq_5000":"CONTROL-GOAL-BABBLING",
          "TOP-DOWN-CMA-Tree-seq_250":"TOP-DOWN-GUIDANCE",     
          }

plot_explo = True
plot_explo_comp = True
plot_comp = False


iterations = 5000
eval_at = range(1, iterations, iterations / 10) + [iterations]

if plot_explo:
    fig_explo, ax_explo = plt.subplots()
    fig_explo.canvas.set_window_title('Exploration comparison')

if plot_explo_comp:
    fig_explo_comp, ax_explo_comp = plt.subplots()
    fig_explo_comp.canvas.set_window_title('ExploComp comparison')   
     
if plot_comp:
    fig_comp, ax_comp = plt.subplots()
    fig_comp.canvas.set_window_title('Comp comparison')   


end_explo = {}
end_explo_comp = {}
ns = {}
ns_explo = {}

for xp_dir in conditions_to_plot:
    comp = {}
    explo = {}
    explo_comp = {}
    for log_dir in log_dirs:
        if xp_dir in conditions_to_plot and os.path.isdir(os.path.join(logs, log_dir, xp_dir)):
        
            print os.path.join(logs, log_dir, xp_dir)
            for log_file in os.listdir(os.path.join(logs, log_dir, xp_dir)):
                file_path = os.path.join(logs, log_dir, xp_dir, log_file)
                if file_path.endswith('.pickle'):
                    print file_path
                    try:
                        with open(file_path, 'r') as f:
                            log = cPickle.load(f)
                            f.close()
                            print "explo", log.explo, "explocomp", log.explo_comp
                            
                            if plot_explo:
                                if len(log.explo) == len(eval_at):
                                    explo[log_dir + xp_dir + log_file] = array(log.explo)
                                    
                            if plot_explo_comp:
                                if len(log.explo_comp) == len(eval_at):
                                    explo_comp[log_dir + xp_dir + log_file] = array(log.explo_comp)
                                    
                            if plot_comp:
                                if len(log.eval_errors) == len(eval_at):
                                    comp[log_dir + xp_dir + log_file] = array(log.eval_errors)
                                    
                            
                            drawer = Drawer(log)
                            #drawer.plot_learning_curve()
                            drawer.plot_scatter2D(plot_testcases=False, show=False, log_dir=os.path.join(logs, log_dir, xp_dir))
                            #drawer.plot_exploration_measure()
                    except ValueError:
                        print "ValueError"
    if plot_explo:
        if len(explo.values()) > 0:
            print explo
            print "Explo", mean(array(explo.values()), axis=0)
            l = len(explo.values()[0])
            print log.config.name
            #print explo.values()
            #print l
            x = eval_at[:l]
            y = mean(array(explo.values()), axis=0)
            error = std(array(explo.values()), axis=0)
            error = error / sqrt(len(explo)) # Standard error of the mean
            color_cycle = ax_explo._get_lines.color_cycle
            next_color = next(color_cycle)
            ax_explo.plot(x, y, label = labels[xp_dir], color=next_color)
            if len(explo.values()) > 1:
                ax_explo.fill_between(x, y-error, y+error, alpha=0.2, label = labels[xp_dir], color = next_color)
            #ax.errorbar(eval_at[:l],, yerr=, label = log.config.name)
            
            ns_explo[xp_dir] = len(explo_comp.values())
            
        end_explo[xp_dir] = [explo[key][-1] for key in explo.keys()]


    if plot_explo_comp:
        if len(explo_comp.values()) > 0:
            print "ExploComp", explo_comp.values()
            l = len(explo_comp.values()[0])
            print log.config.name
            #print explo_comp.values()
            #print l
            x = eval_at[:l]
            y = mean(array(explo_comp.values()), axis=0)
            error = std(array(explo_comp.values()), axis=0)
            error = error / sqrt(len(explo_comp)) # Standard error of the mean
            color_cycle = ax_explo_comp._get_lines.color_cycle
            next_color = next(color_cycle)
            ax_explo_comp.plot(x, y, label = labels[xp_dir], color=next_color)
            if len(explo_comp.values()) > 1:
                ax_explo_comp.fill_between(x, y-error, y+error, alpha=0.2, label = labels[xp_dir], color = next_color)
            #ax_comp.errorbar(eval_at[:l],, yerr=, label = log.config.name)
            
            ns[xp_dir] = len(explo_comp.values())
            
        end_explo_comp[xp_dir] = [explo_comp[key][-1] for key in explo_comp.keys()]
        
    if plot_comp:
        if len(comp.values()) > 0:
            print "Comp", mean(mean(array(comp.values()), axis=2), axis=0)
            l = len(comp.values()[0])
            print log.config.name
            #print mean(array(comp.values()), axis=2)
            #print l
            x = eval_at[:l]
            y = mean(mean(array(comp.values()), axis=2), axis=0)
            error = std(mean(array(comp.values()), axis=2), axis=0)
            error = error / sqrt(len(comp)) # Standard error of the mean
            #print x, y, error
            color_cycle = ax_comp._get_lines.color_cycle
            next_color = next(color_cycle)
            ax_comp.plot(x, y, label = labels[xp_dir], color=next_color)
            if len(comp.values()) > 1:
                ax_comp.fill_between(x, y-error, y+error, alpha=0.2, label = labels[xp_dir], color = next_color)
            #ax_comp.errorbar(eval_at[:l],, yerr=, label = log.config.name)
    
    end_explo[xp_dir] = [explo[key][-1] for key in explo.keys()]
    end_explo_comp[xp_dir] = [explo_comp[key][-1] for key in explo_comp.keys()]

if plot_explo and plot_explo_comp:

    print "EXPLO"
    for xp_dir in end_explo:
        print "    ", xp_dir
        print "        ", end_explo[xp_dir]
    print "EXPLO_COMP"
    for xp_dir in end_explo_comp:
        print "    ", xp_dir
        print "        ", end_explo_comp[xp_dir]
    
    from scipy.stats import mannwhitneyu
    x = end_explo["MS2-seq_5000"]
    y = end_explo["TOP-DOWN-CMA-Tree-seq_250"]
    print "STAT EXPLO MS2/TOP", mannwhitneyu(x, y)
    
    x = end_explo_comp["MS2-seq_5000"]
    y = end_explo_comp["TOP-DOWN-CMA-Tree-seq_250"]
    print "STAT EXPLO_COMP MS2/TOP", mannwhitneyu(x, y)
  
if plot_explo:
    fig_explo.show()
    ax_explo.legend(loc='upper left')
    #fig_explo.set_size_inches(19.2,12)
    plt.xlabel('Iterations', fontsize=18)
    plt.ylabel('Explored cells', fontsize=18)
    plt.show(block=False)
    plt.savefig(logs + 'exploration_comparison.png')


if plot_explo_comp:
    fig_explo_comp.show()
    ax_explo_comp.legend(loc='upper left')
    #fig_explo_comp.set_size_inches(19.2,12)
    
    plt.xlabel('Iterations', fontsize=18)
    plt.ylabel('Reached cells', fontsize=18)
    plt.show(block=False)
    plt.savefig(logs + 'explocomp_comparison.png')
    
if plot_comp:
    fig_comp.show()
    ax_comp.legend(loc='upper left')
    #fig_comp.set_size_inches(19.2,12)
    plt.xlabel('Iterations', fontsize=18)
    plt.ylabel('Competence error', fontsize=18)
    plt.show(block=False)
    plt.savefig(logs + 'comp-' + log_dir[0:-1] + '.png')
    

print ns
print ns_explo

plt.show()
