import numpy as np
from math import pi
from typing import cast
import gymnasium as gym
from typing import Optional

from Karana.Frame import FrameContainer
from Karana.Core import LogLevel, discard, allReady, allDestroyed, MsgLogger
from Karana.Dynamics import (
    Multibody,
    HingeType,
    StatePropagator,
    PhysicalBody,
    PhysicalHinge,
    LinearSubhinge,
    PinSubhinge,
)
from Karana.Math import UnitQuaternion, HomTran, SpatialInertia
from Karana.Scene import (
    BoxGeometry,
    CylinderGeometry,
    Color,
    PhysicalMaterialInfo,
    PhysicalMaterial,
    LAYER_PHYSICAL_GRAPHICS,
    ScenePartSpec,
)
from Karana.Models import Gravity, OutputUpdateType, UniformGravity, UpdateProxyScene, SyncRealTime
from Karana.Integrators import IntegratorType


# create a cart and a pole multibody system
def createMultibody(
    fc: FrameContainer, render_mode: str | None
) -> tuple[LinearSubhinge, PinSubhinge, Multibody]:
    """Create the Mutlibody.

    Parameters
    ----------
    fc : FrameContainer
        The FrameContainer the Multibody will use.

    Returns
    -------
    tuple[LinearSubhinge, PinSubhinge, Multibody]
        The subhinges associated with the cart pole system and the Multibody itself.
    """
    mb = Multibody("mb", fc)

    # add a cart of mass 1.0kg and 0.4 x 0.2 x 0.15 (L, W, H)
    cart_height = 0.15
    cart = PhysicalBody("cart", mb)
    cart.setSpatialInertia(SpatialInertia(1.0, np.zeros(3), np.eye(3)))
    slider_hinge = PhysicalHinge(mb.virtualRoot(), cart, HingeType.SLIDER)
    slider_subhinge = cast(LinearSubhinge, slider_hinge.subhinge(0))
    slider_subhinge.setUnitAxis([1, 0, 0])
    cart_to_slider = HomTran()
    cart.setBodyToJointTransform(cart_to_slider)

    # add a pole of mass 0.1 kg, and 0.02 x 0.5 (R, L)
    pole_length = 0.5
    pole = PhysicalBody("pole", mb)
    inertia_matrix = np.diag([0.0020933, 0.0020933, 0.00002])
    pole.setSpatialInertia(SpatialInertia(0.1, np.zeros(3), inertia_matrix))
    pin_hinge = PhysicalHinge(cart, pole, HingeType.REVOLUTE)
    pole_to_pin = HomTran(np.array([0.0, 0.0, -(pole_length / 2.0)]))
    pole.setBodyToJointTransform(pole_to_pin)
    cart_to_pin = HomTran(np.array([0.0, 0.0, (cart_height / 2.0)]))
    pin_hinge.onode().setBodyToNodeTransform(cart_to_pin)
    pin_subhinge = cast(PinSubhinge, pin_hinge.subhinge(0))
    pin_subhinge.setUnitAxis([0, 1, 0])

    # Add parts if appropriate
    if render_mode == "human":
        box_geom = BoxGeometry(0.4, 0.2, 0.15)
        cylinder_geom = CylinderGeometry(0.02, 0.5)
        mat_info = PhysicalMaterialInfo()
        mat_info.color = Color.FIREBRICK
        firebrick = PhysicalMaterial(mat_info)
        mat_info.color = Color.GOLD
        gold = PhysicalMaterial(mat_info)
        mat_info.color = Color.WHITE
        white = PhysicalMaterial(mat_info)

        # add geometry to bodies
        cart.addScenePartSpec(
            ScenePartSpec(
                "cart_body", geometry=box_geom, material=firebrick, layers=LAYER_PHYSICAL_GRAPHICS
            )
        )
        pole.addScenePartSpec(
            ScenePartSpec(
                "pole_body",
                geometry=cylinder_geom,
                material=gold,
                transform=HomTran(UnitQuaternion(pi / 2, [1, 0, 0])),
                layers=LAYER_PHYSICAL_GRAPHICS,
            )
        )

        # add a floor
        ground_geom = BoxGeometry(20, 20, 0.5)
        mb.virtualRoot().addScenePartSpec(
            ScenePartSpec(
                "ground",
                geometry=ground_geom,
                material=white,
                transform=HomTran(np.array([0, 0, -0.325])),
                layers=LAYER_PHYSICAL_GRAPHICS,
            )
        )

    # check
    mb.ensureHealthy()
    mb.resetData()
    assert allReady()
    return slider_subhinge, pin_subhinge, mb


