# -*- coding: iso-8859-1 -*-
"""
This code is used to weigh the UV radiances we compute by biological action spectra.
"""
########################
###Import useful libraries
########################

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import pdb
from matplotlib.pyplot import cm
from scipy import interpolate as interp
import scipy.integrate

########################
###Set physical constants
########################
hc=1.98645e-9 #value of h*c in erg*nm

def cm2inch(cm): #function to convert cm to inches; useful for complying with Astrobiology size guidelines
	return cm/2.54

########################
###Decide which bits of the calculation will be run
########################
plotactionspec=False #if true, plots the action spectra we are using.
plotactionspec_talk=False #if true, plots the action spectra we are using...but, optimized for a talk instead of a paper
calculatealbaz=False #if true, generates the table for the albedo and zenith angle study
calculateco2=False #if true, generates the table for the co2 study
calculatealtgas=True #if true, generates the table for the alternate gas study

########################
###Helper functions: I/O
########################
def get_UV(filename):
	"""
	Input: filename (including path)
	Output: (wave_leftedges, wav_rightedges, surface radiance) in units of (nm, nm, photons/cm2/sec/nm)
	"""
	wav_leftedges, wav_rightedges, wav, toa_intensity, surface_flux,  surface_intensity, surface_intensity_diffuse, surface_intensity_direct=np.genfromtxt(filename, skip_header=1, skip_footer=0, usecols=(0, 1, 2,3,4,6,7,8), unpack=True)

	surface_intensity_photons=surface_intensity*(wav/(hc))

	return wav_leftedges, wav_rightedges, surface_intensity_photons

########################
###Helper functions: UV Dosimeters
########################

def integrated_radiance(wav_left, wav_right, surf_int, leftlim, rightlim):
	"""
	Computes the surface radiance integrated from leftlim to rightlim. Does this by doing a trapezoid sum. NOTE: The method I have chosen works only so long as the limits line up with the bin edges!

	wav_left: left edge of wavelength bin, in nm
	wav_right: right edge of wavelength bin, in nm
	surf_int: total surface intensity (radiance, hemispherically-integrated) in photons/cm2/s/nm, in bin defined by wav_left and wav_right
	produceplots: if True, shows plots of what it is computing
	returnxy: if True, returns x,y for action spectrum.
	"""

	allowed_inds=np.where((wav_left>=leftlim) & (wav_right<=rightlim))
	delta_wav=wav_right[allowed_inds]-wav_left[allowed_inds]

	surf_int_integrated=np.sum(surf_int[allowed_inds]*delta_wav) #integration converts from photons/cm2/s/nm to photons/cm2/s

	return surf_int_integrated

def tricyano_aqe_prodrate(wav_left, wav_right, surf_int, lambda0, produceplots, returnxy):
	"""
	Weights the input surface intensities by the action spectrum for the photoproduction of aquated electrons from Ritson+2012 and Patel+2015, i.e. irradiation of tricyano cuprate. The action spectrum is composed of the absorption spectrum multiplied by an assumed quantum yield function. We assume the QY function to be a step function, stepping from 0 at wavelengths longer than lambda0 to 0.06 at wavelengths shorter than lambda0. We choose  0.06 for the step function to match the estimate found by Horvath+1984; we note this value may be pH sensitive. Empirically, we know that lambda0>254 nm, but that's about it.

	This process is an eustressor for abiogenesis.

	wav_left: left edge of wavelength bin, in nm
	wav_right: right edge of wavelength bin, in nm
	surf_int: total surface intensity (radiance, hemispherically-integrated) in photons/cm2/s/nm, in bin defined by wav_left and wav_right
	lambda0: value assume for lambda0.
	produceplots: if True, shows plots of what it is computing
	returnxy: if True, returns x,y for action spectrum.
	"""
	####Step 1: reduce input spectrum to match bounds of available dataset.
	int_min=190.0 #This lower limit of integration is set by the limits of the cucn3 absorption dataset (left edge of bin)
	int_max=351.0 #This upper limit of integration is set by the limits of the cucn3 absorption dataset (right edge of bin)
	allowed_inds=np.where((wav_left>=int_min) & (wav_right<=int_max)) #indices that correspond to included data

	wav_left=wav_left[allowed_inds]
	wav_right=wav_right[allowed_inds]
	surf_int=surf_int[allowed_inds]

	delta_wav=wav_right-wav_left #size of wavelength bins in nm

	####Step 2: form the action spectrum from the absorption spectrum and QY curve.
	#Import the tricyanocuprate absorption spectrum
	importeddata=np.genfromtxt('./Raw_Data/Magnani_Data/CuCN3_XC.dat', skip_header=2)
	cucn3_wav=importeddata[:,0] #wav in nm
	cucn3_molabs=importeddata[:,1] #molar absorptivities in  L/(mol*cm), decadic
	cucn3_molabs_func=interp.interp1d(cucn3_wav, cucn3_molabs, kind='linear') #functionalized form of cucn3 molar absorption
	#does not matter if you use decadic or natural logarithmic as constant factors normalize out anyway

	#Formulate the step-function quantum yield curve
	def qy_stepfunc(wav, lambda0): #step function, for the photoionization model
		"""Returns 1 for wav<=lambda0 and 0 for wav>lambda0"""
		qy=np.zeros(np.size(wav))# initialize all to zero
		inds=np.where(wav<=lambda0) #indices where the wavelength is below the threshold
		qy[inds]=qy[inds]+0.06 #increase the QE to 1 at the indices where the wavelength is below the threshold
		return qy

	#Integrate these quantities to match the input spectral resolution
	qy_dist=np.zeros(np.shape(wav_left))#initialize variable to hold the QY integrated over the surface intensity wavelength bins
	cucn3_molabs_dist=np.zeros(np.shape(wav_left))#initialize variable to hold the QY integrated over the surface intensity wavelength bins
	for ind in range(0, len(wav_left)):
		leftedge=wav_left[ind]
		rightedge=wav_right[ind]

		cucn3_molabs_dist[ind]=scipy.integrate.quad(cucn3_molabs_func, leftedge, rightedge, epsabs=0, epsrel=1e-5)[0]/(rightedge-leftedge)
		qy_dist[ind]=scipy.integrate.quad(qy_stepfunc, leftedge, rightedge, args=(lambda0), epsabs=0, epsrel=1e-5)[0]/(rightedge-leftedge)

	action_spectrum=cucn3_molabs_dist*qy_dist

	#Normalize action spectrum to 1 at 195 (arbitrary)
	action_spectrum=action_spectrum*(1./(np.interp(190., 0.5*(wav_left+wav_right), action_spectrum)))

	####Step 3: Compute action-spectrum weighted total intensity
	weighted_surface_intensity=surf_int*action_spectrum

	total_weighted_radiance=np.sum(weighted_surface_intensity*delta_wav) #units: photons/cm2/s

	####Step 4 (Optional): Plot various components of action spectrum to show the multiplication
	if produceplots:
		legendfontsize=12
		axisfontsize=12

		##Plot ribonucleotide absorption and interpolation
		fig1, axarr=plt.subplots(3,2,sharex=True, figsize=(8., 10.5)) #specify figure size (width, height) in inches

		axarr[0,0].bar(wav_left, surf_int,width=delta_wav, color='black', alpha=0.5, log=True)
		axarr[0,0].set_ylim([1e10,1e16])
		axarr[0,0].legend(loc=2, prop={'size':legendfontsize})
		axarr[0,0].yaxis.grid(True)
		axarr[0,0].xaxis.grid(True)
		axarr[0,0].set_ylabel('Surface Radiance \n(photons cm$^{-2}$s$^{-1}$nm$^{-1}$)', fontsize=axisfontsize)
		#axarr[0,0].title.set_position([0.5, 1.11])
		#axarr[0,0].text(0.5, 1.1, r'a(i)', transform=axarr[0].transAxes, va='top')

		axarr[1,0].bar(wav_left, cucn3_molabs_dist,width=delta_wav, color='black', alpha=0.5, log=True)
		#axarr[1,0].set_ylim([-0.1, 1.1])
		axarr[1,0].legend(loc=6, prop={'size':legendfontsize})
		axarr[1,0].yaxis.grid(True)
		axarr[1,0].xaxis.grid(True)
		axarr[1,0].set_ylabel('CuCN3 Molar Absorptivity\n(M$^{-1}$cm$^{-1}$)', fontsize=axisfontsize)
		#axarr[1,0].text(0.5, 1.10, r'b(i)', fontsize=12, transform=axarr[1].transAxes, va='top')

		axarr[2,0].bar(wav_left, qy_dist,width=delta_wav, color='black', alpha=0.5)
		axarr[2,0].set_ylim([-0.01, 0.06])
		axarr[2,0].legend(loc=6, prop={'size':legendfontsize})
		axarr[2,0].yaxis.grid(True)
		axarr[2,0].xaxis.grid(True)
		axarr[2,0].set_ylabel('Quantum Efficiency \n(reductions absorption$^{-1}$)', fontsize=axisfontsize)
		#axarr[2,0].text(0.5, 1.10, r'c(i)', fontsize=12,transform=axarr[2].transAxes, va='top')

		axarr[0,1].bar(wav_left, action_spectrum,width=delta_wav, color='black', alpha=0.5)
		#axarr[0,1].set_ylim([-0.1, 1.1])
		axarr[0,1].legend(loc=6, prop={'size':legendfontsize})
		axarr[0,1].yaxis.grid(True)
		axarr[0,1].xaxis.grid(True)
		axarr[0,1].set_ylabel('Action Spectrum', fontsize=axisfontsize)
		#axarr[0,1].text(0.5, 1.10, r'b(i)', fontsize=12, transform=axarr[1].transAxes, va='top')

		axarr[1,1].bar(wav_left, weighted_surface_intensity,width=delta_wav, color='black', alpha=0.5)
		#axarr[1,1].set_ylim([-0.1, 1.1])
		axarr[1,1].legend(loc=6, prop={'size':legendfontsize})
		axarr[1,1].yaxis.grid(True)
		axarr[1,1].xaxis.grid(True)
		axarr[1,1].set_ylabel('Weighted Surface Radiance', fontsize=axisfontsize)
		#axarr[1,1].text(0.5, 1.10, r'b(i)', fontsize=12, transform=axarr[1].transAxes, va='top')
		#plt.savefig('/home/sranjan/Python/UV/Plots/ritson_assumed_qe_v3.pdf', orientation='portrait',papertype='letter', format='pdf')
		plt.show()
	if returnxy:
		return 0.5*(wav_left+wav_right), action_spectrum
	else:
		return total_weighted_radiance

