# This scripts is doing...
# Patrick Kuehl, date
import time
from pylab import *
from datetime import datetime
import matplotlib.pyplot as plt
import numpy as np
import numpy.ma as ma
from matplotlib.colors import LogNorm
from scipy.optimize import leastsq
from matplotlib import dates
import time
import datetime as date
from datetime import datetime
from datetime import timedelta
import os
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning) 
from matplotlib.backends.backend_pdf import PdfPages


# load some data
year,doy,total_num,int_num,prio_num,int_prio_num=np.loadtxt("/home/pacifix/kuehl/work/analysis/priority_ephin_checks/10_penetrating_total_priority_phas/count_pha_stuff.dat",unpack=True)
datearray=[(datetime(int(year[date]),1,1,0)+timedelta(int(doy[date])-1)) for date in range(len(year))]
x,y,z=total_num,int_num,prio_num


#general plot stuff
pp= PdfPages('plot_3d_scatter.pdf')
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

ax.set_xlim3d(x.min(),x.max())
ax.set_ylim3d(y.min(),y.max())
ax.set_zlim3d(z.min(),z.max())

ax.set_xlabel("Total # of PHAs")
ax.set_ylabel("# of PHAs in Int.Ch.")
ax.set_zlabel("# of priority PHAs")

# contourf 
import colormaps as cmaps
plt.register_cmap(name='viridis', cmap=cmaps.viridis)
plt.set_cmap(cmaps.viridis)
X,Y,Z=total_num,int_num,prio_num
mybins=50 #14
lvls=np.linspace(-.55,2.8,20)
hist_offset_for_logs=0.3
mynchunks=2 # to get rid of artifacts
dist_from_ax=0.0 # 0 to match contours together, >0 for spacing between them

# contour on xy plane (z direction)
hist, binx, biny = np.histogram2d( X, Y,bins=mybins)
x = np.linspace(X.min(), X.max(), hist.shape[0])
y = np.linspace(Y.min(), Y.max(), hist.shape[1])
y, x = np.meshgrid(y, x)
(minz,maxz)= ax.get_zlim3d()
myoffset=minz-dist_from_ax*(maxz-minz)
phist=np.log10(hist+hist_offset_for_logs)
ax.contourf(x, y, phist, zdir='z', offset=myoffset,levels=lvls,nchunk=mynchunks,rasterized=True)
print phist.min(),phist.max()

# contour on xz plane (y direction)
hist, binx, biny = np.histogram2d( X, Z,bins=mybins)
x = np.linspace(X.min(), X.max(), hist.shape[0])
z = np.linspace(Z.min(), Z.max(), hist.shape[1])
z, x = np.meshgrid(z,x)
(miny,maxy)= ax.get_ylim3d()
myoffset=maxy+dist_from_ax*(maxy-miny)
phist=np.log10(hist+hist_offset_for_logs)
ax.contourf(x, phist, z, zdir='y', offset=myoffset,levels=lvls,nchunk=mynchunks,rasterized=True)
print phist.min(),phist.max()

# contour on yz plane (x direction)
hist, binx, biny = np.histogram2d( Y, Z,bins=mybins)
y = np.linspace(Y.min(), Y.max(), hist.shape[0])
z = np.linspace(Z.min(), Z.max(), hist.shape[1])
z, y = np.meshgrid(z, y)
(minx,maxx)= ax.get_xlim3d()
myoffset=minx-dist_from_ax*(maxx-minx)
phist=np.log10(hist+hist_offset_for_logs)
ax.contourf(phist, y, z, zdir='x', offset=myoffset,levels=lvls,nchunk=mynchunks,rasterized=True)
print phist.min(),phist.max()
 
# some general plto stuff 
plt.axis('tight') 
ax.view_init(elev=30., azim=315)
ax.dist=9
pp.savefig()

# add scatter to histoplot and save again :)
x,y,z=total_num,int_num,prio_num
ax.scatter(x,y,z,c='r',rasterized=True,alpha=0.7,edgecolor="r",s=1,zorder=1) # c='z' can also be used

pp.savefig()
plt.close()
pp.close()