# -*- coding: utf-8 -*-
# <nbformat>3.0</nbformat>

# <codecell>

import os
import datetime
import re
import shutil
import numpy as n
from glob import glob
from scipy.io import netcdf

# <codecell>

def create_atts(ncfile, newfile, datecode):
    now = datetime.datetime.utcnow().replace(microsecond=0).strftime("%Y-%m-%d %H:%M:%S")
    newfile.title                = ncfile.title               
    newfile.institution          = ncfile.institution         
    newfile.institution_url      = ncfile.institution_url     
    newfile.institution_dods_url = ncfile.institution_dods_url
    newfile.metadata_url         = ncfile.metadata_url        
    newfile.references           = ncfile.references          
    newfile.contact              = ncfile.contact             
    newfile.source               = ncfile.source              
    newfile.history              = ncfile.history             
    newfile.comment              = ncfile.comment             
    newfile.Conventions          = ncfile.Conventions         
    newfile.format_category_code = ncfile.format_category_code
    newfile.institution_code     = ncfile.institution_code    
    newfile.platform_code        = ncfile.platform_code       
    newfile.package_code         = ncfile.package_code        
    newfile.project              = ncfile.project             
    newfile.project_url          = ncfile.project_url         
    newfile.start_date           = ncfile.start_date          
    newfile.end_date             = ncfile.end_date            
    newfile.release_date         = now       
    newfile.creation_date        = now
    newfile.modification_date    = now
    newfile.process_level        = 'level2'
    if datecode < '2013_02':
        newfile.level2_flag_elevation_correction = True
    if '2011_09' < datecode < '2012_05':
        newfile.level2_flag_bottom_z_nan_padded = True
        newfile.level2_flag_top_z_nan_padded = True
    if '2011_05' < datecode < '2011_11':
        newfile.level2_flag_upper_z_cropped = True
    newfile._FillValue           = ncfile._FillValue          
   

# <codecell>

def create_dims(ncfile, newfile):
    newfile.createDimension('ntime', ncfile.dimensions['ntime'])
    newfile.createDimension('nlat', ncfile.dimensions['nlat'])
    newfile.createDimension('nlon', ncfile.dimensions['nlon'])
    newfile.createDimension('nz', 39)
    

# <codecell>