def ump_glycosidic_photol(wav_left, wav_right, surf_int, lambda0, produceplots, returnxy):
	"""
	Weights the input surface intensities by the action spectrum for cleavage of the glycosidic bond in UMP (the U-RNA monomer), aka base release. We form this spectrum by convolving the pH=7.6 absorption spectrum for Uridine-3'-(2')-phosporic acid (i.e. uridylic acid, UMP) from Voet et al (1963) with an assumed QY curve. The QY curve is based on the work of Gurzadyan and Gorner (1994); they measure (wavelength, QY) for N-glycosidic bond cleavage in UMP in anoxic aqueous solution (Ar-suffused) to be (193 nm, 4.3e-3) and (254 nm, (2-3)e-5). Specifically, we assume that QY=4.3e-3 for lambda<=lambda_0 and QY=2.5e-5 for lambda>lambda_0. natural choices of lambda_0 are 194, 254, and 230 (first two: empirical limits. Last: end of pi-pi* absorption bad, Sinsheimer+1949 suggest it is onset of irreversible photolytic damage).

	This process is a stressor for abiogenesis.

	wav_left: left edge of wavelength bin, in nm
	wav_right: right edge of wavelength bin, in nm
	surf_int: total surface intensity (radiance, hemispherically-integrated) in photons/cm2/s/nm, in bin defined by wav_left and wav_right
	lambda0: value assume for lambda0.
	produceplots: if True, shows plots of what it is computing
	returnxy: if True, returns x,y for action spectrum.
	"""
	####Step 1: reduce input spectrum to match bounds of available dataset (absorption).
	int_min=184.0 #This lower limit of integration is set by the limits of the cucn3 absorption dataset (left edge of bin)
	int_max=299.0 #This upper limit of integration is set by the limits of the cucn3 absorption dataset (right edge of bin)
	allowed_inds=np.where((wav_left>=int_min) & (wav_right<=int_max)) #indices that correspond to included data

	wav_left=wav_left[allowed_inds]
	wav_right=wav_right[allowed_inds]
	surf_int=surf_int[allowed_inds]

	delta_wav=wav_right-wav_left #size of wavelength bins in nm

	####Step 2: form the action spectrum from the absorption spectrum and QY curve.
	#Import the UMP absorption spectrum from Voet et al 1963
	importeddata=np.genfromtxt('./Raw_Data/Voet_Data/ribouridine_pH_7.3_v2.txt', skip_header=0, delimiter=',')
	ump_wav=importeddata[:,0] #wav in nm
	ump_molabs=importeddata[:,1] #molar absorptivities\times 10^{3}, i.e. in units of 10^{-3} L/(mol*cm), decadic (I think -- unit scheme unclear in paper. Not important since normalized out)
	ump_molabs_func=interp.interp1d(ump_wav, ump_molabs, kind='linear') #functionalized form of molar absorption
	#does not matter if you use decadic or natural logarithmic as constant factors normalize out anyway

	#Formulate the step-function quantum yield curve
	def qy_stepfunc(wav, lambda0): #step function, for the photoionization model
		"""QY based on work of  Gurzadyan and Gorner 1994"""
		qy=np.zeros(np.size(wav))# initialize all to zero
		inds1=np.where(wav<=lambda0) #indices where the wavelength is below the threshold
		inds2=np.where(wav>lambda0) #indices where the wavelength is below the threshold
		qy[inds1]=qy[inds1]+4.3e-3 #High QY for lambda<=lambda0
		qy[inds2]=qy[inds2]+2.5e-5 #Low QY for lambda>lambda0
		return qy

	#Integrate these quantities to match the input spectral resolution
	qy_dist=np.zeros(np.shape(wav_left))#initialize variable to hold the QY integrated over the surface intensity wavelength bins
	ump_molabs_dist=np.zeros(np.shape(wav_left))#initialize variable to hold the UMP absorption integrated over the surface intensity wavelength bins
	for ind in range(0, len(wav_left)):
		leftedge=wav_left[ind]
		rightedge=wav_right[ind]

		ump_molabs_dist[ind]=scipy.integrate.quad(ump_molabs_func, leftedge, rightedge, epsabs=0, epsrel=1e-5)[0]/(rightedge-leftedge)
		qy_dist[ind]=scipy.integrate.quad(qy_stepfunc, leftedge, rightedge, args=(lambda0),epsabs=0, epsrel=1e-5)[0]/(rightedge-leftedge)

	action_spectrum=ump_molabs_dist*qy_dist

	#Normalize action spectrum to 1 at 195 (arbitrary)
	action_spectrum=action_spectrum*(1./(np.interp(190., 0.5*(wav_left+wav_right), action_spectrum)))

	####Step 3: Compute action-spectrum weighted total intensity
	weighted_surface_intensity=surf_int*action_spectrum

	total_weighted_radiance=np.sum(weighted_surface_intensity*delta_wav) #units: photons/cm2/s

	####Step 4 (Optional): Plot various components of action spectrum to show the multiplication
	if produceplots:
		legendfontsize=12
		axisfontsize=12

		##Plot ribonucleotide absorption and interpolation
		fig1, axarr=plt.subplots(3,2,sharex=True, figsize=(8., 10.5)) #specify figure size (width, height) in inches

		axarr[0,0].bar(wav_left, surf_int,width=delta_wav, color='black', alpha=0.5, log=True)
		axarr[0,0].set_ylim([1e10,1e16])
		axarr[0,0].legend(loc=2, prop={'size':legendfontsize})
		axarr[0,0].yaxis.grid(True)
		axarr[0,0].xaxis.grid(True)
		axarr[0,0].set_ylabel('Surface Radiance \n(photons cm$^{-2}$s$^{-1}$nm$^{-1}$)', fontsize=axisfontsize)
		#axarr[0,0].title.set_position([0.5, 1.11])
		#axarr[0,0].text(0.5, 1.1, r'a(i)', transform=axarr[0].transAxes, va='top')

		axarr[1,0].bar(wav_left, ump_molabs_dist,width=delta_wav, color='black', alpha=0.5, log=False)
		#axarr[1,0].set_ylim([-0.1, 1.1])
		axarr[1,0].legend(loc=6, prop={'size':legendfontsize})
		axarr[1,0].yaxis.grid(True)
		axarr[1,0].xaxis.grid(True)
		axarr[1,0].set_ylabel('UMP Molar Absorptivity\n(M$^{-1}$cm$^{-1}$)', fontsize=axisfontsize)
		#axarr[1,0].text(0.5, 1.10, r'b(i)', fontsize=12, transform=axarr[1].transAxes, va='top')

		axarr[2,0].bar(wav_left, qy_dist,width=delta_wav, color='black', alpha=0.5, log=True)
		axarr[2,0].set_ylim([1e-5, 1e-2])
		axarr[2,0].legend(loc=6, prop={'size':legendfontsize})
		axarr[2,0].yaxis.grid(True)
		axarr[2,0].xaxis.grid(True)
		axarr[2,0].set_ylabel('Quantum Efficiency \n(reductions absorption$^{-1}$)', fontsize=axisfontsize)
		#axarr[2,0].text(0.5, 1.10, r'c(i)', fontsize=12,transform=axarr[2].transAxes, va='top')

		axarr[0,1].bar(wav_left, action_spectrum,width=delta_wav, color='black', alpha=0.5)
		#axarr[0,1].set_ylim([-0.1, 1.1])
		axarr[0,1].legend(loc=6, prop={'size':legendfontsize})
		axarr[0,1].yaxis.grid(True)
		axarr[0,1].xaxis.grid(True)
		axarr[0,1].set_ylabel('Action Spectrum', fontsize=axisfontsize)
		#axarr[0,1].text(0.5, 1.10, r'b(i)', fontsize=12, transform=axarr[1].transAxes, va='top')

		axarr[1,1].bar(wav_left, weighted_surface_intensity,width=delta_wav, color='black', alpha=0.5)
		#axarr[1,1].set_ylim([-0.1, 1.1])
		axarr[1,1].legend(loc=6, prop={'size':legendfontsize})
		axarr[1,1].yaxis.grid(True)
		axarr[1,1].xaxis.grid(True)
		axarr[1,1].set_ylabel('Weighted Surface Radiance', fontsize=axisfontsize)
		#axarr[1,1].text(0.5, 1.10, r'b(i)', fontsize=12, transform=axarr[1].transAxes, va='top')
		#plt.savefig('/home/sranjan/Python/UV/Plots/ritson_assumed_qe_v3.pdf', orientation='portrait',papertype='letter', format='pdf')
		plt.show()
	if returnxy:
		return 0.5*(wav_left+wav_right), action_spectrum
	else:
		return total_weighted_radiance

########################
###Plot UV Dosimeters
########################
if plotactionspec:
	#Set up wavelength scale
	wave_left=np.arange(100., 500.)
	wave_right=np.arange(101., 501.)
	wave_centers=0.5*(wave_left+wave_right)
	surf_int=np.ones(np.shape(wave_centers)) #for our purposes here, this is a thunk.

	#Extract action spectra
	wav_gly_193, actspec_gly_193=ump_glycosidic_photol(wave_left, wave_right, surf_int, 193., False, True)
	wav_gly_230, actspec_gly_230=ump_glycosidic_photol(wave_left, wave_right, surf_int, 230., False, True)
	wav_gly_254, actspec_gly_254=ump_glycosidic_photol(wave_left, wave_right, surf_int, 254., False, True)
	wav_aqe_254, actspec_aqe_254=tricyano_aqe_prodrate(wave_left, wave_right, surf_int, 254., False, True)
	wav_aqe_300, actspec_aqe_300=tricyano_aqe_prodrate(wave_left, wave_right, surf_int, 300., False, True)

	#####Plot action spectra
	#Initialize Figure
	fig, (ax1)=plt.subplots(1, figsize=(cm2inch(16.5),6), sharex=True)
	colorseq=iter(cm.rainbow(np.linspace(0,1,5)))
	#Plot Data
	ax1.plot(wav_gly_193,actspec_gly_193, linestyle='-',linewidth=2, color=next(colorseq), label=r'UMP Gly Bond Cleavage ($\lambda_0=193$)')
	ax1.plot(wav_gly_230,actspec_gly_230, linestyle='-',linewidth=2, color=next(colorseq), label=r'UMP Gly Bond Cleavage ($\lambda_0=230$)')
	ax1.plot(wav_gly_254,actspec_gly_254, linestyle='-',linewidth=2, color=next(colorseq), label=r'UMP Gly Bond Cleavage ($\lambda_0=254$)')
	ax1.plot(wav_aqe_254,actspec_aqe_254, linestyle='-',linewidth=2, color=next(colorseq), label=r'CuCN$_{3}$$^{2-}$ Photoionization ($\lambda_0=254$)')
	ax1.plot(wav_aqe_300,actspec_aqe_300, linestyle='--',linewidth=2, color=next(colorseq), label=r'CuCN$_{3}$$^{2-}$ Photoionization ($\lambda_0=300$)')

	#####Finalize and save figure
	ax1.set_title(r'Action Spectra')
	ax1.set_xlim([180.,360.])
	ax1.set_xlabel('nm')
	ax1.set_ylabel(r'Relative Sensitivity')
	ax1.set_yscale('log')
	ax1.set_ylim([1e-6, 1e2])
	#ax1.legend(bbox_to_anchor=[0, 1.1, 1,1], loc=3, ncol=2, mode='expand', borderaxespad=0., fontsize=10)
	ax1.legend(loc='upper right', ncol=1, fontsize=10)
	plt.tight_layout(rect=(0,0,1,1))
	plt.savefig('./Plots/actionspectra.eps', orientation='portrait',papertype='letter', format='eps')

if plotactionspec_talk:
	#Set up wavelength scale
	wave_left=np.arange(100., 500.)
	wave_right=np.arange(101., 501.)
	wave_centers=0.5*(wave_left+wave_right)
	surf_int=np.ones(np.shape(wave_centers)) #for our purposes here, this is a thunk.

	#Extract action spectra
	wav_gly_193, actspec_gly_193=ump_glycosidic_photol(wave_left, wave_right, surf_int, 193., False, True)
	wav_gly_230, actspec_gly_230=ump_glycosidic_photol(wave_left, wave_right, surf_int, 230., False, True)
	wav_gly_254, actspec_gly_254=ump_glycosidic_photol(wave_left, wave_right, surf_int, 254., False, True)
	wav_aqe_254, actspec_aqe_254=tricyano_aqe_prodrate(wave_left, wave_right, surf_int, 254., False, True)
	wav_aqe_300, actspec_aqe_300=tricyano_aqe_prodrate(wave_left, wave_right, surf_int, 300., False, True)

	#####Plot action spectra
	#Initialize Figure
	fig, (ax1)=plt.subplots(1, figsize=(10,9), sharex=True)
	colorseq=iter(cm.rainbow(np.linspace(0,1,5)))
	#Plot Data
	ax1.plot(wav_gly_193,actspec_gly_193, linestyle='-',linewidth=3, color=next(colorseq), label=r'UMP-193')
	ax1.plot(wav_gly_230,actspec_gly_230, linestyle='-',linewidth=3, color=next(colorseq), label=r'UMP-230')
	ax1.plot(wav_gly_254,actspec_gly_254, linestyle='-',linewidth=3, color=next(colorseq), label=r'UMP-254')
	ax1.plot(wav_aqe_254,actspec_aqe_254, linestyle='-',linewidth=3, color=next(colorseq), label=r'CuCN3-254')
	ax1.plot(wav_aqe_300,actspec_aqe_300, linestyle='--',linewidth=3, color=next(colorseq), label=r'CuCN3-300')

	#####Finalize and save figure
	ax1.set_title(r'Action Spectra', fontsize=24)
	ax1.set_xlim([180.,360.])
	ax1.set_xlabel('nm',fontsize=24)
	ax1.set_ylabel(r'Relative Sensitivity', fontsize=24)
	ax1.set_yscale('log')
	ax1.set_ylim([1e-6, 1e2])
	ax1.legend(bbox_to_anchor=[0, 1.1, 1,0.5], loc=3, ncol=2, mode='expand', borderaxespad=0., fontsize=24)
	#ax1.legend(loc='upper right', ncol=1, fontsize=16)
	ax1.xaxis.set_tick_params(labelsize=24)
	ax1.yaxis.set_tick_params(labelsize=24)
	plt.tight_layout(rect=(0,0,1,0.75))
	plt.savefig('./TalkFigs/actionspectra.pdf', orientation='portrait',papertype='letter', format='pdf')
########################
###Set "base" values to normalize the alb-zen, co2, and alt-gas dosimeters by
########################

#Use the TOA flux in order to get a good, physically understandable denominator.
wav_leftedges, wav_rightedges, wav, toa_intensity=np.genfromtxt('./TwoStreamOutput/AlbZen/rugheimer_earth_epoch0_a=0.2_z=60.dat', skip_header=1, skip_footer=0, usecols=(0, 1,2, 3), unpack=True)
toa_intensity_photons=toa_intensity*(wav/(hc))

#Compute base doses
intrad100_165_base=integrated_radiance(wav_leftedges, wav_rightedges, toa_intensity_photons, 100, 165.) #This measures the flux vulnerable to activity
intrad200_300_base=integrated_radiance(wav_leftedges, wav_rightedges, toa_intensity_photons, 200., 300.) #This is just an empirical gauge.
umpgly_193_base=ump_glycosidic_photol(wav_leftedges, wav_rightedges, toa_intensity_photons, 193., False, False)
umpgly_230_base=ump_glycosidic_photol(wav_leftedges, wav_rightedges, toa_intensity_photons,230.,  False, False)
umpgly_254_base=ump_glycosidic_photol(wav_leftedges, wav_rightedges, toa_intensity_photons, 254., False, False)
tricyano254_base=tricyano_aqe_prodrate(wav_leftedges, wav_rightedges, toa_intensity_photons, 254., False, False)
tricyano300_base=tricyano_aqe_prodrate(wav_leftedges, wav_rightedges, toa_intensity_photons, 300., False, False)
########################
###Run code for albedo, zenith angle
########################
if calculatealbaz:
	#Evaluate only two zenith angles (to show range of variation)
	zenithangles=['66.5', '0']
	albedos=['tundra', 'ocean', 'desert', 'oldsnow', 'newsnow']

	for zenind in range(0, len(zenithangles)):
		zenithangle=zenithangles[zenind]

		for albind in range(0, len(albedos)):
			albedo=albedos[albind]

			datafile='./TwoStreamOutput/AlbZen/rugheimer_earth_epoch0_a='+albedo+'_z='+zenithangle+'.dat'
			left, right, surface_int=get_UV(datafile)

			intrad100_165=integrated_radiance(left, right, surface_int, 100, 165.) #This measures the flux vulnerable to activity
			intrad200_300=integrated_radiance(left, right, surface_int, 200., 300.) #This is just an empirical gauge.
			umpgly_193=ump_glycosidic_photol(left, right, surface_int, 193., False, False)
			umpgly_230=ump_glycosidic_photol(left, right, surface_int,230.,  False, False)
			umpgly_254=ump_glycosidic_photol(left, right, surface_int, 254., False, False)
			tricyano254=tricyano_aqe_prodrate(left, right, surface_int, 254., False, False)
			tricyano300=tricyano_aqe_prodrate(left, right, surface_int, 300., False, False)

			line=np.array([zenithangle, albedo, intrad100_165/intrad100_165_base,intrad200_300/intrad200_300_base, umpgly_193/umpgly_193_base, umpgly_230/umpgly_230_base, umpgly_254/umpgly_254_base, tricyano254/tricyano254_base, tricyano300/tricyano300_base])

			if (albind==0 and zenind==0):
				albzentable=line #need to initialize in this case
			else:
				albzentable=np.vstack((albzentable, line))

	#Save output
	f=open('./Doses/albzen_uv_doses.dat','w')
	f.write('All Dosimeters Normalized to Space Radiation Case\n')
	np.savetxt(f, albzentable, delimiter='	', fmt='%s', newline='\n', header='Zenith Angle &  Albedo & Radiance (100-165 nm) & Radiance (200-300 nm) & UMP Gly Cleavage (lambda0=193nm) & UMP Gly Cleavage (lambda0=230nm) & UMP Gly Cleavage (lambda0=254nm) &  CuCN3 Photoionization (lambda0=254 nm) & CuCN3 Photoionization (lambda0=300 nm)\n')
	f.close()

