"""Generate forward kinematic functions for a two-link arm.

Copyright (c) 2013-2024 Garth Zeglin. All rights reserved. Licensed under the
terms of the BSD 3-clause license as included in LICENSE.

N.B. These have not been checked in detail.
"""

import os, sys, time, argparse

import sympy as sym

# import a utility class for symbolic manipulation of homogeneous transforms
import transform3d

#================================================================
class TwoLinkKinematicsGenerator(object):
    """Generator to derive kinematic functions for the two link arm using sympy.

    The arm joint angles are specified in radians.
    """
    def __init__(self, verbose=False):

        self.verbose = verbose

        # Generate lists of the needed symbols.
        self._joint_symbols = sym.symbols('j1 j2')
        self._parameter_symbols = sym.symbols('link1len link2len')

        # Generate kinematics functions using sympy.
        self._generate_endpoint_frame_transform()

        return

    #================================================================
    def _generate_endpoint_frame_transform(self):
        """Generate a function to compute the homogeneous transform representing the
        default endpoint frame parameterized by kinematic state.
        """
        if self.verbose: print("Generating endpoint frame transforms...")
        j1, j2 = self._joint_symbols
        link1, link2 = self._parameter_symbols

        m = transform3d.Transform3D()

        # apply Z rotation for joint1
        m.rotate_z(j1)

        # apply the X offset between joint1 and joint2
        m.translate(link1, 0, 0)

        # apply Z rotation for joint2
        m.rotate_z(j2)

        # apply the X offset between joint2 and the end
        m.translate(link2, 0, 0)

        # find the homogeneous transform representing the end frame
        end_transform = sym.trigsimp(m.ctm)
        # print("end frame transform:", end_transform)

        # find an expression for the origin of the end frame
        end_vector = end_transform[0:2, 3]

        print("end vector as expression:", end_vector)
        print("\nend vector location as code:", sym.pycode(end_vector))

        return

################################################################
if __name__ == "__main__":
    # process command line arguments
    parser = argparse.ArgumentParser(description = """Kinematic function code generators.""")
    parser.add_argument( '-v', '--verbose', action='store_true', help='Enable more detailed output.' )
    args = parser.parse_args()

    if args.verbose:
        print("Beginning derivation of kinematics functions...")
        start = time.time()

    kine = TwoLinkKinematicsGenerator(verbose=args.verbose)

    if args.verbose:
        end = time.time()
        print("Completed in %f seconds." % (end - start))