# -*- coding: utf-8 -*-
import skimage.io
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import selectivesearch
import numpy as np
import skimage.transform
import os
import shutil
import caffe
from PIL import Image

candidates = set()
merged_candidates = set()
refined = set()
final = set()
final_extended = set()
text_boxes = set()
text=set()
text_cut = set()
no_text = set()
text_cut_final = set()

def getClass(FileList):
    caffe.set_mode_gpu()
    classifier = caffe.Classifier("../ROIs_Indus/deploy.prototxt","../ROIs_Indus/Models/bvlc_googlenet_indusnet_iter_20000.caffemodel" ,
            image_dims=[224,224], raw_scale=255.0, channel_swap = [2,1,0])

    inputs = [caffe.io.load_image(im_f) for im_f in FileList]
    print("Classifying %d inputs." % len(inputs))

    predictions = classifier.predict(inputs)

    return predictions

def texbox_ext():
    global text
    global both_text
    global text_cut_final
    for x, y, w, h in text:
        A = {'x1': x, 'y1': y, 'x2': x+w, 'y2': y+h, 'w': w, 'h': h}
        for x1, y1, w1, h1 in both_text:
            B = {'x1': x1, 'y1': y1, 'x2': x1+w1, 'y2': y1+h1, 'w': w1, 'h': h1}

            # overlap between A and B
            SA = A['w']*A['h']
            SB = B['w']*B['h']
            SI = np.max([ 0, np.min([A['x2'],B['x2']]) - np.max([A['x1'],B['x1']]) ]) * np.max([ 0, np.min([A['y2'],B['y2']]) - np.max([A['y1'],B['y1']]) ])
            SU = SA + SB - SI
            overlap_AB = float(SI) / float(SU)

            overf = 0
            ax1,ay1,aw,ah = A['x1'],A['y1'],A['w'],A['h']
            if overlap_AB > 0.0:
                if A['x1'] > B['x1'] and abs(B['x1']+B['w'] - A['x1']) < A['w']*0.20: # B is left to A
                    ax1 = B['x1']
                    aw =  A['x1'] + A['w'] - B['x1']
                    overf = 1
                # if A['y1'] < B['y1'] and abs(A['y1']-B['y1']) > A['h']*0.70: # B is bottom to A
                #     ah = A['h'] - (A['y1']+A['h'] - B['y1'])
                #     overf = 1
                # if A['y1'] > B['y1']: # B is top to A
                #     ay1 = B['y1'] + B['h']
                if A['x1'] < B['x1']: # B is right to A
                    aw = B['x1']+B['w'] - A['x1']
                    overf = 1
                # if A['y1'] < B['y1']: # B is bottom to A
                #     ah = A['h'] - (A['y1']+A['h'] - B['y1'])
                # REPLACE by Cohen Suderland algo

                A['x1'],A['y1'],A['w'],A['h'] = ax1,ay1,aw,ah
                text_cut_final.add((A['x1'],A['y1'],A['w'],A['h']))
            if overf == 1:
                break
        text_cut_final.add((A['x1'],A['y1'],A['w'],A['h']))
    text_cut_final = text_cut_final - both_text # CHANGE THIS LINE

def texbox_cut():
    global no_text
    no_text = no_text.union(both_text)

    for x, y, w, h in text:
        A = {'x1': x, 'y1': y, 'x2': x+w, 'y2': y+h, 'w': w, 'h': h}
        for x1, y1, w1, h1 in no_text:
            B = {'x1': x1, 'y1': y1, 'x2': x1+w1, 'y2': y1+h1, 'w': w1, 'h': h1}

            # overlap between A and B
            SA = A['w']*A['h']
            SB = B['w']*B['h']
            SI = np.max([ 0, np.min([A['x2'],B['x2']]) - np.max([A['x1'],B['x1']]) ]) * np.max([ 0, np.min([A['y2'],B['y2']]) - np.max([A['y1'],B['y1']]) ])
            SU = SA + SB - SI
            overlap_AB = float(SI) / float(SU)

            overf = 0
            ax1,ay1,aw,ah = A['x1'],A['y1'],A['w'],A['h']
            if overlap_AB > 0.0:
                if A['x1'] > B['x1'] and abs(B['x1']+B['w'] - A['x1']) < A['w']*0.20: # B is left to A
                    ax1 = B['x1'] + B['w']
                    overf = 1
                if A['y1'] < B['y1'] and abs(A['y1']-B['y1']) > A['h']*0.70: # B is bottom to A
                    ah = A['h'] - (A['y1']+A['h'] - B['y1'])
                    overf = 1
                # if A['y1'] > B['y1']: # B is top to A
                #     ay1 = B['y1'] + B['h']
                # if A['x1'] < B['x1']: # B is right to A
                #     aw = A['w'] - (A['x1']+A['w'] - B['x1'])
                # if A['y1'] < B['y1']: # B is bottom to A
                #     ah = A['h'] - (A['y1']+A['h'] - B['y1'])
                # REPLACE by Cohen Suderland algo

                A['x1'],A['y1'],A['w'],A['h'] = ax1,ay1,aw,ah
                text_cut.add((A['x1'],A['y1'],A['w'],A['h']))
            if overf == 1:
                break
        text_cut.add((A['x1'],A['y1'],A['w'],A['h']))

def extend_text_rect(l):
    return (min([i[0] for i in l]), min([i[1] for i in l]), max([i[0]+i[2] for i in l]) - min([i[0] for i in l]), max([i[3] for i in l]))

def draw_textbox():
    global width, height
    thresh = ((width+height)/2)*(0.25)
    tempc = set()
    for x, y, w, h in text_boxes:
        if (x, y, w, h) in tempc: continue
        temp = set()
        temp.add((x, y, w, h))
        f = 0
        for x1, y1, w1, h1 in text_boxes:
            if abs(y1-y) <= thresh and abs(h1-h) <= thresh:
                temp.add((x1, y1, w1, h1))
                tempc.add((x1, y1, w1, h1))
                f = 1
        if f == 0:
            text.add((x, y, w, h))
        text.add(extend_text_rect(temp))

def contains():
    x1, y1, w1, h1 = p
    for x, y, w, h in candidates:
        if x1>=x and y1 >= y and x1+w1 <= x+w and y1+h1 <= y+h:
            return True
        if x1<=x and y1 <= y and x1+w1 >= x+w and y1+h1 >= y+h:
            candidates.remove((x, y, w, h))
            return False
    return False

def extend_rect(l):
    return (min([i[0] for i in l]), min([i[1] for i in l]), max([i[0]+i[2] for i in l]) - min([i[0] for i in l]), max([i[3] for i in l]))

def extend_superbox():
    global width, height
    thresh = ((width+height)/2)*(0.06)

    tempc = set()
    for x, y, w, h in final:
        if (x, y, w, h) in tempc: continue
        temp = set()
        temp.add((x, y, w, h))
        for x1, y1, w1, h1 in final:
            if abs(y1-y) <= thresh and abs(h1-h) <= thresh:
                temp.add((x1, y1, w1, h1))
                tempc.add((x1, y1, w1, h1))
        final_extended.add(extend_rect(temp))

def draw_superbox(finals=[]):
    noover = []
    refinedT = []

    global final
    final = set()

    # (x1,y1) top-left coord, (x2,y2) bottom-right coord, (w,h) size
    if finals != []:
        refinedT = finals
    else:
        refinedT = refined
    remp = set(refinedT)
    ref = list(refinedT)

    while len(ref) > 0:
        x1, y1, w1, h1 = ref[0]

        if len(ref) == 1: # final box
            final.add((x1, y1, w1, h1))
            ref.remove((x1, y1, w1, h1))
            remp.remove((x1, y1, w1, h1))
        else:
            ref.remove((x1, y1, w1, h1))
            remp.remove((x1, y1, w1, h1))

        over = set()
        for x2, y2, w2, h2 in remp:
            A = {'x1': x1, 'y1': y1, 'x2': x1+w1, 'y2': y1+h1, 'w': w1, 'h': h1}
            B = {'x1': x2, 'y1': y2, 'x2': x2+w2, 'y2': y2+h2, 'w': w2, 'h': h2}

            # overlap between A and B
            SA = A['w']*A['h']
            SB = B['w']*B['h']
            SI = np.max([ 0, np.min([A['x2'],B['x2']]) - np.max([A['x1'],B['x1']]) ]) * np.max([ 0, np.min([A['y2'],B['y2']]) - np.max([A['y1'],B['y1']]) ])
            SU = SA + SB - SI
            overlap_AB = float(SI) / float(SU)
            overlap_A = float(SI) / float(SA)
            overlap_B = float(SI) / float(SB)
            # print(overlap_AB)
            #

            if overlap_A >= 0.40 or overlap_B >= 0.40:
                over.add((B['x1'],B['y1'],B['w'],B['h']))
        # print(len(over))
        if len(over) != 0: #Overlap
            remp = remp - over
            for i in over: ref.remove(i)
            over.add((A['x1'],A['y1'],A['w'],A['h']))
            # print(over)
            final.add((min([i[0] for i in over]), min([i[1] for i in over]), max([i[0]+i[2] for i in over]) - min([i[0] for i in over]), max([i[1]+i[3] for i in over]) - min([i[1] for i in over])))
            # final.add((np.mean([i[0] for i in over]), np.mean([i[1] for i in over]), np.mean([i[2] for i in over]), np.mean([i[3] for i in over])))
            noover.append(False)
        else:   #No overlap
            final.add((x1,y1,w1,h1))
            noover.append(True)

    if all(noover):
        return
    else:
        draw_superbox(final)
        return

def contains_remove():

    for x, y, w, h in merged_candidates:
        f = False
        temp = set(merged_candidates)
        temp.remove((x, y, w, h))
        for x1, y1, w1, h1 in temp:
            if x1>=x and y1 >= y and x1+w1 <= x+w and y1+h1 <= y+h:
                f = False
                break
            # if x1<=x and y1 <= y and x1+w1 >= x+w and y1+h1 >= y+h:
            else:
                f = True
        if f == True:
            refined.add((x, y, w, h))

