from myosuite.utils import gym

import mujoco

import numpy as np

class CumulativeFatigue():

    # 3CC-r model, adapted from https://dl.acm.org/doi/pdf/10.1145/3313831.3376701 for muscles 

    # based on the implementation from Aleksi Ikkala and Florian Fischer https://github.com/aikkala/user-in-the-box/blob/main/uitb/bm_models/effort_models.py

    def __init__(self, mj_model, frame_skip=1, seed=None):

        self._r = 10 * 15 # Recovery time multiplier i.e. how many times more than during rest intervals https://www.ncbi.nlm.nih.gov/pmc/articles/PMC6092960/ (factor 10 to compensate for 0.1 below)

        self._F = 0.00912  # Fatigue coefficients (default parameter was identified for elbow torque https://pubmed.ncbi.nlm.nih.gov/22579269/)

        self._R = 0.1 * 0.00094  # Recivery coefficients (default parameter was identified for elbow torque https://pubmed.ncbi.nlm.nih.gov/22579269/; factor 0.1 to get an approx. 1% R/F ratio)

        # self.na = mj_model.na

        self._dt = mj_model.opt.timestep * frame_skip # dt might be different from model dt because it might include a frame skip

        muscle_act_ind = mj_model.actuator_dyntype == mujoco.mjtDyn.mjDYN_MUSCLE

        self.na = sum(muscle_act_ind)  #WARNING: here, self.na denotes the number of muscle actuators only!

        self._tauact = np.array([mj_model.actuator_dynprm[i][0] for i in range(len(muscle_act_ind)) if muscle_act_ind[i]])

        self._taudeact = np.array([mj_model.actuator_dynprm[i][1] for i in range(len(muscle_act_ind)) if muscle_act_ind[i]])

        self._MA = np.zeros((self.na,))  # Muscle Active

        self._MR = np.ones((self.na,))   # Muscle Resting

        self._MF = np.zeros((self.na,))  # Muscle Fatigue

        self.TL  = np.zeros((self.na,))  # Target Load

        self.seed(seed)  # Create own Random Number Generator (RNG) used when reset is called with fatigue_reset_random=True

        ### NOTE: the seed from CumulativeFatigue is not synchronised with the seed used for the rest of MujocoEnv!

    def set_FatigueCoefficient(self, F):

        # Set Fatigue coefficients

        self._F = F

    def set_RecoveryCoefficient(self, R):

        # Set Recovery coefficients

        self._R = R

    def set_RecoveryMultiplier(self, r):

        # Set Recovery time multiplier

        self._r = r

    def compute_act(self, act):

        # Get target load (actual activation, which might be reached only with some "effort", 

        # depending on how many muscles can be activated (fast enough) and how many are in fatigue state)

        self.TL = act.copy()

        # Calculate effective time constant tau (see https://mujoco.readthedocs.io/en/stable/modeling.html#muscles)

        self._LD = 1/self._tauact*(0.5 + 1.5*self._MA)

        self._LR = (0.5 + 1.5*self._MA)/self._taudeact

        ## TODO: account for smooth transition phase of length tausmooth (tausmooth = mj_model.actuator_dynprm[i][2])?

        # Calculate C(t) -- transfer rate between MR and MA

        C = np.zeros_like(self._MA)

        idxs = (self._MA < self.TL) & (self._MR > (self.TL - self._MA))

        C[idxs] = self._LD[idxs] * (self.TL[idxs] - self._MA[idxs])

        idxs = (self._MA < self.TL) & (self._MR <= (self.TL - self._MA))

        C[idxs] = self._LD[idxs] * self._MR[idxs]

        idxs = self._MA >= self.TL

        C[idxs] = self._LR[idxs] * (self.TL[idxs] - self._MA[idxs])

        # Calculate rR

        rR = np.zeros_like(self._MA)

        idxs = self._MA >= self.TL

        rR[idxs] = self._r*self._R

        idxs = self._MA < self.TL

        rR[idxs] = self._R

        # Clip C(t) if needed, to ensure that MA, MR, and MF remain between 0 and 1

        C = np.clip(C, np.maximum(-self._MA/self._dt + self._F*self._MA, (self._MR - 1)/self._dt + rR*self._MF),

                    np.minimum((1 - self._MA)/self._dt + self._F*self._MA, self._MR/self._dt + rR*self._MF))

        # Update MA, MR, MF

        dMA = (C - self._F*self._MA)*self._dt

        dMR = (-C + rR*self._MF)*self._dt

        dMF = (self._F*self._MA - rR*self._MF)*self._dt

        self._MA += dMA

        self._MR += dMR

        self._MF += dMF

        return self._MA, self._MR, self._MF

    def get_effort(self):

        # Calculate effort

        return np.linalg.norm(self._MA - self.TL)

    def reset(self, fatigue_reset_vec=None, fatigue_reset_random=False):

        if fatigue_reset_random:

            assert fatigue_reset_vec is None, "Cannot use 'fatigue_reset_vec' if fatigue_reset_random=False."

            non_fatigued_muscles = self.np_random.random(size=(self.na,))

            active_percentage = self.np_random.random(size=(self.na,))

            self._MA = non_fatigued_muscles * active_percentage         # Muscle Active

            self._MR = non_fatigued_muscles * (1 - active_percentage)   # Muscle Resting

            self._MF = 1 - non_fatigued_muscles                         # Muscle Fatigue

        else:

            if fatigue_reset_vec is not None:

                assert len(fatigue_reset_vec) == self.na, f"Invalid length of initial/reset fatigue vector (expected {self.na}, but obtained {len(fatigue_reset_vec)})."

                self._MF = fatigue_reset_vec     # Muscle Fatigue

                self._MR = 1 - fatigue_reset_vec # Muscle Resting

                self._MA = np.zeros((self.na,))  # Muscle Active

            else:

                self._MA = np.zeros((self.na,))  # Muscle Active

                self._MR = np.ones((self.na,))   # Muscle Resting

                self._MF = np.zeros((self.na,))  # Muscle Fatigue

    def seed(self, seed=None):

        """

        Set random number seed

        """

        self.input_seed = seed

        self.np_random, seed = gym.utils.seeding.np_random(seed)

        return [seed]

    @property

    def MF(self):

        return self._MF

    @property

    def MR(self):

        return self._MR

    @property

    def MA(self):

        return self._MA

    @property

    def F(self):

        return self._F

    @property

    def R(self):

        return self._R

    @property

    def r(self):

        return self._r