from gymfc_nf.envs import *
import os.path
import time
import datetime
import subprocess
import numpy as np
np.seterr('ignore')
import gym
import argparse

def sample_noise(inst):
    # Experiementally derived for MatekF7 FC, see Chapter 5 of "Flight
    # Controller Synthesis Via Deep Reinforcement Learning" for methodology.
    r_noise = inst.np_random.normal(-0.25465, 1.3373)
    p_noise = inst.np_random.normal(0.241961, 0.9990)
    y_noise = inst.np_random.normal(0.07906, 1.45168)
    return np.array([r_noise, p_noise, y_noise])

class StepCallback:

    def __init__(self, total_timesteps, log_freq=1):
        """
        Args:
            total_timesteps: Total timesteps for training
            log_freq: Number of episodes until an update log message is printed
        """
        self.timesteps = total_timesteps
        self.steps_taken = 0
        self.es = []
        self.sps = []
        self.ep = 1
        self.rewards = []
        self.log_freq = log_freq
        self.log_header = ["Ep",
                           "Done",
                           "Steps",

                           "r",
                           "-ydelta",
                           "+ymin",
                           "+/-e",
                           "-ahigh",
                           "-nothing",

                           "score",

                           "pMAE",
                           "qMAE",
                           "rMAE"]

        header_format = ["{:<5}",
                         "{:<7}",
                         "{:<15}",

                         "{:<15}",
                         "{:<15}",
                         "{:<15}",
                         "{:<15}",
                         "{:<15}",
                         "{:<15}",

                         "{:<10}",

                         "{:<7}",
                         "{:<7}",
                         "{:<7}"]
        self.header_format = "".join(header_format)

        log_format_entries = ["{:<5}",
                              "{:<7.0%}",
                              "{:<15}",

                              "{:<15.0f}",
                              "{:<15.0f}",
                              "{:<15.0f}",
                              "{:<15.0f}",
                              "{:<15.0f}",
                              "{:<15.0f}",

                              "{:<10.2f}",

                              "{:<7.0f}",
                              "{:<7.0f}",
                              "{:<7.0f}"]

        self.log_format = "".join(log_format_entries)

    def callback(self, local, state, reward, done):

        self.es.append(local.true_error)
        self.sps.append(local.angular_rate_sp)

        assert local.ind_rewards[0] <= 0 # oscillation penalty
        assert local.ind_rewards[1] >= 0 # min output reward
        assert local.ind_rewards[3] <= 0 # over saturation penalty
        assert local.ind_rewards[4] <= 0 # do nothing penalty

        self.rewards.append(local.ind_rewards)

        if done:
            if self.ep == 1:
                print(self.header_format.format(*self.log_header))
            # XXX (wfk) Try this new score, we need something normalized to handle the
            # random setpoints. Scale by the setpoint, larger setpoints incur
            # more error. +1 prevents divide by zero
            mae = np.mean(np.abs(self.es))
            mae_pqr = np.mean(np.abs(self.es), axis=0)
            e_score = mae / (1 + np.mean(np.abs(self.sps)))
            self.steps_taken += local.step_counter

            if self.ep % self.log_freq == 0:
                ave_ind_rewards = np.mean(self.rewards, axis=0)
                ind_rewards = ""
                for r in ave_ind_rewards:
                    ind_rewards += "{:<15.2f} ".format(r)

                log_data = [
                    self.ep,
                    self.steps_taken/self.timesteps,
                    self.steps_taken,

                    np.mean(self.rewards),
                    ave_ind_rewards[0],
                    ave_ind_rewards[1],
                    ave_ind_rewards[2],
                    ave_ind_rewards[3],
                    ave_ind_rewards[4],

                    e_score,
                    mae_pqr[0],
                    mae_pqr[1],
                    mae_pqr[2]
                ]
                print (self.log_format.format(*log_data))

            self.ep += 1
            self.es = []
            self.sps = []
            self.rewards = []

if __name__=='__main__':
    env_id = 'gymfc_nf-step-v1'
    env = gym.make(env_id)
    env.sample_noise = sample_noise
    env.set_aircraft_model("/remote-home/zzq/12-drone-dreamer/gymfc/examples/gymfc_nf/twins/nf1/model.sdf")
    cb = StepCallback(10e6)
    env.step_callback = cb.callback
    
    ob = env.reset()

    a = 1