# def contains_remove():
#     for x, y, w, h in merged_candidates:
#         temp = set(merged_candidates)
#         temp.remove((x, y, w, h))
#         test = []
#         for x1, y1, w1, h1 in temp:
#             A = {'x1': x, 'y1': y, 'x2': x+w, 'y2': y+h, 'w': w, 'h': h}
#             B = {'x1': x1, 'y1': y1, 'x2': x1+w1, 'y2': y1+h1, 'w': w1, 'h': h1}
#             # overlap between A and B
#             SA = A['w']*A['h']
#             SB = B['w']*B['h']
#             SI = np.max([ 0, np.min([A['x2'],B['x2']]) - np.max([A['x1'],B['x1']]) ]) * np.max([ 0, np.min([A['y2'],B['y2']]) - np.max([A['y1'],B['y1']]) ])
#             SU = SA + SB - SI
#             overlap_AB = float(SI) / float(SU)
#             if overlap_AB > 0.0:
#                 # if x1>=x and y1 >= y and x1+w1 <= x+w and y1+h1 <= y+h:
#                 if x1<=x and y1 <= y and x1+w1 >= x+w and y1+h1 >= y+h:
#                     test.append(False)
#                 else:
#                     test.append(True)
#             else:
#                 test.append(True)
#         if all(test):
#             refined.add((x, y, w, h))

def mean_rect(l):
    return (min([i[0] for i in l]), min([i[1] for i in l]), max([i[0]+i[2] for i in l]) - min([i[0] for i in l]), max([i[1]+i[3] for i in l]) - min([i[1] for i in l]))

def merge():
    global width, height
    thresh = int(((width+height)/2)*(0.14))
    tempc = set()
    for x, y, w, h in candidates:
        if (x, y, w, h) in tempc: continue
        temp = set()
        temp.add((x, y, w, h))
        for x1, y1, w1, h1 in candidates:
            if abs(x1-x) <= thresh and abs(y1-y) <= thresh and abs(w1-w) <= thresh and abs(h1-h) <= thresh:
                temp.add((x1, y1, w1, h1))
                tempc.add((x1, y1, w1, h1))
        merged_candidates.add(mean_rect(temp))
    contains_remove()

for name in os.listdir("./Images"):
	candidates = set()
	merged_candidates = set()
	refined = set()
	final = set()
	final_extended = set()
	text_boxes = set()
	text=set()
	text_cut = set()
	no_text = set()

	print("Processing Image " + name.split(".")[0])
	fname = "./Images/" + name
	print(fname)
	img = skimage.io.imread(fname)
	width = len(img[0])
	height = len(img)
	# new_size = 256
	# height = int(new_size * height / width)
	# width  = new_size

	if width*height < 256*256*(0.95) and abs(width-height) <= 3 :
	    new_size  = 512
	    height = int(new_size * height / width)
	    width  = new_size
	    print("A")
	elif width*height < 220*220*(1.11):
	    new_size  = 256
	    height = int(new_size * height / width)
	    width  = new_size
	    print("B")
	elif width*height < 256*256:
	    new_size  = 256
	    height = int(new_size * height / width)
	    width  = new_size
	    print("B1")
	elif width*height > 512*512*(0.99) and width < 800 and height < 800:
	    new_size  = 512
	    height = int(new_size * height / width)
	    width  = new_size
	    print("C")
	elif width*height < 512*512*(0.95) and width*height > 256*256*(1.15):
	    new_size  = 512
	    height = int(new_size * height / width)
	    width  = new_size
	    print("D")
	tried = []
	while True:
		tried.append(width)
		candidates = set()
		merged_candidates = set()
		refined = set()
		final = set()
		final_extended = set()
		text_boxes = set()
		text=set()
		text_cut = set()
		no_text = set()
		stage = 1
		text_cut_final = set()

		for sc in [350,450,500]:
			for sig in [0.8]:
				for mins in [30,60,120]: # important
					img = skimage.io.imread(fname)[:,:,:3]
					if height == len(img) and width == len(img[0]):
						pass
					else:
						img = skimage.transform.resize(img, (height, width))

					img_lbl, regions = selectivesearch.selective_search(
					img, scale=sc, sigma= sig,min_size = mins)

					for r in regions:
						# excluding same rectangle (with different segments)
						if r['rect'] in candidates:
							continue
						# excluding regions smaller than 2000 pixels
						if r['size'] < 2000:
							continue
						# distorted rects
						x, y, w, h = r['rect']
						if w / h > 1.2 or h / w > 1.2:
							continue
						if w >= (img.shape[0]-1)*(0.7) and h >= (img.shape[1]-1)*(0.7):
							continue
						candidates.add(r['rect'])
			print("Stage " + str(stage) + " Complete.")
			stage+=1

		print(candidates)
		merge()
		print(refined)
		draw_superbox()
		print(final)
		extend_superbox()
		print(final_extended)

		os.makedirs("Regions/"+name.split(".")[0])

		# draw rectangles on the original image
		fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(6, 6))
		ax.imshow(img)
		for x, y, w, h in final_extended:
			rect = mpatches.Rectangle((x, y), w, h, fill=False, edgecolor='red', linewidth=1)
			ax.add_patch(rect)
		plt.savefig("Regions/"+name.split(".")[0]+"/FinalRegions.png")
		plt.close('all')

		img1 = skimage.io.imread(fname)[:,:,:3]
		if height == len(img1) and width == len(img1[0]): pass
		else: img1 = skimage.transform.resize(img1, (height, width))
		# imgT = Image.open(fname).convert('L')
		# w, h = imgT.size
		# if height == h and width == w:
		# 	pass
		# else:
		# 	# img1 = skimage.transform.resize(img1, (height, width))
		# 	imgT = imgT.resize((width,height), Image.ANTIALIAS)

		ij = 1
		fList = []
		box_list = []
		for x, y, w, h in final_extended:
			skimage.io.imsave("Regions/"+name.split(".")[0]+"/"+str(ij)+"_sub.jpg", img1[y:y+h,x:x+w])
			# imgT.crop((x,y,x+w,y+h)).save("Regions/"+name.split(".")[0]+"/"+str(ij)+"_sub_b.png")
			# imgT = Image.open("Regions/"+name.split(".")[0]+"/"+str(ij)+"_sub.png").convert('L')
			# imgT.save("Regions/"+name.split(".")[0]+"/"+str(ij)+"_sub_b.png")
			fList.append("Regions/"+name.split(".")[0]+"/"+str(ij)+"_sub.jpg")
			box_list.append((x, y, w, h))
			ij+=1

		# classify text no text
		text_boxes=set()
		text = set()
		no_text = set()
		both_text = set()
		text_cut_final = set()

		i = 0
		try:
			a = getClass(fList)
			l = np.array([0,1,2])
			for pred in a:
				idx = list((-pred).argsort())
				pred = l[np.array(idx)]
				if pred[0] == 1 or pred[0] == 2:
					text_boxes.add(box_list[i])
				elif pred[0] == 0:
					no_text.add(box_list[i])
				if pred[0] == 2:
					both_text.add(box_list[i])
				print(pred)
				i+=1
		except:
			print("No Text Regions")

		draw_textbox()
		print(text)
		texbox_cut()
		print(text_cut)
		texbox_ext()
		print(text_cut_final)

		# draw rectangles on the original image
		img = skimage.io.imread(fname)[:,:,:3]
		if height == len(img) and width == len(img[0]): pass
		else: img = skimage.transform.resize(img, (height, width))

		fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(6, 6))
		ax.imshow(img)
		for x, y, w, h in text_cut_final:
			rect = mpatches.Rectangle((x, y), w, h, fill=False, edgecolor='red', linewidth=1)
			ax.add_patch(rect)

		plt.savefig("Result/final_"+name.split(".")[0]+".png")
		plt.close('all')

		ij = 1
		for x, y, w, h in text_cut_final:
			skimage.io.imsave("Regions/"+name.split(".")[0]+"/"+str(ij)+"_text.png", img[y:y+h,x:x+w])
			ij+=1

		# min area check
		minf = 0
		for x, y, w, h in text_cut_final:
			if w*h < width*height*0.20  and (w < width*0.20 or h < height*0.20):
				minf = 1

		if (len(text_cut_final) == 0 or minf == 1) and len(tried) < 3:
			print(tried)
			print("New size being tried.")
			shutil.rmtree("Regions/"+name.split(".")[0]+"/")
			img = skimage.io.imread(fname)
			twidth = len(img[0])
			theight = len(img)

			new_size = list(set([256,512,twidth]) - set(tried))[0]
			height = int(new_size * theight / twidth)
			width = new_size

		else:
			break


#coding:utf8

import numpy as np, scipy
import pylab as pl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import math
from matplotlib import cm
from matplotlib import mlab
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from itertools import *
import collections
from multiprocessing import Pool
import random
from scipy.optimize import leastsq
from adasurf import AdaSurfConfig, adasurf, paint_surfs, identifysurf, point_normalize, Surface

ELAPSE_SEG = 0
class SurfSegConfig:
    def __init__(self):
        self.slice_count = 4
        self.origin_points = 5
        self.most_combination_points = 20
        self.same_threshold = 0.1 # the smaller, the more accurate when judging two surfaces are identical, more surfaces can be generated
        self.pointsame_threshold = 1.0
        self.filter_rate = 0.08
        self.filter_count = 50
        self.ori_adarate = 2.0
        self.step_adarate = 1.0
        self.max_adarate = 2.0
        self.split_by_count = True
        self.weak_abort = 45

def paint_points(points, show = True, title = '', xlim = None, ylim = None, zlim = None):
    fig = pl.figure()
    ax = fig.add_subplot(111, projection='3d')
    if xlim == None:
        xlim = (np.min(points[:, 0]), np.max(points[:, 0]))
    if ylim == None:
        ylim = (np.min(points[:, 1]), np.max(points[:, 1]))
    if zlim == None:
        zlim = (np.min(points[:, 2]), np.max(points[:, 2]))
    x1 = points[:, 0]
    y1 = points[:, 1]
    z1 = points[:, 2]
    ax.scatter(x1, y1, z1, c='r')

    ax.set_zlim(zlim[0], zlim[1])
    ax.set_ylim(ylim[0], ylim[1])
    ax.set_xlim(xlim[0], xlim[1])
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")
    ax.zaxis.set_major_locator(LinearLocator(10))
    ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
    pl.title(title)
    if show:
        pl.show()
    return fig

