"""
dist function return mat f(y,x),
where f(y,x) is the pixel distance in echo pulse RF signal between point(0,0) and point(x,y).
"""
import numpy as np
from scipy.optimize import fsolve

from ds.Config import DeviceConfig
from utils.filename import DS


def refraction_dist(dev_cfg=DeviceConfig()):
    v1 = dev_cfg.v1
    v1x = dev_cfg.v1x
    v2 = dev_cfg.v2
    y1 = dev_cfg.y1
    SECONDS_PER_Y_PIX = dev_cfg.second_per_y_pix
    METER_PER_X_PIX = dev_cfg.meter_per_x_pix

    save_path = DS / f'{y1}_{v1x}_{v2}_fslove.npy'
    if save_path.exists():
        return np.load(str(save_path))

    def gf(x, y2):
        def f(arg):
            θ1 = arg[0]
            return [
                y1 * np.tan(θ1) + v2 * SECONDS_PER_Y_PIX * y2 * np.tan(
                    np.arcsin(v2 / v1x * np.sin(θ1))) - METER_PER_X_PIX * x]

        return f

    iter_start = np.arcsin(v1x / v2 * np.sin(np.deg2rad(90)))
    iter_start = iter_start // 1e-7 / 1e+7
    m = np.zeros((dev_cfg.rows, dev_cfg.cols - 1))
    for y2_ in range(m.shape[0]):
        for x_ in range(m.shape[1]):
            m[y2_, x_] = fsolve(gf(x_ + 1, y2_ + 1), np.array([iter_start]))[0]

    dmat = np.zeros((dev_cfg.rows, dev_cfg.cols))
    t1 = (np.arange(dev_cfg.rows)) * SECONDS_PER_Y_PIX
    y2_mat = (t1 * v2)[:, np.newaxis]
    x1_mat = y1 * np.tan(m)
    x2_mat = y2_mat * np.tan(np.arcsin((v2 / v1x) * np.sin(m)))
    d1 = np.hypot(x1_mat, y1)
    d2 = np.hypot(x2_mat, y2_mat)
    dmat[:, 1:] = (d1 / v1 + d2 / v2)
    dmat[:, 0] = y1 / v1 + t1
    dmat /= SECONDS_PER_Y_PIX
    np.save(save_path.__str__(), dmat)
    return dmat


def direct_dist(dev_cfg=DeviceConfig(), p=np):
    y, x = p.ogrid[:dev_cfg.rows, :dev_cfg.cols]
    y1_meter = dev_cfg.y1
    y2_meter = y * dev_cfg.second_per_y_pix * dev_cfg.v2
    y_meter = y1_meter + y2_meter
    x_meter = x * dev_cfg.meter_per_x_pix
    xy_meter = p.hypot(y_meter, x_meter)
    xy1_meter = xy_meter * (y1_meter / y_meter)
    xy2_meter = xy_meter * (y2_meter / y_meter)
    m = xy1_meter / dev_cfg.v1 + xy2_meter / dev_cfg.v2
    return m / dev_cfg.second_per_y_pix


if __name__ == '__main__':
    print(refraction_dist())
    print(direct_dist())