class CartPoleEnv(gym.Env):
    """Environment for the cart pole system."""

    metadata = {"render_modes": ["human", "none"]}

    def __init__(
        self,
        render_mode: Optional[str] = None,
        port: int = 0,
        sync_real_time: bool = False,
        **kwargs,
    ):
        """Create an instance of CartPoleEnv.

        Parameters
        ----------
        render_mode : Optional[str]
            The render mode for the ProxyScene.
        port : int
            The port to use for ProxyScene.
        sync_real_time : bool
            If True, then add a SyncRealTime model. If False, then do not add it.
        **kwargs
            Extra keyword arguments.
        """
        self.render_mode = "none"
        if render_mode and render_mode in self.metadata["render_modes"]:
            self.render_mode = render_mode

        # init simulation
        self.init_sim(port=port, sync_real_time=sync_real_time)

        # observation space of [cart pos, cart vel, pin pos, pin vel]
        low = np.array([-1.5, -2, -pi / 3, -6.0], dtype=np.float32)
        high = np.array([1.5, 2, pi / 3, 6.0], dtype=np.float32)
        self.observation_space = gym.spaces.Box(low=low, high=high, dtype=np.float32, shape=(4,))

        # action space of negative or positive force on cart
        self.action_space = gym.spaces.Discrete(2)
        self._action_to_force = {0: -10.0, 1: 10.0}

        # initialize values
        self._cart_pos: float = 0.0
        self._cart_vel: float = 0.0
        self._pin_pos: float = 0.0
        self._pin_vel: float = 0.0
        self.np_random = None
        self._force = 0.0
        self._terminate = False

    def init_sim(self, port: int = 0, sync_real_time: bool = False):
        """Initialize the simulation.

        Parameters
        ----------
        port : int
            The port to use for the ProxyScene.
        sync_real_time : bool
            If True, then add a SyncRealTime model. If False, then do not add it.
        """
        self._fc = FrameContainer("root")
        self._slider, self._pin, self._mb = createMultibody(self._fc, self.render_mode)

        self._sp = StatePropagator(self._mb, integrator_type=IntegratorType.RK4)

        # add visuals
        if self.render_mode == "human":
            self._cleanup_graphics, web_scene = self._mb.setupGraphics(port=port, axes=0)
            web_scene.defaultCamera().pointCameraAt(
                np.array([0, 3, 1]), np.zeros(3), np.array([0, 0, 1])
            )

            # update proxy scene model
            UpdateProxyScene("update_proxy_scene", self._sp, self._mb.getScene())

            # for visualization
            if sync_real_time:
                SyncRealTime("sync_real_time", self._sp, 1.0)

        # add models
        # gravity model
        ug = Gravity("grav_model", self._sp, UniformGravity("uniform_gravity"), self._mb)
        ug.getGravityInterface().setGravity(np.array([0, 0, -9.81]), 0.0, OutputUpdateType.PRE_HOP)
        del ug

        # apply slider force to move the cart
        def apply_force(t, x):
            self._slider.setT(self._force)

        # check if pole or cart fall past a threshold
        # max cart position is 1.5, max pole angle is ~12 (0.21)
        def check_terminate(t, x):
            cart_pos = self._slider.getQ()
            pin_pos = self._pin.getQ()
            if abs(cart_pos) > 1.5 or abs(pin_pos) > 0.21:
                self._terminate = True
                return True
            return False

        self._sp.fns.pre_deriv_fns["apply_force"] = apply_force

        self._sp.setTime(0.0)
        self._sp.setMaxStepSize(np.timedelta64(int(1e7), "ns"))

    # what the model sees
    def _get_obs(self):
        return np.array(
            [
                self._cart_pos,
                self._cart_vel,
                self._pin_pos,
                self._pin_vel,
            ],
            dtype=np.float32,
        )

    # for user debugging
    def _get_info(self):
        return {
            "cart_pos": self._cart_pos,
            "cart_vel": self._cart_vel,
            "pin_pos": self._pin_pos,
            "pin_vel": self._pin_vel,
        }

    def close(self):
        """Close down the environment and remove the simulation."""
        del self._slider, self._pin
        discard(self._sp)
        if self.render_mode == "human":
            self._cleanup_graphics()
        discard(self._mb)
        discard(self._fc)
        assert allDestroyed()

    def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
        """Reset the cart pole system.

        Parameters
        ----------
        seed : Optional[int]
            The seed used for random values.
        options : Optional[dict]
            Extra options used for the reset.
        """
        super().reset(seed=seed)
        if not self.np_random:
            self.np_random, _ = gym.utils.seeding.np_random(seed)

        # init values
        self._cart_pos = self.np_random.uniform(-0.1, 0.1)
        self._cart_vel = self.np_random.uniform(-0.1, 0.1)
        self._pin_pos = self.np_random.uniform(-0.1, 0.1)
        self._pin_vel = self.np_random.uniform(-0.1, 0.1)
        self._force = 0.0
        self._terminate = False

        # set our state
        self._mb.resetData()
        self._slider.setQ(self._cart_pos)
        self._slider.setU(self._cart_vel)
        self._pin.setQ(self._pin_pos)
        self._pin.setU(self._pin_vel)
        t_init = np.timedelta64(0, "ns")
        MsgLogger.changeVerbosity("stdout", LogLevel.ERROR)
        self._sp.setTime(t_init)
        MsgLogger.changeVerbosity("stdout", LogLevel.WARN)
        self._sp.setState(self._sp.assembleState())

        # check if pole or cart fall past a threshold
        # max cart position is 1.5, max pole angle is ~12 (0.21)
        def check_terminate(t, x):
            cart_pos = self._slider.getQ().m
            pin_pos = self._pin.getQ().m
            if abs(cart_pos) > 1.5 or abs(pin_pos) > 0.21:
                self._terminate = True
                return True
            return False

        self._sp.fns.terminate_advance_to_fns["check_terminate"] = check_terminate

        # setup timed event
        self._mb.getScene().update()
        return self._get_obs(), self._get_info()

    def step(self, action):
        """Step the simulation.

        Parameters
        ----------
        action :
            The action to take.
        """
        self._force = self._action_to_force[action]
        self._sp.advanceBy(0.02)
        self._cart_pos = self._slider.getQ().m[0]
        self._cart_vel = self._slider.getU().m[0]
        self._pin_pos = self._pin.getQ().m[0]
        self._pin_vel = self._pin.getU().m[0]
        # observation, reward, terminated, truncated, info
        return self._get_obs(), 1.0, self._terminate, False, self._get_info()