import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import click
import numba


def prepare_data(data_pd, parameter):
    lon_set = set(data_pd["lon"])
    lat_set = set(data_pd["lat"])
    dep_set = set(data_pd["dep"])

    lon_list = sorted(lon_set)
    lat_list = sorted(lat_set)
    dep_list = sorted(dep_set)

    lon_mesh, lat_mesh, dep_mesh = np.meshgrid(
        lon_list, lat_list, dep_list, indexing="ij")
    dx, dy, dz = np.shape(lon_mesh)
    value_mesh = np.zeros_like(lon_mesh)
    x_mesh = np.zeros_like(lon_mesh)
    y_mesh = np.zeros_like(lon_mesh)
    z_mesh = np.zeros_like(lon_mesh)
    r_mesh = np.zeros_like(lon_mesh)
    for i in range(dx):
        for j in range(dy):
            for k in range(dz):
                x_mesh[i, j, k], y_mesh[i, j, k], z_mesh[i, j, k], r_mesh[i, j, k] = lld2xyzr(
                    lat_mesh[i, j, k], lon_mesh[i, j, k], dep_mesh[i, j, k])

    for index, row in data_pd.iterrows():
        i = int(round((row.lon-lon_list[0])/(lon_list[1]-lon_list[0]), 0))
        j = int(round((row.lat-lat_list[0])/(lat_list[1]-lat_list[0]), 0))
        k = int(round((row.dep-dep_list[0])/(dep_list[1]-dep_list[0]), 0))
        value_mesh[i, j, k] = row[parameter]

    return x_mesh, y_mesh, z_mesh, value_mesh


def get_value(data_pd, lat, lon, dep, parameter):
    return data_pd.loc[(data_pd.lat == lat) & (data_pd.lon == lon) & (data_pd.dep == dep)][parameter].values[0]


@numba.njit()
def lld2xyzr(lat, lon, dep):
    R_EARTH_KM = 6371.0
    r = (R_EARTH_KM-dep)/R_EARTH_KM
    theta = 90-lat
    phi = lon

    z = r*cosd(theta)
    h = r*sind(theta)
    x = h*cosd(phi)
    y = h*sind(phi)

    return (x, y, z, r)


@numba.njit()
def cosd(x):
    return np.cos(np.deg2rad(x))


@numba.njit()
def sind(x):
    return np.sin(np.deg2rad(x))


# def get_value_func(x_mesh, y_mesh, z_mesh, value_mesh):
#     value_func = RegularGridInterpolator(
#         (x_mesh, y_mesh, z_mesh), value_mesh, method="nearest")
#     return value_func


@numba.njit()
def interp_value(lat, lon, dep, x_mesh, y_mesh, z_mesh, value_mesh):
    x, y, z, _ = lld2xyzr(lat, lon, dep)
    distance2 = (x_mesh-x)**2+(y_mesh-y)**2+(z_mesh-z)**2
    mindistance2 = np.min(distance2)
    coors = np.where(distance2 == mindistance2)
    value = value_mesh[coors[0][0], coors[1][0], coors[2][0]]
    return value


def generate_vertical_profile_grids(lon_list, lat_list, dep_list, hnpts, vnpts):
    lons = np.linspace(lon_list[0], lon_list[1], hnpts)
    lats = np.linspace(lat_list[0], lat_list[1], hnpts)
    deps = np.linspace(dep_list[0], dep_list[1], vnpts)
    return lons, lats, deps


