import matplotlib.pyplot as plt
import numpy as np
import ncdiag as ncd
import matplotlib  as mpl
import math
from mpl_toolkits.basemap import Basemap

plt.rcParams.update({'font.size': 8})

exp = 'x46d_wgu'
exp2 = 'x46d_crtm'
sat = 'CrIS-FSR_N20'

filename = './data/'+exp+'/'+exp+'.diag_cris-fsr_n20_ges.20211215_00z.nc4'
filename2 = './data/'+exp2+'/'+exp2+'.diag_cris-fsr_n20_ges.20211215_00z.nc4'

fsr = ncd.obs(filename,verbose=True)
fsr2 = ncd.obs(filename2,verbose=True)

print filename
print filename2

#startdate = 2019122800
#enddate   = 2019122818
#exp = 'x042_hp2_wg'
#exp = 'x042_hp3_wg'
#sat = 'iasi_metop-a'
#fsr_tmpl = ncd.obs_template(exp+'.diag_'+sat+'_ges.$yyyy$mm${dd}_${hh}z.nc4',startdate=startdate,enddate=enddate)

stype = 'ice'
stype = 'snow'
stype = 'land'
stype = 'sea'
stype = 'glb'
print stype

maskplus = ' & ('+stype+' >=0.99) & (used ==1)'
outname = stype+'_used'
maskplus = ' & ('+stype+' >=0.99)'
maskplus = ''
outname = stype

#fig, (ax1,ax2) = plt.subplots(nrows=2,ncols=1, figsize=(6,5))
fig, (ax1) = plt.subplots(nrows=1,ncols=1, figsize=(7,4))
#fig, ax = plt.subplots(2,2,sharex=True, sharey=True)
#fig, ax = plt.subplots(1,1,figsize=(10, 6))

tbmin = 200
tbmax = 500
font1 = {'family': 'serif', 'color':  'Black', 'weight': 'normal', 'size': 9, 'horizontalalignment': 'left'}
font2 = {'family': 'serif', 'color':  'Red', 'weight': 'normal', 'size': 9, 'horizontalalignment': 'left'}

ichstr = 219
chindex = 501
chstr = str(chindex)
print "ichstr=",ichstr, "chindex= ", chindex
mask = '(ichan == 219)'+maskplus
fsr.set_mask(mask)
fsr.use_mask(True)
fsr2.set_mask(mask)
fsr2.use_mask(True)
#m1 = Basemap(projection='mill',lon_0=180)
#m1.drawcoastlines()
nobs = len(fsr.v('omfnbc'))
nobs2 = len(fsr2.v('omfnbc'))
print "nobs= ",nobs," nobs2= ", nobs2

counts = np.zeros(nobs)
jproj = np.zeros(nobs)
ntot = 7
n_count = np.zeros(ntot)
val = np.zeros(ntot)
val[0]=-40
val[1]=-30
val[2]=-20
val[3]=-10
val[4]=0
val[5]=10
val[6]=20

for i in range(nobs):
  j0 = -1
  if fsr.v('omfnbc')[i] >= fsr2.v('omfnbc')[i] :
    for j in range(ntot):
      if j == 0 :
        if fsr.v('omfnbc')[i] <= val[j] :
          n_count[j] = n_count[j] + 1
          j0 = 0
      else :
        if fsr.v('omfnbc')[i] <= val[j] and fsr.v('omfnbc')[i] > val[j-1] :
          n_count[j] = n_count[j] + 1
          j0 = j
          break
  else :
    for j in range(ntot):
      if j == 0 :
        if fsr2.v('omfnbc')[i] <= val[j] :
          n_count[j] = n_count[j] + 1
          j0 = 0
      else :
        if fsr2.v('omfnbc')[i] <= val[j] and fsr2.v('omfnbc')[i] > val[j-1] :
          n_count[j] = n_count[j] + 1
          j0 = j
  if j0 == -1 :
     print " could not found proper subdomain i, fsr(omfnbc),fsr2(omfnbc),j,val[j]= ", i,fsr.v('omfnbc')[i],fsr2.v('omfnbc')[i]
  jproj[i] = j0

print "n_count=", n_count
print "jproj[0:20]: ", jproj[0:20]
print fsr.v('omfnbc')[0:20]
print fsr2.v('omfnbc')[0:20]

for j in range(ntot):
  n_count[j] = math.log(n_count[j],10)

print "n_count=", n_count
print "jproj[0:20]: ", jproj[0:20]

for i in range(nobs):
   i0 = jproj[i]
   counts[i] = n_count[int(i0)]

print counts[0:10]

#mycm   = mpl.colors.LinearSegmentedColormap.from_list('mycolors',['cyan','green','yellow','magenta','red'])
mycm = plt.cm.get_cmap('RdYlBu')
#mycm = plt.cm.get_cmap('uBlYdR')
sc1 = ax1.scatter(fsr.v('omfnbc'),fsr2.v('omfnbc'),label='CH'+chstr,s=4,marker="o",c=counts,cmap=mycm,edgecolors='face')
ax1.legend(loc='upper left',fontsize=8)
cax1 = plt.axes([0.92, 0.15, 0.03, 0.7])
plt.colorbar(sc1,ax=ax1,cax=cax1)

ax1.set_xlabel(sat+' TB 2.3.0',fontsize=8, fontweight='normal')
ax1.set_ylabel(sat+' TB 2.4.0',fontsize=8, fontweight='normal')


#n_count = np.zeros(ntot)
#for i in range(nobs):
#    for j in range(ntot):
#        if fsr.v('omfnbc')[i] <= val[j] and fsr2.v('omfnbc')[i] <= val[j] :
#          n_count[j] = n_count[j] + 1
#
#print "n_count=", n_count

#ax.set_title('Channel 22 OMF Histogram, Used Obs')
#ax.set_title('Channel 219 OMF Histogram, Used Obs')
#ax1.set_title(exp+': '+sat+' CH 92 OMF Used Obs')
#ax.set_title(exp+': '+sat+' Ch 93 OMF Used Obs')
#ax.set_title(exp+': '+sat+' Ch 95 OMF Used Obs')

#ax2.set_xlim(-2.5,2.51)
#ax2.xaxis.set_ticks(np.arange(-2.5, 2.51, 0.5))
#ax2.set_ylim(0,2000.1)
#ax2.yaxis.set_ticks(np.arange(0, 2000.1, 500))
#ax2.hist(fsr_tmpl.v('omfbc'),bins=41, label="W/ BC")
#ax2.hist(fsr.v('omfbc'),bins=41, label="W/ BC")
#ax2.legend(loc='upper left',fontsize=8)

fig.savefig(exp+'_'+exp2+'_'+sat+'_scatter_omf_'+outname+'_ch501.png')
#fig.savefig(exp+'_'+sat+'_ch92_histogram_omf.png')
#fig.savefig(exp+'_'+sat+'_ch93_histogram_omf.png')
#fig.savefig(exp+'_'+sat+'_ch95_histogram_omf.png')

plt.show()