#!/usr/bin/env python3

# a simple Robotics Toolbox "shell", runs Python3 and loads in NumPy, RTB, SMTB
#
# Run it from the shell
#  % rtb.py
#
# or setup an alias
#
#  alias rtb=PATH/rtb.py   # sh/bash
#  alias rtb PATH/rtb.py   # csh/tcsh
#
# % rtb

# import stuff
import os
from pathlib import Path
import sys
from importlib.metadata import version
import argparse

from pygments.token import Token
from IPython.terminal.prompts import Prompts
from IPython.terminal.prompts import ClassicPrompts
from traitlets.config import Config
import IPython

try:
    from colored import fg, bg, attr

    _colored = True
    # print('using colored output')
except ImportError:
    # print('colored not found')
    _colored = False

parser = argparse.ArgumentParser("RVC book console")
parser.add_argument("script", default=None, nargs="?", help="specify script to run")
parser.add_argument("--backend", "-B", default=None, help="specify graphics backend")
parser.add_argument(
    "--color",
    "-C",
    default="neutral",
    help="specify terminal color scheme (neutral, lightbg, nocolor, linux), linux is for dark mode",
)
parser.add_argument("--confirmexit", "-x", default=False,
    help="confirm exit")
parser.add_argument("--prompt", "-p", default=None,
    help="input prompt")
parser.add_argument(
    "--resultprefix",
    "-r",
    default=None,
    help="execution result prefix, include {} for execution count number",
)
parser.add_argument(
    "--nobanner", "-b",
    dest="banner",
    default=True, action="store_false",
    help="suppress startup banner",
)
parser.add_argument(
    "--nocwd", "-c",
    dest="cwd",
    default=True, action="store_false",
    help="suppress cwd to RVC3 folder",
)
parser.add_argument(
    "--showassign",
    "-a",
    default=False,
    action="store_true",
    help="do not display the result of assignments",
)
parser.add_argument(
    "--normal", "-n",
    dest="book",
    default=True, action="store_false",
    help="use normal ipython settings for prompts and display on assignment"
)
parser.add_argument(
    "--norobot", "-R",
    dest="robot",
    default=True, action="store_false",
    help="do not import robotics toolbox (RTB-P)"
)
parser.add_argument(
    "--novision", "-V",
    dest="vision",
    default=True, action="store_false",
    help="do not import vision toolbox (MVTB-P)"
)
parser.add_argument(
    "--ansi",
    default=False,
    action="store_true",
    help="use ANSImatrix to display matrices",
)
parser.add_argument(
    "--examples",
    "-e",
    default=False,
    action="store_true",
    help="change working directory to shipped examples",
)
parser.add_argument(
    "--swift",
    "-s",
    default=False,
    action="store_true",
    help="use Swift as default backend",
)
args = parser.parse_args()

# TODO more options
# color scheme, light/dark
# silent startup

sys.argv = [sys.argv[0]]

# http://patorjk.com/software/taag/#p=display&f=Standard&t=RVC%203
banner = fg("yellow")
banner += r""" ____       _           _   _             __     ___     _                ___      ____            _             _   _____ 
|  _ \ ___ | |__   ___ | |_(_) ___ ___    \ \   / (_)___(_) ___  _ __    ( _ )    / ___|___  _ __ | |_ _ __ ___ | | |___ / 
| |_) / _ \| '_ \ / _ \| __| |/ __/ __|    \ \ / /| / __| |/ _ \| '_ \   / _ \/\ | |   / _ \| '_ \| __| '__/ _ \| |   |_ \ 
|  _ < (_) | |_) | (_) | |_| | (__\__ \_    \ V / | \__ \ | (_) | | | | | (_>  < | |__| (_) | | | | |_| | | (_) | |  ___) |
|_| \_\___/|_.__/ \___/ \__|_|\___|___( )    \_/  |_|___/_|\___/|_| |_|  \___/\/  \____\___/|_| |_|\__|_|  \___/|_| |____/ 
                                      |/                                                                                   
for Python"""

if args.book:
    # set book options
    args.resultprefix = ""
    args.prompt = ">>> "
    args.showassign = True
    args.ansi = False
    args.examples = True

from math import pi  # lgtm [py/unused-import]

