"""
Geometric calculation functions for tilt sensors.

Includes axis transformations, rotations, and quaternion operations.
"""

import numpy as np
import logging
from typing import Tuple

logger = logging.getLogger(__name__)


def asse_a(
    ax: np.ndarray,
    angle: float,
    spe_tl: np.ndarray,
    i: int,
    j: int
) -> Tuple[float, float, float]:
    """
    Calculate axis A displacement components.

    Converts MATLAB ASSEa.m function.

    Args:
        ax: Acceleration/inclination data for axis X
        angle: Installation angle in degrees
        spe_tl: Sensor spacing/length array
        i: Sensor index
        j: Time index

    Returns:
        Tuple of (North component, East component, Vertical component)
    """
    # Convert angle to radians
    angle_rad = angle * 2 * np.pi / 360

    if ax[i, j] >= 0:
        na = spe_tl[i] * ax[i, j] * np.cos(angle_rad)
        ea = -spe_tl[i] * ax[i, j] * np.sin(angle_rad)
    else:
        na = -spe_tl[i] * ax[i, j] * np.cos(angle_rad)
        ea = spe_tl[i] * ax[i, j] * np.sin(angle_rad)

    # Calculate cosine of inclination angle
    cos_beta = np.sqrt(1 - ax[i, j]**2)
    z = spe_tl[i] * cos_beta
    za = spe_tl[i] - z  # Lowering is POSITIVE

    return na, ea, za


def asse_a_hr(
    ax: np.ndarray,
    angle: float,
    spe_tl: np.ndarray,
    i: int,
    j: int
) -> Tuple[float, float, float]:
    """
    Calculate axis A displacement components for high-resolution sensors.

    Converts MATLAB ASSEa_HR.m function.

    Args:
        ax: Angle data for axis X (in degrees)
        angle: Installation angle in degrees
        spe_tl: Sensor spacing/length array
        i: Sensor index
        j: Time index

    Returns:
        Tuple of (North component, East component, Vertical component)
    """
    # Convert angles to radians
    angle_rad = angle * np.pi / 180
    ax_rad = ax[i, j] * np.pi / 180

    # Calculate displacement components
    na = spe_tl[i] * np.sin(ax_rad) * np.cos(angle_rad)
    ea = -spe_tl[i] * np.sin(ax_rad) * np.sin(angle_rad)

    # Vertical component
    za = spe_tl[i] * (1 - np.cos(ax_rad))

    return na, ea, za


def asse_b(
    ay: np.ndarray,
    angle: float,
    spe_tl: np.ndarray,
    i: int,
    j: int
) -> Tuple[float, float, float]:
    """
    Calculate axis B displacement components.

    Converts MATLAB ASSEb.m function.

    Args:
        ay: Acceleration/inclination data for axis Y
        angle: Installation angle in degrees
        spe_tl: Sensor spacing/length array
        i: Sensor index
        j: Time index

    Returns:
        Tuple of (North component, East component, Vertical component)
    """
    # Convert angle to radians
    angle_rad = angle * 2 * np.pi / 360

    if ay[i, j] >= 0:
        nb = -spe_tl[i] * ay[i, j] * np.sin(angle_rad)
        eb = -spe_tl[i] * ay[i, j] * np.cos(angle_rad)
    else:
        nb = spe_tl[i] * ay[i, j] * np.sin(angle_rad)
        eb = spe_tl[i] * ay[i, j] * np.cos(angle_rad)

    # Calculate cosine of inclination angle
    cos_beta = np.sqrt(1 - ay[i, j]**2)
    z = spe_tl[i] * cos_beta
    zb = spe_tl[i] - z  # Lowering is POSITIVE

    return nb, eb, zb


def asse_b_hr(
    ay: np.ndarray,
    angle: float,
    spe_tl: np.ndarray,
    i: int,
    j: int
) -> Tuple[float, float, float]:
    """
    Calculate axis B displacement components for high-resolution sensors.

    Converts MATLAB ASSEb_HR.m function.

    Args:
        ay: Angle data for axis Y (in degrees)
        angle: Installation angle in degrees
        spe_tl: Sensor spacing/length array
        i: Sensor index
        j: Time index

    Returns:
        Tuple of (North component, East component, Vertical component)
    """
    # Convert angles to radians
    angle_rad = angle * np.pi / 180
    ay_rad = ay[i, j] * np.pi / 180

    # Calculate displacement components
    nb = -spe_tl[i] * np.sin(ay_rad) * np.sin(angle_rad)
    eb = -spe_tl[i] * np.sin(ay_rad) * np.cos(angle_rad)

    # Vertical component
    zb = spe_tl[i] * (1 - np.cos(ay_rad))

    return nb, eb, zb


