#!/usr/bin/env python

from __future__ import print_function

import sys, os
import numpy as np
import netCDF4 as nc

# https://trac.nci.org.au/svn/access_tools/umfile_utils/branches/umfile-operators
# installed locally.
from umfile_utils import umfile

MASK_ITEM_CODE = 30
LANDFRAC_ITEM_CODE = 505

def find_closest_unmasked(mask, row, col):
    """
    Find the index of the closest unmasked value.
    
    Do a spiral search.
    """

    def circle_around(x, y):

        # First try the centre
        yield 0, (x, y)

        # Now spiral outwards. 
        r = 1
        i, j = x-1, y-1
        while True:
            while i < x+r:
                i += 1
                yield r, (i, j)
            while j < y+r:
                j += 1
                yield r, (i, j)
            while i > x-r:
                i -= 1
                yield r, (i, j)
            while j > y-r:
                j -= 1
                yield r, (i, j)
            r += 1
            j -= 1
            yield r, (i, j)

    def normalise(i, j, max_i, max_j):
        """
        The spiral search may have returned indices which are
        out-of-bounds. Fix up. 
        """

        # Negative latitude index is not allowed. Note:
        # negative longitude index is OK, it just wraps
        # around.
        i = abs(i)

        # Set maximum latitude index.                    
        if i >= max_i:
            i = max_i - 1

        # Longitude index wraps around in positive direction.
        if j >= max_j:
            j = j - max_j
            
        return i, j

    for r, (i, j) in circle_around(row, col):
        i, j = normalise(i, j, mask.shape[0], mask.shape[1])
        if not mask[i, j]:
            return r, (i, j)


def remask_var(new_var, old_var, new_mask, old_mask):
    """
    Make new_var use the new_mask. Data is taken from old_var on old_mask. For
    multi-level variables each level is handled separately.

    Firstly any land points missing in new_var (according to new_mask) need to
    be set. This is done with a nearest-neighbour approach. 

    Any land points that var has but new_mask doesn't are unset.
    """

    # New var is presently 'broken' because it uses a new mask, and the mask
    # determines how data is packed. Reset to previous value.
    new_var[:] = old_var[:]

    # Build list of references to levels so that each can be treated
    # separately. Include a copy of the original level data to read from.
    levels = []
    if len(new_var[:].shape) == 3:
        for l in range(new_var[:].shape[0]):
            levels.append((new_var[l, :, :], old_var[l, :, :],
                           new_var[l, :, :].mask))
    else:
        levels.append((new_var[:], old_var[:], new_var[:].mask))

    # Fix up values in new_var which are missing (according to new_mask).
    for (level_data, level_data_orig, level_mask) in levels:
        land_missing = np.where((new_mask == False) & (old_mask == True))

        for i, j in zip(land_missing[0], land_missing[1]):
            _, (closest_i, closest_j) = find_closest_unmasked(old_mask, i, j)
            level_data[i, j] = level_data_orig[closest_i, closest_j]

            # Need to fix up the mask also, changing numpy array views do not
            # change the mask of the original!
            level_mask[i, j] = False

    # Mask out land points which are set in new_var but not set in new_mask.
    if len(new_var[:].shape) == 3:
        for i in range(new_var[:].shape[0]):
            tmp = new_var[i, :, :].mask
            tmp[np.where(new_mask == True)] = True
    else:
        new_var[np.where(new_mask == True)].mask = True


def get_land_mask(um_filename):
    """
    Read out the old land mask.
    """

    um_f = umfile.UMFile(um_filename)
    tmp = um_f.variables[MASK_ITEM_CODE][:]

    # The UM uses -1 to mean unmasked and 0 for masked. So we need to do some
    # swapping. 
    old_mask = np.zeros_like(tmp, dtype='bool') 
    old_mask[np.where(tmp == -1)] = False
    old_mask[np.where(tmp == 0)] = True

    return old_mask


def set_land_mask(um_filename, new_mask):
    """
    Put in new land mask.

    This will effectively 'break' all data fields that use the land mask for
    packing. After calling this all data fields need to be remasked.
    """

    um_f = umfile.UMFile(um_filename, 'r+')
    tmp = um_f.variables[MASK_ITEM_CODE][:]

    # The UM dump uses -1 to mean unmasked and 0 for masked. So we need to do
    # some swapping. 
    tmp[np.where(new_mask == False)] = -1
    tmp[np.where(new_mask == True)] = 0

    um_f.close()


def copy_masked_vars(f, mask):
    """
    Guess which vars in f are masked with a land-sea mask and copy them. These
    are the ones which need to be updated when the mask changes.

    The masked variables are defined on land but not on sea.
    """

    var_copies = []
    for v in f.variables.itervalues():
        if len(v[:].shape) == 3:
            var = v[0,:,:]
        else:
            var = v[:]

        if np.array_equal(np.where(var.mask == True), np.where(mask == True)):
            var_copies.append(v)
        else:
            # This means that the field contains missing values but not due
            # to land-sea masking. Broken data? Mismatching land-sea mask?

            # Either all masked or no mask. Something funny with the
            # river fields, they have a different grid, ignore for now.  
            if 'RIVER' not in v.__str__():
                assert(np.all(var.mask) or var.mask == False)

    return var_copies


def remask(um_restart_filename, new_mask, new_landfrac=None):
    """
    Update all fields with a new mask. Can also update land fraction.
    """

    old_mask = get_land_mask(um_restart_filename)

    # Copy out the variables. Changing the masks will corrupt these so we need
    # to keep a copy around to do the remasking.
    var_copies = []
    with umfile.UMFile(um_restart_filename) as um_f:
        var_copies = copy_masked_vars(um_f, old_mask)

    # Change the land mask. The file format uses this for compression, so if
    # it's not changed first then updates on other fields may not work.
    set_land_mask(um_restart_filename, new_mask)

    # Now update all variables.
    um_f = umfile.UMFile(um_restart_filename, 'r+')
    for v in var_copies:
        if v.item_code != MASK_ITEM_CODE and v.item_code != LANDFRAC_ITEM_CODE:
            remask_var(um_f.variables[v.item_code], v, new_mask, old_mask) 

    if new_landfrac is not None:
        um_f.variables[LANDFRAC_ITEM_CODE][:] = new_landfrac[:]

    um_f.close()

    # Return item codes of all remasked variables.
    item_codes = [v.item_code for v in var_copies]
    item_codes.append(MASK_ITEM_CODE)
    if new_landfrac is not None:
        item_codes.append(LANDFRAC_ITEM_CODE)

    return item_codes