########################
###Run code for varying CO2 levels
########################
if calculateco2:
	N_co2_rugh=2.09e24 #column density of CO2 in Rugheimer base model (cm**-2)
	co2multiples=np.array([0., 1.e-6,1.e-5, 1.e-4, 1.e-3, 0.00893, 1.e-2, 1.e-1, 0.6, 1., 1.33, 1.e1, 46.6, 1.e2, 470., 1.e3])
	zenithangles=['0', '66.5']
	albedos=['newsnow', 'tundra']

	for surfind in range(0, len(zenithangles)):
		albedo=albedos[surfind]
		zenithangle=zenithangles[surfind]

		for multind in range(0, len(co2multiples)):
			multiple=co2multiples[multind]
			colden_co2=N_co2_rugh*multiple

			datafile='./TwoStreamOutput/CO2lim/surface_intensities_co2limits_co2multiple='+str(multiple)+'_a='+albedo+'_z='+zenithangle+'.dat'
			left, right, surface_int=get_UV(datafile)

			intrad100_165=integrated_radiance(left, right, surface_int, 100, 165.) #This measures the flux vulnerable to activity
			intrad200_300=integrated_radiance(left, right, surface_int, 200., 300.) #This is just an empirical gauge.
			umpgly_193=ump_glycosidic_photol(left, right, surface_int, 193., False, False)
			umpgly_230=ump_glycosidic_photol(left, right, surface_int,230.,  False, False)
			umpgly_254=ump_glycosidic_photol(left, right, surface_int, 254., False, False)
			tricyano254=tricyano_aqe_prodrate(left, right, surface_int, 254., False, False)
			tricyano300=tricyano_aqe_prodrate(left, right, surface_int, 300., False, False)
			#print intrad200_300
			#pdb.set_trace()
			line=np.array([zenithangle, albedo, colden_co2, intrad100_165/intrad100_165_base,intrad200_300/intrad200_300_base, umpgly_193/umpgly_193_base, umpgly_230/umpgly_230_base, umpgly_254/umpgly_254_base, tricyano254/tricyano254_base, tricyano300/tricyano300_base])

			if (multind==0 and surfind==0):
				co2table=line #need to initialize in this case
			else:
				co2table=np.vstack((co2table, line))

	#Save Output
	f=open('./Doses/co2_uv_doses.dat','w')
	f.write('All Dosimeters Normalized to Space Radiation Case\n')
	np.savetxt(f, co2table, delimiter='	', fmt='%s', newline='\n', header='Zenith Angle & Albedo & Radiance (100-165 nm) & Radiance (200-300 nm) & UMP Gly Cleavage (lambda0=193nm) & UMP Gly Cleavage (lambda0=230nm) & UMP Gly Cleavage (lambda0=254nm) &  CuCN3 Photoionization (lambda0=254 nm) & CuCN3 Photoionization (lambda0=300 nm)\n')
	f.close()

########################
###Run code for alternate gas absorption.
########################
if calculatealtgas:
	#####Set up info about the files to extract # All are the maximum possible natural surface radiance case (z=0, albedo=fresh snow) aka "max"
	N_tot=2.0925e25#total column density of Rugheimer+2015 model in cm**-2
	gaslist=['h2o', 'ch4', 'so2', 'o2', 'o3', 'h2s'] #list of gases we are doing this for
	base_abundances=np.array([4.657e-3, 1.647e-6, 3.548e-11, 2.241e-6, 8.846e-11, 7.097e-11]) #molar concentration of each of these gases in the Rugheimer model.

	gasmultiples={}#dict holding the multiples of the molar concentration we are using
	gasmultiples['h2o']=np.array([1.e-5, 1.e-4, 1.e-3, 1.e-2, 1.e-1, 1., 1.e1, 1.e2, 1.e3, 1.e4, 1.e5])
	gasmultiples['ch4']=np.array([1.e-5, 1.e-4, 1.e-3, 1.e-2, 1.e-1, 1., 1.e1, 1.e2, 1.e3, 1.e4, 1.e5])
	gasmultiples['so2']=np.array([1.e-5, 1.e-4, 1.e-3, 1.e-2, 1.e-1, 1., 1.e1, 1.e2, 1.e3, 1.e4, 1.e5, 1.e6, 1.e7])
	gasmultiples['o2']=np.array([1.e-5, 1.e-4, 1.e-3, 1.e-2, 1.e-1, 1., 1.e1, 1.e2, 1.e3, 1.e4, 1.e5])
	gasmultiples['o3']=np.array([1.e-5, 1.e-4, 1.e-3, 1.e-2, 1.e-1, 1., 1.e1, 1.e2, 1.e3, 1.e4, 1.e5])
	gasmultiples['h2s']=np.array([1.e-5, 1.e-4, 1.e-3, 1.e-2, 1.e-1, 1., 1.e1, 1.e2, 1.e3, 1.e4, 1.e5, 1.e6, 1.e7])

	#####In a loop, extract the files and compute the statistics
	for gasind in range(0, len(gaslist)):
		gas=gaslist[gasind]
		base_abundance=base_abundances[gasind]
		multiples=gasmultiples[gas]

		for multind in range(0, len(multiples)):
			multiple=multiples[multind]
			colden_X=base_abundance*multiple*N_tot #total column density of gas X

			datafile='./TwoStreamOutput/gaslim/surface_intensities_'+gas+'limits_'+gas+'multiple='+str(multiple)+'_a=newsnow_z=0.dat'
			left, right, surface_int=get_UV(datafile)

			intrad100_165=integrated_radiance(left, right, surface_int, 100, 165.) #This measures the flux vulnerable to activity
			intrad200_300=integrated_radiance(left, right, surface_int, 200., 300.) #This is just an empirical gauge.
			umpgly_193=ump_glycosidic_photol(left, right, surface_int, 193., False, False)
			umpgly_230=ump_glycosidic_photol(left, right, surface_int,230.,  False, False)
			umpgly_254=ump_glycosidic_photol(left, right, surface_int, 254., False, False)
			tricyano254=tricyano_aqe_prodrate(left, right, surface_int, 254., False, False)
			tricyano300=tricyano_aqe_prodrate(left, right, surface_int, 300., False, False)

			line=np.array([gas, colden_X, intrad100_165/intrad100_165_base,intrad200_300/intrad200_300_base, umpgly_193/umpgly_193_base, umpgly_230/umpgly_230_base, umpgly_254/umpgly_254_base, tricyano254/tricyano254_base, tricyano300/tricyano300_base])

			if (multind==0):
				altgastable=line #need to initialize in this case
			else:
				altgastable=np.vstack((altgastable, line))
			f=open('./Doses/'+gas+'_uv_doses.dat','w')
			f.write('All Dosimeters Normalized to Space Radiation Case\n')
			np.savetxt(f, altgastable, delimiter=' & ', fmt='%s', newline='\n', header='Gas & Column Density (cm-2) & Radiance (100-165 nm) & Radiance (200-300 nm) & UMP Gly Cleavage (lambda0=193nm) & UMP Gly Cleavage (lambda0=230nm) & UMP Gly Cleavage (lambda0=254nm) &  CuCN3 Photoionization (lambda0=254 nm) & CuCN3 Photoionization (lambda0=300 nm)\n')
			f.close()

	#Wrap Up

########################
###Wrap Up
########################
plt.show()


# Copyright (c) 2014 Google Inc. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""A simple shell that uses the IPython messaging system."""

# Override platform information.
import platform
platform.system = lambda: "pnacl"
platform.release = lambda: "chrome"

import time
import json
import logging
import sys
import Queue
import thread

stdin_input = Queue.Queue()
shell_input = Queue.Queue()
stdin_output = Queue.Queue()
shell_output = Queue.Queue()
iopub_output = Queue.Queue()

sys_stdout = sys.stdout
sys_stderr = sys.stderr

def emit(s):
    print >> sys_stderr, "EMITTING: %s" % (s)
    time.sleep(1)

import IPython
from IPython.core.interactiveshell import InteractiveShell, InteractiveShellABC
from IPython.utils.traitlets import Type, Dict, Instance
from IPython.core.displayhook import DisplayHook
from IPython.utils import py3compat
from IPython.utils.py3compat import builtin_mod
from IPython.utils.jsonutil import json_clean, encode_images
from IPython.core.displaypub import DisplayPublisher
from IPython.config.configurable import Configurable

# module defined in shell.cc for communicating via pepper API
from pyppapi import nacl_instance

def CreateMessage(msg_type, parent_header=None, content=None):
  if parent_header is None:
    parent_header = {}
  if content is None:
    content = {}
  return {
      'header': {'msg_type': msg_type},
      'parent_header': parent_header,
      'content': content,
      'msg_type': msg_type,
  }

class MsgOutStream(object):
  """Class to overrides stderr and stdout."""

  def __init__(self, stream_name):
    self._stream_name = stream_name
    self._parent_header = {}

  def SetParentHeader(self, parent_header):
    self._parent_header = parent_header

  def close(self):
    pass

  def flush(self):
    pass

  def write(self, string):
    iopub_output.put(CreateMessage('stream', parent_header=self._parent_header,
                content={'name': self._stream_name, 'data': string}))

  def writelines(self, sequence):
    for string in sequence:
      self.write(string)

# override sys.stdout and sys.stderr to broadcast on iopub
stdout_stream = MsgOutStream('stdout')
stderr_stream = MsgOutStream('stderr')
sys.stdout = stdout_stream
sys.stderr = stderr_stream

class PepperShellDisplayHook(DisplayHook):
  parent_header = Dict({})

  def set_parent_header(self, parent_header):
    """Set the parent for outbound messages."""
    self.parent_header = parent_header

  def start_displayhook(self):
    self.content = {}

  def write_output_prompt(self):
    self.content['execution_count'] = self.prompt_count

  def write_format_data(self, format_dict, md_dict=None):
    self.content['data'] = encode_images(format_dict)
    self.content['metadata'] = md_dict

  def finish_displayhook(self):
    sys.stdout.flush()
    sys.stderr.flush()
    iopub_output.put(CreateMessage('pyout', parent_header=self.parent_header,
                content=self.content))
    self.content = None

class PepperDisplayPublisher(DisplayPublisher):
  parent_header = Dict({})

  def set_parent_header(self, parent_header):
    self.parent_header = parent_header

  def _flush_streams(self):
    """flush IO Streams prior to display"""
    sys.stdout.flush()
    sys.stderr.flush()

  def publish(self, source, data, metadata=None):
    self._flush_streams()
    if metadata is None:
      metadata = {}
    self._validate_data(source, data, metadata)
    content = {}
    content['source'] = source
    content['data'] = encode_images(data)
    content['metadata'] = metadata
    iopub_output.put(CreateMessage('display_data', content=json_clean(content),
                parent_header=self.parent_header))

  def clear_output(self, stdout=True, stderr=True, other=True):
    content = dict(stdout=stdout, stderr=stderr, other=other)

    if stdout:
      sys.stdout.write('\r')
    if stderr:
      sys.stderr.write('\r')

    self._flush_streams()
    iopub_output.put(CreateMessage('clear_output', content=content,
                parent_header=self.parent_header))

class PepperInteractiveShell(InteractiveShell):
    """A subclass of InteractiveShell for the Pepper Messagin API."""
    displayhook_class = Type(PepperShellDisplayHook)
    display_pub_class = Type(PepperDisplayPublisher)
    @staticmethod
    def enable_gui(gui):
        pass

InteractiveShellABC.register(PepperInteractiveShell)

class PepperKernel(Configurable):
    shell = Instance('IPython.core.interactiveshell.InteractiveShellABC')
    shell_class = Type(PepperInteractiveShell)

    def __init__(self):
        self.shell = self.shell_class.instance(parent=self)
        self.shell.run_cell("""
import os
matplotlib_config_dir = '/mplconfigdir'
os.environ['XDG_CONFIG_HOME'] = matplotlib_config_dir
os.environ['TMP'] = ''
import matplotlib
import matplotlib.cbook
""")