import numpy as np
import scipy as sp
import matplotlib.pyplot as plt  # lgtm [py/unused-import]

# setup defaults
np.set_printoptions(
    linewidth=120,
    formatter={"float": lambda x: f"{x:8.4g}" if abs(x) > 1e-10 else f"{0:8.4g}"},
)

if args.robot:
    from roboticstoolbox import *  # lgtm [py/unused-import]

if args.vision:
    from machinevisiontoolbox import *
    import machinevisiontoolbox.base as mvbase

from spatialmath import *  # lgtm [py/polluting-import]
from spatialgeometry import *  # lgtm [py/polluting-import]
from spatialmath.base import *

# set matrix printing mode for spatialmath
SE3._ansimatrix = args.ansi

# set default matplotlib backend
if args.backend is not None:
    print(f"Using matplotlib backend {args.backend}")
    mpl.use(args.backend)

# load some robot models
puma = models.DH.Puma560()
panda = models.DH.Panda()

# print the banner: standard
# https://patorjk.com/software/taag/#p=display&f=Standard&t=Robotics%2C%20Vision%20%26%20Control%203

# set default backend for Robot.plot
if args.swift:
    Robot.default_backend = "swift"

versions = []
if args.robot:
    versions.append(f"RTB=={version('roboticstoolbox-python')}")
if args.vision:
    versions.append(f"MVTB=={version('machinevision-toolbox-python')}")
versions.append(f"SMTB=={version('spatialmath-python')}")
banner += " (" + ", ".join(versions) + ")"
banner += r"""

import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
from math import pi
from spatialmath import *
from spatialmath.base import *
"""
if args.robot:
    banner += "from roboticstoolbox import *\n"
if args.vision:
    banner += """from machinevisiontoolbox import *
import machinevisiontoolbox.base as mvbase
"""

banner += r"""
func/object?       - show brief help
help(func/object)  - show detailed help
func/object??      - show source code

"""
banner += attr(0)

if args.banner:
    print(banner)

# append to the module path
# - RVC3 models and examples
# - RTB examples
root = Path(__file__).absolute().parent.parent
sys.path.append(str(root / "models"))
sys.path.append(str(root / "examples"))

import roboticstoolbox as rtb
sys.path.append(str(Path(rtb.__path__[0]) / "examples"))

if args.cwd:
    os.chdir(root)

if args.showassign and args.banner:
    print(
        fg("red")
        + """Results of assignments will be displayed, use trailing ; to suppress

""",
        attr(0),
    )

# drop into IPython

# def my_file_finder(*args, **kwargs):
#     print("FF", *args, **kwargs)
# IPython.core.magics.execution.ExecutionMagics.run.im_func.func_defaults[2] = my_file_finder


class MyPrompt(Prompts):
    def in_prompt_tokens(self, cli=None):
        if args.prompt is None:
            return super().in_prompt_tokens()
        else:
            return [(Token.Prompt, args.prompt)]

    def out_prompt_tokens(self, cli=None):
        if args.resultprefix is None:
            # traditional behaviour
            return super().out_prompt_tokens()
        else:
            return [
                (Token.Prompt, args.resultprefix.format(self.shell.execution_count))
            ]

# set configuration options, there are lots, see
# https://ipython.readthedocs.io/en/stable/config/options/terminal.html
c = Config()
c.InteractiveShellEmbed.colors = args.color
c.InteractiveShell.confirm_exit = args.confirmexit
# c.InteractiveShell.prompts_class = ClassicPrompts
c.InteractiveShell.prompts_class = MyPrompt
if args.showassign:
    c.InteractiveShell.ast_node_interactivity = "last_expr_or_assign"
c.TerminalIPythonApp.display_banner = args.banner

# set precision, same as %precision
c.PlainTextFormatter.float_precision = "%.3f"

# set up a script to be executed by IPython when we get there
code = None
if args.script is not None:
    path = Path(args.script)
    print(path)
    if not path.exists():
        raise ValueError(f"script does not exist: {args.script}")
    code = path.open("r").readlines()
if code is None:
    code = [
        "%precision %.3g;",
        "plt.ion();",
    ]

else:
    code.append("plt.ion()")
c.InteractiveShellApp.exec_lines = code
IPython.start_ipython(config=c, user_ns=globals())
