flag_noloop

#!/usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt
from netCDF4 import Dataset
import os as os
#import flagy as flagy

def listevar(fichier):
	datago = Dataset(fichier, 'r')
	x = datago.variables.keys()
	print x

def extraction(fichier):
	datago = Dataset(fichier, 'r')
	lg = datago.variables['longitude'][:]
	lt = datago.variables['latitude'][:]
	ti = datago.variables['time'][:]
	atb = datago.variables['ATB'][:,:]
	SR = datago.variables['instant_SR'][:,:]
	sol_alt = datago.variables['SE'][:]
	alt = datago.variables['alt_bound'][:,:]
	u, v = np.shape(alt)
	alt_bas = alt[0,:]
	alt_haut = alt[1,:]
	alt = np.zeros((v,1))
	for i in range(0,v):
		alt[i] = alt_bas[i]
	#alt[v] = alt_haut[v-1]
	alt = alt[:,0]
	return lg, lt, ti, atb,alt, sol_alt, SR


def affiche(n1,n2,lg,lt,ti,atb,alt,sol_alt):
	ti = ti[n1:n2]
	sol_alt = sol_alt[n1:n2]
	u,v = np.shape(atb[n1:n2,:])
	trait = ti[int (((n2-n1)/2))] 
	plt.plot ([trait,trait] , [alt[0],alt[-1]],'w--',linewidth=2)
	plt.plot(ti, sol_alt,'w',linewidth=2)
	plt.pcolormesh(ti,alt,atb[n1:n2,:].T)
	plt.ylim(-1, 16)
	plt.xlim(ti[0], ti[-1])
	plt.clim(0, 0.03)
	plt.xlabel('Time')
	plt.ylabel('Altitude [km]')
	plt.colorbar()	  
	plt.show() 

def afficheSR(lg,lt,ti,atb,alt,sol_alt):
	n1 = 0
	n2,x = np.shape(atb)
	lt = lt[n1:n2]
	sol_alt = sol_alt[n1:n2]
	u,v = np.shape(atb[n1:n2,:])	
 	sol_alt = sol_alt[n1:n2]
	fig = plt.figure(figsize=(15, 10), dpi=300)
	plt.plot(lt, sol_alt,'w',linewidth=2)
	plt.pcolormesh(lt,alt,atb[n1:n2,:].T)
	plt.ylim(-1, 16)
	plt.clim(1e-5, 1e-2)
	plt.xlabel('Latitude')
	plt.ylabel('Altitude [km]')
	plt.colorbar()	  
	plt.show() 

def lonlat(n1, n2,lg,lt,ti,atb):
	from mpl_toolkits.basemap import Basemap
	m = Basemap()
	m.plot(lg, lt)
	lg = lg[n1:n2]
	lt = lt[n1:n2]
	m.plot(lg, lt, 'r')
	m.drawcoastlines()
	m.fillcontinents()
	plt.show()


def profil(n,lg,lt,ti,atb,alt):
	plt.semilogx(atb[n,:], alt)
	#plt.xlim(1e-7, 1e-1)
	plt.ylim(-20,20)
	plt.xlabel('ATB [km-1 sr-1]')
	plt.ylabel('Altitude [km]')
	plt.title('Profil %d' % (n))
	plt.show()

def profil_SR(n,lg,lt,ti,atb,alt,fichier):
	plt.plot(atb[n,:], alt)
	plt.xlim(-0.1, 30)
	plt.ylim(-20,20)
	plt.xlabel('ATB [km-1 sr-1]')
	plt.ylabel('Altitude [km]')
	plt.title('Profil %d' % (n))
	plt.show()
	path = fichier
	path += '.png'
	#plt.savefig(path)
	#plt.close()
	

def coord_pole(lg,lt,ti,atb, sol_alt):
	idx = ( lt < -60 ) & ( lt >= -90 )
	lt6090 = lt[idx]
	lg6090 = lg[idx]
	ti6090 = ti[idx]
	sol_alt = sol_alt[idx]
	atb = atb [idx,:]
	n = len(lt6090)
	return n, lt6090, lg6090, ti6090,atb, sol_alt

def alt_max(SR, alt, n):
	u, v = np.shape(SR)
	j = - 1
	i = 0
	while SR[n,i] < 0.01 and i < v and SR[n,i] > -10 :
		if (i >2) and (SR[n,i-1] > -1 and SR[n,i-2]> -10) and (SR[n,i-1]<0.01  and SR[n,i-2]<0.01) :
			j = i
		elif i <=2 :
			j = i
		i = i + 1
	return j