def surf_segmentation(points, config, paint_when_end = False):
    global ELAPSE_SEG
    config.slice_count = min(int(len(points) / config.origin_points), config.slice_count)
    assert len(points) / config.slice_count >= config.origin_points
    adasurconfig = AdaSurfConfig({'origin_points': config.origin_points
        , 'most_combination_points': config.most_combination_points
        , 'same_threshold': config.same_threshold
        , 'filter_rate': config.filter_rate
        , 'ori_adarate': config.ori_adarate
        , 'step_adarate': config.step_adarate
        , 'max_adarate': config.max_adarate
        , 'pointsame_threshold': config.pointsame_threshold
        , 'filter_count' : config.filter_count
        , 'weak_abort' : config.weak_abort
        })
    surfs = []
    slice_fig = []
    npoints = point_normalize(points)
    starttime = time.clock()
    xlim = (np.min(npoints[:, 0]), np.max(npoints[:, 0]))
    ylim = (np.min(npoints[:, 1]), np.max(npoints[:, 1]))
    zlim = (np.min(npoints[:, 2]), np.max(npoints[:, 2]))

    pca_md = mlab.PCA(np.copy(npoints))

    projection0_direction = None

    # projection0_direction = pca_md.Y[0]
    # projection0 = np.inner(projection0_direction, npoints)
    projection0 = npoints[:, 0]
    if config.split_by_count:
        step_count = len(projection0) / config.slice_count
        pointsets = [np.array([]).reshape(0,3)] * config.slice_count
        sorted_projection0_index = np.argsort(projection0)
        current_slot_count, ptsetid = 0, 0
        for index in sorted_projection0_index:
            pointsets[ptsetid] = np.vstack((pointsets[ptsetid], npoints[index, :]))
            current_slot_count += 1
            if current_slot_count > step_count:
                current_slot_count = 0
                ptsetid += 1
    else:
        projection0min, projection0max = np.min(projection0), np.max(projection0)
        step_len = (projection0max - projection0min) / config.slice_count
        pointsets = [np.array([]).reshape(0,3)] * config.slice_count
        for i in xrange(len(projection0)):
            if projection0[i] == projection0max:
                ptsetid = config.slice_count - 1
            else:
                ptsetid = int((projection0[i] - projection0min) / step_len)
            pointsets[ptsetid] = np.vstack((pointsets[ptsetid], npoints[i]))

    # random.shuffle(pointsets)

    partial_surfs, fail = [], np.array([]).reshape(0,3)
    # for (ptset, ptsetindex) in zip(pointsets, range(len(pointsets))):
    #     print "slice", len(ptset), xlim, ylim, zlim
        # paint_points(ptset, xlim = xlim, ylim = ylim, zlim = zlim)
    for (ptset, ptsetindex) in zip(pointsets, range(len(pointsets))):
        print "--------------------------------------"
        print "before segment", ptsetindex, '/', len(pointsets)
        print 'derived surfs:'
        # print '---000', ptset.shape, np.array(fail).shape, np.array(fail), fail
        if fail == None:
            allptfortest = np.array(ptset)
        else:
            allptfortest = np.vstack((ptset, np.array(fail).reshape(-1,3)))
        print "len of surf is: ", len(partial_surfs), ", len of points is: ", len(allptfortest)
        if allptfortest != None and len(allptfortest) > 0 :
            partial_surfs, _, fail, extradata = identifysurf(allptfortest, adasurconfig, donorm = False, surfs = partial_surfs, title = str(ptsetindex)
                , paint_when_end = paint_when_end, current_direction = projection0_direction)
            if paint_when_end:
                slice_fig.append(extradata[0])
        if fail == None:
            print "after segment", ptsetindex, "len of surf", len(partial_surfs), "fail is None", fail
        else:
            print "after segment", ptsetindex, "len of surf", len(partial_surfs), "len of fail", len(fail)

        for x in partial_surfs:
            x.printf()

    surfs.extend(partial_surfs)

    # fig = pl.figure()
    # ax = fig.add_subplot(111, projection='3d')
    # ax.scatter(npoints[:, 0], npoints[:, 1], npoints[:, 2], c='r')
    # x = np.linspace(0, pca_md.Wt[0, 0] * 100, 300)
    # y = np.linspace(0, pca_md.Wt[0, 1] * 100, 300)
    # z = np.linspace(0, pca_md.Wt[0, 2] * 100, 300)
    # ax.plot(x, y, z, c='k')
    # x = np.linspace(0, pca_md.Wt[1, 0] * 100, 300)
    # y = np.linspace(0, pca_md.Wt[1, 1] * 100, 300)
    # z = np.linspace(0, pca_md.Wt[1, 2] * 100, 300)
    # ax.plot(x, y, z, c='g')
    # pl.show()

    return surfs, npoints, (slice_fig, )

if __name__ == '__main__':
    c = np.loadtxt('5.py', comments='#')
    config = SurfSegConfig()
    print 'config', config.__dict__
    import time
    starttime = time.clock()
    surfs, npoints, extradata = surf_segmentation(c, config, paint_when_end = True)

    print "----------BELOW ARE SURFACES---------- count:", len(surfs)
    print 'TOTAL: ', time.clock() - starttime
    print 'ELAPSE_SEG: ', ELAPSE_SEG
    ALL_POINT = 0
    for s,i in zip(surfs, range(len(surfs))):
        print "SURFACE ", i
        print s.args # surface args
        print s.residuals # MSE
        print len(s.points)
        ALL_POINT += len(s.points)
        # print s[2] # npoints
        print '**************************************'
    print 'ALL_POINT: ', ALL_POINT
    print '----------BELOW ARE POINTS----------'
    # for s,i in zip(surfs, range(len(surfs))):
    #     print "SURFACE ", i
    #     print s.points
    paint_surfs(surfs, npoints, 'all')
    print extradata
    for slice_fig in extradata[0]:
        slice_fig.show()


#    This file is part of DEAP.
#
#    DEAP is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Lesser General Public License as
#    published by the Free Software Foundation, either version 3 of
#    the License, or (at your option) any later version.
#
#    DEAP is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
#    GNU Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public
#    License along with DEAP. If not, see <http://www.gnu.org/licenses/>.

import numpy

from deap import algorithms
from deap import base
from deap import benchmarks
from deap import cma
from deap import creator
from deap import tools

import matplotlib.pyplot as plt

# Problem size
N = 10
NGEN = 125

creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
creator.create("Individual", list, fitness=creator.FitnessMin)

toolbox = base.Toolbox()
toolbox.register("evaluate", benchmarks.rastrigin)

def main(verbose=True):
    # The cma module uses the numpy random number generator
    numpy.random.seed(64)

    # The CMA-ES algorithm takes a population of one individual as argument
    # The centroid is set to a vector of 5.0 see http://www.lri.fr/~hansen/cmaes_inmatlab.html
    # for more details about the rastrigin and other tests for CMA-ES
    strategy = cma.Strategy(centroid=[5.0]*N, sigma=5.0, lambda_=20*N)
    toolbox.register("generate", strategy.generate, creator.Individual)
    toolbox.register("update", strategy.update)

    halloffame = tools.HallOfFame(1)
    stats = tools.Statistics(lambda ind: ind.fitness.values)
    stats.register("avg", numpy.mean)
    stats.register("std", numpy.std)
    stats.register("min", numpy.min)
    stats.register("max", numpy.max)

    logbook = tools.Logbook()
    logbook.header = "gen", "evals", "std", "min", "avg", "max"

    # Objects that will compile the data
    sigma = numpy.ndarray((NGEN,1))
    axis_ratio = numpy.ndarray((NGEN,1))
    diagD = numpy.ndarray((NGEN,N))
    fbest = numpy.ndarray((NGEN,1))
    best = numpy.ndarray((NGEN,N))
    std = numpy.ndarray((NGEN,N))

    for gen in range(NGEN):
        # Generate a new population
        population = toolbox.generate()
        # Evaluate the individuals
        fitnesses = toolbox.map(toolbox.evaluate, population)
        for ind, fit in zip(population, fitnesses):
            ind.fitness.values = fit

        # Update the strategy with the evaluated individuals
        toolbox.update(population)

        # Update the hall of fame and the statistics with the
        # currently evaluated population
        halloffame.update(population)
        record = stats.compile(population)
        logbook.record(evals=len(population), gen=gen, **record)

        if verbose:
            print(logbook.stream)

        # Save more data along the evolution for latter plotting
        # diagD is sorted and sqrooted in the update method
        sigma[gen] = strategy.sigma
        axis_ratio[gen] = max(strategy.diagD)**2/min(strategy.diagD)**2
        diagD[gen, :N] = strategy.diagD**2
        fbest[gen] = halloffame[0].fitness.values
        best[gen, :N] = halloffame[0]
        std[gen, :N] = numpy.std(population, axis=0)

    # The x-axis will be the number of evaluations
    x = list(range(0, strategy.lambda_ * NGEN, strategy.lambda_))
    avg, max_, min_ = logbook.select("avg", "max", "min")
    plt.figure()
    plt.subplot(2, 2, 1)
    plt.semilogy(x, avg, "--b")
    plt.semilogy(x, max_, "--b")
    plt.semilogy(x, min_, "-b")
    plt.semilogy(x, fbest, "-c")
    plt.semilogy(x, sigma, "-g")
    plt.semilogy(x, axis_ratio, "-r")
    plt.grid(True)
    plt.title("blue: f-values, green: sigma, red: axis ratio")

    plt.subplot(2, 2, 2)
    plt.plot(x, best)
    plt.grid(True)
    plt.title("Object Variables")

    plt.subplot(2, 2, 3)
    plt.semilogy(x, diagD)
    plt.grid(True)
    plt.title("Scaling (All Main Axes)")

    plt.subplot(2, 2, 4)
    plt.semilogy(x, std)
    plt.grid(True)
    plt.title("Standard Deviations in All Coordinates")

    plt.show()

if __name__ == "__main__":
    main(False)


import numpy as np
from sklearn.model_selection import cross_val_predict, cross_val_score, StratifiedKFold
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.svm import SVC
from ..metrics import METRICS