shell = PepperKernel().shell

# Taken from IPython 2.x branch, IPython/kernel/zmq/ipykernel.py
def _complete(msg):
  c = msg['content']
  try:
    cpos = int(c['cursor_pos'])
  except:
    # If we don't get something that we can convert to an integer, at
    # least attempt the completion guessing the cursor is at the end of
    # the text, if there's any, and otherwise of the line
    cpos = len(c['text'])
    if cpos==0:
      cpos = len(c['line'])
  return shell.complete(c['text'], c['line'], cpos)

# Special message to indicate the NaCl kernel is ready.
iopub_output.put(CreateMessage('status', content={'execution_state': 'nacl_ready'}))

def _no_raw_input(self):
  """Raise StdinNotImplentedError if active frontend doesn't support
  stdin."""
  raise StdinNotImplementedError("raw_input was called, but this "
                                 "frontend does not support stdin.")

def _raw_input(prompt, parent_header):
    # Flush output before making the request.
    sys.stderr.flush()
    sys.stdout.flush()
    # flush the stdin socket, to purge stale replies
    while True:
        try:
            stdin_input.get_nowait()
        except Queue.Empty:
            break

    # Send the input request.
    content = json_clean(dict(prompt=prompt))
    stdin_output.put(CreateMessage('input_request', content=content,
      parent_header=parent_header))

    # Await a response.
    while True:
        try:
            reply = stdin_input.get()
        except Exception:
            print "Invalid Message"
        except KeyboardInterrupt:
            # re-raise KeyboardInterrupt, to truncate traceback
            raise KeyboardInterrupt
        else:
            break
    try:
        value = py3compat.unicode_to_str(reply['content']['value'])
    except:
        print "Got bad raw_input reply: "
        print reply
        value = ''
    if value == '\x04':
        # EOF
        raise EOFError
    return value

def main_loop():
  execution_count = 1
  while 1:
    iopub_output.put(CreateMessage('status', content={'execution_state': 'idle'}))
    msg = shell_input.get()
    iopub_output.put(CreateMessage('status', content={'execution_state': 'busy'}))

    if not 'header' in msg:
      continue
    request_header = msg['header']
    if not 'msg_type' in request_header:
      continue
    msg_type = request_header['msg_type']
    if msg_type == 'execute_request':
      try:
        content = msg[u'content']
        code = content[u'code']
        silent = content[u'silent']
        store_history = content.get(u'store_history', not silent)
      except:
        self.log.error("Got bad msg: ")
        self.log.error("%s", msg)
        continue

      # Replace raw_input. Note that is not sufficient to replace
      # raw_input in the user namespace.
      if content.get('allow_stdin', False):
        raw_input = lambda prompt='': _raw_input(prompt, request_header)
        input = lambda prompt='': eval(raw_input(prompt))
      else:
        raw_input = input = lambda prompt='' : _no_raw_input()

      if py3compat.PY3:
        _sys_raw_input = builtin_mod.input
        builtin_mod.input = raw_input
      else:
        _sys_raw_input = builtin_mod.raw_input
        _sys_eval_input = builtin_mod.input
        builtin_mod.raw_input = raw_input
        builtin_mod.input = input

      # Let output streams know which message the output is for
      stdout_stream.SetParentHeader(request_header)
      stderr_stream.SetParentHeader(request_header)
      shell.displayhook.set_parent_header(request_header)
      shell.display_pub.set_parent_header(request_header)

      status = 'ok'
      content = {}
      try:
        shell.run_cell(msg['content']['code'],
                       store_history=store_history,
                       silent=silent)
      except Exception, ex:
        status = 'error'
        logging.exception('Exception occured while running cell')
      finally:
        # Restore raw_input.
        if py3compat.PY3:
          builtin_mod.input = _sys_raw_input
        else:
          builtin_mod.raw_input = _sys_raw_input
          builtin_mod.input = _sys_eval_input

      content = {'status': status,
               'execution_count': execution_count}

      if status == 'ok':
        content['payload'] = []
        content['user_variables'] = {}
        content['user_expressions'] = {}
      elif status == 'error':
        content['ename'] = type(ex).__name__
        content['evalue'] = str(ex)
        content['traceback'] = []

      execution_count += 1
      if status == 'error':
        iopub_output.put(CreateMessage('pyerr', parent_header=request_header,
                    content={
                        'execution_count': execution_count,
                        'ename': type(ex).__name__,
                        'evalue': str(ex),
                        'traceback': []
                        }
                    ))
      shell_output.put(CreateMessage('execute_reply', parent_header=request_header,
                  content=content))
    elif msg_type == 'complete_request':
      # Taken from IPython 2.x branch, IPython/kernel/zmq/ipykernel.py
      txt, matches = _complete(msg)
      matches = {'matches' : matches,
                 'matched_text' : txt,
                 'status' : 'ok'}
      matches = json_clean(matches)
      shell_output.put(CreateMessage('complete_reply',
                  parent_header = request_header,
                  content = matches))
    elif msg_type == 'object_info_request':
      # Taken from IPython 2.x branch, IPython/kernel/zmq/ipykernel.py
      content = msg['content']
      object_info = shell.object_inspect(content['oname'],
                      detail_level = content.get('detail_level', 0))
      # Before we send this object over, we scrub it for JSON usage
      oinfo = json_clean(object_info)
      shell_output.put(CreateMessage('object_info_reply',
                  parent_header = request_header,
                  content = oinfo))
    elif msg_type == 'restart':
      # break out of this loop, ending this program.
      # The main event loop in shell.cc will then
      # run this program again.
      break
    elif msg_type == 'kill':
      # Raise an exception so that the function
      # running this script will return -1, resulting
      # in no restart of this script.
      raise RuntimeError

thread.start_new_thread(main_loop, ())

def deal_message(msg):
  channel = msg['stream']
  content = json.loads(msg['json'])

  queues = {'shell': shell_input, 'stdin': stdin_input}
  queue = queues[channel]

  queue.put(content)

def send_message(stream, msg):
  nacl_instance.send_raw_object({
    'stream': stream,
    'json': json.dumps(msg)
  })

while 1:
  msg = nacl_instance.wait_for_message(timeout=1, sleeptime=10000)
  try:
    deal_message(msg)
  except:
    pass

  output_streams = [
    (stdin_output, 'stdin'),
    (shell_output, 'shell'),
    (iopub_output, 'iopub')
  ]
  for msg_queue, stream in output_streams:
    msg = None
    try:
      msg = msg_queue.get_nowait()
      send_message(stream, msg)
    except Queue.Empty:
      pass


# Run h264dec on a single file compiled with PRIOR_STATS and then run this script
# Outputs timeseries plot at /tmp/misses.pdf
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
import os

def temporal_misses(key):
    values = data[key]
    numbins = 100
    binsize = len(values) // numbins
    bins = [[]]
    for v in values:
        if len(bins[-1]) >= binsize:
            bins.append([])
        bins[-1].append(v)

    x = range(len(bins))
    total_misses = float(sum(values))
    y = [100 * float(sum(b)) / total_misses for b in bins]
    return plt.plot(x, y, label=key)[0]

paths = filter(lambda s: 'misses.log' in s, os.listdir('/tmp/'))
data = {p.split('_misses.')[0]: map(lambda c: c == '0', open('/tmp/' + p).read()) for p in paths}
handles = []
plt.figure(figsize=(20,10))
keys = data.keys()
for k in keys:
    handles.append(temporal_misses(k))
plt.axis((0, 100, 0, 2))
plt.xlabel('temporal %')
plt.ylabel('% total misses')
plt.legend(handles, keys, bbox_to_anchor=(1, 1), bbox_transform=plt.gcf().transFigure)

out = PdfPages('/tmp/misses.pdf')
out.savefig()
out.close()


import numpy as np
from optparse import OptionParser
import scipy.linalg as la
import scipy.stats as stats
import scipy.linalg.blas as blas
import pandas as pd
import csv
import time
import fastlmm.util.VertexCut as vc
from pysnptools.snpreader.bed import Bed
import pysnptools.util as pstutil
import pysnptools.util.pheno as phenoUtils
np.set_printoptions(precision=3, linewidth=200)

def loadData(bfile, extractSim, phenoFile, missingPhenotype='-9', loadSNPs=False, standardize=True):
	bed = Bed(bfile, count_A1=True)

	if (extractSim is not None):
		f = open(extractSim)
		csvReader = csv.reader(f)
		extractSnpsSet = set([])
		for l in csvReader: extractSnpsSet.add(l[0])
		f.close()
		keepSnpsInds = [i for i in range(bed.sid.shape[0]) if bed.sid[i] in extractSnpsSet]
		bed = bed[:, keepSnpsInds]

	phe = None
	if (phenoFile is not None):	bed, phe = loadPheno(bed, phenoFile, missingPhenotype)

	if (loadSNPs):
		bed = bed.read()
		if (standardize): bed = bed.standardize()

	return bed, phe

def loadPheno(bed, phenoFile, missingPhenotype='-9', keepDict=False):
	pheno = phenoUtils.loadOnePhen(phenoFile, missing=missingPhenotype, vectorize=True)
	checkIntersection(bed, pheno, 'phenotypes')
	bed, pheno = pstutil.intersect_apply([bed, pheno])
	if (not keepDict): pheno = pheno['vals']
	return bed, pheno

def checkIntersection(bed, fileDict, fileStr, checkSuperSet=False):
	bedSet = set((b[0], b[1]) for b in bed.iid)
	fileSet = set((b[0], b[1]) for b in fileDict['iid'])

	if checkSuperSet:
		if (not fileSet.issuperset(bedSet)): raise Exception(fileStr + " file does not include all individuals in the bfile")

	intersectSet = bedSet.intersection(fileSet)
	if (len(intersectSet) != len (bedSet)):
		print(len(intersectSet), 'individuals appear in both the plink file and the', fileStr, 'file')

def symmetrize(a):
    return a + a.T - np.diag(a.diagonal())

def loadRelatedFile(bed, relFile):
	relatedDict = phenoUtils.loadOnePhen(relFile, vectorize=True)
	checkIntersection(bed, relatedDict, 'relatedness', checkSuperSet=True)
	_, relatedDict = pstutil.intersect_apply([bed, relatedDict])
	related = relatedDict['vals']
	keepArr = (related < 0.5)
	print(np.sum(~keepArr), 'individuals will be removed due to high relatedness')
	return keepArr

def findRelated(bed, cutoff, kinshipFile=None):

	if (kinshipFile is None):
		print('Computing kinship matrix...')
		t0 = time.time()
		XXT = symmetrize(blas.dsyrk(1.0, bed.val, lower=1) / bed.val.shape[1])
		print('Done in %0.2f'%(time.time()-t0), 'seconds')
	else:
		XXT = np.loadtxt(kinshipFile)

	#Find related individuals
	removeSet = set(np.sort(vc.VertexCut().work(XXT, cutoff))) #These are the indexes of the IIDs to remove
	print('Marking', len(removeSet), 'individuals to be removed due to high relatedness')

	#keepArr = np.array([(1 if iid in keepSet else 0) for iid in bed.iid], dtype=bool)
	keepArr = np.ones(bed.iid.shape[0], dtype=bool)
	for i in removeSet: keepArr[i] = False
	return keepArr

def eigenDecompose(XXT, ignore_neig=False):
	t0 = time.time()
	print('Computing eigendecomposition...')
	s,U = la.eigh(XXT)
	if (not ignore_neig and (np.min(s) < -1e-4)): raise Exception('Negative eigenvalues found')
	s[s<0]=0
	ind = np.argsort(s)
	ind = ind[s>1e-12]
	U = U[:, ind]
	s = s[ind]
	print('Done in %0.2f'%(time.time()-t0), 'seconds')
	return s,U

def loadCovars(bed, covarFile):
	covarsDict = phenoUtils.loadPhen(covarFile)
	checkIntersection(bed, covarsDict, 'covariates', checkSuperSet=True)
	_, covarsDict = pstutil.intersect_apply([bed, covarsDict])
	covar = covarsDict['vals']
	return covar

def getSNPCovarsMatrix(bed, resfile, pthresh, mindist):
	snpNameToNumDict = dict([])
	for i,s in enumerate(bed.sid): snpNameToNumDict[s] = i

	f = open(resfile)
	csvReader = csv.reader(f, delimiter="\t")
	next(csvReader)
	significantSNPs = []
	significantSNPNames = []
	lastPval = 0
	featuresPosList = []
	for l in csvReader:
		snpName, pVal = l[0], float(l[4])
		if (pVal < lastPval): raise Exception('P-values are not sorted in descending order: ' + str(pVal) + ">" + str(lastPval))
		lastPval = pVal
		if (pVal > pthresh): break
		if (snpName not in snpNameToNumDict): continue
		significantSNPNames.append(snpName)
		if (mindist == 0):
			significantSNPs.append(snpNameToNumDict[snpName])
			print('Using SNP', snpName, 'with p<%0.2e'%pVal, 'as a fixed effect')
		else:
			posArr = bed.pos[snpNameToNumDict[snpName]]
			chrom, pos = posArr[0], int(posArr[2])
			addSNP = True
			for (c,p) in featuresPosList:
				if (chrom == c and abs(pos-p) < mindist):
					addSNP = False
					break
			if addSNP:
				significantSNPs.append(snpNameToNumDict[snpName])
				featuresPosList.append((chrom, pos))
				print('Using SNP', snpName, '('+str(int(chrom))+':'+str(pos)+') with p<%0.2e'%pVal, 'as a fixed effect')
	f.close()

	snpCovarsMat = bed.val[:, significantSNPs]
	return snpCovarsMat

