#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-

import os, sys
import numpy as np
from optparse import OptionParser

from pywrf.post import Wrfout
from pywrf.post.utils import choose, create_ncfile, create_ncvar

def compute_iso0c(wrfout, ncout_filename):

    # Compute T°C
    tc  = wrfout["tc"]

    # Compute height
    z   = wrfout["z"]
    z0c = np.ma.masked_array(np.zeros_like(z[:,0]))
    z0c.set_fill_value(-999.)

    # Compute level of first inversion
    ntimes  = tc.shape[0]
    nlevels = tc.shape[1]

    # Interpolation linéaire de (z0, z1, z2) avec les poids calculés par (t05,t15)
    #  on a : 0 ------------------------------------> zmax
    #              z0  z05    z1     z*   z15     z2
    #              |   t05>0  |      0    t15<0   |
    #         0 ------------------------------------> zmax
    #              z0  z05      z*     z1 z15     z2
    #    ou        |   t05>0    0      |  t15<0   |
    #
    #    ie:       z05      z*     z15
    #              t05>0    0      t15<0
    #
    #  z05 = (z0+z1)/2
    #  z15 = (z1+z2)/2
    #  t(z*)  := 0.0
    #  t(z05) := t05 > 0
    #  t(z15) := t15 < 0
    #  (z*-z05) / (z15-z05) = (t(z*)-t(z05))/(t(z15)-t(z05))
    # => z* = z05 - (z15-z05) * t(z05)/(t(z15)-t(z05))

    for n in range(ntimes):

        # Restrict data to time step
        t3d = np.ma.masked_array(tc[n])
        z3d = np.ma.masked_array(z[n])

        # Compute level of first inversion
        tsup = t3d > 0.
        amin = nlevels - tsup[::-1].argmax(axis=0)
        amask = amin == nlevels
        amin[amask] = 0

        t05 = choose(amin-1, t3d) ; t05[amask] = np.ma.masked     # t05 = last  positive temp. (bottom->up)
        t15 = choose(amin,   t3d) ; t15[amask] = np.ma.masked     # t15 = first negative temp.

        z0 = choose(amin-1, z3d) ; z0[amask] = np.ma.masked
        z1 = choose(amin,   z3d) ; z1[amask] = np.ma.masked
        z2 = choose(amin+1, z3d) ; z2[amask] = np.ma.masked

        z05  = 0.5*(z0+z1)
        z15  = 0.5*(z1+z2)
        z0c[n] = z05 - (z15-z05)*t05/(t15-t05)

    dimensions = wrfout.variables["T2"].dimensions
    attributes = wrfout.variables["T2"].get_attributes()
    attributes["units"] = "m"
    attributes["description"] = "0C isotherm level height"
    attributes["_FillValue"] = -999.

    fout = create_ncfile(ncout_filename, model=wrfout)
    var_z0c = create_ncvar(fout, "ISO0C", z0c.dtype, dimensions, attributes)
    var_z0c[:] = z0c
    fout.close()

def get_options():

    usage = "usage: %prog [options] WRF_FILE"
    parser = OptionParser(usage)

    (options, args) = parser.parse_args()

    if len(args) != 1:
        parser.print_help()
        sys.exit("\nERROR: incorrect number of arguments")

    wrf_filename = args[0]

    return wrf_filename, options

if __name__=="__main__":

    wrf_filename, options = get_options()
    wrfile = Wrfout(wrf_filename)
    compute_iso0c(wrfile, "iso0c.nc")