from smtenv import SMTEnv
from smtenv.envs.atesim import ATESim

import numpy as np
from collections import defaultdict

import matplotlib.pyplot as plt

class QLearnAgent(object):
    def __init__(self, action_space, alpha = 0.6, gamma = 0.1, argmax=np.argmax, state_parse=None):
        self.action_space = action_space
        
        self.q = defaultdict(lambda: np.zeros(action_space.n)) 
        self.alpha = alpha
        self.gamma = gamma

        self.na = self.action_space.n
        self.last_action = None

        self.state = None
        self.state_parse = state_parse

        self.argmax = argmax
    
    def reset(self):
        self.q = defaultdict(lambda: np.zeros(self.action_space.n)) 
    
    def policy(self, state, epsilon = 0.00):
        self.action_dist = np.ones(self.na, dtype=float) * (epsilon / self.na)
        opt_a = self.argmax(self.q[state])
        # print(self.q[state])
        self.action_dist[opt_a] += (1.0 - epsilon)
        return self.action_dist
    
    def prop_policy(self, ob):
        states = self.state_parse(ob)

        positives = np.zeros(self.na) 
        tot = 0

        # for state in states:
        #     dist = self.q[state].clip(0)
        #     positives += dist
        #     tot += np.sum(dist)
        for state in states:
            dist = self.q[state]
            positives += dist
        
        positives = positives.clip(0)
        tot = np.sum(positives)
        
        if tot != 0:
            self.action_dist = positives / tot
        else:
            self.action_dist = np.ones(self.na, dtype=float) / self.na
        return self.action_dist

    def act(self, observation, reward, done):
        return np.random.choice(np.arange(self.na), p=self.prop_policy(observation)) 
        # next_action = self.policy(observation)

        # if self.last_action != None:
        #     td = reward + self.gamma * self.q[observation][next_action] - self.q[self.state][self.last_action]
        #     self.q[self.state][self.last_action] += self.alpha * td

        # self.state = observation
        # self.last_action = next_action
        # return self.last_action

def ate_state_parser(state):
    states = []
    for i in range(2, len(state), 2):
        pair = state[i], state[i + 1]
        states.append(pair)

    return states


def train(agent, env, num_episodes = 100, episode_len = 100):
    episode_results = dict()
    episode_results["total_rwd"] = []
    episode_results["length"] = []

    for episode_i in range(num_episodes):
        ob = env.reset()

        reward = 0
        total_rwd = 0

        done = False

        episode_log = {
            "rwd": list()
        }

        steps = 0

        print("Starting Episode {0}".format(episode_i))
        for _ in range(episode_len):
            action = agent.act(ob, reward, done)
            next_ob, reward, done, _ = env.step(action)

            episode_log["rwd"].append(reward)
            total_rwd += reward

            for last_state in agent.state_parse(ob):
                for next_state in agent.state_parse(next_ob):
                    next_action = np.argmax(agent.q[next_state])
                    td = reward + agent.gamma * agent.q[next_state][next_action] - agent.q[last_state][action]
                    agent.q[last_state][action] += agent.alpha * td

            if done:
                break

            steps += 1

            ob = next_ob
        
        episode_log["total_rwd"] = total_rwd
        episode_results[episode_i] = episode_log
        episode_results["total_rwd"].append(total_rwd)
        episode_results["length"].append(steps)

    return episode_results


# NOTE TODO:    instead of treating the state as seeing three enemies at once, treat it as three separate states, seeing one enemy at a time, and 
#               sum the three action distributions generated by applying the policy to the three states, and renormalize. MUCH smaller state space
#               to consider, and actually boosts the information gain as well.
if __name__ == '__main__':
    scenario = {
        "width": 4,
        "height": 4, 
        "cellsize": 100,
        "player": {
            "x": 0,
            "y": 0,
            "vision": 2
        },
        "enemies": [
            {
                "id": "A",
                "x": 2,
                "y": 2,
                "behavior": "chase",
                "action_rate": 6
            }
        ]
    }

    env = SMTEnv.register(
        ATESim, 
        display_screen=False, 
        state_preprocessor=ATESim.state_flatten, 
        config=scenario,
        fps=10
    )

    agent = QLearnAgent(env.action_space, state_parse=ate_state_parser)
    results = train(agent, env, num_episodes=1000)

    env.set_display(True)
    SMTEnv.run(env, agent, episode_count=5, debug=True)

    env.close()

    # ep = np.linspace(0, len(results["total_rwd"]), num=len(results["total_rwd"]))
    # plt.plot(ep, results["total_rwd"])
    # plt.show()