def getExcludedChromosome(bfile, chrom):
	bed = Bed(bfile, count_A1=True)
	indsToKeep = (bed.pos[:,0] != chrom)
	bed = bed[:, indsToKeep]
	return bed.read().standardize()

def getChromosome(bfile, chrom):
	bed = Bed(bfile, count_A1=True)
	indsToKeep = (bed.pos[:,0] == chrom)
	bed = bed[:, indsToKeep]
	return bed.read().standardize()

def _fixupBedAndPheno(bed, pheno, missingPhenotype='-9'):
	bed = _fixupBed(bed)
	bed, pheno = _fixup_pheno(pheno, bed, missingPhenotype)
	return bed, pheno

def _fixupBed(bed):
	if isinstance(bed, str):
		return Bed(bed, count_A1=True).read().standardize()
	else: return bed

def _fixup_pheno(pheno, bed=None, missingPhenotype='-9'):
	if (isinstance(pheno, str)):
		if (bed is not None):
			bed, pheno = loadPheno(bed, pheno, missingPhenotype, keepDict=True)
			return bed, pheno
		else:
			phenoDict = phenoUtils.loadOnePhen(pheno, missing=missingPhenotype, vectorize=True)
			return phenoDict
	else:
		if (bed is not None): return bed, pheno
		else: return pheno

def linreg(bed, pheno):

	#Extract snps and phenotype
	bed, pheno = _fixupBedAndPheno(bed, pheno)
	if isinstance(pheno, dict):	phe = pheno['vals']
	else: phe = pheno
	if (len(phe.shape)==2):
		if (phe.shape[1]==1): phe=phe[:,0]
		else: raise Exception('More than one phenotype found')

	#Normalize y. We assume X is already normalized.
	y = phe - phe.mean(); y /= y.std()

	#Compute p-values
	Xy = bed.val.T.dot(y) / y.shape[0]
	Xy[Xy>1.0] = 1.0
	Xy[Xy<-1.0] = -1.0
	df = y.shape[0]-2
	TINY = 1.0e-20
	t = Xy * np.sqrt(df / ((1.0-Xy+TINY) * (1.0+Xy+TINY)))
	pValT = stats.t.sf(np.abs(t), df)*2

	#Create pandas data frame
	items = [
		('SNP', bed.sid),
		('Chr', bed.pos[:,0]),
		('GenDist', bed.pos[:,1]),
		('ChrPos', bed.pos[:,2]),
		('PValue', pValT),
	]
	frame = pd.DataFrame.from_items(items)
	frame.sort("PValue", inplace=True)
	frame.index = np.arange(len(frame))
	return frame

def powerPlot(df, causalSNPs, title=''):
	import pylab
	causalSNPs = set(causalSNPs)
	csnpPvals = df[df['SNP'].isin(causalSNPs)]["PValue"]
	pvalPoints = np.logspace(-6, -2, num=1000)
	power = [np.mean(csnpPvals < p ) for p in list(pvalPoints)]
	pylab.plot(-np.log10(pvalPoints), power)
	pylab.xlabel("-log10(Significance Threshold)")
	pylab.ylabel("Power")
	pylab.title(title)

def computeCovar(bed, shrinkMethod, fitIndividuals):
	eigen = dict([])

	if (shrinkMethod in ['lw', 'oas', 'l1', 'cv']):
		import sklearn.covariance as cov
		t0 = time.time()
		print('Estimating shrunk covariance using', shrinkMethod, 'estimator...')

		if (shrinkMethod == 'lw'): covEstimator = cov.LedoitWolf(assume_centered=True, block_size = 5*bed.val.shape[0])
		elif (shrinkMethod == 'oas'): covEstimator = cov.OAS(assume_centered=True)
		elif (shrinkMethod == 'l1'): covEstimator = cov.GraphLassoCV(assume_centered=True, verbose=True)
		elif (shrinkMethod == 'cv'):
			shrunkEstimator = cov.ShrunkCovariance(assume_centered=True)
			param_grid = {'shrinkage': [0.01, 0.1, 0.3, 0.5, 0.7, 0.9, 0.99]}
			covEstimator = sklearn.grid_search.GridSearchCV(shrunkEstimator, param_grid)
		else: raise Exception('unknown covariance regularizer')

		covEstimator.fit(bed.val[fitIndividuals, :].T)
		if (shrinkMethod == 'l1'):
			alpha = covEstimator.alpha_
			print('l1 alpha chosen:', alpha)
			covEstimator2 = cov.GraphLasso(alpha=alpha, assume_centered=True, verbose=True)
		else:
			if (shrinkMethod == 'cv'): shrinkEstimator = clf.best_params_['shrinkage']
			else: shrinkEstimator = covEstimator.shrinkage_
			print('shrinkage estimator:', shrinkEstimator)
			covEstimator2 = cov.ShrunkCovariance(shrinkage=shrinkEstimator, assume_centered=True)
		covEstimator2.fit(bed.val.T)
		XXT = covEstimator2.covariance_ * bed.val.shape[1]
		print('Done in %0.2f'%(time.time()-t0), 'seconds')

	else:
		print('Computing kinship matrix...')
		t0 = time.time()
		XXT = symmetrize(blas.dsyrk(1.0, bed.val, lower=1))
		print('Done in %0.2f'%(time.time()-t0), 'seconds')
		try: shrinkParam = float(shrinkMethod)
		except: shrinkParam = -1
		if (shrinkMethod == 'mylw'):
			XXT_fit = XXT[np.ix_(fitIndividuals, fitIndividuals)]
			sE2R = (np.sum(XXT_fit**2) - np.sum(np.diag(XXT_fit)**2)) / (bed.val.shape[1]**2)
			#temp = (bed.val**2).dot((bed.val.T)**2)
			temp = symmetrize(blas.dsyrk(1.0, bed.val[fitIndividuals, :]**2, lower=1))
			sER2 = (temp.sum() - np.diag(temp).sum()) / bed.val.shape[1]
			shrinkParam = (sER2 - sE2R) / (sE2R * (bed.val.shape[1]-1))
		if (shrinkParam > 0):
			print('shrinkage estimator:', 1-shrinkParam)
			XXT = (1-shrinkParam)*XXT + bed.val.shape[1]*shrinkParam*np.eye(XXT.shape[0])

	return XXT

def standardize(X, method, optionsDict):
	fitIndividuals = np.ones(X.shape[0], dtype=np.bool)
	if (method == 'frq'):
		empMean = X.mean(axis=0) / 2.0
		X[:, empMean>0.5] = 2 - X[:, empMean>0.5]
		print('regularizng SNPs according to frq file...')
		frqFile = (optionsDict['bfilesim']+'.frq' if (optionsDict['frq'] is None) else optionsDict['frq'])
		mafs = np.loadtxt(frqFile, usecols=[1,2]).mean(axis=1)
		snpsMean = 2*mafs
		snpsStd = np.sqrt(2*mafs*(1-mafs))
	elif (method == 'related'):
		if (optionsDict['related'] is None): raise Exception('related file not supplied')
		print('regularizng SNPs according to non-related individuals...')
		relLines = np.loadtxt(optionsDict['related'], usecols=[2])
		keepArr = (relLines != 1)
		print('Excluding', np.sum(~keepArr), 'from the covariance matrix standardization')
		snpsMean = X[keepArr, :].mean(axis=0)
		snpsStd = X[keepArr, :].std(axis=0)
		fitIndividuals = keepArr
	elif (method == 'controls'):
		phe = optionsDict['pheno']
		pheThreshold = phe.mean()
		controls = (phe<pheThreshold)
		print('regularizng SNPs according to controls...')
		snpsMean = X[controls, :].mean(axis=0)
		snpsStd = X[controls, :].std(axis=0)
		fitIndividuals = controls
	elif (method is None):
		snpsMean = X.mean(axis=0)
		snpsStd = X.std(axis=0)
	else:
		raise Exception('unknown SNP standardization option: ' + method)

	X -= snpsMean,
	X /= snpsStd
	return X, fitIndividuals


from __future__ import division  #brings in Python 3.0 mixed type calculation rules
import logging
import json
import requests
import math
import pandas as pd
import os

rest_url_varroapop = os.environ.get('OPENCPU_REST_SERVER')
#rest_url_varroapop = 'http://localhost'

if not os.environ.get('OPENCPU_REST_SERVER'):
    rest_url_varroapop = 'http://172.20.100.18:5656'