class Blender:

    def make_greedy_blend(self, X, y, models, scoring=None, cv=3,
                          proba=True, random_state=42, verbose=False):
        """
        This func makes greedy blend for many models.

        Attributes
        ----------
        X : array-like
            The data to fit.
        y : array-like
            The target variable to try to predict.
        models : list
            List of models to blend.
        scoring : string or callable, optional
            Scoring function from sklearn.
        cv : int, cross validation generator, optional, default: 3
            Cross validation from sklearn or number of folds.
        proba : bool, optional, default: True
            If true than probabilities were predicted else labels.
        random_state : int, optional, default: 42

        Returns
        -------

        """
        try:
            metric = METRICS[scoring]
        except KeyError:
            metrics = [metric for metric in METRICS]
            raise ValueError('%r is not a valid scoring value. '
                             'Valid options are %s'
                             % (scoring, sorted(metrics)))

        if isinstance(cv, int):
            cv = StratifiedKFold(n_splits=cv, shuffle=True, random_state=random_state)

        scores, preds = self._evaluate_models(X, y, models, metric, cv)

        models, scores = np.array(models), np.array(scores)
        isorted = np.argsort(scores)[::-1]

        blend = BlendClassifier((models[isorted][0],), (1,))
        best_scores, best_pred = np.array([scores[isorted][0]]), preds[isorted][0]

        if verbose:
            print('First blending model:\n{}\nScore: {}'
                  .format(models[isorted][0], scores[isorted][0]))

        for model, pred in zip(models[isorted][1:], preds[isorted][1:]):
            score, alpha = self._blend_two_preds(y, best_pred, pred, metric, cv)
            if alpha != 1 and score > best_scores[-1]:
                blend = blend.get_updated_classifier((model, ), (1 - alpha, ))
                best_scores = np.append(best_scores, score)
                best_pred = alpha * best_pred + (1 - alpha) * pred
                if verbose:
                    print('The model added to blending:\n{}\nCoef: {}\nNew score: {}'
                          .format(model, 1 - alpha, score))
            elif verbose:
                print('The model is not added to blending:\n{}'
                      .format(model))

        return blend, best_scores

    def blend_two_models(self, X, y, est1, est2, scoring, cv,
                         proba=True, random_state=42):
        """
        This func blends two estimators as the following combination:
            alpha * est1_prediction + (1 - alpha) * est2_prediction
        and finds the best combination.

        Attributes
        ----------
        X : array-like
            The data to fit.
        y : array-like
            The target variable to try to predict.
        est1 : estimator
            The first estimator to blend.
        est2 : estimator
            The second estimator to blend.
        scoring : string or callable, optional
            Scoring function from sklearn.
        cv : int, cross validation generator, optional, default: 3
            Cross validation from sklearn or number of folds.
        proba : bool, optional, default: True
            If true than probabilities were predicted else labels.
        random_state : int, optional, default: 42

        Returns
        -------
        best_score : float
            The best score of blending.
        best_alpha : float
            The alpha parameter of best blending combination.
        """
        try:
            metric = METRICS[scoring]
        except KeyError:
            metrics = [metric for metric in METRICS]
            raise ValueError('%r is not a valid scoring value. '
                             'Valid options are %s'
                             % (scoring, sorted(metrics)))

        weights = np.linspace(0, 1, 101)
        method = 'predict_proba' if proba else 'predict'

        if isinstance(cv, int):
            cv = StratifiedKFold(n_splits=cv, random_state=random_state)

        preds1 = cross_val_predict(est1, X, y, cv=cv, method=method)
        preds2 = cross_val_predict(est2, X, y, cv=cv, method=method)
        if proba:
            preds1, preds2 = preds1[:, 1], preds2[:, 1]

        best_score, best_alpha = metric(y, preds1), 1
        for idx, alpha in enumerate(weights):
            preds = alpha * preds1 + (1 - alpha) * preds2
            score = metric(y, preds)
            if score > best_score:
                best_score = score
                best_alpha = alpha

        return best_score, best_alpha

    def _blend_two_preds(self, y, pred1, pred2, metric, cv):
        weights = np.linspace(0, 1, 101)
        best_score, best_alpha = metric(y, pred1), 1
        for idx, alpha in enumerate(weights):
            preds = alpha * pred1 + (1 - alpha) * pred2
            score = metric(y, preds)
            if score > best_score:
                best_score = score
                best_alpha = alpha

        return best_score, best_alpha

    def _evaluate_models(self, X, y, models, metric, cv, proba=True):
        scores = []
        preds  = []
        method = 'predict_proba' if proba else 'predict'
        for model in models:
            preds.append(cross_val_predict(model, X, y, cv=cv, method=method))
            scores.append(metric(y, preds[-1]))
        return scores, preds

    def _get_n_best_estimators_from_each_class(self, task_manager, n, classes):
        if classes is None:
            classes = task_manager.get_done_model_classes()

        models = []
        for cls in classes:
            models.extend(task_manager.get_best_models(cls, n))

        return list(reversed(sorted(models, key=lambda x: x[0])))

    def _get_models_with_scores(self, models, scores):
        return np.array(list(reversed(sorted(zip(scores, models)))))

class BlendClassifier(BaseEstimator, ClassifierMixin):
    def __init__(self, estimators=(SVC(),), coefs=(1,)):
        """
        Estimator which result is blending of many models.

        Parameters
        ----------
        estimators : tuple of classifiers, optional, default: (SVC(),)
            The tuple of classifiers to blend.
        coefs : tuple of numbers (float, int), optional, default: (1,)
            The tuple of coefficients for classifiers blending.
        """

        self._check_params(estimators, coefs)
        self.estimators = estimators
        self.coefs = coefs

    def __repr__(self):
        output = ''
        for idx, (est, coef) in enumerate(zip(self.estimators, self.coefs)):
            output += 'step {}. {} : {}\n'.format(idx + 1, est, coef)
        return output

    @property
    def n_estimators(self):
        return len(self.estimators)

    def fit(self, X, y):
        for est in self.estimators:
            est.fit(X, y)
        return self

    def predict(self, X):
        result = np.zeros(len(X))
        for coef, est in zip(self.coefs, self.estimators):
            result += coef * est.predict(X)
        return result

    def predict_proba(self, X):
        result = np.zeros((len(X), 2))
        for coef, est in zip(self.coefs, self.estimators):
            result += coef * est.predict_proba(X)
        return result

    def get_updated_classifier(self, new_estimators, new_coefs):
        """
        This func makes updated blending classifier
        and returns new blending classifier.

        Parameters
        ----------
        new_estimators : tuple of estimators
            The tuple of new models to blend.
        new_coefs : tuple of numbers (float, int)
            The tuple of coefficients of new models
            If sum of coefficients were equal to 1 then all models will be added
            with saving that rule. I.e. models will be added one by one and coefs
            will be updated by the following algorithm:
                - we have coefs array and new_coef to be added
                - coefs = coefs * (1 - new_coef)
                - coefs.append(new_coef)

        Returns
        -------
        blend_clf : BlendClassifier
            Updated blend classifier
        """

        _estimators = self.estimators + new_estimators
        _coefs = self.coefs
        if self._check_coefs_sum(verbose=True):
            # Append each coefficient saving whole sum equal to 1
            for coef in new_coefs:
                _coefs = tuple(map(lambda x: x * (1 - coef), self.coefs))
                _coefs = _coefs + (coef,)
        else:
            _coefs += new_coefs

        blend_clf = BlendClassifier(_estimators, _coefs)
        return blend_clf

    def _check_coefs_sum(self, verbose=True):
        if len(self.coefs) > 0 and sum(self.coefs) != 1:
            if verbose:
                print('WARNING: the sum of coefficients is not equal to 1.')
            return False
        return True

    def _check_array_elems_type(self, arr, types):
        for elem in arr:
            if not isinstance(elem, types):
                return False
        return True

    def _get_models_and_coefs(self, *args):
        if len(args) % 2:
            raise ValueError('Number of model and number of coefficients must be equal.')

        if len(args) == 2 and isinstance(args[0], (list, tuple)):
            models, coefs = args[0], args[1]
        else:
            models, coefs = args[:len(args) / 2], args[len(args) / 2:]

        if not self._check_array_elems_type(models, BaseEstimator):
            raise ValueError('All models must have some of estimator type.')
        if not self._check_array_elems_type(coefs, (float, int)):
            raise ValueError('All coefficients must be float.')

        return models, coefs

    def _check_params(self, estimators, coefs):
        if len(estimators) != len(coefs):
            raise ValueError('Number of estimators and number of coefficients must be the same.\n'
                             'Given estimators parameter has len {} and coefs parameter has len {}'
                             .format(len(estimators), len(coefs)))
        if not isinstance(estimators, tuple):
            raise ValueError('The estimators parameter must be a tuple type.\n'
                             'Given estimators parameter have {} type'.format(type(estimators)))
        if not isinstance(coefs, tuple):
            raise ValueError('The coefs must be a tuple type.\n'
                             'Given coefs parameter have {} type'.format(type(coefs)))


"""
========================================
Lasso and Elastic Net for Sparse Signals
========================================

Estimates Lasso and Elastic-Net regression models on a manually generated
sparse signal corrupted with an additive noise. Estimated coefficients are
compared with the ground-truth.

"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt

from sklearn.metrics import r2_score

###############################################################################
# generate some sparse data to play with
np.random.seed(42)

n_samples, n_features = 50, 200
X = np.random.randn(n_samples, n_features)
coef = 3 * np.random.randn(n_features)
inds = np.arange(n_features)
np.random.shuffle(inds)
coef[inds[10:]] = 0  # sparsify coef
y = np.dot(X, coef)

# add noise
y += 0.01 * np.random.normal((n_samples,))

# Split data in train set and test set
n_samples = X.shape[0]
X_train, y_train = X[:n_samples / 2], y[:n_samples / 2]
X_test, y_test = X[n_samples / 2:], y[n_samples / 2:]

###############################################################################
# Lasso
from sklearn.linear_model import Lasso

alpha = 0.1
lasso = Lasso(alpha=alpha)

y_pred_lasso = lasso.fit(X_train, y_train).predict(X_test)
r2_score_lasso = r2_score(y_test, y_pred_lasso)
print(lasso)
print("r^2 on test data : %f" % r2_score_lasso)

###############################################################################
# ElasticNet
from sklearn.linear_model import ElasticNet

enet = ElasticNet(alpha=alpha, l1_ratio=0.7)

y_pred_enet = enet.fit(X_train, y_train).predict(X_test)
r2_score_enet = r2_score(y_test, y_pred_enet)
print(enet)
print("r^2 on test data : %f" % r2_score_enet)

plt.plot(enet.coef_, color='lightgreen', linewidth=2,
         label='Elastic net coefficients')
plt.plot(lasso.coef_, color='gold', linewidth=2,
         label='Lasso coefficients')
plt.plot(coef, '--', color='navy', label='original coefficients')
plt.legend(loc='best')
plt.title("Lasso R^2: %f, Elastic Net R^2: %f"
          % (r2_score_lasso, r2_score_enet))
plt.show()


"""Analyze results.

