import pyproj
import netCDF4 as nc
import os
import numpy as np
from mpl_toolkits.basemap import Basemap

os.chdir("C:\\Users\\bnordgren\\Documents\\Orchidee\\Snow")

d = nc.Dataset('nhtsw100e2_20030107_20030113.nc')

print d

print d.variables['time']

print d.variables['time'][0]

full_years = 365.25 * (2002-1966)
oct_to_jan = np.sum([(31-3),30,31]) # rest of oct, nov, dec
print full_years, full_years + oct_to_jan
print d.variables['time'][0] - (full_years + oct_to_jan)

print d.variables['latitude']

print d.variables['latitude'][0,0]

lats = d.variables['latitude'][:]
print lats.shape

print "The mask is: ", lats.mask[0,0]
print "The data value is: ", lats.data[0,0]

lats[0,0] = 42.
print lats[0,0]
print "The mask is: ", lats.mask[0,0]
print "The data value is: ", lats.data[0,0]

# ... and set it back
lats[0,0] = np.ma.masked
print lats[0,0]
print "The mask is: ", lats.mask[0,0]
print "The data value is: ", lats.data[0,0]

lats.data[0,0] = -999.0
print lats[0,0]
print "The mask is: ", lats.mask[0,0]
print "The data value is: ", lats.data[0,0]

snow = d.variables['merged_snow_cover_extent']
print snow

print snow[0,0,0]

snow[0,0,0] = 42

coords = d.variables['coord_system']
print coords

unmasked_ij = np.where(np.logical_not(lats.mask))
indices = zip(*unmasked_ij)
first = indices[0]
last = indices[-1]
print "First: ", first
print "Last: ", last

lons = d.variables['longitude']

print "First: ", lats[first], lons[first]
print "Last: ", lats[last], lons[last]

#remember that coords is the python variable referring to the "coord_system" netcdf variable.
lat_0 = coords.latitude_of_projection_origin 
lon_0 = coords.longitude_of_projection_origin
false_e = coords.false_easting
false_n = coords.false_northing
s_major = coords.semimajor_axis 
s_minor = coords.semiminor_axis

the_proj = pyproj.Proj(proj="laea", lat_0=lat_0, lon_0=lon_0, x_0=false_e, y_0=false_n, a=s_major, b=s_minor)
geo      = pyproj.Proj(proj="latlong", a=s_major, b=s_minor)

x1,y1 = pyproj.transform(geo,the_proj, lons[first], lats[first])
x2,y2 = pyproj.transform(geo,the_proj, lons[last], lats[last])

# it's (rows,cols), so index[1]==column 
m_x = (x2-x1)/(last[1]-first[1])
b_x = x1 - (m_x*first[1])
print x2, last[1]*m_x+b_x

# it's (rows,cols) so index[0]==rows
m_y = (y2-y1)/(last[0]-first[0])
b_y = y1 -(m_y*first[0])
print y2, last[0]*m_y+b_y

print m_x, m_y

cv_rows = m_x * np.arange(len(d.dimensions['rows'])) + b_x
cv_cols = m_y * np.arange(len(d.dimensions['cols'])) + b_y

# first, close the current file
d.close()
# now open again for writing
d = nc.Dataset('nhtsw100e2_20030107_20030113.nc', 'r+')

# make the rows coordinate variable
# note the variable name is the same as the dimension name
rows_var = d.createVariable('rows',np.float32, ('rows',))

# give it the expected attributes
rows_var.standard_name = 'projection_y_coordinate'
rows_var.axis          = 'Y'
rows_var.units         = 'meters'

# write the values to the variable
rows_var[:] = cv_rows

# make the cols coordinate variable
# note the variable name is the same as the dimension name
cols_var = d.createVariable('cols',np.float32, ('cols',))

# give it the expected attributes
cols_var.standard_name = 'projection_x_coordinate'
cols_var.axis          = 'X'
cols_var.units         = 'meters'

# write the values to the variable
cols_var[:] = cv_cols

#close the file to make sure everything is written
d.close()

print m_y

d = nc.Dataset('nhtsw100e2_20030107_20030113.nc')
snow = d.variables['merged_snow_cover_extent'][:]
lats = d.variables['latitude'][:]
lons = d.variables['longitude'][:]
imshow(snow[0,...])

def init_basemap() : 

    m = Basemap(width=100000*180, height=100000*180,
            #llcrnrx=cv_cols[0], llcrnry=cv_rows[-1], urcrnrx=cv_cols[-1], urcrnry=cv_rows[0],
            resolution='l', projection='laea',
            lon_0=lon_0, lat_0=lat_0)
    m.drawcoastlines()
    m.drawcountries()
    m.drawmeridians(np.arange(-180.,180.,20.),labels=[False,False,False,True])
    m.drawparallels(np.arange(10.,80.,20.), labels=[True,False,False,False])
    return m

figsize(12,10)
m = init_basemap()
c = m.contourf(lons[:], lats[:],
               snow[0,...], 20, latlon=True)
cb = m.colorbar(c)
title("Snow cover, 13 Jan 2003")