class VarroapopFunctions(object):
    """
    Function class for Stir.
    """

    def __init__(self):
        """Class representing the functions for VarroaPop"""
        super(VarroapopFunctions, self).__init__()

    def call_varroapop_api(self):
        logging.info("=========== formatting Varroapop JSON payload")
        input_json = self.format_varroapop_payload()
        logging.info("=========== calling Varroapop windows REST API")
        called_endpoint = (rest_url_varroapop + '/ocpu/apps/quanted/VarroaPopWrapper/R/RunVarroaPop/json')
        logging.info(called_endpoint)
        http_headers = {'Content-Type': 'application/json'}
        logging.info("JSON payload:")
        print(input_json)
        return requests.post(called_endpoint, headers=http_headers, data=input_json, timeout=60)

    def fill_model_out_attr(self, output_json):
        outputs = json.loads(json.loads(output_json)[0])
        self.out_date = self.out_date.append(pd.Series(outputs.get('Date')))
        self.out_colony_size = self.out_colony_size.append(pd.Series(outputs.get('Colony.Size')))
        self.out_adult_drones = self.out_adult_drones.append(pd.Series(outputs.get('Adult.Drones')))
        self.out_adult_workers = self.out_adult_workers.append(pd.Series(outputs.get('Adult.Workers')))
        self.out_foragers = self.out_foragers.append(pd.Series(outputs.get('Foragers')))
        self.out_capped_drone_brood = self.out_capped_drone_brood.append(pd.Series(outputs.get('Capped.Drone.Brood')))
        self.out_capped_worker_brood = self.out_capped_worker_brood.append(pd.Series(outputs.get('Capped.Worker.Brood')))
        self.out_drone_larvae = self.out_drone_larvae.append(pd.Series(outputs.get('Drone.Larvae')))
        self.out_worker_larvae =self.out_worker_larvae.append(pd.Series(outputs.get('Worker.Larvae')))
        self.out_drone_eggs = self.out_drone_eggs.append(pd.Series(outputs.get('Drone.Eggs')))
        self.out_worker_eggs = self.out_worker_eggs.append(pd.Series(outputs.get('Worker.Eggs')))
        self.out_free_mites = self.out_free_mites.append(pd.Series(outputs.get('Free.Mites')))
        self.out_drone_brood_mites =self.out_drone_brood_mites.append(pd.Series(outputs.get('Drone.Brood.Mites')))
        self.out_worker_brood_mites =self.out_worker_brood_mites.append(pd.Series(outputs.get('Worker.Brood.Mites')))
        self.out_drone_mites_per_cell = self.out_drone_mites_per_cell.append(pd.Series(outputs.get('Mites.Drone.Cell')))
        self.out_worker_mites_per_cell = self.out_worker_mites_per_cell.append(pd.Series(outputs.get('Mites.Worker.Cell')))
        self.out_mites_dying = self.out_mites_dying.append(pd.Series(outputs.get('Mites.Dying')))
        self.out_proportion_mites_dying =self.out_proportion_mites_dying.append(pd.Series(outputs.get('Proportion.Mites.Dying')))
        self.out_colony_pollen = self.out_colony_pollen.append(pd.Series(outputs.get('Colony.Pollen..g.')))
        self.out_chemical_conc_pollen =self.out_chemical_conc_pollen.append(pd.Series(outputs.get('Pollen.Pesticide.Concentration')))
        self.out_colony_nectar = self.out_colony_nectar.append(pd.Series(outputs.get('Colony.Nectar')))
        self.out_chemical_conc_nectar =self.out_chemical_conc_nectar.append(pd.Series(outputs.get('Nectar.Pesticide.Concentration')))
        self.out_dead_drone_larvae = self.out_dead_drone_larvae.append(pd.Series(outputs.get('Dead.Drone.Larvae')))
        self.out_dead_worker_larvae =self.out_dead_worker_larvae.append(pd.Series(outputs.get('Dead.Worker.Larvae')))
        self.out_dead_drone_adults = self.out_dead_drone_adults.append(pd.Series(outputs.get('Dead.Drone.Adults')))
        self.out_dead_worker_adults = self.out_dead_worker_adults.append(pd.Series(outputs.get('Dead.Worker.Adults')))
        self.out_dead_foragers = self.out_dead_foragers.append(pd.Series(outputs.get('Dead.Foragers')))
        self.out_queen_strength = self.out_queen_strength.append(pd.Series(outputs.get('Queen.Strength')))
        self.out_average_temp_c = self.out_average_temp_c.append(pd.Series(outputs.get('Average.Temperature..celsius.')))
        self.out_rain_inch = self.out_rain_inch.append(pd.Series(outputs.get('Rain')))

    def fill_summary_stats(self):
        self.out_mean_colony_size = self.out_mean_colony_size.append(pd.Series(self.out_colony_size.mean()))
        self.out_max_colony_size = self.out_max_colony_size.append(pd.Series(self.out_colony_size.max()))
        self.out_min_colony_size = self.out_min_colony_size.append(pd.Series(self.out_colony_size.min()))
        self.out_total_bee_mortality = self.out_total_bee_mortality.append(pd.Series(sum([self.out_dead_drone_adults.sum(),
                                                                                self.out_dead_drone_larvae.sum(),
                                                                                self.out_dead_worker_adults.sum(),
                                                                                self.out_dead_worker_larvae.sum(),
                                                                                self.out_dead_foragers.sum()])))
        self.out_max_chemical_conc_pollen = self.out_max_chemical_conc_pollen.append(pd.Series(self.out_chemical_conc_pollen.max()))
        self.out_max_chemical_conc_nectar = self.out_max_chemical_conc_nectar.append(pd.Series(self.out_chemical_conc_nectar.max()))

    def fill_sessionid(self, sessionid):
        self.out_api_sessionid = self.out_api_sessionid.append(pd.Series(sessionid))

    def format_varroapop_payload(self):
        input_dict = self.pd_obj.to_dict('records')[0]
        weather_loc = input_dict.pop('weather_location')
        print('Weather location: '+ weather_loc )
        input_dict = self.collapse_dates(input_dict)
        input_dict = self.rename_inputs(input_dict)
        input_dict = self.remove_unused_inputs(input_dict)
        data = json.dumps({'parameters':input_dict, 'weather_file':weather_loc})
        return data

    def collapse_dates(self, input_dict):
        sim_start_keys = ['SimStart_month', 'SimStart_day', 'SimStart_year']
        input_dict['SimStart'] = "/".join([str(int(input_dict.get(key))) for key in sim_start_keys])
        sim_end_keys = ['SimEnd_month', 'SimEnd_day', 'SimEnd_year']
        input_dict['SimEnd'] = "/".join([str(int(input_dict.get(key))) for key in sim_end_keys])
        requeen_date_keys = ['RQReQueenDate_month', 'RQReQueenDate_day', 'RQReQueenDate_year']
        input_dict['RQReQueenDate'] = "/".join([str(int(input_dict.get(key))) for key in requeen_date_keys])
        imm_start_keys = ['ImmStart_month', 'ImmStart_day', 'ImmStart_year']
        input_dict['ImmStart'] = "/".join([str(int(input_dict.get(key))) for key in imm_start_keys])
        imm_end_keys = ['ImmEnd_month', 'ImmEnd_day', 'ImmEnd_year']
        input_dict['ImmEnd'] = "/".join([str(int(input_dict.get(key))) for key in imm_end_keys])
        vt_treatment_start_keys = ['VTTreatmentStart_month', 'VTTreatmentStart_day', 'VTTreatmentStart_year']
        input_dict['VTTreatmentStart'] = "/".join([str(int(input_dict.get(key))) for key in vt_treatment_start_keys])
        foliar_app_date_keys = ['FoliarAppDate_month', 'FoliarAppDate_day', 'FoliarAppDate_year']
        input_dict['FoliarAppDate'] = "/".join([str(int(input_dict.get(key))) for key in foliar_app_date_keys])
        foliar_forage_begin_keys = ['FoliarForageBegin_month', 'FoliarForageBegin_day', 'FoliarForageBegin_year']
        input_dict['FoliarForageBegin'] = "/".join([str(int(input_dict.get(key))) for key in foliar_forage_begin_keys])
        foliar_forage_end_keys = ['FoliarForageEnd_month', 'FoliarForageEnd_day', 'FoliarForageEnd_year']
        input_dict['FoliarForageEnd'] = "/".join([str(int(input_dict.get(key))) for key in foliar_forage_end_keys])
        soil_forage_begin_keys = ['SoilForageBegin_month', 'SoilForageBegin_day', 'SoilForageBegin_year']
        input_dict['SoilForageBegin'] = "/".join([str(int(input_dict.get(key))) for key in soil_forage_begin_keys])
        soil_forage_end_keys = ['SoilForageEnd_month', 'SoilForageEnd_day', 'SoilForageEnd_year']
        input_dict['SoilForageEnd'] = "/".join([str(int(input_dict.get(key))) for key in soil_forage_end_keys])
        seed_forage_begin_keys = ['SeedForageBegin_month', 'SeedForageBegin_day', 'SeedForageBegin_year']
        input_dict['SeedForageBegin'] = "/".join([str(int(input_dict.get(key))) for key in seed_forage_begin_keys])
        seed_forage_end_keys = ['SeedForageEnd_month', 'SeedForageEnd_day', 'SeedForageEnd_year']
        input_dict['SeedForageEnd'] = "/".join([str(int(input_dict.get(key))) for key in seed_forage_end_keys])
        sup_pollen_begin_keys = ['SupPollenBegin_month', 'SupPollenBegin_day', 'SupPollenBegin_year']
        input_dict['SupPollenBegin'] = "/".join([str(int(input_dict.get(key))) for key in sup_pollen_begin_keys])
        sup_pollen_end_keys = ['SupPollenEnd_month', 'SupPollenEnd_day', 'SupPollenEnd_year']
        input_dict['SupPollenEnd'] = "/".join([str(int(input_dict.get(key))) for key in sup_pollen_end_keys])
        sup_nectar_begin_keys = ['SupNectarBegin_month', 'SupNectarBegin_day', 'SupNectarBegin_year']
        input_dict['SupNectarBegin'] = "/".join([str(int(input_dict.get(key))) for key in sup_nectar_begin_keys])
        sup_nectar_end_keys = ['SupNectarEnd_month', 'SupNectarEnd_day', 'SupNectarEnd_year']
        input_dict['SupNectarEnd'] = "/".join([str(int(input_dict.get(key))) for key in sup_nectar_end_keys])
        inputs_to_remove = sum([sim_start_keys,sim_end_keys,requeen_date_keys,imm_start_keys,
                                         imm_end_keys,vt_treatment_start_keys,foliar_app_date_keys,
                                         foliar_forage_begin_keys, foliar_forage_end_keys,soil_forage_begin_keys,
                                         soil_forage_end_keys, seed_forage_begin_keys, seed_forage_end_keys,
                                         sup_pollen_begin_keys, sup_pollen_end_keys, sup_nectar_begin_keys, sup_nectar_end_keys], [])
        [input_dict.pop(k, None) for k in inputs_to_remove]
        return input_dict

    def rename_inputs(self, input_dict):
        input_dict['EAppRate'] = input_dict.pop('ar_lb')
        input_dict['AIKOW'] = math.exp(input_dict.pop('l_kow'))
        input_dict['AIKOC'] = input_dict.pop('k_oc')
        return input_dict

    def remove_unused_inputs(self, input_dict):
        keys = list(input_dict.keys())
        to_remove = [i for i in keys if i[0].islower()]
        for k in to_remove:
         input_dict.pop(k, None)
        return input_dict

    def get_input_file(self, api_sessionid):
        file_endpoint =  (rest_url_varroapop + '/ocpu/tmp/'  + api_sessionid + '/files/')
        return requests.get(file_endpoint+'vp_input.txt')

    def get_log_file(self, api_sessionid):
        file_endpoint =  (rest_url_varroapop + '/ocpu/tmp/'  + api_sessionid + '/files/')
        return requests.get(file_endpoint+'vp_log.txt')

    def get_results_file(self, api_sessionid):
        file_endpoint =  (rest_url_varroapop + '/ocpu/tmp/'  + api_sessionid + '/files/')
        return requests.get(file_endpoint+'vp_results.txt')


from pyutil.mongo.engine.strategy import Strategy, strategies, configuration
from pyutil.mongo.engine.symbol import Symbol, Group
from pyutil.performance.drawdown import drawdown
from pyutil.performance.month import monthlytable
from pyutil.performance.return_series import from_nav
from pyutil.portfolio.portfolio import  similar
import pandas.testing as pt
from test.config import *

@pytest.fixture()
def group():
    Group.objects.delete()
    return Group(name="US Equity").save()

@pytest.fixture()
def symbols(group, portfolio):
    Symbol.objects.delete()
    # add the symbols to database
    for symbol in portfolio.assets:
        Symbol(name=symbol, group=group).save()

def test_strategy(symbols, portfolio):
    Strategy.objects.delete()

    s = Strategy(name="mdt", type="mdt", active=True, source="AAA")

    assert s.source == "AAA"
    assert s.type == "mdt"
    assert s.active

    assert s.portfolio is None
    assert s.last_valid_index is None

    # empty dictionary as portfolio hasn't been set
    assert Strategy.portfolios(strategies=[s]) == {}

    s.save()

    frame = Strategy.reference_frame()
    assert frame.index.name == "strategy"

    s.portfolio = portfolio
    pt.assert_frame_equal(s.portfolio.weights, portfolio.weights)
    pt.assert_frame_equal(s.portfolio.prices, portfolio.prices)

    s.save()

    similar(Strategy.portfolios(strategies=[s])["mdt"], portfolio)

    navs = Strategy.navs()
    assert not navs["mdt"].empty

    frame = Strategy.sectors(strategies=[s])
    assert frame.index.name == "Portfolio"
    assert set(frame.keys()) == {"US Equity", "Total"}
    assert frame.loc["mdt"]["US Equity"] == pytest.approx(0.308974, abs=1e-5)

def test_source(portfolio):
    with open(resource("source.py"), "r") as f:
        s = Strategy(name="Peter", source=f.read(), active=True, type="wild")
        # construct the configuration based on the strategy (and it's source code)
        c = configuration(strategy=s)
        # verify the names of the configuration
        assert c.names == portfolio.assets
        # also possible to ask the strategy directly
        assert s.assets == portfolio.assets

def test_last_valid(portfolio):
    s = Strategy(name="Maffay", source="AAA", active=True, type="wild2")
    s.portfolio = portfolio
    assert s.last_valid_index == portfolio.prices.last_valid_index()
    assert similar(s.portfolio, portfolio)

def test_strategies():
    folder = resource(name="strat")
    for name, source in strategies(folder=folder):
        assert name in {"P1", "P2"}

def test_active():
    Strategy.objects.delete()
    Strategy(name="A", active=False).save()
    Strategy(name="B", active=True).save()

    assert len(Strategy.active_strategies()) == 1
    assert len(Strategy.objects) == 2

def test_drawdown(portfolio):
    Strategy.objects.delete()
    s = Strategy(name="Maffay", source="")
    s.portfolio = portfolio
    pt.assert_series_equal(drawdown(portfolio.nav), s.drawdown)

def test_volatility(portfolio):
    Strategy.objects.delete()
    s = Strategy(name="Maffay", source="")
    s.portfolio = portfolio
    pt.assert_series_equal(from_nav(portfolio.nav).ewm_volatility().dropna(), s.ewm_volatility())

def test_monthlytable(portfolio):
    Strategy.objects.delete()
    s = Strategy(name="Maffay", source="")
    s.portfolio = portfolio
    pt.assert_frame_equal(monthlytable(portfolio.nav.pct_change()), s.monthlytable)


""" This is the Bokeh charts testing interface.

"""
#-----------------------------------------------------------------------------
# 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 collections import OrderedDict
import unittest

import numpy as np
from numpy.testing import assert_array_equal
import pandas as pd

from bokeh.charts import Line

from bokeh.charts.builder.tests._utils import create_chart

#-----------------------------------------------------------------------------
# Classes and functions
#-----------------------------------------------------------------------------

class TestLine(unittest.TestCase):
    def test_supported_input(self):
        xyvalues = OrderedDict()
        y_python = xyvalues['python'] = [2, 3, 7, 5, 26]
        y_pypy = xyvalues['pypy'] = [12, 33, 47, 15, 126]
        y_jython = xyvalues['jython'] = [22, 43, 10, 25, 26]

        xyvaluesdf = pd.DataFrame(xyvalues)

        for i, _xy in enumerate([xyvalues, xyvaluesdf]):
            hm = create_chart(Line, _xy)
            builder = hm._builders[0]
            self.assertEqual(sorted(builder._groups), sorted(list(xyvalues.keys())))
            assert_array_equal(builder._data['x'], [0, 1, 2, 3, 4])
            assert_array_equal(builder._data['y_python'], y_python)
            assert_array_equal(builder._data['y_pypy'], y_pypy)
            assert_array_equal(builder._data['y_jython'], y_jython)

        lvalues = [[2, 3, 7, 5, 26], [12, 33, 47, 15, 126], [22, 43, 10, 25, 26]]
        for _xy in [lvalues, np.array(lvalues)]:
            hm = create_chart(Line, _xy)
            builder = hm._builders[0]
            self.assertEqual(builder._groups, ['0', '1', '2'])
            assert_array_equal(builder._data['x'], [0, 1, 2, 3, 4])
            assert_array_equal(builder._data['y_0'], y_python)
            assert_array_equal(builder._data['y_1'], y_pypy)
            assert_array_equal(builder._data['y_2'], y_jython)


# -*- coding: utf-8 -*-
# Comunicacao com a placa agilent U2531A
#
# UFC - Universidade de Federal do Ceará
#
# Responsáveis:
# Felipe Bandeira da Silva
# Francisco Alexander
#

from __future__ import division

import platform

#if platform.system() == 'Windows':
#    import visa
#else:
#    import visa_linux_emulation as visa

try:
    import visa
except:
    # Durante o processo de instalação normal usando o NSIS, o path do windows
    # não estava atualizado com o Python, portanto não era possível, durante a instalação,
    # a execução do pip para instalar o "pyvisa" que requer por natureza, várias
    # dependências que são simplesmene tratadas pelo pip. Portanto para a primeira
    # utilização do programa é necessário a utilização da internet.
    #
    # Para que tudo funcione corretamente e necessario pyvisa 1.4
    #import pip
    #pip.main(['install', 'pyvisa'])

    import subprocess
    print u"aviso: instalando o PyVISA 1.4"
    subprocess.call(['pip', 'install', 'PyVISA==1.4'])
    print u"aviso: instalacao finalizada"
    import visa

import matplotlib.pyplot as plt
from time import sleep, time, asctime, localtime
import numpy as np

###############################################################################
# Constantes para correçao. As mesmas usadas pelo programa feito no LabView
###############################################################################
FATOR_CORRECAO_TENSAO = 100
FATOR_CORRECAO_CORRENTE = 2.71