Use the saved model output to calculate AUC and other metrics.
"""

import collections
import cPickle as pickle
import gflags as flags
import gzip
import logging
import numpy as np
import os
import pandas as pd
from sklearn import metrics
from statsmodels.stats import proportion
import sys

flags.DEFINE_string('root', None, 'Root directory containing model results.')
flags.DEFINE_string('dataset_file', None, 'Filename containing datasets.')
flags.DEFINE_string('prefix', None, 'Dataset prefix.')
flags.DEFINE_boolean('tversky', False, 'If True, use Tversky features.')
flags.DEFINE_integer('num_folds', 5, 'Number of cross-validation folds.')
flags.DEFINE_boolean('cycle', False,
                     'If True, expect multiple query molecules.')
flags.DEFINE_string('reload', None, 'Load previously analyzed results.')
flags.DEFINE_string('subset', None, 'Subset.')
FLAGS = flags.FLAGS

logging.getLogger().setLevel(logging.INFO)

FEATURES_MAP = {
    'rocs': 'TanimotoCombo',
    'shape_color': 'ST-CT',
    'shape_color_components': 'ST-CCT',
    'shape_color_overlaps': 'ST-CAO',
    'shape_color_components_overlaps': 'ST-CCT-CAO',
    'rocs_tversky': 'TverskyCombo',
    'shape_color_tversky': 'STv-CTv',
    'shape_color_components_tversky': 'STv-CCTv',
    'shape_color_components_tversky_overlaps': 'STv-CCTv-CAO',
}

MODEL_MAP = {
    'logistic': 'LR',
    'random_forest': 'RF',
    'svm': 'SVM',
}

def roc_enrichment(fpr, tpr, target_fpr):
    """Get ROC enrichment."""
    assert fpr[0] == 0
    assert fpr[-1] == 1
    assert np.all(np.diff(fpr) >= 0)
    return np.true_divide(np.interp(target_fpr, fpr, tpr), target_fpr)

def get_cv_metrics(y_true, y_pred):
    """Get 5-fold mean AUC."""
    assert len(y_true) == len(y_pred)
    fold_metrics = collections.defaultdict(list)
    for yt, yp in zip(y_true, y_pred):
        assert len(yt) == len(yp)
        fold_metrics['auc'].append(metrics.roc_auc_score(yt, yp))
        fpr, tpr, _ = metrics.roc_curve(yt, yp)
        for x in [0.005, 0.01, 0.02, 0.05, 0.1, 0.2]:
            fold_metrics['e-%g' % x].append(roc_enrichment(fpr, tpr, x))
    return fold_metrics

def add_rows(features, scores, rows, dataset, index=None):
    """Record per-fold and averaged cross-validation results."""
    for fold in range(len(scores['auc'])):
        row = {'dataset': dataset, 'features': features, 'fold': fold}
        if index is not None:
            row['index'] = index
        for key, values in scores.iteritems():
            row[key] = values[fold]
        rows.append(row)

    # Averages
    row = {'dataset': dataset, 'features': features, 'fold': 'all'}
    if index is not None:
        row['index'] = index
    for key, values in scores.iteritems():
        row[key] = np.mean(values)
    rows.append(row)

def load_output_and_calculate_metrics(model, subset):
    """Calculate metrics using saved model output.

    Args:
        model: String model type (e.g. logistic).
        subset: String query subset (e.g. omega1).

    Returns:
        DataFrame containing calculated metrics for each model/subset, including
        per-fold and average values for each reference molecule.
    """
    with open(FLAGS.dataset_file) as f:
        datasets = [line.strip() for line in f]
    rows = []
    for dataset in datasets:
        ref_idx = 0
        while True:  # Cycle through reference molecules.
            ref_idx_exists = get_ref_rows(model, subset, dataset, ref_idx, rows)
            if not FLAGS.cycle or not ref_idx_exists:
                break
            ref_idx += 1
        logging.info('%s\t%d', dataset, ref_idx)
    return pd.DataFrame(rows)

def get_ref_rows(model, subset, dataset, ref_idx, rows):
    logging.debug('ref_idx %d', ref_idx)
    for features in FEATURES_MAP.keys():
        logging.debug('Features: %s', features)
        fold_y_true = []
        fold_y_pred = []
        for fold_idx in range(FLAGS.num_folds):
            filename = get_output_filename(dataset, model, subset, features,
                                           fold_idx, ref_idx)
            if not os.path.exists(filename):
                return False
            logging.debug(filename)
            with gzip.open(filename) as f:
                df = pickle.load(f)
            fold_y_true.append(df['y_true'].values)
            fold_y_pred.append(df['y_pred'].values)
        scores = get_cv_metrics(fold_y_true, fold_y_pred)
        add_rows(features, scores, rows, dataset, index=ref_idx)
    return True

def get_output_filename(dataset, model, subset, features, fold_idx, ref_idx):
    if FLAGS.cycle:
        filename = os.path.join(
            '%s-%s' % (FLAGS.root, subset),
            dataset,
            'fold-%d' % fold_idx,
            '%s-%s-%s-%s-%s-fold-%d-ref-%d-output.pkl.gz' % (
                FLAGS.prefix, dataset, model, subset, features,
                fold_idx, ref_idx))
    else:
        assert ref_idx == 0
        filename = os.path.join(
            '%s-%s' % (FLAGS.root, subset),
            dataset,
            'fold-%d' % fold_idx,
            '%s-%s-%s-%s-%s-fold-%d-output.pkl.gz' % (
                FLAGS.prefix, dataset, model, subset, features,
                fold_idx))
    return filename

def load_data(model, subset):
    data = []
    with open(FLAGS.dataset_file) as f:
        for line in f:
            dataset = line.strip()
            filename = os.path.join(FLAGS.root, '%s-%s-%s-%s.pkl.gz' % (
                                    FLAGS.prefix, dataset, model, subset))
            assert os.path.exists(filename)
            logging.info(filename)
            with gzip.open(filename) as g:
                df = pickle.load(g)
            data.append(df)
    return pd.concat(data)

def confidence_interval(delta, metric):
    """Calculate a two-sided 95% confidence interval for differences."""
    # Wilson score interval for sign test.
    num_successes = np.count_nonzero(delta > 0)
    num_trials = np.count_nonzero(delta != 0)  # Exclude zero differences.
    lower, upper = proportion.proportion_confint(
        num_successes, num_trials, alpha=0.05, method='wilson')
    median_delta = delta.median()
    if metric == 'auc':
        median = r'%.3f' % median_delta
        ci = r'(%.2f, %.2f)' % (lower, upper)
    else:
        median = r'%.0f' % median_delta
        ci = r'(%.2f, %.2f)' % (lower, upper)
    if lower < 0.5 and upper < 0.5:
        median = r'\bfseries \color{red} ' + median
        ci = r'\bfseries \color{red} ' + ci
    elif lower > 0.5 and upper > 0.5:
        median = r'\bfseries ' + median
        ci = r'\bfseries ' + ci
    return median, ci

def data_table(data, subsets, models, kind=None, tversky=False):
    """Get medians and compare everything to ROCS.

    Args:
        data: DataFrame containing model performance.
        subsets: List of query subsets.
        models: List of models to include in the table.
        kind: List of metrics to report. Defaults to ['auc'].
        tversky: Boolean whether to use Tversky features. If False, use Tanimoto
            features.
    """
    if kind is None:
        kind = ['auc']
    if tversky:
        rocs_baseline = 'rocs_tversky'
        features_order = ['shape_color_tversky',
                          'shape_color_components_tversky',
                          'shape_color_overlaps',
                          'shape_color_components_tversky_overlaps']
    else:
        rocs_baseline = 'rocs'
        features_order = ['shape_color', 'shape_color_components',
                          'shape_color_overlaps',
                          'shape_color_components_overlaps']
    table = []

    # Get ROCS row.
    row = [r'\cellcolor{white} ROCS', FEATURES_MAP[rocs_baseline]]

    for subset in subsets:
        rocs_mask = ((data['features'] == rocs_baseline) &
                     (data['subset'] == subset) &
                     (data['model'] == models[0]))
        rocs_df = data[rocs_mask]
        logging.info('Confidence interval N = %d', len(rocs_df))
        logging.info('Number of datasets = %d',
                     len(pd.unique(rocs_df['dataset'])))
        for metric in kind:
            if metric == 'auc':
                number = '%.3f'
            else:
                number = '%.0f'
            row.extend([number % rocs_df[metric].median(), '', ''])
    table.append(' & '.join(row))

    # Get model rows.
    for model in models:
        for features in features_order:
            if features == features_order[-1]:
                row = [r'\multirow{-%d}{*}{\cellcolor{white} %s}' % (
                    len(features_order), MODEL_MAP[model])]
            else:
                row = [r'\cellcolor{white}']
            row.append(FEATURES_MAP[features])
            for subset in subsets:
                mask = ((data['features'] == features) &
                        (data['subset'] == subset) &
                        (data['model'] == model))
                df = data[mask]
                rocs_mask = ((data['features'] == rocs_baseline) &
                             (data['subset'] == subset) &
                             (data['model'] == model))
                rocs_df = data[rocs_mask]
                for metric in kind:
                    if metric == 'auc':
                        number = '%.3f'
                    else:
                        number = '%.0f'
                    row.append(number % df[metric].median())
                    if features == rocs_baseline:
                        row.append('')
                        row.append('')
                    else:
                        assert np.array_equal(df['dataset'].values,
                                              rocs_df['dataset'].values)
                        if 'index' in df.columns:
                            assert np.array_equal(df['index'].values,
                                                  rocs_df['index'].values)
                        delta = df.copy()
                        delta[metric] -= rocs_df[metric].values
                        row.extend(confidence_interval(delta[metric], metric))
            table.append(' & '.join(row))
    print ' \\\\\n'.join(table)

def main():
    if FLAGS.prefix == 'muv':
        subsets = ['omega1']
        assert FLAGS.cycle
    elif FLAGS.prefix == 'dude':
        subsets = ['xtal', 'omega1']
    elif FLAGS.prefix == 'chembl':
        subsets = ['omega1']
        assert FLAGS.cycle
    else:
        raise ValueError(FLAGS.prefix)
    if FLAGS.subset is not None:
        subsets = [FLAGS.subset]

    # Load data from output or previously processed.
    models = ['logistic', 'random_forest', 'svm']
    if FLAGS.reload is not None:
        logging.info('Loading processed data from %s', FLAGS.reload)
        data = pd.read_pickle(FLAGS.reload)
    else:
        data = []
        for model in models:
            for subset in subsets:
                logging.info('%s\t%s', model, subset)
                df = load_output_and_calculate_metrics(model, subset)
                df['model'] = model
                df['subset'] = subset
                data.append(df)
        data = pd.concat(data)

        # Save processed data.
        filename = '%s-processed.pkl.gz' % FLAGS.prefix
        logging.info('Saving processed data to %s', filename)
        with gzip.open(filename, 'wb') as f:
            pickle.dump(data, f, pickle.HIGHEST_PROTOCOL)

    # Only keep 5-fold mean information.
    mask = data['fold'] == 'all'
    data = data[mask]

    # AUC tables.
    # Combine subsets into a single table here.
    logging.info('AUC table')
    data_table(data, subsets, models, kind=['auc'], tversky=FLAGS.tversky)

    # Enrichment tables.
    # One per FPR.
    for metric in ['e-0.005', 'e-0.01', 'e-0.02', 'e-0.05']:
        logging.info('Metric: %s', metric)
        logging.info('Enrichment table')
        data_table(data, subsets, models, kind=[metric], tversky=FLAGS.tversky)

if __name__ == '__main__':
    flags.MarkFlagAsRequired('root')
    flags.MarkFlagAsRequired('dataset_file')
    flags.MarkFlagAsRequired('prefix')
    FLAGS(sys.argv)
    main()


import os
import vtkAll as vtk
import math
import time
import re
import numpy as np

from director.timercallback import TimerCallback
from director import objectmodel as om
from director.simpletimer import SimpleTimer
from director.utime import getUtime
from director import robotstate

import copy
import pickle
import scipy.interpolate

def asRobotPlan(msg):
    '''
    If the given message is a robot_plan_with_supports_t then this function returns
    the plan message contained within it.  For any other message type, this function
    just returns its input argument.
    '''
    try:
        import drc as lcmdrc
    except ImportError:
        pass
    else:
        if isinstance(msg, lcmdrc.robot_plan_with_supports_t):
            return msg.plan
    return msg

class PlanPlayback(object):

    def __init__(self):
        self.animationCallback = None
        self.animationTimer = None
        self.interpolationMethod = 'slinear'
        self.playbackSpeed = 1.0
        self.jointNameRegex = ''

    @staticmethod
    def getPlanPoses(msgOrList):

        if isinstance(msgOrList, list):
            messages = msgOrList
            allPoseTimes, allPoses = PlanPlayback.getPlanPoses(messages[0])

            for msg in messages[1:]:
                poseTimes, poses = PlanPlayback.getPlanPoses(msg)
                poseTimes += allPoseTimes[-1]
                allPoseTimes = np.hstack((allPoseTimes, poseTimes[1:]))
                allPoses += poses[1:]
            return allPoseTimes, allPoses

        else:
            msg = asRobotPlan(msgOrList)

            poses = []
            poseTimes = []
            for plan in msg.plan:
                pose = robotstate.convertStateMessageToDrakePose(plan)
                poseTimes.append(plan.utime / 1e6)
                poses.append(pose)
            return np.array(poseTimes), poses

    @staticmethod
    def getPlanElapsedTime(msg):
        msg = asRobotPlan(msg)
        startTime = msg.plan[0].utime
        endTime = msg.plan[-1].utime
        return (endTime - startTime) / 1e6

    @staticmethod
    def mergePlanMessages(plans):

        msg = copy.deepcopy(plans[0])

        for plan in plans[1:]:
            plan = copy.deepcopy(plan)

            lastTime = msg.plan[-1].utime
            for state in plan.plan:
                state.utime += lastTime

            msg.plan_info += plan.plan_info
            msg.plan += plan.plan

        msg.num_states = len(msg.plan)
        return msg

    @staticmethod
    def isPlanInfoFeasible(info):
        return 0 <= info < 10

    @staticmethod
    def isPlanFeasible(plan):
        plan = asRobotPlan(plan)
        return plan is not None and (max(plan.plan_info) < 10 and min(plan.plan_info) >= 0)

    def stopAnimation(self):
        if self.animationTimer:
            self.animationTimer.stop()

    def setInterpolationMethod(method):
        self.interpolationMethod = method

    def playPlan(self, msg, jointController):
        self.playPlans([msg], jointController)

    def playPlans(self, messages, jointController):

        assert len(messages)

        poseTimes, poses = self.getPlanPoses(messages)
        self.playPoses(poseTimes, poses, jointController)

    def getPoseInterpolatorFromPlan(self, message):
        poseTimes, poses = self.getPlanPoses(message)
        return self.getPoseInterpolator(poseTimes, poses)

    def getPoseInterpolator(self, poseTimes, poses, unwrap_rpy=True):
        if unwrap_rpy:
            poses = np.array(poses, copy=True)
            poses[:,3:6] = np.unwrap(poses[:,3:6],axis=0)

        if self.interpolationMethod in ['slinear', 'quadratic', 'cubic']:
            f = scipy.interpolate.interp1d(poseTimes, poses, axis=0, kind=self.interpolationMethod)
        elif self.interpolationMethod == 'pchip':
            f = scipy.interpolate.PchipInterpolator(poseTimes, poses, axis=0)
        return f

    def getPlanPoseMeshes(self, messages, jointController, robotModel, numberOfSamples):

        poseTimes, poses = self.getPlanPoses(messages)
        f = self.getPoseInterpolator(poseTimes, poses)
        sampleTimes = np.linspace(poseTimes[0], poseTimes[-1], numberOfSamples)
        meshes = []

        for sampleTime in sampleTimes:

            pose = f(sampleTime)
            jointController.setPose('plan_playback', pose)
            polyData = vtk.vtkPolyData()
            robotModel.model.getModelMesh(polyData)
            meshes.append(polyData)

        return meshes

    def showPoseAtTime(self, time, jointController, poseInterpolator):
        pose = poseInterpolator(time)
        jointController.setPose('plan_playback', pose)

    def playPoses(self, poseTimes, poses, jointController):

        f = self.getPoseInterpolator(poseTimes, poses)

        timer = SimpleTimer()

        def updateAnimation():

            tNow = timer.elapsed() * self.playbackSpeed

            if tNow > poseTimes[-1]:
                pose = poses[-1]
                jointController.setPose('plan_playback', pose)

                if self.animationCallback:
                    self.animationCallback()

                return False

            pose = f(tNow)
            jointController.setPose('plan_playback', pose)

            if self.animationCallback:
                self.animationCallback()

        self.animationTimer = TimerCallback()
        self.animationTimer.targetFps = 60
        self.animationTimer.callback = updateAnimation
        self.animationTimer.start()
        updateAnimation()

    def picklePlan(self, filename, msg):
        poseTimes, poses = self.getPlanPoses(msg)
        pickle.dump((poseTimes, poses), open(filename, 'w'))

    def getMovingJointNames(self, msg):
        poseTimes, poses = self.getPlanPoses(msg)
        diffs = np.diff(poses, axis=0)
        jointIds =  np.unique(np.where(diffs != 0.0)[1])
        jointNames = [robotstate.getDrakePoseJointNames()[jointId] for jointId in jointIds]
        return jointNames

    def plotPlan(self, msg):

        poseTimes, poses = self.getPlanPoses(msg)
        self.plotPoses(poseTimes, poses)

    def plotPoses(self, poseTimes, poses):

        import matplotlib.pyplot as plt

        poses = np.array(poses)

        if self.jointNameRegex:
            jointIds = range(poses.shape[1])
        else:
            diffs = np.diff(poses, axis=0)
            jointIds = np.unique(np.where(diffs != 0.0)[1])

        jointNames = [robotstate.getDrakePoseJointNames()[jointId] for jointId in jointIds]
        jointTrajectories = [poses[:,jointId] for jointId in jointIds]

        seriesNames = []

        sampleResolutionInSeconds = 0.01
        numberOfSamples = (poseTimes[-1] - poseTimes[0]) / sampleResolutionInSeconds
        xnew = np.linspace(poseTimes[0], poseTimes[-1], numberOfSamples)

        fig = plt.figure()
        ax = fig.add_subplot(111)

        for jointId, jointName, jointTrajectory in zip(jointIds, jointNames, jointTrajectories):

            if self.jointNameRegex and not re.match(self.jointNameRegex, jointName):
                continue

            x = poseTimes
            y = jointTrajectory

            y = np.rad2deg(y)

            if self.interpolationMethod in ['slinear', 'quadratic', 'cubic']:
                f = scipy.interpolate.interp1d(x, y, kind=self.interpolationMethod)
            elif self.interpolationMethod == 'pchip':
                f = scipy.interpolate.PchipInterpolator(x, y)

            ax.plot(x, y, 'ko')
            seriesNames.append(jointName + ' points')

            ax.plot(xnew, f(xnew), '-')
            seriesNames.append(jointName + ' ' + self.interpolationMethod)

        ax.legend(seriesNames, loc='upper right').draggable()
        ax.set_xlabel('time (s)')
        ax.set_ylabel('joint angle (deg)')
        ax.set_title('joint trajectories')
        plt.show()


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

import unittest
import nose

import datetime
from pandas.msgpack import packb, unpackb

class DummyException(Exception):
    pass

class TestExceptions(unittest.TestCase):

    def test_raise_on_find_unsupported_value(self):
        import datetime
        self.assertRaises(TypeError, packb, datetime.datetime.now())

    def test_raise_from_object_hook(self):
        def hook(obj):
            raise DummyException
        self.assertRaises(DummyException, unpackb, packb({}), object_hook=hook)
        self.assertRaises(DummyException, unpackb, packb({'fizz': 'buzz'}), object_hook=hook)
        self.assertRaises(DummyException, unpackb, packb({'fizz': 'buzz'}), object_pairs_hook=hook)
        self.assertRaises(DummyException, unpackb, packb({'fizz': {'buzz': 'spam'}}), object_hook=hook)
        self.assertRaises(DummyException, unpackb, packb({'fizz': {'buzz': 'spam'}}), object_pairs_hook=hook)

    def test_invalidvalue(self):
        self.assertRaises(ValueError, unpackb, b'\xd9\x97#DL_')


import numpy as np
import pandas as pd
from pandas import DataFrame

############################################################
# Life tables
# https://www.ssa.gov/oact/STATS/table4c6.html
############################################################

############################################################
# Male life table
############################################################

# survivors from 100000 births

MlivesArray = [100000, 99348, 99302, 99273, 99252, 99235, 99219, 99205, 99192, 99180,
               99170, 99161, 99151, 99138, 99119, 99091, 99052, 99003, 98943, 98870,
               98785, 98685, 98572, 98449, 98321, 98191, 98060, 97928, 97795, 97659,
               97519, 97376, 97230, 97080, 96927, 96772, 96612, 96448, 96277, 96097,
               95908, 95708, 95493, 95262, 95012, 94739, 94441, 94115, 93759, 93368,
               92940, 92472, 91961, 91406, 90804, 90153, 89450, 88693, 87883, 87022,
               86112, 85147, 84125, 83042, 81899, 80691, 79412, 78054, 76613, 75084,
               73461, 71732, 69889, 67930, 65853, 63657, 61329, 58859, 56249, 53504,
               50629, 47621, 44484, 41233, 37890, 34482, 31040, 27598, 24201, 20896,
               17735, 14768, 12043, 9599, 7463, 5647, 4157, 2977, 2075, 1410,
               935, 605, 380, 232, 137, 78, 43, 23, 11, 5,
               2, 1, 0, 0, 0, 0, 0, 0, 0, 0, ]

MlivesSeries = pd.Series(MlivesArray)

# life expectancy
MLEarray = [76.28, 75.78, 74.82, 73.84, 72.85, 71.87, 70.88, 69.89, 68.9, 67.9,
            66.91, 65.92, 64.92, 63.93, 62.94, 61.96, 60.99, 60.02, 59.05, 58.09,
            57.14, 56.2, 55.27, 54.33, 53.4, 52.47, 51.54, 50.61, 49.68, 48.75,
            47.82, 46.89, 45.96, 45.03, 44.1, 43.17, 42.24, 41.31, 40.38, 39.46,
            38.53, 37.61, 36.7, 35.78, 34.88, 33.98, 33.08, 32.19, 31.32, 30.44,
            29.58, 28.73, 27.89, 27.05, 26.23, 25.41, 24.61, 23.82, 23.03, 22.25,
            21.48, 20.72, 19.97, 19.22, 18.48, 17.75, 17.03, 16.32, 15.61, 14.92,
            14.24, 13.57, 12.92, 12.27, 11.65, 11.03, 10.43, 9.85, 9.28, 8.73,
            8.2, 7.68, 7.19, 6.72, 6.27, 5.84, 5.43, 5.04, 4.68, 4.34,
            4.03, 3.74, 3.47, 3.23, 3.01, 2.82, 2.64, 2.49, 2.36, 2.24,
            2.12, 2.01, 1.9, 1.8, 1.7, 1.6, 1.51, 1.42, 1.34, 1.26,
            1.18, 1.11, 1.04, 0.97, 0.9, 0.84, 0.78, 0.72, 0.67, 0.61,
            ]

MLEseries = pd.Series(MLEarray)

# death probability
MdeathrateArray = [0.006519, 0.000462, 0.000291, 0.000209, 0.000176, 0.000159, 0.000146, 0.000133, 0.000118, 0.000102,
                   0.000091, 0.000096, 0.000128, 0.000195, 0.000288, 0.000389, 0.000492, 0.000607, 0.000735, 0.000869,
                   0.001011, 0.001145, 0.001246, 0.001301, 0.001321, 0.00133, 0.001345, 0.001363, 0.001391, 0.001427,
                   0.001467, 0.001505, 0.001541, 0.001573, 0.001606, 0.001648, 0.001704, 0.001774, 0.001861, 0.001967,
                   0.002092, 0.00224, 0.002418, 0.002629, 0.002873, 0.003146, 0.003447, 0.003787, 0.004167, 0.004586,
                   0.005038, 0.00552, 0.006036, 0.006587, 0.00717, 0.007801, 0.008466, 0.009133, 0.009792, 0.010462,
                   0.011197, 0.012009, 0.012867, 0.013772, 0.014749, 0.015852, 0.017097, 0.018463, 0.019959, 0.021616,
                   0.023528, 0.025693, 0.028041, 0.030567, 0.033347, 0.036572, 0.040276, 0.044348, 0.048797, 0.053739,
                   0.059403, 0.065873, 0.073082, 0.08107, 0.089947, 0.099842, 0.110863, 0.123088, 0.136563, 0.151299,
                   0.167291, 0.18452, 0.202954, 0.222555, 0.243272, 0.263821, 0.283833, 0.302916, 0.320672, 0.336706,
                   0.353541, 0.371218, 0.389779, 0.409268, 0.429732, 0.451218, 0.473779, 0.497468, 0.522341, 0.548458,
                   0.575881, 0.604675, 0.634909, 0.666655, 0.699987, 0.734987, 0.771736, 0.810323, 0.850839, 0.893381,
                   ]
MdeathrateSeries = pd.Series(MdeathrateArray)

MlivesSeries.to_csv('MLivesSeries.csv', index_label='index')
MLEseries.to_csv('MLEseries.csv', index_label='index')
MdeathrateSeries.to_csv('MdeathrateSeries.csv', index_label='index')

############################################################
# Female life table
############################################################

FlivesArray = [100000, 99462, 99425, 99403, 99387, 99373, 99361, 99351, 99341, 99331,
               99322, 99312, 99303, 99291, 99278, 99262, 99243, 99220, 99194, 99165,
               99132, 99095, 99054, 99010, 98963, 98915, 98864, 98811, 98755, 98697,
               98635, 98569, 98500, 98426, 98348, 98265, 98176, 98081, 97979, 97870,
               97753, 97627, 97491, 97343, 97182, 97004, 96810, 96597, 96364, 96109,
               95829, 95524, 95193, 94834, 94449, 94038, 93598, 93126, 92623, 92090,
               91526, 90927, 90287, 89600, 88858, 88054, 87177, 86223, 85187, 84069,
               82864, 81561, 80147, 78616, 76961, 75177, 73244, 71148, 68888, 66467,
               63880, 61114, 58159, 55016, 51694, 48205, 44565, 40796, 36933, 33017,
               29104, 25257, 21542, 18027, 14775, 11839, 9267, 7083, 5285, 3852,
               2745, 1909, 1292, 850, 541, 333, 197, 112, 61, 31,
               15, 7, 3, 1, 0, 0, 0, 0, 0, 0,]

FlivesSeries = pd.Series(FlivesArray)

FLEarray = [81.05, 80.49, 79.52, 78.54, 77.55, 76.56, 75.57, 74.58, 73.58, 72.59,
            71.6, 70.6, 69.61, 68.62, 67.63, 66.64, 65.65, 64.67, 63.68, 62.7,
            61.72, 60.75, 59.77, 58.8, 57.82, 56.85, 55.88, 54.91, 53.94, 52.97,
            52.01, 51.04, 50.08, 49.11, 48.15, 47.19, 46.23, 45.28, 44.33, 43.37,
            42.43, 41.48, 40.54, 39.6, 38.66, 37.73, 36.81, 35.89, 34.97, 34.06,
            33.16, 32.27, 31.38, 30.49, 29.62, 28.74, 27.88, 27.01, 26.16, 25.31,
            24.46, 23.62, 22.78, 21.95, 21.13, 20.32, 19.52, 18.73, 17.95, 17.18,
            16.43, 15.68, 14.95, 14.23, 13.53, 12.83, 12.16, 11.5, 10.86, 10.24,
            9.64, 9.05, 8.48, 7.94, 7.42, 6.92, 6.44, 5.99, 5.57, 5.17,
            4.8, 4.45, 4.13, 3.84, 3.57, 3.34, 3.12, 2.93, 2.76, 2.6,
            2.45, 2.3, 2.17, 2.03, 1.91, 1.78, 1.67, 1.56, 1.45, 1.35,
            1.26, 1.17, 1.08, 1, 0.92, 0.85, 0.78, 0.72, 0.67, 0.61,]

FLEseries = pd.Series(FLEarray)

FdeathrateArray =[0.005377, 0.000379, 0.000221, 0.000162, 0.000133, 0.000119, 0.000109, 0.000101, 0.000096, 0.000093,
                  0.000094, 0.0001, 0.000112, 0.000134, 0.000162, 0.000194, 0.000226, 0.000261, 0.000297, 0.000334,
                  0.000373, 0.000412, 0.000446, 0.000472, 0.000493, 0.000513, 0.000537, 0.000563, 0.000593, 0.000627,
                  0.000664, 0.000705, 0.000748, 0.000794, 0.000845, 0.000903, 0.000968, 0.001038, 0.001113, 0.001196,
                  0.001287, 0.001393, 0.001517, 0.001662, 0.001827, 0.002005, 0.002198, 0.002412, 0.002648, 0.002904,
                  0.003182, 0.003473, 0.003767, 0.004058, 0.004352, 0.004681, 0.00504, 0.0054, 0.005756, 0.006128,
                  0.006545, 0.007034, 0.007607, 0.008281, 0.009057, 0.009953, 0.01095, 0.01201, 0.013124, 0.01433,
                  0.015728, 0.017338, 0.019108, 0.021041, 0.023191, 0.025713, 0.028609, 0.03176, 0.035157, 0.03892,
                  0.043289, 0.048356, 0.054041, 0.060384, 0.067498, 0.075516, 0.084556, 0.094703, 0.106014, 0.118513,
                  0.132206, 0.147092, 0.163154, 0.180371, 0.198714, 0.217264, 0.235735, 0.25381, 0.271155, 0.287424,
                  0.30467, 0.32295, 0.342327, 0.362867, 0.384639, 0.407717, 0.43218, 0.458111, 0.485597, 0.514733,
                  0.545617, 0.578354, 0.613055, 0.649839, 0.688829, 0.730159, 0.771736, 0.810323, 0.850839, 0.893381,]

FdeathrateSeries = pd.Series(FdeathrateArray)

def genLifetable(livesSeries, leSeries, ret_age, ret_length):
    """Create frame with life expectancies, lives"""

    # check bounds of lifetable
    # assert start age, end age between the two

    # 2nd version - take DataFrame where everything lines up by age

    end_age = ret_age + ret_length
    survival = livesSeries[ret_age:end_age]
    survival = survival / float(survival[ret_age])
    deathrate = survival - survival.shift(-1)
    deathrate.ix[end_age-1] = 1 - np.sum(deathrate)
    lifetable = DataFrame(survival, columns=['survival'])
    LE = leSeries[ret_age:end_age]
    lifetable['life_expectancy'] = LE
    lifetable['deathrate'] = deathrate
    return lifetable


import logging
pvl_logger = logging.getLogger('pvlib')

import numpy as np
import pandas as pd

from nose.tools import raises
from numpy.testing import assert_almost_equal
from pandas.util.testing import assert_frame_equal, assert_series_equal

from pvlib.location import Location
from pvlib import clearsky
from pvlib import solarposition

from . import requires_scipy

# setup times and location to be tested.
tus = Location(32.2, -111, 'US/Arizona', 700)
times = pd.date_range(start='2014-06-24', end='2014-06-25', freq='3h')
times_localized = times.tz_localize(tus.tz)

ephem_data = solarposition.get_solarposition(times_localized, tus.latitude,
                                             tus.longitude)

@requires_scipy
def test_ineichen_required():
    # the clearsky function should call lookup_linke_turbidity by default
    expected = pd.DataFrame(
        np.array([[    0.        ,     0.        ,     0.        ],
                  [    0.        ,     0.        ,     0.        ],
                  [   51.47811191,   265.33462162,    84.48262202],
                  [  105.008507  ,   832.29100407,   682.67761951],
                  [  121.97988054,   901.31821834,  1008.02102657],
                  [  112.57957512,   867.76297247,   824.61702926],
                  [   76.69672675,   588.8462898 ,   254.5808329 ],
                  [    0.        ,     0.        ,     0.        ],
                  [    0.        ,     0.        ,     0.        ]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times_localized)
    out = clearsky.ineichen(times_localized, tus.latitude, tus.longitude)
    assert_frame_equal(expected, out)

def test_ineichen_supply_linke():
    expected = pd.DataFrame(np.array(
        [[    0.        ,     0.        ,     0.        ],
         [    0.        ,     0.        ,     0.        ],
         [   40.16490879,   321.71856556,    80.12815294],
         [   95.14336873,   876.49252839,   703.47605855],
         [  118.4587024 ,   939.81646535,  1042.34480815],
         [  105.36645492,   909.11265773,   851.32459694],
         [   61.91187639,   647.35889938,   257.42691896],
         [    0.        ,     0.        ,     0.        ],
         [    0.        ,     0.        ,     0.        ]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times_localized)
    out = clearsky.ineichen(times_localized, tus.latitude, tus.longitude,
                            altitude=tus.altitude,
                            linke_turbidity=3)
    assert_frame_equal(expected, out)

def test_ineichen_solpos():
    clearsky.ineichen(times_localized, tus.latitude, tus.longitude,
                      linke_turbidity=3,
                      solarposition_method='ephemeris')

def test_ineichen_airmass():
    expected = pd.DataFrame(
        np.array([[    0.        ,     0.        ,     0.        ],
                  [    0.        ,     0.        ,     0.        ],
                  [   53.90422388,   257.01655613,    85.87406435],
                  [  101.34055688,   842.92925705,   686.39337307],
                  [  117.7573735 ,   909.70367947,  1012.04184961],
                  [  108.6233401 ,   877.30589626,   828.49118038],
                  [   75.23108133,   602.06895546,   257.10961202],
                  [    0.        ,     0.        ,     0.        ],
                  [    0.        ,     0.        ,     0.        ]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times_localized)
    out = clearsky.ineichen(times_localized, tus.latitude, tus.longitude,
                            linke_turbidity=3,
                            airmass_model='simple')
    assert_frame_equal(expected, out)

@requires_scipy
def test_lookup_linke_turbidity():
    times = pd.date_range(start='2014-06-24', end='2014-06-25',
                          freq='12h', tz=tus.tz)
    # expect same value on 2014-06-24 0000 and 1200, and
    # diff value on 2014-06-25
    expected = pd.Series(np.array([3.10126582, 3.10126582, 3.11443038]),
                         index=times)
    out = clearsky.lookup_linke_turbidity(times, tus.latitude, tus.longitude)
    assert_series_equal(expected, out)

@requires_scipy
def test_lookup_linke_turbidity_nointerp():
    times = pd.date_range(start='2014-06-24', end='2014-06-25',
                          freq='12h', tz=tus.tz)
    # expect same value for all days
    expected = pd.Series(np.array([3., 3., 3.]), index=times)
    out = clearsky.lookup_linke_turbidity(times, tus.latitude, tus.longitude,
                                          interp_turbidity=False)
    assert_series_equal(expected, out)

@requires_scipy
def test_lookup_linke_turbidity_months():
    times = pd.date_range(start='2014-04-01', end='2014-07-01',
                          freq='1M', tz=tus.tz)
    expected = pd.Series(np.array([2.8943038, 2.97316456, 3.18025316]),
                         index=times)
    out = clearsky.lookup_linke_turbidity(times, tus.latitude,
                                          tus.longitude)
    assert_series_equal(expected, out)

@requires_scipy
def test_lookup_linke_turbidity_nointerp_months():
    times = pd.date_range(start='2014-04-10', end='2014-07-10',
                          freq='1M', tz=tus.tz)
    expected = pd.Series(np.array([2.85, 2.95, 3.]), index=times)
    out = clearsky.lookup_linke_turbidity(times, tus.latitude, tus.longitude,
                                          interp_turbidity=False)
    assert_series_equal(expected, out)
    # changing the dates shouldn't matter if interp=False
    times = pd.date_range(start='2014-04-05', end='2014-07-05',
                          freq='1M', tz=tus.tz)
    out = clearsky.lookup_linke_turbidity(times, tus.latitude, tus.longitude,
                                          interp_turbidity=False)
    assert_series_equal(expected, out)

def test_haurwitz():
    expected = pd.DataFrame(np.array([[0.],
                                      [0.],
                                      [82.85934048],
                                      [699.74514735],
                                      [1016.50198354],
                                      [838.32103769],
                                      [271.90853863],
                                      [0.],
                                      [0.]]),
                             columns=['ghi'], index=times_localized)
    out = clearsky.haurwitz(ephem_data['zenith'])
    assert_frame_equal(expected, out)


"""This is the Bokeh charts interface. It gives you a high level API to build
complex plot is a simple way.

