#!/usr/bin/env python
"""
Mesh and CVM extraction
"""
import numpy as np
import os, pyproj, cst

# parameters
if 1:
    origin = (-117.8, -117.7), (34.6, 34.6)
    rotate = 45.0
    translate = 405000.0, 202500.0
    xsmooth = 64000.0, 114000.0
    delta =   40.0,   40.0,  40.0; shape = 20250, 10125, 2125; nproc = 65536
    delta =  640.0,  640.0,  20.0; shape =  1266,  633,   550; nproc = 256
    delta =  800.0,  800.0,  20.0; shape =  1013,  507,     1; nproc = 2
    delta =  640.0,  640.0,  20.0; shape =  1266,  633,     1; nproc = 2
    delta =   80.0,   80.0,  80.0; shape = 10125, 5062,     1; nproc = 4
    offset = 0.0, 0.0
    offset = 20.0, 20.0
else:
    origin = -122.4, 35.8
    rotate = 46.0
    translate = 0.0, 0.0
    offset = 0.0, 0.0
    xsmooth = 54000.0, 104000.0
    delta = 1000.0, 1000.0, 100.0; shape =   800,  400,     1; nproc = 1

# smoothing region
jsmooth = int(xsmooth[0] / delta[0]), int(xsmooth[1] / delta[0])

# projection
projection = dict(proj='utm', zone=11, ellps='WGS84')
transform = dict(origin=origin, rotate=rotate, translate=translate)
proj = pyproj.Proj(**projection)
proj = cst.coord.Transform(proj, **transform)

# bounds
x = 0.0, delta[0] * (shape[0] + 1)
y = 0.0, delta[1] * (shape[1] + 1)
bounds = x, y

# corners
x = x[0], x[1], x[1], x[0]
y = y[0], y[0], y[1], y[1]
x, y = proj(x, y, inverse=True)
corners = x, y

# box
x = np.arange(shape[0] + 1) * delta[0]
y = np.arange(shape[1] + 1) * delta[1]
y, x = np.meshgrid(y, x)
x = np.concatenate([ x[:,0], x[-1,1:], x[-2::-1,-1], x[0,-2::-1] ])
y = np.concatenate([ y[:,0], y[-1,1:], y[-2::-1,-1], y[0,-2::-1] ])
x, y = proj(x, y, inverse=True)
extent = (x.min(), x.max()), (y.min(), y.max())

# topography
topo, extent = cst.data.topo(extent)
z = cst.coord.interp2(extent, topo, (x, y))
box = x, y, z

# trace
x, y = np.loadtxt('src/trace.txt').T
x, y = cst.data.densify(x, y, delta=(0.00001*delta[0]))
z = cst.coord.interp2(extent, topo, (x, y))
trace = x, y, z
x, y = proj(x, y)
trace_xyz = 0.001 * x, 0.001 * y, 0.001 * z

# mesh
x = np.arange(shape[0]) * delta[0] + offset[0]
y = np.arange(shape[1]) * delta[1] + offset[1]
y, x = np.meshgrid(y, x)
x, y = proj(x, y, inverse=True)
z = cst.coord.interp2(extent, topo, (x, y))
grid = x, y, z

# metadata
meta = dict(
    delta = delta,
    shape = shape,
    extent = extent,
    projection = projection,
    transform = transform,
    jsmooth = jsmooth,
    dtype = np.dtype('f').str,
)

# generate mesh if not importing
if __name__ == '__main__':

    # save data
    path = 'data/%.0f/' % delta[0]
    cst.util.save(path + 'meta.py', meta)
    np.savetxt(path + 'corners.txt', np.array(corners, 'f').T)
    np.savetxt(path + 'box.txt', np.array(box, 'f').T)
    np.savetxt(path + 'trace.txt', np.array(trace, 'f').T)
    np.savetxt(path + 'trace-xyz.txt', np.array(trace_xyz, 'f').T)
    x.T.astype('f').tofile(path + 'lon.bin')
    y.T.astype('f').tofile(path + 'lat.bin')
    z.T.astype('f').tofile(path + 'topo.bin')

    # mapdata
    for kind in 'coastlines', 'borders':
        x, y = cst.data.mapdata(kind, 'high', extent, 50.0)
        x -= 360.0
        z = cst.coord.interp2(extent, topo, (x, y))
        x, y = proj(x, y)
        x, y, i = cst.data.clipdata(x, y, bounds)
        z = z[i]
        np.savetxt(path + kind + '-xyz.txt', 0.001 * np.array([x,y,z]).T)

    if 0:
        # setup CVM
        n = shape[0] * shape[1] * shape[2]
        job = cst.cvm.stage(nsample=n, nproc=nproc)
        path1 = job.rundir + os.sep

        # write CVM input files
        x, y = np.array(grid[:2], 'f')
        z = np.arange(0, shape[2]) * delta[2]
        f1 = open(path1 + 'lon.bin', 'wb')
        f2 = open(path1 + 'lat.bin', 'wb')
        f3 = open(path1 + 'dep.bin', 'wb')
        for i in range(z.size):
            x.T.tofile(f1)
            y.T.tofile(f2)
        for i in range(z.size):
            x.fill(z[i])
            x.T.tofile(f3)
        f1.close()
        f2.close()
        f3.close()

        # run job and smooth
        if shape[-1] == 1:
            cst.cvm.launch(job)
            j1, j2 = jsmooth
            n = j2 - j1 + 1
            w = np.arange(n)[:,None] / (n-1.0)
            nn = shape[:2]
            n = nn[0] * nn[1]
            for f in 'vs.bin', 'vp.bin', 'rho.bin':
                s = np.fromfile(path1 + f, 'f', n).reshape(nn[::-1]).T
                s[j1:j2+1] = (1.0-w) * s[j1:j1+1] + w * s[j2:j2+1]
                s.T.tofile(path + f)