# 0 - nao mostra as mensagens
# 1 - mostras as mensagens para debug
DEBUG = 0

# um pequeno pulso inicial é visto no inicio da
# aquisição, puro ruido. Para que o sinal seja
# visualizado corretamento foi necessário aumentar
# o número de aquisições. Isso implica em uma
# aquisição mais demorada.

#QUANTIDADE_PONTOS = 50000
QUANTIDADE_PONTOS = 800000

###############################################################################

# testBit() returns a nonzero result, 2**offset, if the bit at 'offset' is one.
def testBit(int_type, offset):
    mask = 1 << offset
    return(int_type & mask)

# setBit() returns an integer with the bit at 'offset' set to 1.
def setBit(int_type, offset):
    mask = 1 << offset
    return(int_type | mask)

# clearBit() returns an integer with the bit at 'offset' cleared.
def clearBit(int_type, offset):
    mask = ~(1 << offset)
    return(int_type & mask)

# toggleBit() returns an integer with the bit at 'offset' inverted, 0 -> 1 and 1 -> 0.
def toggleBit(int_type, offset):
    mask = 1 << offset
    return(int_type ^ mask)

def lerEndian(data):
    """
    Converte um sequencia de dados em valores de 2 bytes
    A sequencia de entrada é dada no formato little-endian
    com entrada do 13 bit para o carry.
    Entrada:
        data = string pura com informacoes do bloco de bytes
    Saída:
        t = tamanho do vetor de bytes
        v = valores em um vetor
    """
    raw = data[10:]

    valores = []
    passo = 0
    for i in raw:
        if passo == 0:
            lsb = i
            passo = 1
        elif passo == 1:
            msb = i
            passo = 0
            num = ((ord(msb)<<8)+(ord(lsb)))>>2
            #print hex(num)
            valores.append(num)

    return [len(valores), valores]

def ler2Endian(data):
    """
    Ler um bloco de bytes composto por duas leitura simultaneas do canal.
    """

    raw = data[10:]
    A = []
    B = []
    passo = 0

    for i in raw:
        if passo == 0:
            lsb = i
            passo = 1
        elif passo == 1:
            msb = i
            passo = 2
            A.append(((ord(msb)<<8)+(ord(lsb)))>>2)
        elif passo == 2:
            lsb = i
            passo = 3
        elif passo == 3:
            msb = i
            passo = 0
            B.append(((ord(msb)<<8)+(ord(lsb)))>>2)

    return [len(A), A, B]

def convBIP(raw, range_ad=10, resolution=14):
    v = []
    for i in raw:
        v.append( (2*i)/(2**resolution)  * range_ad )
    return v

def convUNI(raw, range_ad=10, resolution=14):
    v = []
    for i in raw:
        # se o 13 bit do byte for 1 então o número é "negativo"
        # a conversão unipolar é dada por
        # MAX = 1FFF
        # MAX/2 = 0000
        # 0 = 2000
        if testBit(i, 13) > 0:
            valor = clearBit(i, 13) - (2**14)/2
            v.append( (valor/(2**resolution) + 0.5)*range_ad )
        else:
            v.append( (i/(2**resolution) + 0.5)*range_ad )
    return v

def lerTensaoCorrente(ag):
    """
    Faz a leitura de dois canais de forma simultanea
    Canal 101(corrente) e 102(tensão)
    """

    # reseta a placa a de aquisição
    ag.write("*CLS")
    ag.write("*RST")

    ag.write("ROUT:ENAB 0,(@103, 104)") # desabilita os canais 103 e 104
    ag.write("ROUT:ENAB 1,(@101, 102)") # habilita os canais 101 e 102
    ag.write("ROUT:CHAN:RANG 10,(@101, 102)")   # coloca no mesmo nivel que o programa da National
    ag.write("ROUT:CHAN:POL UNIP,(@101, 102)")  # unipolar
    ag.write("ACQ:SRAT 2000000")    # frequencia de amostragem
    #ag.write("ACQ:POIN 2000000")
    #ag.write("ACQ:POIN 50000")  # número de pontos para aquisição
    ag.write("ACQ:POIN %d" % QUANTIDADE_PONTOS)

    ####################
    # inicia aquisicao #
    ####################
    ag.write("DIG")

    disparaTensao(ag)

    #ag.write("DIG")

    while True:
        ag.write("WAV:COMP?")
        if ag.read() == 'YES':
            break
        sleep(0.2)  # espera um tempo até que amostra fique pronta

    # Uma pequena mudança no capacitor do primeiro 555
    # faz com que o set e reset necessitem de um tempo
    # maior para que ambos acontecam.
    sleep(.2)

    retiraTensao(ag)

    ag.write("WAV:DATA?")
    dados = ag.read()

    t, I, V = ler2Endian(dados)

    V = convUNI(V, 10)
    I = convUNI(I, 10)

    return [dados, V, I]

def lerTensao(ag):
    """
    Ler apenas o canal de tensão da fonte. Canal 102
    Com toda a sequencia de acionamento do set e reset.
    """
    # reset
    ag.write("*CLS")
    ag.write("*RST")

    # inicia a leitura do canal 102 tensao
    ag.write("ROUT:ENAB 0,(@103, 101, 104)")
    ag.write("ROUT:ENAB 1,(@102)")
    ag.write("ROUT:CHAN:RANG 10,(@102)")        # coloca no mesmo nivel que o programa da National
    ag.write("ROUT:CHAN:POL UNIP,(@102)")
    ag.write("ACQ:SRAT 2000000")
    #ag.write("ACQ:POIN 2000000")
    #ag.write("ACQ:POIN 50000")

    # um pequeno pulso inicial é visto no inicio da
    # aquisição, puro ruido. Para que o sinal seja
    # visualizado corretamento foi necessário aumentar
    # o número de aquisições. Isso implica em uma
    # aquisição mais demorada.
    ag.write("ACQ:POIN %d" % (QUANTIDADE_PONTOS))

    # inicia aquisicao
    ag.write("DIG")

    disparaTensao(ag)

    while True:
        ag.write("WAV:COMP?")
        if ag.read() == 'YES':
            break
        sleep(0.5)

    ag.write("WAV:DATA?")
    dados = ag.read()

    sleep(.2)

    retiraTensao(ag)

    #print dados

    t, R = lerEndian(dados)
    V = convUNI(R, 10)
    plt.grid()
    plt.plot(range(0, t), V)
    plt.show()

    return t, V

def lerCorrente(ag):
    """
    Ler apenas o canal de corrente da fonte. Canal 101
    Com toda a sequencia de acionamento do set e reset.
    """
    # reset
    ag.write("*CLS")
    ag.write("*RST")

    # inicia a leitura do canal 101 corrente
    ag.write("ROUT:ENAB 0,(@103, 102, 104)")
    ag.write("ROUT:ENAB 1,(@101)")
    ag.write("ROUT:CHAN:RANG 10,(@101)")
    ag.write("ROUT:CHAN:POL UNIP,(@101)")
    ag.write("ACQ:SRAT 2000000")
    ag.write("ACQ:POIN 2000000")

    # inicia aquisicao
    ag.write("DIG")

    disparaTensao(ag)

    while True:
        ag.write("WAV:COMP?")
        if ag.read() == 'YES':
            break
        sleep(0.5)

    ag.write("WAV:DATA?")
    dados = ag.read()

    sleep(.2)
    retiraTensao(ag)

    #print dados

    t, R = lerEndian(dados)
    V = convUNI(R, 10)
    plt.grid()
    plt.plot(range(0, t), V)
    plt.show()

    return t, V

def lerCanal103(ag):
    """
    Este canal foi usado para os testes iniciais da conversão
    do análogico digital. Não sendo mais necessário.
    As funçoes para leitura de tensão e corrente são identicas
    a esta funçao. Mudando apenas o canal.
    """
    # reset
    ag.write("*CLS")
    ag.write("*RST")

    # inicia a leitura do canal 103
    ag.write("ROUT:ENAB 0,(@101, 102, 104)")
    ag.write("ROUT:ENAB 1,(@103)")
    ag.write("ROUT:CHAN:RANG 10,(@103)")
    #ag.write("ROUT:CHAN:POL BIP,(@103)")
    ag.write("ROUT:CHAN:POL UNIP,(@103)")
    ag.write("ACQ:SRAT 2000000")
    ag.write("ACQ:POIN 2000000")

    # inicia aquisicao
    ag.write("DIG")
    # espera o fim

    disparaTensao(ag)

    while True:
        ag.write("WAV:COMP?")
        if ag.read() == 'YES':
            break
        sleep(0.1)

    ag.write("WAV:DATA?")
    dados = ag.read()

    sleep(.2)
    retiraTensao(ag)

    #print dados

    t, R = lerEndian(dados)
    V = convUNI(R)
    plt.grid()
    plt.plot(range(0, t), V)

    return t, V

def disparaTensao(ag):
    """
    Envia um pulso de alta tensão para o sistema de aterramento.
    Acionando para isto o primeiro 555.
    Os pulso não deve ser enviando em um curto intervalo de tempo
    já que a fonte não foi projetada para tal situaçao.
    Portanto deve-se tormar cuidado no acionamento sequencia.
    SET - Pino 68 na placa U2901-60602
    RESET - Pino 34 na placa U2901-60602
    """
    ag.write("CONF:DIG:DIR OUTP,(@501)")
    ag.write("SOUR:DIG:DATA 1,(@501)")

    return 0

def retiraTensao(ag):
    """
    Reseta a fonte. Habilitando a mesma para um novo envio
    de um pulso de alta tensão.
    """
    ag.write("CONF:DIG:DIR OUTP,(@501)")
    ag.write("SOUR:DIG:DATA 0,(@501)")  # desabilita o set
    sleep(0.1)                          # espera um tempo para resetar
    ag.write("SOUR:DIG:DATA 2,(@501)")  # reseta a fonte
    sleep(0.1)                          # espera um tempo para entrar em repouso
    ag.write("SOUR:DIG:DATA 0,(@501)")  # entra em repouso

    return 0

def pltTensaoCorrente(V, I):
    t1 = np.arange(0, len(V))
    plt.figure(1)
    plt.title("Leitura do U2531A")
    plt.subplot(211)
    plt.plot(t1, V)
    plt.subplot(212)
    plt.plot(t1, I)
    plt.show()

def aplicaCorrecoes(V, I):
    V = np.array(V)
    V = FATOR_CORRECAO_TENSAO * V
    I = np.array(I)
    I = FATOR_CORRECAO_CORRENTE * I

    return [V, I]

def sequenciaAquisicoes(ag, quantidade, local="C:\\Temp", rotulo = '0'):
    """
    Faz um aquisiçao sequencial dos canais de tensão e corrente.
    ag = objeto usada para o controle da placa
    """

    print "Iniciando aquisicao sequencial"
    print "Equipamento = ", ag
    print "quantidade = ", quantidade
    print "Tempo de inicio = ", asctime()

    tempoInicio = time()
    contagem = quantidade

    plt.figure(1)

    while quantidade > 0:
        print "Atual = ", quantidade
        tempoIndividual = time()

        # inicia aquisição
        raw, V, I = lerTensaoCorrente(ag)
        V, I = aplicaCorrecoes(V, I)

        # não é uma boa ideia plotar desta forma
        #pltTensaoCorrente(V, I)

        plt.subplot(211)
        plt.plot(np.arange(0, len(V)), V)

        plt.subplot(212)
        plt.plot(np.arange(0, len(I)), I)

        salvaTensaoTXT(local, rotulo, contagem-quantidade+1, V)
        salvaCorrenteTXT(local, rotulo, contagem-quantidade+1, I)

        print "Individual = ", time()-tempoIndividual
        quantidade -=1

    total = time()-tempoInicio
    print 'Completo em [seg]: ', total
    plt.show()
    return 0

def salvaTensaoTXT(local, rotulo, posicao, V):
    """
    Salva o vetor tensão em um arquivo com nome formatado para isso
    """
    nomeCompleto = local+"\\"+rotulo+"V"+str(posicao)+".txt"
    return salvaTXT(nomeCompleto, V)

def salvaCorrenteTXT(local, rotulo, posicao, I):
    """
    Salva o vetor corrente em um arquivo com nome formatado para isso
    """
    nomeCompleto = local+"\\"+rotulo+"I"+str(posicao)+".txt"
    return salvaTXT(nomeCompleto, I)

def salvaTXT(caminhoCompleto, vetor):
    """
    Salva em um arquivo txt os valores de um vetor
    onde a primeira coluna informa o indice e a segunda
    coluna informa o valor para o indice.
    """
    try:
        arquivo = open(caminhoCompleto, 'w')
    except:
        print 'erro: nao foi possivel escrever no arquivo'
        print '    : ', caminhoCompleto
        return -1

    #for i in range(len(vetor)):
    #    string = "%d %f\n" % (i, float(vetor[i]))
    #    arquivo.write(string)

    for i in vetor:
        arquivo.write(i)

    arquivo.close()

    # escrita finalizada com sucesso
    return 0

def buscaAgilent():
    """
    Busca o equipamento conectado a porta usb do computador
    Retornando o objeto a ser usado pelas funções de controle
    da placa de aquisiçao da agilent.
    """

    listaInstrumentos = visa.get_instruments_list() # adquiri a lista de equipamentos conectados ao computador
    listaAgilent = listaInstrumentos[0]     # pega o primeiro equipamento

    print 'Lista de instrumentos:'
    print listaAgilent      # espera-se que o equipamento seja da agilent

    ag = visa.instrument(listaAgilent)  # cria um objeto a ser manipulado e passado para as outras funções

    identificacao = ag.ask("*IDN?")
    print identificacao

    return ag

