#!/usr/bin/python
"""
Simple utility to scan through a set of ROMS or HOPS data, and compute the
depth-adjusted mean and standard deviation over all timesteps, which are then
written to stdout.
The inputs can be a list of single-timestep netcdf files, or 4D datasets,
as long as the time dimension is named 'time', and is the first dimension
for the variable.
"""

import os
import sys
from optparse import OptionParser

from Scientific.IO.NetCDF import *
import Numeric as Nm
import MA

#from MBARI_utils import *


default_missing_value = 1e30


def get_timesteps(filenames,varname):
    """
    get_timesteps is a generator which will yield successive timesteps
    from a list of files. The result type is a 3D masked array.
    The value to mask is read from the variable's 'missing_value' attribute,
    if one exists.
    """
    for fname in filenames:
        ncf=NetCDFFile(fname)
        ds=ncf.variables[varname]
        if hasattr(ds,"missing_value"):
            missing_value=ds.missing_value.astype(ds.typecode())
        else:
            missing_value=default_missing_value
            
        if "time" in ds.dimensions and len(ds.shape)==4:
            assert "time"==ds.dimensions[0]
            for ts in range(ds.shape[0]):
                yield MA.masked_values(ds[ts,:,:,:],missing_value)
        else:
            yield MA.masked_values(ds.getValue(),missing_value)
        del ncf
        del ds

def compute_timestep_means(ds):
    nlevels=ds.shape[0]
    # Initialize the result.
    profile=Nm.zeros(nlevels,Nm.Float64)

    for lev in range(nlevels):
        # Gather all the unmasked data at this level (or depth)
        valid_data=ds[lev,:,:].compressed()
        # Sum up all the data and find the average
        if valid_data.size()>0:
            profile[lev]=Nm.add.reduce(valid_data)/valid_data.size()
    return profile
        
def compute_global_means(filenames,varname):
    # Find the average of all the means for all the input datasets
    # This assumes that the number of unmasked datapoints is equal for
    # all timesteps
    total=None
    count=0
    for ds in get_timesteps(filenames,varname):
        profile=compute_timestep_means(ds)
        if total:
            total+=profile
        else:
            total=profile
        count+=1
    return total/count


def compute_global_stddevs(filenames,varname,means):
    nlevels=means.shape[0]
    accum=Nm.zeros(nlevels,Nm.Float64)
    counts=Nm.zeros(nlevels,Nm.Int)
    
    for dat in get_timesteps(filenames,varname):
        for lev in range(nlevels):
            # Gather all the unmasked data at this level
            valid_data=dat[lev,:,:].compressed()
            # Sum up the squared deviation from the mean
            accum[lev]+=Nm.add.reduce((valid_data-means[lev])**2)
            # Keep track of the number of data points at this level
            counts[lev]+=valid_data.size()
    return Nm.sqrt(accum/(counts-1))
            
        


def usage():
    print >>sys.stderr,"Usage: vertical_profiler <variable> <files...>"
    sys.exit(0)
    
if __name__=="__main__":

    # Basic no-frills command line.
    if len(sys.argv)<3: usage()
    varname=sys.argv[1]
    filenames=sys.argv[2:]

    # Extract the depth variable. ROMS and HOPS have different naming
    # conventions here.
    ncf=NetCDFFile(filenames[0])
    if "depth" in ncf.variables:
        depths=ncf.variables['depth'].getValue()
    else:
        depths=ncf.variables['ocean_depth'].getValue()

    # Do the computation.
    means=compute_global_means(filenames,varname)
    devs=compute_global_stddevs(filenames,varname,means)

    # Print results to stdout.
    print "# depth\t  mean\t  stddev"
    print "# -----\t  ----\t  ------"
    for depth,mean,dev in zip(depths,means,devs):
        print "  %4g\t  %0.4g\t  %0.4g" %(depth,mean,dev)