from pylab import *

N=30

def fvary(g, gmax):
    return where(g>gmax, (gmax -g)/(gmax-1), 0)

g = linspace(0,1,N)
tau = ones(N)
w0 = ones(N) * .90

dtau = lambda tau,w0,f: tau * (1. - w0 * f)
dw0 = lambda w0, f: w0 * (1.-f) / (1.-f * w0)
dg = lambda g,f: (g-f) / (1.-f)

figure(1,figsize=(2*9,2*3))
clf()
for i in arange(1,16):
    if i==1:
        ng = g
        nw0 = w0
        ntau = tau
    f = ng**2
    ntau = dtau(ntau, nw0, f)
    nw0 = dw0(nw0, f)
    ng = dg(ng, f)

    if i<3 or i%5==0:
        subplot(131); title('tau'); plot(g, ntau, label='iter {} (f=g**2)'.format(i))
        subplot(132); title('w0'); plot(g, nw0, label='iter {} (f=g**2)'.format(i))
        subplot(133); title('g'); plot(g, ng, label='iter {} (f=g**2)'.format(i))

subplot(131); plot(g, dtau(tau, w0, fvary(g, .65)), 'o-', label='gmax .65'.format(i))
subplot(132); plot(g, dw0 (w0 , fvary(g, .65)    ), 'o-', label='gmax .65'.format(i))
subplot(133); plot(g, dg  (g  , fvary(g, .65)    ), 'o-', label='gmax .65'.format(i))

subplot(131); plot(g, dtau(tau, w0, fvary(g, .5)), 'o-', label='gmax .5'.format(i))
subplot(132); plot(g, dw0 (w0 , fvary(g, .5)    ), 'o-', label='gmax .5'.format(i))
subplot(133); plot(g, dg  (g  , fvary(g, .5)    ), 'o-', label='gmax .5'.format(i))

subplot(131); plot(g, dtau(tau, w0, fvary(g, .0)), 'o-', label='gmax .0'.format(i))
subplot(132); plot(g, dw0 (w0 , fvary(g, .0)    ), 'o-', label='gmax .0'.format(i))
subplot(133); plot(g, dg  (g  , fvary(g, .0)    ), 'o-', label='gmax .0'.format(i))

for i in arange(1,4):
    sub = subplot(1,3,i)
    sub.legend(loc='best', framealpha=.5)
    sub.set_xlabel('original g')

plt.savefig('iterative_classic_delta_scaling.pdf')
plt.savefig('iterative_classic_delta_scaling.png')


import numpy as np
import glob
search_files = lambda ident: zip(glob.glob("/project/meteo/work/Fabian.Jakub/uvspec/i3rc1/{}/output/*.flx.spc.npy".format(ident)), glob.glob("/project/meteo/work/Fabian.Jakub/uvspec/i3rc1/{}/output/*.abs.spc.npy".format(ident)))
ffile3d, afile3d = search_files('i3rc1.sza60.pprts_3d')[0]
f3d = np.load(ffile3d)
a3d = np.load(afile3d)

others = search_files('i3rc1.sza60.pprts_3_10_gmax*')

import load_librad as N
edir  = lambda f: f[0]
edn   = lambda f: f[1]
eup   = lambda f: f[2]

edir0 = lambda f: edir(f)[0]
edn0  = lambda f: edn(f)[0]
eup0  = lambda f: eup(f)[0]

ident = lambda f: f.replace("/project/meteo/work/Fabian.Jakub/uvspec/i3rc1/","").split("/")[0]

gmax_list = []
err_edir_by_gmax = []
err_edn_by_gmax = []
err_eup_by_gmax = []
err_abso_by_gmax = []

for fnames in others:
    ffile, afile = fnames
    if ffile == ffile3d:
        continue
    fo = np.load(ffile)
    ao = np.load(afile)
    err_edir = N.rmse(edir0(fo), edir0(f3d))
    err_edn  = N.rmse(edn0 (fo), edn0 (f3d))
    err_eup  = N.rmse(eup0 (fo), eup0 (f3d))
    err_abso = N.rmse(ao, a3d)
    bias_srfc = N.bias(edir0(fo)+edn0 (fo), edir0(f3d)+edn0 (f3d))
    bias_atm = N.bias(ao, a3d)

    ident3d, idento = [ ident(f) for f in (ffile3d, ffile) ]

    gmax = float(idento[idento.find('gmax')+4:])
    print(idento, gmax, bias_srfc)

    gmax_list.append(gmax)
    err_edir_by_gmax.append(err_edir[1])
    err_edn_by_gmax.append(err_edn[1])
    err_eup_by_gmax.append(err_eup[1])
    err_abso_by_gmax.append(err_abso[1])

gmax_list, err_edir_by_gmax, err_edn_by_gmax, err_eup_by_gmax, err_abso_by_gmax = zip(*sorted(zip(gmax_list, err_edir_by_gmax, err_edn_by_gmax, err_eup_by_gmax, err_abso_by_gmax)))

figure(2,figsize=(14,7))
clf()
title("Dependence of error on gmax for i3rc scene with Tenstream 3_10 sza 60deg")
plt.plot(gmax_list, err_edir_by_gmax, 'o-', label='edir')
plt.plot(gmax_list, err_edn_by_gmax, 'o-', label='edn')
plt.plot(gmax_list, err_eup_by_gmax, 'o-', label='eup')
plt.plot(gmax_list, err_abso_by_gmax, 'o-', label='abso')
legend(loc="best", framealpha=.5)
xlabel("value of `g_max`")
ylabel("rel. RMSE [%]")
ylim(ymin=10)
xlim(-.05, .8)
plt.savefig('err_as_function_of_gmax.pdf')
plt.savefig('err_as_function_of_gmax.png')