#!/usr/bin/env python
import sys
import os
import glob
from scipy.io.netcdf import netcdf_file
import numpy as np
from darwin import iofmt
from subprocess import call

#srcdir = '/home/jahn/darwin3/Darwin2/darwin2/verification'
ftmpl1 = '{exp}/results/var_trcstat_ptracer{itr}_mean.txt'

args = sys.argv[1:]
exp, = args
fname = exp
dir,name = os.path.split(fname)
dir,_ = os.path.split(dir)
srcdir,exp = os.path.split(dir)
_,ext = os.path.splitext(name)
_,ext2 = os.path.splitext(_)
ext = ext2 + ext

names0 = np.genfromtxt('{}/{}/names'.format(srcdir, exp), object)
names = np.genfromtxt('{}/names'.format(exp), object)
namel0 = names0.tolist()
namel = names.tolist()

ii = []
ff = []
for name in namel:
    nameorig = name
    if name in namel0:
        ii.append(namel0.index(name))
        ff.append(1.)
        continue
    if name[-1] == 'C':
        name = name[:-1] + 'P'
        if name in namel0:
            ii.append(namel0.index(name))
            ff.append(120.)
            continue
    if name.startswith('c'):
        name = 'Phy' + str(int(name[1:]))
        if name in namel0:
            ii.append(namel0.index(name))
            ff.append(120.)
            continue
    if name.startswith('Phy'):
        name = 'Phy{:02d}'.format(int(name[3:]))
        if name in namel0:
            ii.append(namel0.index(name))
            ff.append(120.)
            continue
    if name.startswith('zc'):
        name = 'ZOO' + str(int(name[2:])) + 'P'
        if name in namel0:
            ii.append(namel0.index(name))
            ff.append(120.)
            continue
    ii.append(0)
    ff.append(0.)
    sys.stderr.write('Tracer not found: {}\n'.format(nameorig))

for i,name in enumerate(namel):
    print '{i:02d} {0:<6s} <- {1:6.2f}*{2}'.format(name, ff[i], ff[i] and namel0[ii[i]] or '', i=i+1)

i0s = ii
f0s = ff
n = len(namel)

#fname = '{}/{}/run/ptracers.nc'.format(srcdir, exp)
#if os.path.exists(fname):
#    ext = '.nc'
#    fnames = [fname]
#else:
#    ext = '.t001.nc'
#    fnames = glob.glob('{}/{}/run/ptracers.*{}'.format(srcdir, exp, ext))
#    print '\n'.join(fnames)
#    if len(fnames) != 1:
#        sys.exit('Found {} ptracers files.'.format(len(fnames)))
#    fname, = fnames

print exp

gridname = '{}/{}/run/grid{}'.format(srcdir, exp, ext)
print gridname
gridf = netcdf_file('{}/{}/run/grid{}'.format(srcdir, exp, ext))
hFacC = gridf.variables['HFacC'][:]
Ac = gridf.variables['rA'][:]
drf = gridf.variables['drF'][:][:,None,None]
vol = (hFacC*Ac*drf).sum()

print fname
nc = netcdf_file(fname)
for i,i0 in enumerate(i0s):
    p = nc.variables[namel0[i0]][:]
    av = (p*Ac*hFacC*drf).sum(-1).sum(-1).sum(-1)/vol
    av *= f0s[i]
    with open(ftmpl1.format(exp=exp, itr=iofmt(i+1)), 'w') as f:
        for t,v in enumerate(av):
            f.write('{:6d}\t{:25.16E}\n'.format(t+1, v))