def create_vars(ncfile, newfile, zs, newz):
    ncfiles = None
    if type(ncfile) is list:
        ncfiles = ncfile
        ncfile = ncfile[0]
    
    time = newfile.createVariable('time', 'i', ('ntime',))
    if ncfiles:
        time[:] = n.concatenate([ncf.variables['time'].data for ncf in ncfiles])
    else:
        time[:] = ncfile.variables['time'].data
    time.short_name = ncfile.variables['time'].short_name
    time.long_name = ncfile.variables['time'].long_name
    time.standard_name = ncfile.variables['time'].standard_name
    time.units = ncfile.variables['time'].units
    time.axis = ncfile.variables['time'].axis

    lat = newfile.createVariable('lat', 'f', ('nlat',))
    lat[:] = ncfile.variables['lat'].data
    lat.short_name = ncfile.variables['lat'].short_name
    lat.long_name = ncfile.variables['lat'].long_name
    lat.standard_name = ncfile.variables['lat'].standard_name
    lat.reference = ncfile.variables['lat'].reference
    lat.units = ncfile.variables['lat'].units
    lat.valid_range = ncfile.variables['lat'].valid_range
    lat.axis = ncfile.variables['lat'].axis

    lon = newfile.createVariable('lon', 'f', ('nlon',))
    lon[:] = ncfile.variables['lon'].data
    lon.short_name = ncfile.variables['lon'].short_name
    lon.long_name = ncfile.variables['lon'].long_name
    lon.standard_name = ncfile.variables['lon'].standard_name
    lon.reference = ncfile.variables['lon'].reference
    lon.units = ncfile.variables['lon'].units
    lon.valid_range = ncfile.variables['lon'].valid_range
    lon.axis = ncfile.variables['lon'].axis

    z = newfile.createVariable('z', 'f', ('nz',))
    if newz is not None:
        z[:] = newz
    else:
        z[:] = ncfile.variables['z'].data
    z.short_name = ncfile.variables['z'].short_name
    z.long_name = ncfile.variables['z'].long_name
    z.standard_name = ncfile.variables['z'].standard_name
    z.reference = ncfile.variables['z'].reference
    z.positive = ncfile.variables['z'].positive
    z.units = ncfile.variables['z'].units
    z.axis = ncfile.variables['z'].axis

    u = newfile.createVariable('u', 'f', ('ntime', 'nz'))
    if ncfiles:
        u[:,:] = n.concatenate([ncf.variables['u'].data for ncf in ncfiles])
    else:
        if zs < 39:
            u[:,:] = ones(u.shape, dtype=float) * nan
            u[:,1:38] = ncfile.variables['u'].data
        elif zs == 39:
            u[:,:] = ncfile.variables['u'].data
        else:
            u[:,:] = ncfile.variables['u'].data[:,0:39]
    u.short_name = ncfile.variables['u'].short_name
    u.long_name = ncfile.variables['u'].long_name
    u.standard_name = ncfile.variables['u'].standard_name
    u.positive = ncfile.variables['u'].positive
    u.units = ncfile.variables['u'].units

    v = newfile.createVariable('v', 'f', ('ntime', 'nz'))
    if ncfiles:
        v[:,:] = n.concatenate([ncf.variables['v'].data for ncf in ncfiles])
    else:
        if zs < 39:
            v[:,:] = ones(u.shape, dtype=float) * nan
            v[:,1:38] = ncfile.variables['v'].data
        elif zs == 39:
            v[:,:] = ncfile.variables['v'].data
        else:
            v[:,:] = ncfile.variables['v'].data[:,0:39]
    v.short_name = ncfile.variables['v'].short_name
    v.long_name = ncfile.variables['v'].long_name
    v.standard_name = ncfile.variables['v'].standard_name
    v.positive = ncfile.variables['v'].positive
    v.units = ncfile.variables['v'].units

    w = newfile.createVariable('w', 'f', ('ntime', 'nz'))
    if ncfiles:
        w[:,:] = n.concatenate([ncf.variables['w'].data for ncf in ncfiles])
    else:
        if zs < 39:
            w[:,:] = ones(v.shape, dtype=float) * nan
            w[:,1:38] = ncfile.variables['w'].data
        elif zs == 39:
            w[:,:] = ncfile.variables['w'].data
        else:
            w[:,:] = ncfile.variables['w'].data[:,0:39]
    w.short_name = ncfile.variables['w'].short_name
    w.long_name = ncfile.variables['w'].long_name
    w.standard_name = ncfile.variables['w'].standard_name
    w.positive = ncfile.variables['w'].positive
    w.units = ncfile.variables['w'].units

    sigw = newfile.createVariable('sigw', 'f', ('ntime', 'nz'))
    if ncfiles:
        sigw[:,:] = n.concatenate([ncf.variables['sigw'].data for ncf in ncfiles])
    else:
        if zs < 39:
            sigw[:,:] = ones(sigw.shape, dtype=float) * nan
            sigw[:,1:38] = ncfile.variables['sigw'].data
        elif zs == 39:
            sigw[:,:] = ncfile.variables['sigw'].data
        else:
            sigw[:,:] = ncfile.variables['sigw'].data[:,0:39]
    sigw.short_name = ncfile.variables['sigw'].short_name
    sigw.long_name = ncfile.variables['sigw'].long_name
    sigw.standard_name = ncfile.variables['sigw'].standard_name

    bck = newfile.createVariable('bck', 'f', ('ntime', 'nz'))
    if ncfiles:
        bck[:,:] = n.concatenate([ncf.variables['bck'].data for ncf in ncfiles])
    else:
        if zs < 39:
            bck[:,:] = ones(bck.shape, dtype=float) * nan
            bck[:,1:38] = ncfile.variables['bck'].data
        elif zs == 39:
            bck[:,:] = ncfile.variables['bck'].data
        else:
            bck[:,:] = ncfile.variables['bck'].data[:,0:39]
    bck.short_name = ncfile.variables['bck'].short_name
    bck.long_name = ncfile.variables['bck'].long_name
    bck.standard_name = ncfile.variables['bck'].standard_name

    error = newfile.createVariable('error', 'i', ('ntime', 'nz'))
    if ncfiles:
        error[:,:] = n.concatenate([ncf.variables['error'].data for ncf in ncfiles])
    else:
        if zs < 39:
            error[:,:] = zeros(error.shape, dtype=int)
            error[:,1:38] = ncfile.variables['error'].data
        elif zs == 39:
            error[:,:] = ncfile.variables['error'].data
        else:
            error[:,:] = ncfile.variables['error'].data[:,0:39]
    error.short_name = ncfile.variables['error'].short_name
    error.long_name = ncfile.variables['error'].long_name
    error.standard_name = ncfile.variables['error'].standard_name

