import numpy as np
import argparse

import time

from safe_agents.policies import MLP
from open_safety.envs.balance_bot_env import BalanceBotEnv 
#from open_safety_gym.envs.kart_env import KartEnv
#from open_safety_gym.envs.hoverboard_env import HoverboardEnv

import skimage
import skimage.io as sio

def get_fitness(agent, env, epds, get_cost=True, max_steps=1000, save_frames=False):

    #epd_rewards = []
    #epd_costs = []
    total_steps = 0
    sum_reward = 0
    sum_cost = 0

    exp_id = str(int(time.time()))

    for epd in range(epds):
        steps = 0
        done = False
        obs = env.reset()
        
        while not done and steps < max_steps:

            action = agent.forward(obs)
            if len(action.shape) > 1:
                action = action.squeeze()
            
            obs, reward, done, info = env.step(action)

            sum_reward += reward
            sum_cost += info["cost"]
            steps += 1

            if save_frames: 

                img = env.render()[2]
                sio.imsave("./frames/id{}_epd{}_step{}.png".format(exp_id, epd, steps), img)

        total_steps += steps

    sum_reward /= max_steps * epds
    sum_cost /= max_steps * epds

    return sum_reward, sum_cost, total_steps

def get_elite_mean(population, reward, cost, cost_constraint=2.5, pure_rewards=None, rh=False):

    if not(rh):
        adjusted_cost = [max([cost_constraint, elem]) for elem in cost]
        cost_fitness_agent = [[cost, fit, agent.parameters]
                for a_cost, cost, fit, agent in \
                sorted(zip(adjusted_cost, cost, pure_rewards, population),\
                key = lambda trip: [-trip[0], trip[2]], reverse=True)]
            
        fitness = [elem[1] for elem in cost_fitness_agent]
        cost = [elem[0] for elem in cost_fitness_agent]
        population = [elem[2] for elem in cost_fitness_agent]

    else:
        fitness_agent = [[fit, agent.parameters, my_cost, r] \
                for fit, agent, my_cost, r in \
                sorted(zip(reward, population, cost, pure_rewards),\
                key = lambda trip: [trip[0]], reverse=True) 
                ]

        fitness = [elem[0] for elem in fitness_agent]
        cost = [elem[2] for elem in fitness_agent]
        population = [elem[1] for elem in fitness_agent]
        my_rewards = [elem[3] for elem in fitness_agent]

    keep = int(0.125 * len(population))

    elite_pop = population[:keep] 
    elite_cost = cost[:keep] 
    elite_fitness = fitness[:keep]

    if rh:
        elite_rewards = my_rewards[:keep]
    else:
        elite_rewards = fitness[:keep]


    print("population mean cost, rewards: {:.3e}, {:.3e}".format(\
            np.mean(cost), np.mean(pure_rewards)))

    print("elite mean cost, rewards: {:.3e}, {:.3e}".format(\
            np.mean(elite_cost), np.mean(elite_rewards)))

    param_sum = elite_pop[0]

    for agent_idx in range(1,keep):
        param_sum += elite_pop[agent_idx] 
    

    param_means = param_sum / keep

    return [param_means, np.mean(cost), np.mean(elite_cost), \
            np.mean(reward), np.mean(elite_rewards)]


def enjoy(env, input_dim, output_dim, max_episodes=1, model=None, \
        steps_per_epd=1000, save_frames=False):

    # hard-coded params
    hid_dim = [32]

    exp_id = "eval_{}".format(str(int(time.time())))

    results = {"costs": [],
            "rewards": [],
            "steps": []}
    
    pop_size = 1
    population = [MLP(input_dim, output_dim, hid_dim) \
            for ii in range(pop_size)]

    if model is not None:
        param_means = np.load(model)

        for ll in range(pop_size):
            population[ll].mean = param_means 
            population[ll].covar *= 1e-3
            population[ll].init_params() 

    for epd in range(max_episodes):

        reward, cost, steps = get_fitness(population[0], env, epds=1, get_cost=True,\
                max_steps=steps_per_epd, save_frames=save_frames)
        
        print("episode {} accumulated reward, cost = {:.3e}, {:.3e}".format(epd, reward, cost))

        results["costs"].append(cost)
        results["rewards"].append(reward)
        results["steps"].append(steps)



    mean_epd_cost = np.mean(results["costs"])
    mean_epd_reward = np.mean(results["rewards"])

    print("mean episodic reward and cost over {} episodes".format(max_episodes))
    print("     reward : {:.3e} +/- {:.3e}".format(mean_epd_reward, np.std(results["rewards"])))
    print("     cost   : {:.3e} +/- {:.3e}".format(mean_epd_cost, np.std(results["costs"])))

