# Manual Coarsen
# --------------
# Author : Riley X. Brady
# Date   : 11/14/2018
# --------------
# The pyAMG method of coarsening our particles traverses in 3D, rather than
# just removing full 2D columns. This method simply strides through the
# lat and lon list, removing full columns.
#
# NOTE: Run make_particle_file.py first without any coarsening turned on.
# This will take in that particles.nc file and coarsen it.
import numpy as np
import xarray as xr
import argparse


def xyz_to_lat_lon(x, y, z):
    x, y, z = np.asarray(x), np.asarray(y), np.asarray(z)
    plat = np.arcsin(z / np.sqrt(x**2 + y**2 + z**2))
    plong = np.arctan2(y, x)
    # longitude adjustment
    plong[plong < 0.0] = 2*np.pi + plong[plong < 0.0]
    return plong * 180./np.pi, \
           plat * 180./np.pi


def main(init, stride, output):
    ds = xr.open_dataset(init)
    plong, plat = xyz_to_lat_lon(ds.xParticle, ds.yParticle, ds.zParticle)
    ds.coords['particleLat'] = (('Time', 'nParticles'), plat)
    ds.coords['particleLon'] = (('Time', 'nParticles'), plong)
    un_lon, un_lat = np.unique(ds.particleLon), np.unique(ds.particleLat)
    # lat filter
    latlist = un_lat[0::stride]
    particleList = []
    for i, lat in enumerate(ds.squeeze().particleLat.values):
        if lat in latlist:
            particleList.append(i)
    ds_lat_trimmed = ds.sel(nParticles=particleList)
    # lon filter
    lonlist = un_lon[0::stride]
    particleList = []
    for i, lon in enumerate(ds_lat_trimmed.squeeze().particleLon.values):
        if lon in lonlist:
            particleList.append(i)
    ds_full_trimmed = ds_lat_trimmed.sel(nParticles=particleList)
    ds_full_trimmed.to_netcdf(output)


if __name__ == '__main__':
    ap = argparse.ArgumentParser(
            description="Coarsens particle init file horizontally by a determined stride.")
    ap.add_argument('-i', '--init', required=True, type=str,
            help="Particle init file (netCDF)")
    ap.add_argument('-s', '--stride', required=True, type=int,
            help="Number of particles to skip in lat and lon")
    ap.add_argument('-o', '--output', required=True, type=str,
            help="Coarsened particle output file (.nc at end)")
    args = vars(ap.parse_args())
    init = args['init']
    stride = args['stride']
    output = args['output']
    main(init, stride, output)