#!/usr/bin/env python3

import numpy as np

CACHE_DIR = "analysis_cache"
TPW_DEG = 104.0
AXIS_LON = 31.0
EARTH_RADIUS = 1.0  # unit sphere


def rotation_matrix(axis, theta):
    """
    Rodrigues rotation formula
    """
    axis = axis / np.linalg.norm(axis)
    a = np.cos(theta)
    b = np.sin(theta)

    K = np.array([
        [0, -axis[2], axis[1]],
        [axis[2], 0, -axis[0]],
        [-axis[1], axis[0], 0]
    ])

    I = np.eye(3)

    return a * I + b * K + (1 - a) * np.outer(axis, axis)


def rotation_axis_vector():
    lat = np.radians(90.0)
    lon = np.radians(AXIS_LON)

    return np.array([
        np.cos(lat) * np.cos(lon),
        np.cos(lat) * np.sin(lon),
        np.sin(lat)
    ])


def compute_shear(coords):

    axis = rotation_axis_vector()
    theta = np.radians(TPW_DEG)

    R = rotation_matrix(axis, theta)

    N = len(coords)

    # Displacement field
    displaced = (R @ coords.T).T
    disp_vec = displaced - coords

    shear_vals = np.zeros(N)

    # Use local least squares gradient from nearest neighbours
    dist_matrix = np.load(f"{CACHE_DIR}/dist_matrix.npy")

    neighbour_radius = np.radians(5.0)  # ~550 km

    for i in range(N):

        neighbors = np.where(dist_matrix[i] < neighbour_radius)[0]

        if len(neighbors) < 5:
            continue

        Xi = coords[neighbors] - coords[i]
        Ui = disp_vec[neighbors] - disp_vec[i]

        # Project onto tangent plane
        n = coords[i]
        P = np.eye(3) - np.outer(n, n)

        Xi = (P @ Xi.T).T
        Ui = (P @ Ui.T).T

        # Solve gradient tensor G via least squares
        G, _, _, _ = np.linalg.lstsq(Xi, Ui, rcond=None)

        # Symmetric strain
        E = 0.5 * (G + G.T)

        eigvals = np.linalg.eigvalsh(E)

        shear_vals[i] = np.abs(eigvals[-1] - eigvals[0])

    return shear_vals


if __name__ == "__main__":

    coords = np.load(f"{CACHE_DIR}/coords.npy")

    print("Computing finite-rotation shear field...")

    shear_vals = compute_shear(coords)

    np.save(f"{CACHE_DIR}/shear_vals.npy", shear_vals)

    print("Done.")
    print("Shear range:", shear_vals.min(), shear_vals.max())