@click.command()
@click.option('--lon1', required=True, type=float, help="lon1")
@click.option('--lon2', required=True, type=float, help="lon2")
@click.option('--lat1', required=True, type=float, help="lat1")
@click.option('--lat2', required=True, type=float, help="lat2")
@click.option('--dep1', required=True, type=float, help="dep1")
@click.option('--dep2', required=True, type=float, help="dep2")
@click.option('--data', required=True, type=str, help="the pickle file")
@click.option('--parameter', required=True, type=str, help="physicial parameter to plot")
@click.option('--hnpts', required=True, type=int, help="horizontal npts")
@click.option('--vnpts', required=True, type=int, help="vertical npts")
def main(lon1, lon2, lat1, lat2, dep1, dep2, data, parameter, hnpts, vnpts):
    lon_list = [lon1, lon2]
    lat_list = [lat1, lat2]
    dep_list = [dep1, dep2]
    data_pd_raw = pd.read_pickle(data)

    # data_pd is too big
    minlon = min(lon1, lon2)
    maxlon = max(lon1, lon2)
    minlat = min(lat1, lat2)
    maxlat = max(lat1, lat2)
    mindep = min(dep1, dep2)
    maxdep = max(dep1, dep2)
    data_pd = data_pd_raw.loc[(data_pd_raw.lat <= maxlat) & (
        data_pd_raw.lat >= minlat) & (data_pd_raw.lon < maxlon) & (data_pd_raw.lon > minlon) & (data_pd_raw.dep >= mindep) & (data_pd_raw.dep <= maxdep)]

    x_mesh, y_mesh, z_mesh, value_mesh = prepare_data(data_pd, parameter)
    lons_plot, lats_plot, deps_plot = generate_vertical_profile_grids(
        lon_list, lat_list, dep_list, hnpts, vnpts)
    values = np.zeros((hnpts, vnpts))
    for ih in range(hnpts):
        for iv in range(vnpts):
            values[ih, iv] = interp_value(
                lats_plot[ih], lons_plot[ih], deps_plot[iv], x_mesh, y_mesh, z_mesh, value_mesh)
            # print(lats_plot[ih], lons_plot[ih], deps_plot[iv], values[ih, iv])

    # plotting part
    plt.figure()
    mesh_plot_lat, mesh_plot_dep = np.meshgrid(
        lats_plot,  deps_plot, indexing="ij")

    # get vmin and vmax
    vmin_round = round(np.min(values), 2)
    if(vmin_round < np.min(values)):
        vmin = vmin_round
    else:
        vmin = vmin_round-0.01
    vmax_round = round(np.max(values), 2)
    if(vmax_round > np.max(values)):
        vmax = vmax_round
    else:
        vmax = vmax_round+0.01
    print(vmin, vmax, np.max(values), np.min(values), vmin_round, vmax_round)
    plt.contourf(mesh_plot_lat, mesh_plot_dep,
                 values,  101, cmap=plt.cm.seismic_r)
    v = np.arange(vmin, vmax, 0.01)
    plt.colorbar(ticks=v, label="perturbation")
    plt.gca().invert_yaxis()
    plt.xlabel(
        f"latitude(°) between (lon: {lon1}°, lat: {lat1}°) and (lon: {lon2}°, lat: {lat2}°)")
    plt.ylabel("depth(km)")
    plt.show()


if __name__ == "__main__":
    main()

