#!/usr/bin/env python
# -*- coding: utf-8 -*-

import netCDF4
# Basemap is a library that lets us to draw on a map.
from mpl_toolkits.basemap import Basemap
import numpy as np
import matplotlib.pyplot as plt

# Setting up the map
# We use stereographic projection centered on 65.5 degrees
# north, 12 degrees east. (Approximately the location of
# Brønnøysund).
# With these parameters, and scale = 1, the plot covers
# the entire NorKyst dataset. Reduce scale to zoom in.
scale  = 1
width  = 1800000
height = 2300000
lat_0  = 65.5
lon_0  = 12
print('Creating map')
fig, ax = plt.subplots(figsize = (10, 10))
m = Basemap(width=width*scale,height=height*scale,
    resolution='i',projection='stere',\
    lat_ts=65,lat_0=lat_0,lon_0=lon_0)

# Draw continents, from the database that comes with
# matplotlib as well as country borders and a border
# around the map
m.fillcontinents(color='#aaaaaa',lake_color='#cccccc')
m.drawcountries(linewidth=0.2)
m.drawmapboundary(fill_color='#cccccc')

# draw parallels and meridians for every 5 and 10 degrees.
# Most of these won't show up in the map selection used
# here, but that doesn't matter.
m.drawparallels(np.arange(-80.,81.,5.))
m.drawmeridians(np.arange(-180.,181.,10.))

# met.no provides us with free data (or rather, free as
# in 200 free texts, since we pay for the service with
# our taxes), downloadable through the magic of thredds.
# You can open this url as if it was a local file, and
# the data are only downloaded when needed, and only as
# much as needed. Note that this link will remain valid
# only for a limited time. See the website for
# available files at any given time.
dataurl = 'http://thredds.met.no/thredds/dodsC/fou-hi/norkyst800m-1h/NorKyst-800m_ZDEPTHS_his.an.2016020500.nc'
data = netCDF4.Dataset(dataurl)

print('Dowloading data')
# The full dataset is pretty large, so we will download
# only every second datapoint in either direction for
# a large overview map (change as needed).
dx, dy = 1, 1
# The dimensions of the dataset is time, depth, lon, lat.
# We only download the first timestep, only the top layer
# and only every third datapoint in the horisontal
# directions.
# Downloading vector components of the velocity. Note that
# the components are aligned with the grid, not North-East.
u = data.variables['u'][0,0,::dx,::dy]
v = data.variables['v'][0,0,::dx,::dy]
# Calculate speed
speed = np.sqrt(u**2 + v**2)
# gridlons and gridlats tell us the lon and lat of each
# point in the grid. This is necessary for drawing on
# the map. These are two dimensional grids and we only
# download every dx and dy element in either direction to
# match the temperature data.
gridlons = data.variables['lon'][::dx,::dy]
gridlats = data.variables['lat'][::dx,::dy]

# Here we use information from the map projection to
# convert from lon and lat to coordinates that can be
# used to draw data onto the map.
X, Y = m(gridlons, gridlats)

print('Plotting')
# pcolormesh is a function which plots an array of data
# as a grid of cells with color according to the value.
plt.pcolormesh(X, Y, speed)
# Add a colorbar, so we can read the speed.
plt.colorbar()
# Crop away white space around figure
plt.tight_layout()
# Saving file
plt.savefig('depth.png', dpi = 200)