import pyarrow as pa
import numpy as np
from scipy.interpolate import RBFInterpolator
import os


def rotate_vector(v, k, theta):
    v = np.asarray(v)
    k = np.asarray(k) / np.linalg.norm(k)

    cos_theta = np.cos(theta)

    return v * cos_theta + np.cross(k, v) * np.sin(theta) + k * np.dot(k, v) * (1 - cos_theta)


def main():
    # Accessing an environment variable
    global my_variable
    try:
        my_variable = os.environ['WG_WorkingDirectory']
        print(f"The value of MY_VARIABLE_NAME is: {my_variable}")
    except KeyError:
        print("MY_VARIABLE_NAME is not set.")

    fixe_surface = os.path.join(my_variable, 'fixed_surface.arrow')
    df1 = None
    with pa.memory_map(fixe_surface, 'r') as source:
        data = pa.ipc.open_file(source).read_all()
        df1 = data.to_pandas()

    control_points = os.path.join(my_variable, 'control_points.arrow')
    df2 = None
    with pa.memory_map(control_points, 'r') as source:
        data = pa.ipc.open_file(source).read_all()
        df2 = data.to_pandas()

    n1 = df1.shape[0]
    n2 = df2.shape[0]

    n = n1 + n2
    xyz = np.zeros((n, 3))
    uvw = np.zeros((n, 3))

    for i in range(n1):
        xyz[i, :] = df1.iloc[i][1:]

    for i in range(n2):
        v = df2.iloc[i][1:4]
        k = [0, 0, 1]
        theta = np.pi / 4
        nv = rotate_vector(v, k, theta) - v

        xyz[n1+i, :] = v
        uvw[n1 + i, :] = nv

    deformed_domain = os.path.join(my_variable, 'deformed_domain.arrow')
    df3 = None
    with pa.memory_map(deformed_domain, 'r') as source:
        data = pa.ipc.open_file(source).read_all()
        df3 = data.to_pandas()

    nn = df3.shape[0]
    ids = np.zeros((nn,), dtype=np.int32)
    vert = np.zeros((nn, 3), dtype=np.float64)
    for i in range(nn):
        ids[i] = df3.iloc[i]['id']
        vert[i, :] = df3.iloc[i][1:]

    rbf = RBFInterpolator(xyz, uvw)
    off = rbf(vert)

    Id = pa.array(ids)
    dx = pa.array(off[:, 0])
    dy = pa.array(off[:, 1])
    dz = pa.array(off[:, 2])

    schema = pa.schema([
        pa.field('id', pa.int32(), nullable=False),
        pa.field('dx', pa.float64(), nullable=False),
        pa.field('dy', pa.float64(), nullable=False),
        pa.field('dz', pa.float64(), nullable=False),
    ])

    deformed_mesh = os.path.join(my_variable, 'deformed_mesh.arrow')
    with pa.OSFile(deformed_mesh, 'wb') as sink:
        with pa.ipc.new_file(sink, schema=schema) as writer:
            batch = pa.record_batch([Id, dx, dy, dz], schema=schema)
            writer.write(batch)


if __name__ == '__main__':
    main()