[((1443, 1455), 'numba.njit', 'numba.njit', ([], {}), '()\n', (1453, 1455), False, 'import numba\n'), ((1691, 1703), 'numba.njit', 'numba.njit', ([], {}), '()\n', (1701, 1703), False, 'import numba\n'), ((1753, 1765), 'numba.njit', 'numba.njit', ([], {}), '()\n', (1763, 1765), False, 'import numba\n'), ((2009, 2021), 'numba.njit', 'numba.njit', ([], {}), '()\n', (2019, 2021), False, 'import numba\n'), ((2636, 2651), 'click.command', 'click.command', ([], {}), '()\n', (2649, 2651), False, 'import click\n'), ((2653, 2715), 'click.option', 'click.option', (['"""--lon1"""'], {'required': '(True)', 'type': 'float', 'help': '"""lon1"""'}), "('--lon1', required=True, type=float, help='lon1')\n", (2665, 2715), False, 'import click\n'), ((2717, 2779), 'click.option', 'click.option', (['"""--lon2"""'], {'required': '(True)', 'type': 'float', 'help': '"""lon2"""'}), "('--lon2', required=True, type=float, help='lon2')\n", (2729, 2779), False, 'import click\n'), ((2781, 2843), 'click.option', 'click.option', (['"""--lat1"""'], {'required': '(True)', 'type': 'float', 'help': '"""lat1"""'}), "('--lat1', required=True, type=float, help='lat1')\n", (2793, 2843), False, 'import click\n'), ((2845, 2907), 'click.option', 'click.option', (['"""--lat2"""'], {'required': '(True)', 'type': 'float', 'help': '"""lat2"""'}), "('--lat2', required=True, type=float, help='lat2')\n", (2857, 2907), False, 'import click\n'), ((2909, 2971), 'click.option', 'click.option', (['"""--dep1"""'], {'required': '(True)', 'type': 'float', 'help': '"""dep1"""'}), "('--dep1', required=True, type=float, help='dep1')\n", (2921, 2971), False, 'import click\n'), ((2973, 3035), 'click.option', 'click.option', (['"""--dep2"""'], {'required': '(True)', 'type': 'float', 'help': '"""dep2"""'}), "('--dep2', required=True, type=float, help='dep2')\n", (2985, 3035), False, 'import click\n'), ((3037, 3108), 'click.option', 'click.option', (['"""--data"""'], {'required': '(True)', 'type': 'str', 'help': '"""the pickle file"""'}), "('--data', required=True, type=str, help='the pickle file')\n", (3049, 3108), False, 'import click\n'), ((3110, 3203), 'click.option', 'click.option', (['"""--parameter"""'], {'required': '(True)', 'type': 'str', 'help': '"""physicial parameter to plot"""'}), "('--parameter', required=True, type=str, help=\n    'physicial parameter to plot')\n", (3122, 3203), False, 'import click\n'), ((3200, 3272), 'click.option', 'click.option', (['"""--hnpts"""'], {'required': '(True)', 'type': 'int', 'help': '"""horizontal npts"""'}), "('--hnpts', required=True, type=int, help='horizontal npts')\n", (3212, 3272), False, 'import click\n'), ((3274, 3344), 'click.option', 'click.option', (['"""--vnpts"""'], {'required': '(True)', 'type': 'int', 'help': '"""vertical npts"""'}), "('--vnpts', required=True, type=int, help='vertical npts')\n", (3286, 3344), False, 'import click\n'), ((369, 425), 'numpy.meshgrid', 'np.meshgrid', (['lon_list', 'lat_list', 'dep_list'], {'indexing': '"""ij"""'}), "(lon_list, lat_list, dep_list, indexing='ij')\n", (380, 425), True, 'import numpy as np\n'), ((452, 470), 'numpy.shape', 'np.shape', (['lon_mesh'], {}), '(lon_mesh)\n', (460, 470), True, 'import numpy as np\n'), ((488, 511), 'numpy.zeros_like', 'np.zeros_like', (['lon_mesh'], {}), '(lon_mesh)\n', (501, 511), True, 'import numpy as np\n'), ((525, 548), 'numpy.zeros_like', 'np.zeros_like', (['lon_mesh'], {}), '(lon_mesh)\n', (538, 548), True, 'import numpy as np\n'), ((562, 585), 'numpy.zeros_like', 'np.zeros_like', (['lon_mesh'], {}), '(lon_mesh)\n', (575, 585), True, 'import numpy as np\n'), ((599, 622), 'numpy.zeros_like', 'np.zeros_like', (['lon_mesh'], {}), '(lon_mesh)\n', (612, 622), True, 'import numpy as np\n'), ((636, 659), 'numpy.zeros_like', 'np.zeros_like', (['lon_mesh'], {}), '(lon_mesh)\n', (649, 659), True, 'import numpy as np\n'), ((2209, 2226), 'numpy.min', 'np.min', (['distance2'], {}), '(distance2)\n', (2215, 2226), True, 'import numpy as np\n'), ((2239, 2274), 'numpy.where', 'np.where', (['(distance2 == mindistance2)'], {}), '(distance2 == mindistance2)\n', (2247, 2274), True, 'import numpy as np\n'), ((2448, 2492), 'numpy.linspace', 'np.linspace', (['lon_list[0]', 'lon_list[1]', 'hnpts'], {}), '(lon_list[0], lon_list[1], hnpts)\n', (2459, 2492), True, 'import numpy as np\n'), ((2504, 2548), 'numpy.linspace', 'np.linspace', (['lat_list[0]', 'lat_list[1]', 'hnpts'], {}), '(lat_list[0], lat_list[1], hnpts)\n', (2515, 2548), True, 'import numpy as np\n'), ((2560, 2604), 'numpy.linspace', 'np.linspace', (['dep_list[0]', 'dep_list[1]', 'vnpts'], {}), '(dep_list[0], dep_list[1], vnpts)\n', (2571, 2604), True, 'import numpy as np\n'), ((3524, 3544), 'pandas.read_pickle', 'pd.read_pickle', (['data'], {}), '(data)\n', (3538, 3544), True, 'import pandas as pd\n'), ((4172, 4196), 'numpy.zeros', 'np.zeros', (['(hnpts, vnpts)'], {}), '((hnpts, vnpts))\n', (4180, 4196), True, 'import numpy as np\n'), ((4503, 4515), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (4513, 4515), True, 'import matplotlib.pyplot as plt\n'), ((4551, 4599), 'numpy.meshgrid', 'np.meshgrid', (['lats_plot', 'deps_plot'], {'indexing': '"""ij"""'}), "(lats_plot, deps_plot, indexing='ij')\n", (4562, 4599), True, 'import numpy as np\n'), ((5009, 5087), 'matplotlib.pyplot.contourf', 'plt.contourf', (['mesh_plot_lat', 'mesh_plot_dep', 'values', '(101)'], {'cmap': 'plt.cm.seismic_r'}), '(mesh_plot_lat, mesh_plot_dep, values, 101, cmap=plt.cm.seismic_r)\n', (5021, 5087), True, 'import matplotlib.pyplot as plt\n'), ((5114, 5141), 'numpy.arange', 'np.arange', (['vmin', 'vmax', '(0.01)'], {}), '(vmin, vmax, 0.01)\n', (5123, 5141), True, 'import numpy as np\n'), ((5146, 5189), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ([], {'ticks': 'v', 'label': '"""perturbation"""'}), "(ticks=v, label='perturbation')\n", (5158, 5189), True, 'import matplotlib.pyplot as plt\n'), ((5223, 5329), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['f"""latitude(°) between (lon: {lon1}°, lat: {lat1}°) and (lon: {lon2}°, lat: {lat2}°)"""'], {}), "(\n    f'latitude(°) between (lon: {lon1}°, lat: {lat1}°) and (lon: {lon2}°, lat: {lat2}°)'\n    )\n", (5233, 5329), True, 'import matplotlib.pyplot as plt\n'), ((5333, 5356), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""depth(km)"""'], {}), "('depth(km)')\n", (5343, 5356), True, 'import matplotlib.pyplot as plt\n'), ((5361, 5371), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (5369, 5371), True, 'import matplotlib.pyplot as plt\n'), ((1735, 1748), 'numpy.deg2rad', 'np.deg2rad', (['x'], {}), '(x)\n', (1745, 1748), True, 'import numpy as np\n'), ((1797, 1810), 'numpy.deg2rad', 'np.deg2rad', (['x'], {}), '(x)\n', (1807, 1810), True, 'import numpy as np\n'), ((4658, 4672), 'numpy.min', 'np.min', (['values'], {}), '(values)\n', (4664, 4672), True, 'import numpy as np\n'), ((4697, 4711), 'numpy.min', 'np.min', (['values'], {}), '(values)\n', (4703, 4711), True, 'import numpy as np\n'), ((4804, 4818), 'numpy.max', 'np.max', (['values'], {}), '(values)\n', (4810, 4818), True, 'import numpy as np\n'), ((4843, 4857), 'numpy.max', 'np.max', (['values'], {}), '(values)\n', (4849, 4857), True, 'import numpy as np\n'), ((4949, 4963), 'numpy.max', 'np.max', (['values'], {}), '(values)\n', (4955, 4963), True, 'import numpy as np\n'), ((4965, 4979), 'numpy.min', 'np.min', (['values'], {}), '(values)\n', (4971, 4979), True, 'import numpy as np\n'), ((5194, 5203), 'matplotlib.pyplot.gca', 'plt.gca', ([], {}), '()\n', (5201, 5203), True, 'import matplotlib.pyplot as plt\n')]