import numpy as np
import mbrltools.leaderboard as ml


def theta_1(X_df):
    return np.arctan2(X_df['sin_theta_1'], X_df['cos_theta_1'])
def sin_theta_1(X_df):
    return X_df['sin_theta_1']
def cos_theta_1(X_df):
    return X_df['cos_theta_1']
def theta_2(X_df):
    return np.arctan2(X_df['sin_theta_2'], X_df['cos_theta_2'])
def sin_theta_2(X_df):
    return X_df['sin_theta_2']
def cos_theta_2(X_df):
    return X_df['cos_theta_2']
def theta_dot_1(X_df):
    return X_df['theta_dot_1']
def theta_dot_2(X_df):
    return X_df['theta_dot_2']
def tip_height(X_df):
    reward = 2 - (X_df['cos_theta_1'] +
                  X_df['cos_theta_1'] * X_df['cos_theta_2'] -
                  X_df['sin_theta_1'] * X_df['sin_theta_2'])

    return reward
def tip_width(X_df):
    width = (X_df['sin_theta_1'] +
             X_df['sin_theta_1'] * X_df['cos_theta_2'] +
             X_df['sin_theta_2'] * X_df['cos_theta_1'])
    return width
        
obss = []
obss.append(ml.Observable(
    col='theta_1', funct=theta_1, domain=[-4, 4], is_angle=True))
obss.append(ml.Observable('theta_2', theta_2, [-4, 4], is_angle=True))
obss.append(ml.Observable('theta_dot_1', theta_dot_1, [-10, 10]))
obss.append(ml.Observable('theta_dot_2', theta_dot_2, [-10, 10]))
obss.append(ml.Observable('tip_height', tip_height, [0, 4], is_cummean=True))
obss.append(ml.Observable('tip_width', tip_width, [-2, 2]))

obss.append(ml.Observable('sin_theta_1', sin_theta_1, [-1, 1]))
obss.append(ml.Observable('cos_theta_1', cos_theta_1, [-1, 1]))
obss.append(ml.Observable('sin_theta_2', sin_theta_2, [-1, 1]))
obss.append(ml.Observable('cos_theta_2', cos_theta_2, [-1, 1]))