###############################################################################
# MAIN                                                                        #
###############################################################################
if __name__ == '__main__':
    print 'Agilente U3125A'

    ag = buscaAgilent()

    ##############################
    # leitura de apenas um canal #
    ##############################
    #lerCanal103(ag)
    #lerTensao(ag)
    #lerCorrente(ag)

    ##########################
    # leitura de dois canais #
    ##########################
    raw, V, I = lerTensaoCorrente(ag)
    V, I = aplicaCorrecoes(V, I)
    pltTensaoCorrente(V, I)

    #########################
    # Aquisiçoes sequencial #
    #########################
    # 60 aquisições
    # local onde é salvo "C:\Temp"
    #sequenciaAquisicoes(ag, 10)


"""
============================
Underfitting vs. Overfitting
============================

This example demonstrates the problems of underfitting and overfitting and
how we can use linear regression with polynomial features to approximate
nonlinear functions. The plot shows the function that we want to approximate,
which is a part of the cosine function. In addition, the samples from the
real function and the approximations of different models are displayed. The
models have polynomial features of different degrees. We can see that a
linear function (polynomial with degree 1) is not sufficient to fit the
training samples. This is called **underfitting**. A polynomial of degree 4
approximates the true function almost perfectly. However, for higher degrees
the model will **overfit** the training data, i.e. it learns the noise of the
training data.
We evaluate quantitatively **overfitting** / **underfitting** by using
cross-validation. We calculate the mean squared error (MSE) on the validation
set, the higher, the less likely the model generalizes correctly from the
training data.
"""

print(__doc__)

import numpy as np
import matplotlib.pyplot as plt
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import cross_val_score

np.random.seed(0)

n_samples = 30
degrees = [1, 4, 15]

true_fun = lambda X: np.cos(1.5 * np.pi * X)
X = np.sort(np.random.rand(n_samples))
y = true_fun(X) + np.random.randn(n_samples) * 0.1

plt.figure(figsize=(14, 5))
for i in range(len(degrees)):
    ax = plt.subplot(1, len(degrees), i + 1)
    plt.setp(ax, xticks=(), yticks=())

    polynomial_features = PolynomialFeatures(degree=degrees[i],
                                             include_bias=False)
    linear_regression = LinearRegression()
    pipeline = Pipeline([("polynomial_features", polynomial_features),
                         ("linear_regression", linear_regression)])
    pipeline.fit(X[:, np.newaxis], y)

    # Evaluate the models using crossvalidation
    scores = cross_val_score(pipeline, X[:, np.newaxis], y,
                             scoring="mean_squared_error", cv=10)

    X_test = np.linspace(0, 1, 100)
    plt.plot(X_test, pipeline.predict(X_test[:, np.newaxis]), label="Model")
    plt.plot(X_test, true_fun(X_test), label="True function")
    plt.scatter(X, y, label="Samples")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.xlim((0, 1))
    plt.ylim((-2, 2))
    plt.legend(loc="best")
    plt.title("Degree {}\nMSE = {:.2e}(+/- {:.2e})".format(
        degrees[i], -scores.mean(), scores.std()))
plt.show()


# Script to draw world countries as patches.

import numpy as np
import matplotlib.pyplot as  plt
from matplotlib.path import Path
import matplotlib.patches as patches
from osgeo import ogr

def order_coords(coords, clockwise):
    """Orders coordinates."""
    total = 0
    x1, y1 = coords[0]
    for x, y in coords[1:]:
        total += (x - x1) * (y + y1)
        x1, y1 = x, y
    x, y = coords[0]
    total += (x - x1) * (y + y1)
    is_clockwise = total > 0
    if clockwise != is_clockwise:
        coords.reverse()
    return coords

def make_codes(n):
    """Makes a list of path codes."""
    codes = [Path.LINETO] * n
    codes[0] = Path.MOVETO
    return codes

def plot_polygon_patch(poly, color):
    """Plots a polygon as a patch."""
    # Outer clockwise path.
    coords = poly.GetGeometryRef(0).GetPoints()
    coords = order_coords(coords, True)
    codes = make_codes(len(coords))
    for i in range(1, poly.GetGeometryCount()):

        # Inner counter-clockwise paths.
        coords2 = poly.GetGeometryRef(i).GetPoints()
        coords2 = order_coords(coords2, False)
        codes2 = make_codes(len(coords2))

        # Concatenate the paths.
        coords = np.concatenate((coords, coords2))
        codes = np.concatenate((codes, codes2))

    # Add the patch to the plot
    path = Path(coords, codes)
    patch = patches.PathPatch(path, facecolor=color)
    plt.axes().add_patch(patch)

# Loop through all of the features in the countries layer and create
# patches for the polygons.
ds = ogr.Open(r'D:\osgeopy-data\global\ne_110m_admin_0_countries.shp')
lyr = ds.GetLayer(0)
for row in lyr:
    geom = row.geometry()
    if geom.GetGeometryType() == ogr.wkbPolygon:
        plot_polygon_patch(geom, 'yellow')
    elif geom.GetGeometryType() == ogr.wkbMultiPolygon:
        for i in range(geom.GetGeometryCount()):
            plot_polygon_patch(geom.GetGeometryRef(i), 'yellow')
plt.axis('equal')
plt.show()


# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for histogram_ops."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from tensorflow.contrib.metrics.python.ops import histogram_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test

class Strict1dCumsumTest(test.TestCase):
  """Test this private function."""

  def test_empty_tensor_returns_empty(self):
    with self.test_session():
      tensor = constant_op.constant([])
      result = histogram_ops._strict_1d_cumsum(tensor, 0)
      expected = constant_op.constant([])
      np.testing.assert_array_equal(expected.eval(), result.eval())

  def test_length_1_tensor_works(self):
    with self.test_session():
      tensor = constant_op.constant([3], dtype=dtypes.float32)
      result = histogram_ops._strict_1d_cumsum(tensor, 1)
      expected = constant_op.constant([3], dtype=dtypes.float32)
      np.testing.assert_array_equal(expected.eval(), result.eval())

  def test_length_3_tensor_works(self):
    with self.test_session():
      tensor = constant_op.constant([1, 2, 3], dtype=dtypes.float32)
      result = histogram_ops._strict_1d_cumsum(tensor, 3)
      expected = constant_op.constant([1, 3, 6], dtype=dtypes.float32)
      np.testing.assert_array_equal(expected.eval(), result.eval())

class AUCUsingHistogramTest(test.TestCase):

  def setUp(self):
    self.rng = np.random.RandomState(0)

  def test_empty_labels_and_scores_gives_nan_auc(self):
    with self.test_session():
      labels = constant_op.constant([], shape=[0], dtype=dtypes.bool)
      scores = constant_op.constant([], shape=[0], dtype=dtypes.float32)
      score_range = [0, 1.]
      auc, update_op = histogram_ops.auc_using_histogram(labels, scores,
                                                         score_range)
      variables.local_variables_initializer().run()
      update_op.run()
      self.assertTrue(np.isnan(auc.eval()))

  def test_perfect_scores_gives_auc_1(self):
    self._check_auc(
        nbins=100,
        desired_auc=1.0,
        score_range=[0, 1.],
        num_records=50,
        frac_true=0.5,
        atol=0.05,
        num_updates=1)

  def test_terrible_scores_gives_auc_0(self):
    self._check_auc(
        nbins=100,
        desired_auc=0.0,
        score_range=[0, 1.],
        num_records=50,
        frac_true=0.5,
        atol=0.05,
        num_updates=1)

  def test_many_common_conditions(self):
    for nbins in [50]:
      for desired_auc in [0.3, 0.5, 0.8]:
        for score_range in [[-1, 1], [-10, 0]]:
          for frac_true in [0.3, 0.8]:
            # Tests pass with atol = 0.03.  Moved up to 0.05 to avoid flakes.
            self._check_auc(
                nbins=nbins,
                desired_auc=desired_auc,
                score_range=score_range,
                num_records=100,
                frac_true=frac_true,
                atol=0.05,
                num_updates=50)

  def test_large_class_imbalance_still_ok(self):
    # With probability frac_true ** num_records, each batch contains only True
    # records.  In this case, ~ 95%.
    # Tests pass with atol = 0.02.  Increased to 0.05 to avoid flakes.
    self._check_auc(
        nbins=100,
        desired_auc=0.8,
        score_range=[-1, 1.],
        num_records=10,
        frac_true=0.995,
        atol=0.05,
        num_updates=1000)

  def test_super_accuracy_with_many_bins_and_records(self):
    # Test passes with atol = 0.0005.  Increased atol to avoid flakes.
    self._check_auc(
        nbins=1000,
        desired_auc=0.75,
        score_range=[0, 1.],
        num_records=1000,
        frac_true=0.5,
        atol=0.005,
        num_updates=100)

  def _check_auc(self,
                 nbins=100,
                 desired_auc=0.75,
                 score_range=None,
                 num_records=50,
                 frac_true=0.5,
                 atol=0.05,
                 num_updates=10):
    """Check auc accuracy against synthetic data.

    Args:
      nbins:  nbins arg from contrib.metrics.auc_using_histogram.
      desired_auc:  Number in [0, 1].  The desired auc for synthetic data.
      score_range:  2-tuple, (low, high), giving the range of the resultant
        scores.  Defaults to [0, 1.].
      num_records:  Positive integer.  The number of records to return.
      frac_true:  Number in (0, 1).  Expected fraction of resultant labels that
        will be True.  This is just in expectation...more or less may actually
        be True.
      atol:  Absolute tolerance for final AUC estimate.
      num_updates:  Update internal histograms this many times, each with a new
        batch of synthetic data, before computing final AUC.

    Raises:
      AssertionError: If resultant AUC is not within atol of theoretical AUC
        from synthetic data.
    """
    score_range = [0, 1.] or score_range
    with self.test_session():
      labels = array_ops.placeholder(dtypes.bool, shape=[num_records])
      scores = array_ops.placeholder(dtypes.float32, shape=[num_records])
      auc, update_op = histogram_ops.auc_using_histogram(
          labels, scores, score_range, nbins=nbins)
      variables.local_variables_initializer().run()
      # Updates, then extract auc.
      for _ in range(num_updates):
        labels_a, scores_a = synthetic_data(desired_auc, score_range,
                                            num_records, self.rng, frac_true)
        update_op.run(feed_dict={labels: labels_a, scores: scores_a})
      labels_a, scores_a = synthetic_data(desired_auc, score_range, num_records,
                                          self.rng, frac_true)
      # Fetch current auc, and verify that fetching again doesn't change it.
      auc_eval = auc.eval()
      self.assertAlmostEqual(auc_eval, auc.eval(), places=5)

    msg = ('nbins: %s, desired_auc: %s, score_range: %s, '
           'num_records: %s, frac_true: %s, num_updates: %s') % (nbins,
                                                                 desired_auc,
                                                                 score_range,
                                                                 num_records,
                                                                 frac_true,
                                                                 num_updates)
    np.testing.assert_allclose(desired_auc, auc_eval, atol=atol, err_msg=msg)

def synthetic_data(desired_auc, score_range, num_records, rng, frac_true):
  """Create synthetic boolean_labels and scores with adjustable auc.

  Args:
    desired_auc:  Number in [0, 1], the theoretical AUC of resultant data.
    score_range:  2-tuple, (low, high), giving the range of the resultant scores
    num_records:  Positive integer.  The number of records to return.
    rng:  Initialized np.random.RandomState random number generator
    frac_true:  Number in (0, 1).  Expected fraction of resultant labels that
      will be True.  This is just in expectation...more or less may actually be
      True.

  Returns:
    boolean_labels:  np.array, dtype=bool.
    scores:  np.array, dtype=np.float32
  """
  # We prove here why the method (below) for computing AUC works.  Of course we
  # also checked this against sklearn.metrics.roc_auc_curve.
  #
  # First do this for score_range = [0, 1], then rescale.
  # WLOG assume AUC >= 0.5, otherwise we will solve for AUC >= 0.5 then swap
  # the labels.
  # So for AUC in [0, 1] we create False and True labels
  # and corresponding scores drawn from:
  # F ~ U[0, 1],  T ~ U[x, 1]
  # We have,
  # AUC
  #  = P[T > F]
  #  = P[T > F | F < x] P[F < x] + P[T > F | F > x] P[F > x]
  #  = (1 * x) + (0.5 * (1 - x)).
  # Inverting, we have:
  # x = 2 * AUC - 1, when AUC >= 0.5.
  assert 0 <= desired_auc <= 1
  assert 0 < frac_true < 1

  if desired_auc < 0.5:
    flip_labels = True
    desired_auc = 1 - desired_auc
    frac_true = 1 - frac_true
  else:
    flip_labels = False
  x = 2 * desired_auc - 1

  labels = rng.binomial(1, frac_true, size=num_records).astype(bool)
  num_true = labels.sum()
  num_false = num_records - labels.sum()

  # Draw F ~ U[0, 1], and T ~ U[x, 1]
  false_scores = rng.rand(num_false)
  true_scores = x + rng.rand(num_true) * (1 - x)

  # Reshape [0, 1] to score_range.
  def reshape(scores):
    return score_range[0] + scores * (score_range[1] - score_range[0])

  false_scores = reshape(false_scores)
  true_scores = reshape(true_scores)

  # Place into one array corresponding with the labels.
  scores = np.nan * np.ones(num_records, dtype=np.float32)
  scores[labels] = true_scores
  scores[~labels] = false_scores

  if flip_labels:
    labels = ~labels

  return labels, scores

if __name__ == '__main__':
  test.main()