# <codecell>

def combine(ncfiles, dest):
    print 'Writing', os.path.basename(dest)
    newfile = netcdf.netcdf_file(dest, 'w')
    try:
        pattern = re.compile(r'(\d{4}_\d{2})')
        create_atts(ncfiles[0], newfile, pattern.search(os.path.basename(dest)).groups()[0])
        create_dims(ncfiles[0], newfile)
        create_vars(ncfiles, newfile, None, None)
    finally:
        newfile.close()

# <codecell>

home = os.path.join(os.sep, 'Users', 'cbcoasis', 'Documents', 'sodar')
src = os.path.join(home, 'ncs')
dest = os.path.join(home, 'new')
ncext = 'nc'
ncs = os.path.join(src, '*' + os.extsep + ncext)
ncs = glob(ncs)
ncs = [os.path.basename(nc) for nc in ncs]
ncs.sort()

# <codecell>

for nc in ncs:
    print 'Starting', nc
    ncfile = netcdf.netcdf_file(os.path.join(src, nc))
    try:
        newz = None
        zs = len(ncfile.variables['z'].data)
        z0 = ncfile.variables['z'].data[0]
        if zs < 39:
            newz = ncfile.variables['z'].data
            newz = insert(newz, 0, float32(10.0))
            newz = append(newz, float32(200.0))
            newz = newz + 1.0
        elif zs == 39:
            if z0 != 11.0:
                newz = ncfile.variables['z'].data + 1.0
        else:
            newz = ncfile.variables['z'].data[0:39] + 1.0
        newfile = netcdf.netcdf_file(os.path.join(dest, nc), 'w')
        try:
            pattern = re.compile(r'(\d{4}_\d{2})')
            create_atts(ncfile, newfile, pattern.search(nc).groups()[0])
            create_dims(ncfile, newfile)
            create_vars(ncfile, newfile, zs, newz)
        finally:
            newfile.close()
    finally:
        ncfile.close()
    print 'Finishing', nc

# <codecell>

src = os.path.join(home, 'new')
ncs = os.path.join(src, '*' + os.extsep + ncext)
ncs = glob(ncs)
ncs = [os.path.basename(nc) for nc in ncs]
ncs.sort()
pattern = re.compile(r'(\d{4}_\d{2})[A-Za-z]+')
keys = []
for nc in ncs:
    try:
        key = pattern.search(nc).groups()[0]
        if key not in keys:
            keys.append(key)
    except:
        pass
    

# <codecell>

dest = os.path.join(home, 'newer')
for key in keys:
    ncs = os.path.join(src, '*' + key + '*' + os.extsep + ncext)
    ncs = glob(ncs)
    ncs = [os.path.basename(nc) for nc in ncs]
    ncs.sort
    ncfiles = []
    try:
        for nc in ncs:
            print 'Reading', nc
            ncfiles.append(netcdf.netcdf_file(os.path.join(src, nc)))
        combine(ncfiles,os.path.join(dest, ncs[0]))
    finally:
        for ncfile in ncfiles:
            ncfile.close()
    ncfiles = []; ncfile = None # close doesn't work on weak references?
    

# <codecell>

dest = os.path.join(home, 'newest')
shutil.rmtree(dest, True)
shutil.copytree(src, dest)
for key in keys:
    ncs = os.path.join(dest, '*' + key + '*' + os.extsep + ncext)
    ncs = glob(ncs)
    for nc in ncs:
        os.remove(nc)
src = os.path.join(home, 'newer')
ncs = os.path.join(src, '*' + os.extsep + ncext)
ncs = glob(ncs)
for nc in ncs:
    shutil.copy(nc, dest)