# pylint: disable=C0413
# flake8: noqa: E402

from nimbus.utils.utils import init_env

init_env()

# Patch scipy.spatial.transform module so that `from_quat` and `as_quat`
# accept `scalar_first=` on older scipy (< 1.11).
try:
    from scipy.spatial.transform import Rotation as _R
    _R.from_quat([1, 0, 0, 0], scalar_first=True)
except TypeError:
    import numpy as _np
    import scipy.spatial.transform as _sst
    from scipy.spatial.transform import Rotation as _OrigR

    class Rotation(_OrigR):
        @classmethod
        def from_quat(cls, quat, *args, scalar_first=False, **kwargs):
            quat = _np.asarray(quat, dtype=float)
            if scalar_first:
                if quat.ndim == 1:
                    quat = _np.array([quat[1], quat[2], quat[3], quat[0]])
                else:
                    quat = _np.concatenate([quat[..., 1:], quat[..., :1]], axis=-1)
            return super().from_quat(quat, *args, **kwargs)

        def as_quat(self, *args, scalar_first=False, **kwargs):
            q = super().as_quat(*args, **kwargs)
            if scalar_first:
                if q.ndim == 1:
                    q = _np.array([q[3], q[0], q[1], q[2]])
                else:
                    q = _np.concatenate([q[..., 3:], q[..., :3]], axis=-1)
            return q

    _sst.Rotation = Rotation
    import sys
    sys.modules['scipy.spatial.transform._rotation'].Rotation = Rotation
    sys.modules['scipy.spatial.transform'].Rotation = Rotation

import argparse

from nimbus import run_data_engine
from nimbus.utils.config_processor import ConfigProcessor
from nimbus.utils.flags import set_debug_mode, set_random_seed


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True, help="path to config file")
    parser.add_argument("--random_seed", help="random seed")
    parser.add_argument("--debug", action="store_true", help="enable debug mode: all errors raised immediately")
    args, extras = parser.parse_known_args()

    processor = ConfigProcessor()

    try:
        config = processor.process_config(args.config, cli_args=extras)
    except ValueError as e:
        print(f"Configuration Error: {e}")
        print(f"\n Available configuration paths can be found in: {args.config}")
        print("   Use dot notation to override nested values, e.g.:")
        print("   --stage_pipe.worker_num='[2,4]'")
        print("   --load_stage.layout_random_generator.args.random_num=500")
        return 1

    processor.print_final_config(config)

    if args.debug:
        set_debug_mode(True)

    if args.random_seed is not None:
        set_random_seed(int(args.random_seed))

    try:
        run_data_engine(config, args.random_seed)
    except Exception as e:
        raise e


if __name__ == "__main__":
    main()