import sys
import numpy as np
from matplotlib import pyplot as pl
from rw import WriteGTiff

fn = '../pozo-steep-vegetated-pcl.npy'
pts = np.load(fn)
x, y, z, c = pts[:, 0], pts[:, 1], pts[:, 2], pts[:, 5]
ix = (0.2 * (x - x.min())).astype('int')
iy = (0.2 * (y - y.min())).astype('int')

shape = (100, 100)
xb = np.arange(shape[1]+1)
yb = np.arange(shape[0]+1)
fg, ax = pl.subplots(ncols = 2, nrows = 2,
            figsize = (10.24, 10.24),
            sharex = True, sharey = True)

uc = (2, 5)
for j in range(len(uc)):
    print('Class %i' % uc[j])
    b = c == uc[j]
    cx, cy, cz = ix[b], iy[b], z[b]

    mean = np.zeros(shape)
    stdr = np.zeros(shape)
    for i in range(shape[0]):
        print('% 3d%%' % i)
        for k in range(shape[1]):
            b = (cy == i) * (cx == k)
            mean[i, k] = cz[b].mean()
            stdr[i, k] = cz[b].std()

    fname = 'pozo_5m_dem_mean_cl%i.tif' % uc[j]
    WriteGTiff(fname, mean, x.min(), y.min()+500, step = 5)
    np.save('pozo_5m_dem_mean_cl%i.npy' % uc[j], mean)
    np.save('pozo_5m_dem_stdr_cl%i.npy' % uc[j], stdr)
    ax[0, j].set_title('Class %i' % uc[j])
    im = ax[0, j].pcolormesh(xb, yb,
                  np.ma.masked_invalid(mean),
                  cmap = pl.cm.viridis_r)
    cb = fg.colorbar(im, ax = ax[0, j])
    cb.set_label('Mean elevation [m]')

    im = ax[1, j].pcolormesh(xb, yb,
                  np.ma.masked_invalid(stdr),
                  cmap = pl.cm.magma_r)
    cb = fg.colorbar(im, ax = ax[1, j])
    cb.set_label('Elevation STD')
    ax[0, j].set_aspect('equal')
    ax[1, j].set_aspect('equal')

pl.savefig('%s.png' % sys.argv[0][:-3])

[(7, 'numpy.load', 'np.load', True, 'import numpy as np\n'), (13, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (14, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (15, 'matplotlib.pyplot.subplots', 'pl.subplots', True, 'from matplotlib import pyplot as pl\n'), (53, 'matplotlib.pyplot.savefig', 'pl.savefig', True, 'from matplotlib import pyplot as pl\n'), (25, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (26, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (36, 'numpy.save', 'np.save', True, 'import numpy as np\n'), (37, 'numpy.save', 'np.save', True, 'import numpy as np\n'), (40, 'numpy.ma.masked_invalid', 'np.ma.masked_invalid', True, 'import numpy as np\n'), (46, 'numpy.ma.masked_invalid', 'np.ma.masked_invalid', True, 'import numpy as np\n')]