import numpy as np

unit_x = np.array([1, 0, 0])
unit_y = np.array([0, 1, 0])
unit_z = np.array([0, 0, 1])

# Axis are numbered from 1 to 3 from x to z

def get_axis(axis):
    match axis:
        case 1:
            ax = unit_x
        case 2:
            ax = unit_y
        case 3:
            ax = unit_z
        case _:
            ax = unit_x
    return ax

def proj(vec, axis=1):
    ax = get_axis(axis)

    return np.dot(vec, ax) * ax

def abs_custom(vec):
    l = 0
    for i in range(3):
        l += vec[i] ** 2
    return np.sqrt(l)


def rotate(v, angle=90, axis=1):
    angle = angle/180 * np.pi
    k = get_axis(axis)

    return (
        v * np.cos(angle)
        + np.cross(k, v) * np.sin(angle)
        + k * np.dot(k, v) * (1 - np.cos(angle))
    )
        
def agl(a, b):
    return np.round(np.acos(np.dot(a, b)/(abs_custom(a) * abs_custom(b)))/(2 * np.pi) * 360)

def normalize(vec):
    l = abs_custom(vec)
    return vec/l

def get_angles(source, target):
    source_planar = source - proj(source, 3)
    target_planar = target - proj(target, 3)

    source_phi = agl(source_planar, unit_x)
    target_phi = agl(target_planar, unit_x)

    source_theta = agl(rotate(source, 90 - source_phi, 3), unit_z)
    target_theta = agl(rotate(target, 90 - target_phi, 3), unit_z)

    phi = None
    theta = None

    theta_diff = None
    phi_diff = agl(source_planar, target_planar)
    if source_phi < target_phi:
        rota = rotate(source_planar, phi_diff, 3)
        theta_diff = agl(rota, target)
        phi = source_phi + phi_diff/2
    else:
        rota = rotate(target_planar, phi_diff, 3)
        theta_diff = agl(rota, source)
        phi = target_phi + phi_diff/2

    if source_theta < target_theta:
        theta = target_theta + theta_diff/2
    else:
        theta = source_theta + theta_diff/2

    return (phi, theta)

GRID_SIZE = 10

source_orig = np.array([30, 20, 50])
target_orig = np.array([30, 20, 20])

for x in range(2):
    for y in range(2):
        x_size = x * GRID_SIZE
        y_size = y * GRID_SIZE

        phi, theta = get_angles(source_orig - unit_x * x_size - unit_y * y_size, target_orig - unit_x * x_size - unit_y * y_size)

        print(f"For grid ({x}, {y}), phi = {phi} and theta = {theta}.")