def arot(
    ax: np.ndarray,
    ay: np.ndarray,
    angle: float,
    spe_tl: np.ndarray,
    i: int,
    j: int
) -> Tuple[float, float, float]:
    """
    Calculate combined rotation displacement.

    Converts MATLAB arot.m function.

    Args:
        ax: Acceleration/inclination data for axis X
        ay: Acceleration/inclination data for axis Y
        angle: Installation angle in degrees
        spe_tl: Sensor spacing/length array
        i: Sensor index
        j: Time index

    Returns:
        Tuple of (North displacement, East displacement, Vertical displacement)
    """
    # Calculate components from both axes
    na, ea, za = asse_a(ax, angle, spe_tl, i, j)
    nb, eb, zb = asse_b(ay, angle, spe_tl, i, j)

    # Combine components
    n_total = na + nb
    e_total = ea + eb
    z_total = za + zb

    return n_total, e_total, z_total


def arot_hr(
    ax: np.ndarray,
    ay: np.ndarray,
    angle: float,
    spe_tl: np.ndarray,
    i: int,
    j: int
) -> Tuple[float, float, float]:
    """
    Calculate combined rotation displacement for high-resolution sensors.

    Converts MATLAB arotHR.m function.

    Args:
        ax: Angle data for axis X (in degrees)
        ay: Angle data for axis Y (in degrees)
        angle: Installation angle in degrees
        spe_tl: Sensor spacing/length array
        i: Sensor index
        j: Time index

    Returns:
        Tuple of (North displacement, East displacement, Vertical displacement)
    """
    # Calculate components from both axes
    na, ea, za = asse_a_hr(ax, angle, spe_tl, i, j)
    nb, eb, zb = asse_b_hr(ay, angle, spe_tl, i, j)

    # Combine components
    n_total = na + nb
    e_total = ea + eb
    z_total = za + zb

    return n_total, e_total, z_total


# Quaternion operations
def q_mult2(q1: np.ndarray, q2: np.ndarray) -> np.ndarray:
    """
    Multiply two quaternions.

    Converts MATLAB q_mult2.m function.

    Args:
        q1: First quaternion [w, x, y, z]
        q2: Second quaternion [w, x, y, z]

    Returns:
        Product quaternion
    """
    w1, x1, y1, z1 = q1
    w2, x2, y2, z2 = q2

    w = w1*w2 - x1*x2 - y1*y2 - z1*z2
    x = w1*x2 + x1*w2 + y1*z2 - z1*y2
    y = w1*y2 - x1*z2 + y1*w2 + z1*x2
    z = w1*z2 + x1*y2 - y1*x2 + z1*w2

    return np.array([w, x, y, z])


def rotate_v_by_q(v: np.ndarray, q: np.ndarray) -> np.ndarray:
    """
    Rotate a vector by a quaternion.

    Converts MATLAB rotate_v_by_q.m function.

    Args:
        v: Vector to rotate [x, y, z]
        q: Quaternion [w, x, y, z]

    Returns:
        Rotated vector
    """
    # Convert vector to quaternion form [0, x, y, z]
    v_quat = np.array([0, v[0], v[1], v[2]])

    # Calculate q * v * q_conjugate
    q_conj = np.array([q[0], -q[1], -q[2], -q[3]])

    temp = q_mult2(q, v_quat)
    result = q_mult2(temp, q_conj)

    # Return vector part
    return result[1:]


def fqa(ax: float, ay: float) -> np.ndarray:
    """
    Calculate quaternion from acceleration angles.

    Converts MATLAB fqa.m function.

    Args:
        ax: Acceleration angle X
        ay: Acceleration angle Y

    Returns:
        Quaternion representation
    """
    # Calculate rotation angles
    theta_x = np.arcsin(ax)
    theta_y = np.arcsin(ay)

    # Build quaternion
    qx = np.array([
        np.cos(theta_x/2),
        np.sin(theta_x/2),
        0,
        0
    ])

    qy = np.array([
        np.cos(theta_y/2),
        0,
        np.sin(theta_y/2),
        0
    ])

    # Combine rotations
    q = q_mult2(qx, qy)

    return q