def alt_max_orbite(alt, sol_alt, SR,lt):
	p,v = np.shape(SR)
	tab1 = np.zeros((p))
	tab2 = np.zeros((p))
	for o in range (0,p):
		j = alt_max(SR,alt,o)
		if j != -1 :
			tab1[o] = alt[j]
			#if tab1[o] > 5.: print o
		else : tab1[o] = -10
		tab2[o] = o

	tab1 = np.ma.masked_where(tab1 <= -1, tab1)
	fig = plt.figure(figsize=(15, 10), dpi=300)
	plt.plot(lt,tab1)
	plt.plot(lt, sol_alt,'g',linewidth=2)
	plt.xlim(np.min(lt), np.max(lt))
	plt.ylabel('Altitude')
	plt.xlabel('indice')
	plt.show()

def alt_max_orbite_flag(alt, SR):
	u,v= np.shape(SR)
	tab1 = np.zeros((u,v))
	for o in range (0,u):
		j = alt_max(SR,alt,o)
		if j != -1 :tab1[:][o] = alt[j]
		else : tab1[:][o] = -10
	return tab1

def alt_max_profil(alt,SR,n,fichier):
	j = alt_max(SR,alt,n)
	s = np.shape(alt)
	if j != -1 :alt_ext= alt[j]
	else : alt_ext= -10
	print alt_ext
	profil_SR(n,lg,lt,ti,SR,alt,fichier)
	
def affiche_flag(lg,lt,atb,alt,sol_alt):
	n1 = 0
	n2, x = np.shape(atb)
	lt = lt[n1:n2]
	sol_alt = sol_alt[n1:n2]
	u,v = np.shape(atb[n1:n2,:])
	fig = plt.figure(figsize=(9, 6), dpi=100)
	plt.plot(lt, sol_alt,'w',linewidth=2)
	plt.pcolormesh(lt,alt,atb.T)
	plt.ylim(-1, 20)
	plt.xlim(lt[0], lt[-1])
	plt.clim(0, 4)
	plt.xlabel('Latitude')
	plt.ylabel('Altitude')
	plt.colorbar()	  
	plt.show() 

def flag_orbite(alt,SR):
	u,v = np.shape(SR)
	flag = np.zeros((u,v))
	alt_array = np.zeros((u,v))
	alt_array[:][:] = alt
	alt_max = alt_max_orbite_flag(alt,SR)
	#print np.shape(alt_array), np.shape(alt_max)	
	idx = (SR > 5)
	print 'oh',np.shape(idx)
	flag[idx] = 1 # nuage
	idx4 = (alt_array < alt_max)&(SR < 0.01) & (SR> -10)#vrai eteint
	print 'he',np.shape(idx4)
	flag[idx4] = 4
	idx3 = (alt_array == alt_max)&(SR < 0.01) & (SR > -10)#Nuage opaque
	flag[idx3] = 3
	idx2 = ((alt_array > alt_max)  & (alt_max!= -10)&(SR < 0.01) & (SR> -10)) | (alt_max == -10&(SR < 0.01) & (SR> -10))
	flag[idx2] = 2
	return flag

	
def lecture_dossier(repertoire):
	u = np.shape(os.listdir(repertoire))
	
	return u, os.listdir(repertoire)

def occur(flag,alt):
	clouds = (flag == 1).sum(axis=0)
	faux_eteint = (flag == 2).sum(axis=0) 
	cloud_opa = (flag == 3).sum(axis=0)
	vrai_eteint = (flag == 4).sum(axis=0)
	return clouds, faux_eteint, cloud_opa, vrai_eteint

def freq_occur(flag,alt):
	u,v = np.shape(flag)
	clouds, faux_eteint, cloud_opa, vrai_eteint = occur(flag,alt)
	clouds, faux_eteint, cloud_opa, vrai_eteint = (1.)*clouds / u, (1.)*faux_eteint / u, (1.)*cloud_opa / u, (1.)*vrai_eteint / u
	fig = plt.figure(figsize=(9, 15), dpi=100)
	p1, = plt.plot(vrai_eteint,alt)
	sum1 = vrai_eteint + cloud_opa
	p2, = plt.plot(sum1,alt)
	sum2 = vrai_eteint + cloud_opa + clouds
	p3, = plt.plot(sum2,alt)
	sum3 = vrai_eteint + cloud_opa + clouds + faux_eteint
	p4, = plt.plot(sum3,alt)
	#plt.xlim(0.0001, 0.1)
	plt.legend([p1,p2,p3,p4], ["Nuages","Faux eteint","Nuages opaques","Vrai eteint"], loc=1)
	plt.xlabel("Frequence d'occurence")
	plt.ylabel("Altitude en  km")
	plt.show()
	
def occur_in_region(flag,alt,lt,latrange):
	idx = (lt < latrange[1]) & (lt > latrange[0])
	flag = flag[idx]
	clouds = (flag == 1).sum(axis=0)
	faux_eteint = (flag == 2).sum(axis=0) 
	cloud_opa = (flag == 3).sum(axis=0)
	vrai_eteint = (flag == 4).sum(axis=0)
	return clouds, faux_eteint, cloud_opa, vrai_eteint