This is the ChartObject class, a minimal prototype class to build more chart
types on top of it. It provides the mechanisms to support the shared chained
methods.
"""
#-----------------------------------------------------------------------------
# Copyright (c) 2012 - 2014, Continuum Analytics, Inc. All rights reserved.
#
# Powered by the Bokeh Development Team.
#
# The full license is in the file LICENSE.txt, distributed with this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

from __future__ import absolute_import

from six import string_types
from collections import OrderedDict
from ..properties import bokeh_integer_types, Datetime

try:
    import numpy as np

except ImportError:
    np = None

try:
    import pandas as pd

except ImportError:
    pd = None
try:
    import blaze
except ImportError:
    blaze=None
#-----------------------------------------------------------------------------
# Classes and functions
#-----------------------------------------------------------------------------

DEFAULT_INDEX_ALIASES = list('abcdefghijklmnopqrstuvz1234567890')
DEFAULT_INDEX_ALIASES += list(zip(DEFAULT_INDEX_ALIASES, DEFAULT_INDEX_ALIASES))

class DataAdapter(object):
    """
    Adapter object used to normalize Charts inputs to a common interface.
    Supported inputs are dict, list, tuple, np.ndarray and pd.DataFrame.
    """
    def __init__(self, data, index=None, columns=None, force_alias=True):
        self.__values = data
        self._values = self.validate_values(data)

        self.convert_index_to_int = False
        self._columns_map = {}
        self.convert_items_to_dict = False

        if columns is None and force_alias:
            # no column 'labels' defined for data... in this case we use
            # default names
            keys = getattr(self._values, 'keys', None)
            if callable(keys):
                columns = list(keys())
            elif keys is None:
                columns = list(map(str, range(len(data))))
            else:
                columns = list(keys)

        if columns:
            self._columns = columns

            # define a mapping between the real keys to access data and the aliases
            # we have defined using 'columns'
            self._columns_map = dict(zip(columns, self.keys()))

        if index is not None:
            self._index = index
            self.convert_items_to_dict = True

        elif force_alias:
            _index = getattr(self._values, 'index', None)

            # check because if it is a callable self._values is not a
            # dataframe (probably a list)
            if _index is None:
                indexes = self.index

                if isinstance(indexes[0], int):
                    self._index = DEFAULT_INDEX_ALIASES[:][:len(self.values()[0])]
                    self.convert_items_to_dict = True

            elif not callable(_index):
                self._index = list(_index)
                self.convert_items_to_dict = True

            else:
                self._index = DEFAULT_INDEX_ALIASES[:][:len(self.values()[0])]
                self.convert_items_to_dict = True

    @staticmethod
    def is_number(value):
        numbers = (float, ) + bokeh_integer_types
        return isinstance(value, numbers)

    @staticmethod
    def is_datetime(value):
        try:
            dt = Datetime(value)
            dt # shut up pyflakes
            return True

        except ValueError:
            return False

    @staticmethod
    def validate_values(values):
        if np and isinstance(values, np.ndarray):
            if len(values.shape) == 1:
                return np.array([values])
            else:
                return values

        elif pd and isinstance(values, pd.DataFrame):
            return values

        elif isinstance(values, (dict, OrderedDict)):
            if all(DataAdapter.is_number(x) for x in values.values()):
                return values

            return values

        elif isinstance(values, (list, tuple)):
            if all(DataAdapter.is_number(x) for x in values):
                return [values]

            return values

        elif hasattr(values, '__array__'):
            values = pd.DataFrame(np.asarray(values))
            return values

        # TODO: Improve this error message..
        raise TypeError("Input type not supported! %s" % values)

    def index_converter(self, x):
        key = self._columns_map.get(x, x)
        if self.convert_index_to_int:
            key = int(key)
        return key

    def keys(self):
        # assuming it's a dict or dataframe
        keys = getattr(self._values, "keys", None)

        if callable(keys):
            return list(keys())

        elif keys is None:
            self.convert_index_to_int = True
            indexes = range(len(self._values))
            return list(map(str, indexes))

        else:
            return list(keys)

    def __len__(self):
        return len(self.values())

    def __iter__(self):
        for k in self.keys():
            yield k

    def __getitem__(self, key):
        val = self._values[self.index_converter(key)]

        # if we have "index aliases" we need to remap the values...
        if self.convert_items_to_dict:
            val = dict(zip(self._index, val))

        return val

    def values(self):
        return self.normalize_values(self._values)

    @staticmethod
    def normalize_values(values):
        _values = getattr(values, "values", None)

        if callable(_values):
            return list(_values())

        elif _values is None:
            return values

        else:
            # assuming it's a dataframe, in that case it returns transposed
            # values compared to it's dict equivalent..
            return list(_values.T)

    def items(self):
        return [(key, self[key]) for key in self]

    def iterkeys(self):
        return iter(self)

    def itervalues(self):
        for k in self:
            yield self[k]

    def iteritems(self):
        for k in self:
            yield (k, self[k])

    @property
    def columns(self):
        try:
            return self._columns

        except AttributeError:
            return list(self.keys())

    @property
    def index(self):
        try:
            return self._index

        except AttributeError:
            index = getattr(self._values, "index", None)

            if not callable(index) and index is not None:
                # guess it's a pandas dataframe..
                return index

        # no, it's not. So it's probably a list so let's get the
        # values and check
        values = self.values()

        if isinstance(values, dict):
            return list(values.keys())

        else:
            first_el = self.values()[0]

            if isinstance(first_el, dict):
                indexes = list(first_el.keys())

            else:
                indexes = range(0, len(self.values()[0]))
                self._index = indexes
            return indexes

    #-----------------------------------------------------------------------------
    # Convenience methods
    #-----------------------------------------------------------------------------
    @staticmethod
    def get_index_and_data(values, index=None):
        """Parse values (that must be one of the DataAdapter supported
        input types) and create an separate/create index and data
        depending on values type and index.

        Args:
            values (iterable): container that holds data to be plotted using
                on the Chart classes

        Returns:
            A tuple of (index, values), where: ``index`` is an iterable that
            represents the data index and ``values`` is an iterable containing
            the values to be plotted.

        """
        _values = DataAdapter(values, force_alias=False)
        if hasattr(values, 'keys'):
            if index is not None:
                if isinstance(index, string_types):
                    xs = _values[index]
                else:
                    xs = index
            else:
                try:
                    xs = _values.index
                except AttributeError:
                    xs = values.index
        else:
            if index is None:
                xs = _values.index
            elif isinstance(index, string_types):
                xs = _values[index]
            else:
                xs = index

        return xs, _values