def train_es(env, input_dim, output_dim, pop_size=6, max_gen=100, \
        cost_constraint=0.0, cost_penalty=False,\
        model=None, reward_hypothesis=False, steps_per_epd=1000):


    # hard-coded policy parameters
    hid_dim = [32]
    es_lr = 1.0 #8e-1
    reward_cost_ratio = 5 
    save_results = True
    exp_id = "c{}_rh{}_{}".format(int(cost_constraint*10),\
        bool(reward_hypothesis), str(int(time.time()))[-10:])
    
    max_steps = steps_per_epd #2000 #np.min([2000, 250 + 10* gen])
    difficulty_threshold = 0.01
    # generate a population
    population = [MLP(input_dim, output_dim, hid_dim) \
            for ii in range(pop_size)]
    
    if model is not None:
        param_means = np.load(model)

        for ll in range(pop_size):
            population[ll].mean = param_means 
            population[ll].init_params() 

    results = {"costs": [],
            "combined_rc": [],
            "rewards": [],
            "elite_costs": [],
            "elite_rewards": [], 
            "steps": []}

    t0 = time.time()    
    try:
        for gen in range(max_gen):
            fitnesses = []
            costs = []
            total_steps = []
            rc_fitnesses = []

            t1 = time.time()

            for agent_idx in range(len(population)):
                epds = 4 #np.max([1,int(10 - gen/10)])

                reward, cost, steps = get_fitness(population[agent_idx],\
                        env, epds=epds, max_steps=max_steps)
                fitnesses.append(reward)
                total_steps.append(steps)
                costs.append(cost)
                
                rc_fitnesses.append(reward - cost / reward_cost_ratio)
            
            if gen:
                param_means *= (1 - es_lr)

                cost_constraint = float("Inf") if reward_hypothesis else cost_constraint
                gen_results = get_elite_mean(population, rc_fitnesses, costs, cost_constraint,\
                        pure_rewards=fitnesses, rh=reward_hypothesis)

                param_means += es_lr * gen_results[0] 
            else:

                cost_constraint = float("Inf") if reward_hypothesis else cost_constraint
                gen_results = get_elite_mean(population, rc_fitnesses, costs, cost_constraint,\
                        pure_rewards=fitnesses, rh=reward_hypothesis)

                param_means = gen_results[0] 

            for ll in range(pop_size):
                population[ll].mean = param_means 
                population[ll].init_params() 
            
            results["costs"].append(gen_results[1])
            results["elite_costs"].append(gen_results[2])
            results["rewards"].append(gen_results[3])
            results["elite_rewards"].append(gen_results[4])
            results["steps"].append(np.sum(total_steps))

            if save_results and gen % 50 == 0:
                np.save("./results/means_{}_gen{}.npy".format(\
                    exp_id, gen), param_means)
                np.save("./results/temp_{}_results.npy".format(\
                    exp_id), results)

            elapsed = time.time() - t0
            epd_elapsed = time.time() - t1
            
            if min(costs) < difficulty_threshold:
                # increase the difficulty (episode duration)
                # if any agent manages to fall down for less than a threshold of time steps
                max_steps = min([max_steps+1, env.max_steps])
            
            print("gen. {} total steps: {} elapsed total: {:.3f} this gen. {:.3f} max_steps {}"\
                    .format(\
                    gen, np.sum(total_steps), elapsed, epd_elapsed, max_steps))

            del fitnesses
            del costs
            del total_steps

    except KeyboardInterrupt:
        pass

    if save_results:
        np.save("./results/means_{}_gen{}.npy".format(\
            exp_id, gen), param_means)
        np.save("./results/temp_{}_results.npy".format(\
            exp_id), results)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="args for safe agents")
    parser.add_argument("-n", "--env_name", type=str,\
            help="name of environment", default="BalanceBot-v0")
    parser.add_argument("-a", "--algo", type=str,\
            help="training algo", default="es")
    parser.add_argument("-c", "--constraint", type=float,\
            help="safety constraint", default=9e9)
    parser.add_argument("-m", "--model", type=str,\
            help="resume parameters stored at filepath (default None)", default=None)
    parser.add_argument("-r", "--reward_hypothesis", type=bool,\
            help="combine cost and reward (reward hypothesis)", default=False)
    parser.add_argument("-p", "--pop_size",type=int,\
            help="population size", default=64)
    parser.add_argument("-e", "--enjoy", type=bool,\
            help="render episodes", default = False)
    parser.add_argument("-f", "--save_frames", type=bool,\
            help="save frames from episodes", default = False)
    parser.add_argument("-g", "--generations", type=int,\
            help="number of generations", default=1024)
    parser.add_argument("-s", "--steps_per_episode", type=int,\
            help="steps per episode ;)", default=1000)
    parser.add_argument("-hm", "--how_many_epds", type=int,\
            help="number of episodes to enjoy", default=3)


    args = parser.parse_args()

    constraint = args.constraint
    rh = args.reward_hypothesis
    model = args.model
    pop_size = args.pop_size
    render = args.enjoy
    save_frames = args.save_frames
    steps_per_episode = args.steps_per_episode
    max_generations = args.generations
    env_name = args.env_name


    env = BalanceBotEnv(render=render)
    obs_dim = env.observation_space.sample().shape[0]
    act_dim = env.action_space.sample().shape[0]

    if not(args.enjoy):
        train_es(env, obs_dim, act_dim, cost_constraint=constraint, pop_size=pop_size,\
                max_gen=max_generations, model=model, reward_hypothesis=rh,\
                steps_per_epd=steps_per_episode)
    else:

        enjoy(env, obs_dim, act_dim, max_episodes=args.how_many_epds,\
                model=model, steps_per_epd=1000, save_frames=save_frames)


    print("all oK")