def freq_occur_region(dossier):
	nbr_fichier,fichier = lecture_dossier(dossier)	

	dic_opa_cloud = dict()
	dic_cloud = dict()
	dic_vrai = dict()
	dic_faux = dict()
	#nbr_fichier=30
	nbr_fichier=np.max(nbr_fichier)
	latbands = {'tropique' :[-20,20],'pole_nord' :[60,90],'pole_sud' : [-90,-60],'mid_sud' :[-50,-30],'mid_nord' :[30, 50 ]}
	for i in range(0,nbr_fichier):	
		fichier3 = dossier + fichier[i]
		lg, lt, ti, atb,alt, sol_alt, SR = extraction(fichier3)
		flag = flag_orbite(alt,SR)
		u,v = np.shape(flag)
		for region, tab_reg  in latbands.items():
			print region, tab_reg
			#flag = flag[tab_reg,:]
			clouds, faux_eteint, cloud_opa, vrai_eteint = occur_in_region(flag, alt, lt, latbands[region])
			if region in dic_opa_cloud:
				dic_opa_cloud[region] += 1.* cloud_opa/ (u * nbr_fichier)
				dic_cloud[region] += 1.* clouds/ (u * nbr_fichier)
				dic_vrai[region] += 1.* vrai_eteint/ (u * nbr_fichier)
				dic_faux[region] += 1.* faux_eteint/ (u * nbr_fichier)
			else:
				dic_opa_cloud[region] = 1.* cloud_opa/ (u * nbr_fichier)
				dic_cloud[region] = 1.* clouds/ (u * nbr_fichier)
				dic_vrai[region] = 1.* vrai_eteint/ (u * nbr_fichier)
				dic_faux[region] = 1.* faux_eteint/ (u * nbr_fichier)
	for	region in latbands.keys():
		
		affiche_freq(alt,dic_opa_cloud[region],dic_cloud[region],dic_vrai[region],dic_faux[region],region)

def affiche_freq(alt,op,cl,vr,fa,region):
	fig = plt.figure(figsize=(9, 15), dpi=100)
	p1, = plt.plot(vr,alt)
	sum1 = vr + op
	p2, = plt.plot(sum1,alt)
	sum2 = sum1 + cl
	p3, = plt.plot(sum2,alt)
	sum3 = sum2 + fa
	p4, = plt.plot(sum3,alt)
	plt.legend([p3,p4,p2,p1], ["Nuages","Faux eteint","Nuages opaques","Vrai eteint"], loc=1)
	plt.title(region)
	plt.xlabel("Frequence d'occurence")
	plt.ylabel("Altitude en  km")
	path = region
	path += '.png'
	plt.savefig(path)
	plt.close()

if __name__=='__main__':

	fichier3="/bdd/CFMIP/CFMIP_OBS_LOCAL/GOCCP/instant_SR_CR_DR/grid_L40/2013/201308/night/instant_SR_CR_DR_2013-08-24T23-19-30ZN_night_CFMIP2_2.70.nc"#38418 mid
	#fichier3="/bdd/CFMIP/CFMIP_OBS_LOCAL/GOCCP/instant_SR_CR_DR/grid_L40/2013/201308/night/instant_SR_CR_DR_2013-08-23T19-18-36ZN_night_CFMIP2_2.70.nc" #20632 trop
	listevar(fichier3)
	#n1 , n2 = 0  , 20632
	lg, lt, ti, atb,alt, sol_alt, SR = extraction(fichier3)
 	#alt_max(SR, alt, n1,2*n2)
	#profil_SR(n1,2*n2,lg,lt,ti,SR,alt)
	#affiche(n1,n2,lg, lt, ti, SR,alt,sol_alt)
	'''
	n1 = 0
	n2, lt, lg, ti, atb, sol_alt = coord_pole(lg, lt, ti, atb,sol_alt)
	alt_max(SR, alt, n1,n2)
	profil(n1,n2,lg,lt,ti,atb,alt)	
	affiche(n1,n2,lg, lt, ti, atb,alt,sol_alt)
	lonlat(n1, n2,lg, lt, ti, atb)
	'''

	#lonlat(n1, n2,lg, lt, ti, atb)
	#affiche(n1,n2,lg, lt, ti, SR,alt,sol_alt)
	#alt_max_orbite(alt, sol_alt, SR,lt)
	'''
	# Tableau de flag
	
	#affiche_flag(lg, lt, flag,alt,sol_alt)
	freq_occur(flag,alt)
	'''
	dossier = '/bdd/CFMIP/CFMIP_OBS_LOCAL/GOCCP/instant_SR_CR_DR/grid_L40/2013/201308/night/'
	#freq_occur_region(dossier)
	flag = flag_orbite(alt,SR)
	affiche_flag(lg, lt, flag,alt,sol_alt)