# General utilities for ROMS/HOPS datasets

from  sets import Set

import Numeric as Nm
from math import *

ROMS_dx = 1794  # meters
ROMS_dy = 2219

HOPS_dx = 1789
HOPS_dy = 2219
# Note: HOPS velocity is in cm/s

earth_equat_radius=6378135
earth_polar_radius=6356750

class Settings(object): pass
class SettingsError(Exception): pass


def earth_radius(latitude):
    """Calculate radius of the earth (in meters)at a given latitude.
    Uses an oblate spheroid model.
    """
    a=earth_equat_radius
    b=earth_polar_radius
    c=cos(latitude*pi/180)
    s=sin(latitude*pi/180)
    return sqrt(((a**2 * c)**2 + (b**2 * s)**2)/
                ((a*c)**2 + (b*s)**2))

earth_polar_circumference=2*pi*earth_radius(45.0)

def earth_circumference(latitude):
    """Compute length of a line of latitude (in meters).
    """
    #radius of circle of latitude 
    latrad=earth_radius(latitude)*cos(latitude*pi/180)
    return 2*pi*latrad



def normalize_array(dat,mn=None,mx=None):
    """Normalize the values in the array to lie in [0,1]. If mn or mx are given,
    then the range [mn,mx] is normalized to [0,1], and values outside this range are
    clamped to the nearest endpoint. Otherwise, the range of values is computed.
    """
    if mn is None:
        mn=Nm.minimum.reduce(Nm.ravel(dat))
    # Find maximum, excluding any sentinel values
    if mx is None:
        mx=Nm.maximum.reduce(Nm.ravel(Nm.where(dat<1e30,dat,mn)))
    scaled= (dat-mn)/(mx-mn)
    #clamp to 0-1
    return Nm.where(scaled<0,0,Nm.where(scaled>1,1,scaled))


# Read the dataset settings file
def load_settings(fname):
    
    try:
        f=open(fname,'r')
        # Create empty local and global dictionaries for 'exec' to use
        gdict={}
        ldict={}
        # Evaluate the settings file. settings will be stored in ldict
        exec f in gdict,ldict
    except IOError,err:
        raise SettingsError("Error reading settings file: "+fname)
    except Exception,err:
        raise SettingsError("Error parsing settings file: "+str(err))

    result=Settings()
    # Check that required items exists and are of correct type:
    # Variable map:
    try:
        result.name_to_canon=dict(ldict["variable_map"])
    except KeyError,err:
        raise SettingsError("Missing setting:"+str(err))
    except ValueError:
        raise SettingsError("Setting 'variable_map' must be a dictionary")

    # Clip plane list
    try:
        result.clip_planes=list(ldict["clip_planes"])
    except KeyError,err:
        raise SettingsError("Missing setting:"+str(err))
    except ValueError:
        raise SettingsError("Setting 'clip_planes' must be a list")

    # invert the variable->canonical lookup
    result.canon_to_name=[([k for k in result.name_to_canon.keys()
                            if result.name_to_canon[k]==v],v)
                          for v in Set(result.name_to_canon.values())]

    # Velocity scale (optional)
    try:
        result.velocity_scale=float(ldict.get("velocity_scale",1.0))
    except ValueError:
        raise SettingsError("The setting 'velocity_scale' must be a number.")
    
    return result
        

def load_profile(fname):
	f=open(fname)
	means=[]
	devs=[]
	for line in f:
                # skip comments
		if line[0]!='#':
                        #split line into depth, mean, stddev triples
			vals=map(float,line.split())
			assert len(vals)==3
			means.append(vals[1])
			devs.append(vals[2])
        # Convert means and stddevs to numeric arrays
	means_arr=Nm.array(means)
	devs_arr=Nm.array(devs)
        # The following lines 'broadcast' the 1-dimensional arrays, so that
        # they can be used as 3D arrays
	return means_arr[:,Nm.NewAxis,Nm.NewAxis], \
	       devs_arr[:,Nm.NewAxis,Nm.NewAxis]


def compute_partial(arr,axis,sentinel=1e30):
    if (axis!=0):
        arr=Nm.swapaxes(arr,0,axis)
    # One sided derivative, invalid if either sample is invalid
    # (use array multiplication instead of 'and', because numeric 'and' seems broke
    diff1=Nm.where((arr[1:,:,:]<sentinel) * (arr[:-1,:,:]<sentinel),
                   arr[1:,:,:]-arr[:-1,:,:],
                   sentinel)
    result=Nm.zeros(arr.shape,diff1.typecode())

    # one sided derivatives on faces
    result[0,:,:]=diff1[0,:,:]
    result[-1,:,:]=diff1[-1,:,:]

    # Two sided derivatives if both one sided derivatives are valid,
    # One sided if other derivative is sentinel,
    result[1:-1,:,:]=Nm.where(diff1[1:,:,:]<sentinel,diff1[1:,:,:],0)
    result[1:-1,:,:]+=Nm.where(diff1[:-1,:,:]<sentinel,diff1[:-1,:,:],0)
    result[1:-1,:,:]=Nm.where((diff1[1:,:,:]<sentinel) * (diff1[:-1,:,:]<sentinel),
                              result[1:-1]/2.0, result[1:-1])
    result[1:-1,:,:]=Nm.where((diff1[1:,:,:]>=sentinel) * (diff1[:-1,:,:]>=sentinel),
                              sentinel, result[1:-1])
    result=Nm.swapaxes(result,0,axis)
    return result
                              
                              
def compute_vertical_vorticity(u,v,delta_x=1.0,delta_y=1.0, sentinel=1e30):
    """Assumes roms/hops (z,y,x) ordering"""
    dvdx=compute_partial(v,2,sentinel)
    dudy=compute_partial(u,1,sentinel)
    return Nm.where((dvdx<sentinel) *(dudy<sentinel),
                    dvdx/delta_x-dudy/delta_y, sentinel)