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()

[((44, 1, 44, 13), 'numba.njit', 'numba.njit', ({}, {}), '()', False, 'import numba\n'), ((59, 1, 59, 13), 'numba.njit', 'numba.njit', ({}, {}), '()', False, 'import numba\n'), ((64, 1, 64, 13), 'numba.njit', 'numba.njit', ({}, {}), '()', False, 'import numba\n'), ((75, 1, 75, 13), 'numba.njit', 'numba.njit', ({}, {}), '()', False, 'import numba\n'), ((92, 1, 92, 16), 'click.command', 'click.command', ({}, {}), '()', False, 'import click\n'), ((93, 1, 93, 63), 'click.option', 'click.option', (), '', False, 'import click\n'), ((94, 1, 94, 63), 'click.option', 'click.option', (), '', False, 'import click\n'), ((95, 1, 95, 63), 'click.option', 'click.option', (), '', False, 'import click\n'), ((96, 1, 96, 63), 'click.option', 'click.option', (), '', False, 'import click\n'), ((97, 1, 97, 63), 'click.option', 'click.option', (), '', False, 'import click\n'), ((98, 1, 98, 63), 'click.option', 'click.option', (), '', False, 'import click\n'), ((99, 1, 99, 72), 'click.option', 'click.option', (), '', False, 'import click\n'), ((100, 1, 100, 89), 'click.option', 'click.option', (), '', False, 'import click\n'), ((101, 1, 101, 73), 'click.option', 'click.option', (), '', False, 'import click\n'), ((102, 1, 102, 71), 'click.option', 'click.option', (), '', False, 'import click\n'), ((17, 35, 18, 52), 'numpy.meshgrid', 'np.meshgrid', (), '', True, 'import numpy as np\n'), ((19, 17, 19, 35), 'numpy.shape', 'np.shape', ({(19, 26, 19, 34): 'lon_mesh'}, {}), '(lon_mesh)', True, 'import numpy as np\n'), ((20, 17, 20, 40), 'numpy.zeros_like', 'np.zeros_like', ({(20, 31, 20, 39): 'lon_mesh'}, {}), '(lon_mesh)', True, 'import numpy as np\n'), ((21, 13, 21, 36), 'numpy.zeros_like', 'np.zeros_like', ({(21, 27, 21, 35): 'lon_mesh'}, {}), '(lon_mesh)', True, 'import numpy as np\n'), ((22, 13, 22, 36), 'numpy.zeros_like', 'np.zeros_like', ({(22, 27, 22, 35): 'lon_mesh'}, {}), '(lon_mesh)', True, 'import numpy as np\n'), ((23, 13, 23, 36), 'numpy.zeros_like', 'np.zeros_like', ({(23, 27, 23, 35): 'lon_mesh'}, {}), '(lon_mesh)', True, 'import numpy as np\n'), ((24, 13, 24, 36), 'numpy.zeros_like', 'np.zeros_like', ({(24, 27, 24, 35): 'lon_mesh'}, {}), '(lon_mesh)', True, 'import numpy as np\n'), ((79, 19, 79, 36), 'numpy.min', 'np.min', ({(79, 26, 79, 35): 'distance2'}, {}), '(distance2)', True, 'import numpy as np\n'), ((80, 12, 80, 47), 'numpy.where', 'np.where', ({(80, 21, 80, 46): 'distance2 == mindistance2'}, {}), '(distance2 == mindistance2)', True, 'import numpy as np\n'), ((86, 11, 86, 55), 'numpy.linspace', 'np.linspace', ({(86, 23, 86, 34): 'lon_list[0]', (86, 36, 86, 47): 'lon_list[1]', (86, 49, 86, 54): 'hnpts'}, {}), '(lon_list[0], lon_list[1], hnpts)', True, 'import numpy as np\n'), ((87, 11, 87, 55), 'numpy.linspace', 'np.linspace', ({(87, 23, 87, 34): 'lat_list[0]', (87, 36, 87, 47): 'lat_list[1]', (87, 49, 87, 54): 'hnpts'}, {}), '(lat_list[0], lat_list[1], hnpts)', True, 'import numpy as np\n'), ((88, 11, 88, 55), 'numpy.linspace', 'np.linspace', ({(88, 23, 88, 34): 'dep_list[0]', (88, 36, 88, 47): 'dep_list[1]', (88, 49, 88, 54): 'vnpts'}, {}), '(dep_list[0], dep_list[1], vnpts)', True, 'import numpy as np\n'), ((107, 18, 107, 38), 'pandas.read_pickle', 'pd.read_pickle', ({(107, 33, 107, 37): 'data'}, {}), '(data)', True, 'import pandas as pd\n'), ((122, 13, 122, 37), 'numpy.zeros', 'np.zeros', ({(122, 22, 122, 36): '(hnpts, vnpts)'}, {}), '((hnpts, vnpts))', True, 'import numpy as np\n'), ((130, 4, 130, 16), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((131, 35, 132, 45), 'numpy.meshgrid', 'np.meshgrid', (), '', True, 'import numpy as np\n'), ((146, 4, 147, 53), 'matplotlib.pyplot.contourf', 'plt.contourf', (), '', True, 'import matplotlib.pyplot as plt\n'), ((148, 8, 148, 35), 'numpy.arange', 'np.arange', ({(148, 18, 148, 22): 'vmin', (148, 24, 148, 28): 'vmax', (148, 30, 148, 34): '0.01'}, {}), '(vmin, vmax, 0.01)', True, 'import numpy as np\n'), ((149, 4, 149, 47), 'matplotlib.pyplot.colorbar', 'plt.colorbar', (), '', True, 'import matplotlib.pyplot as plt\n'), ((151, 4, 152, 98), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(152, 8, 152, 97): '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    )", True, 'import matplotlib.pyplot as plt\n'), ((153, 4, 153, 27), 'matplotlib.pyplot.ylabel', 'plt.ylabel', ({(153, 15, 153, 26): '"""depth(km)"""'}, {}), "('depth(km)')", True, 'import matplotlib.pyplot as plt\n'), ((154, 4, 154, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((61, 18, 61, 31), 'numpy.deg2rad', 'np.deg2rad', ({(61, 29, 61, 30): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((66, 18, 66, 31), 'numpy.deg2rad', 'np.deg2rad', ({(66, 29, 66, 30): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((135, 23, 135, 37), 'numpy.min', 'np.min', ({(135, 30, 135, 36): 'values'}, {}), '(values)', True, 'import numpy as np\n'), ((136, 20, 136, 34), 'numpy.min', 'np.min', ({(136, 27, 136, 33): 'values'}, {}), '(values)', True, 'import numpy as np\n'), ((140, 23, 140, 37), 'numpy.max', 'np.max', ({(140, 30, 140, 36): 'values'}, {}), '(values)', True, 'import numpy as np\n'), ((141, 20, 141, 34), 'numpy.max', 'np.max', ({(141, 27, 141, 33): 'values'}, {}), '(values)', True, 'import numpy as np\n'), ((145, 22, 145, 36), 'numpy.max', 'np.max', ({(145, 29, 145, 35): 'values'}, {}), '(values)', True, 'import numpy as np\n'), ((145, 38, 145, 52), 'numpy.min', 'np.min', ({(145, 45, 145, 51): 'values'}, {}), '(values)', True, 'import numpy as np\n'), ((150, 4, 150, 13), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n')]