class HarlowExperimentEnv():
    def __init__(self, reward = 1, punishment = -1):
        """Initialize Harlow Experiment environment."""
        self.reward = reward
        self.punishment = punishment

        self.rewarded_digit = -1
        self.punished_digit = -1

        self.state = np.array([self.rewarded_digit, self.punished_digit])

    def update_state(self):
        """Update state by selecting rewarded hand for random."""
        if np.random.rand() < 0.5:
            self.state = np.array([self.rewarded_digit, self.punished_digit])
        else:
            self.state = np.array([self.punished_digit, self.rewarded_digit])

    def reset(self):
        """Reset environment by updating new rewarded and punished digits as well as create current state of the world (tuple of observations)."""
        self.rewarded_digit, self.punished_digit = np.random.choice(10, 2, replace=False)
        self.update_state()
        return self.state

    def step(self, action):
        """Evaluate agent's performance, return reward and next observation."""
        if self.state[action] == self.rewarded_digit:
            feedback = self.reward
        else:
            feedback = self.punishment
        self.update_state()
        return self.state, feedback

class ActorCritic(nn.Module):
    def __init__(self, hidden_size, num_inputs = 5, num_actions = 2):
        """Initialize Actor-Critic agent."""
        super(ActorCritic, self).__init__()

        #num_actions is 2 because left/right hand
        self.num_actions = num_actions

        #num_inputs is 5 because one-hot encoding of action (2) + reward (1) + previous state (2)
        self.num_inputs = num_inputs

        self.hidden_size = hidden_size

        #hyperparameters involved in training (important to keep assigned to the agent)
        self.learning_rate = 0.00075 #learning rate for optimizer
        self.discount_factor = 0.91 #gamma
        self.state_value_estimate_cost = 0.4 #beta_v
        self.entropy_cost = 0.001 #beta_e

        self.emb = nn.Linear(num_inputs, hidden_size)
        self.rnn = nn.LSTM(hidden_size, hidden_size)
        self.critic_linear = nn.Linear(hidden_size, 1)
        self.actor_linear = nn.Linear(hidden_size, num_actions)


    def forward(self, state, h, c):
        """Implement forward pass through agent."""
        #at first, input goes through embedding
        state = self.emb(state.unsqueeze(0))

        #then through RNNs (observe that we pass hidden states too!)
        state, hidden_states = self.rnn(state.unsqueeze(0), (h, c))
        state = state.squeeze(0)
        h, c = hidden_states
        state = state.squeeze(0)

        #critic -> value
        value = self.critic_linear(state)

        #actor -> policy
        policy_logits = self.actor_linear(state)

        return value, policy_logits, (h, c)

    def get_init_hidden_states(self, batch_size=1, device = device):
        """Initialize hidden state with 0."""
        #initialize hidden state in RNNs
        return (torch.zeros(1, batch_size, self.hidden_size).to(device), torch.zeros(1, batch_size, self.hidden_size).to(device))

def train_evaluate_agent(env, agent, optimizer_func, num_gradient_steps=25, num_trials=6, num_evaluation_trials=20):
    """Training and evaluation for agent in Harlow experiment environment.
    Evaluation goes only after all gradient steps.

    Inputs:
    - env (HarlowExperimentEnv): environment.
    - agent (ActorCritic): particular instance of Actor Critic agent to train.
    - optimizer_func (torch.optim.Optimizer): optimizer to use for training.
    - num_gradient_steps (int, default=25): number of gradient steps to perform.
    - num_trials (int, default=6): number of times the agent is exposed to the environment per gradient step to be trained.
    - num_evaluation_trials (int, default=20): number of times the agent is exposed to the environment to evaluate it (no training happens during this phase).

    Outputs:
    - score (int): cumulative reward over all trials of evaluation.
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    agent.to(device)

    # Training

    # Reset environment
    state = env.reset()

    # Define optimizer
    optimizer = optimizer_func(agent.parameters(), lr=agent.learning_rate, eps=1e-5)

    for _ in range(num_gradient_steps):

        # For storing variables for training
        log_probs = []
        values = []
        rewards = []
        entropy_term = torch.tensor(0., device=device)

        # Start conditions
        h, c = agent.get_init_hidden_states()
        preceding_reward = torch.tensor([0], device=device)
        preceding_action = torch.tensor([0, 0], device=device)

        for trial in range(num_trials):

            # State + reward + one-hot encoding of action
            full_state = torch.cat((torch.from_numpy(state).float().to(device), preceding_reward, preceding_action), dim=0)
            value, policy_logits, step_hidden_states = agent(full_state, h, c)
            h, c = step_hidden_states
            value = value.squeeze(0)

            # Sample action from policy
            dist = torch.distributions.Categorical(logits=policy_logits.squeeze(0))
            action = dist.sample()

            # Perform action to get reward and new state
            new_state, reward = env.step(action)

            # Update preceding variables
            preceding_reward = torch.tensor([reward], device=device)
            preceding_action = F.one_hot(action, num_classes=2).float().to(device)
            state = new_state

            # For training
            log_prob = dist.log_prob(action)
            entropy = dist.entropy()
            rewards.append(reward)
            values.append(value)
            log_probs.append(log_prob)
            entropy_term += entropy

        # Calculating loss
        Qval = 0
        Qvals = torch.zeros(len(rewards), device=device)
        for t in reversed(range(len(rewards))):
            Qval = rewards[t] + agent.discount_factor * Qval
            Qvals[t] = Qval
        values = torch.stack(values)
        log_probs = torch.stack(log_probs)
        advantage = Qvals - values
        actor_loss = (-log_probs * advantage.detach()).mean()
        critic_loss = advantage.pow(2).mean()
        entropy_term = entropy_term / num_trials

        # Loss incorporates actor/critic terms + entropy
        loss = actor_loss + agent.state_value_estimate_cost * critic_loss - agent.entropy_cost * entropy_term

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # Evaluation (on the same task after all gradient steps!)
    score = 0

    # Start conditions
    h, c = agent.get_init_hidden_states()
    preceding_reward = torch.tensor([0], device=device)
    preceding_action = torch.tensor([0, 0], device=device)

    for _ in range(num_evaluation_trials):

        # State + reward + one-hot encoding of action
        full_state = torch.cat((torch.from_numpy(state).float().to(device), preceding_reward, preceding_action), dim=0)
        value, policy_logits, step_hidden_states = agent(full_state, h, c)
        h, c = step_hidden_states
        value = value.squeeze(0)

        # Sample action from policy
        dist = torch.distributions.Categorical(logits=policy_logits.squeeze(0))
        action = dist.sample()

        # Perform action to get reward and new state
        new_state, reward = env.step(action)

        # Update preceding variables
        preceding_reward = torch.tensor([reward], device=device)
        preceding_action = F.one_hot(action, num_classes=2).float().to(device)
        state = new_state

        # Add reward to the score of agent
        score += reward

    return score

set_seed(42)

#define environment
env = HarlowExperimentEnv()

#define agent and optimizer
agent = ActorCritic(hidden_size=20).to(device)  # Move the agent to the device
optimizer_func = optim.RMSprop

#calculate score
score = train_evaluate_agent(env, agent, optimizer_func)
print(f"Score is {score}.")

def evaluate_individual(env, agent, optimizer_func, num_tasks = 10, num_gradient_steps = 25, num_trials = 6, num_evaluation_trials = 20):
    """Training and evaluation for agent in Harlow experiment environment for the bunch of tasks (thus measuring overall potential for agent's generalization across the tasks).
    Evaluation goes only after all gradient steps.

    Inputs:
    - env (HarlowExperimentEnv): environment.
    - agent (ActorCritic): particular instance of Actor Critic agent to train.
    - optimizer_func (torch.Optim): optimizer to use for training.
    - num_tasks (int, default = 10): number of tasks to evaluate agent on.
    - num_gradient_steps (int, default = 25): number of gradient steps to perform.
    - num_trials (int, default = 6): number of times the agent is exposed to the environment per gradient step to be trained .
    - num_evaluation_trials (int, default = 20): number of times the agent is exposed to the environment to evaluate it (no training happened during this phase).

    Outputs:
    - score (int): total score.
    - scores (list): list of scores obtained during evaluation on the specific tasks.
    """
    scores = []
    for _ in range(num_tasks): #lines[2-3]; notice that environment resets inside `train_evaluate_agent`
      agent_copy = copy.deepcopy(agent) #line[4]; remember that we don't want to change agent's parameters!
      score = train_evaluate_agent(env, agent_copy, optimizer_func, num_gradient_steps, num_trials, num_evaluation_trials)
      scores.append(score)
    return np.sum(scores), scores

set_seed(42)

#define environment
env = HarlowExperimentEnv()

#define agent and optimizer
agent = ActorCritic(hidden_size = 20)
optimizer_func = optim.RMSprop

#calculate score
total_score, scores = evaluate_individual(env, agent, optimizer_func)
print(f"Total score is {total_score}.")
plot_boxplot_scores(scores)

def create_initial_population(env, optimizer_func, population_size=50, hidden_size=20):
    """
    Creates an initial population of agents.

    Inputs:
    - env (HarlowExperimentEnv): environment.
    - optimizer_func (torch.Optim): optimizer to use for training.
    - population_size (int, default = 50): the size of the initial population.
    - hidden_size (int, default = 20): the size of LSTM layer in A2C agent.

    Outputs:
    - population (list): initial population which consists of tuples (agent, score).
    - best_score (int): the best score for the individual in the population registered so far.
    """
    population = []
    total_score = 0
    best_score = 0

    for _ in tqdm(range(population_size), desc="Creating Initial Population"):
        agent = ActorCritic(hidden_size)
        score, _ = evaluate_individual(env, agent, optimizer_func)
        best_score = max(best_score, score)
        total_score += score
        population.append((agent, score))

    print(f"Generation: 0, mean population score: {total_score / population_size}, best score: {best_score}")

    return population, best_score

#for mutation of (hyper)parameters
parameters_noise = 0.02
learning_rate_noise = 0.00005
discount_factor_noise = 0.01
state_value_estimate_cost_noise = 0.05
entropy_cost_noise = 0.001

def create_new_agent(agent1, agent2):
    """
    Creates new agent using crossing over technique over layers of network and mutation of the parameters with Gaussian noise.

    Inputs:
    - agent1 (ActorCritic): first parent agent.
    - agent2 (ActorCritic): second parent agent.

    Outputs:
    - new_agent (ActorCritic): new agent which is offspring of the given two.
    """
    #creates agent as copy of the first one
    new_agent = copy.deepcopy(agent1)

    #evolving network parameters with crossing over (over separate layes) & mutating (Gaussian noise)
    for name, module in new_agent.named_modules():
        if isinstance(module, nn.Linear):
            if random.random() < 0.5:
                module.weight.data = agent2._modules[name].weight.data
                module.bias.data = agent2._modules[name].bias.data
            #add noise
            module.weight.data += torch.randn_like(module.weight.data) * parameters_noise
            module.bias.data += torch.randn_like(module.bias.data) * parameters_noise
        elif isinstance(module, nn.LSTM):
            if random.random() < 0.5:
                module.weight_ih_l0.data = agent2._modules[name].weight_ih_l0.data
                module.weight_hh_l0.data = agent2._modules[name].weight_hh_l0.data
                module.bias_ih_l0.data = agent2._modules[name].bias_ih_l0.data
                module.bias_hh_l0.data = agent2._modules[name].bias_hh_l0.data
            #add noise
            module.weight_ih_l0.data += torch.randn_like(module.weight_ih_l0.data) * parameters_noise
            module.weight_hh_l0.data += torch.randn_like(module.weight_hh_l0.data) * parameters_noise
            module.bias_ih_l0.data += torch.randn_like(module.bias_ih_l0.data) * parameters_noise
            module.bias_hh_l0.data += torch.randn_like(module.bias_hh_l0.data) * parameters_noise

    #evolving & mutating hyperparameters
    if random.random() < 0.5:
        new_agent.learning_rate = agent2.learning_rate
    new_agent.learning_rate += np.random.normal(size = 1).item() * learning_rate_noise
    new_agent.learning_rate = min(max(new_agent.learning_rate, 0.0001), 0.01)
    if random.random() < 0.5:
        new_agent.discount_factor = agent2.discount_factor
    new_agent.discount_factor += np.random.normal(size = 1).item() * discount_factor_noise
    new_agent.discount_factor = min(max(new_agent.discount_factor, 0.6), 0.99)
    if random.random() < 0.5:
        new_agent.state_value_estimate_cost = agent2.state_value_estimate_cost
    new_agent.state_value_estimate_cost += np.random.normal(size = 1).item() * state_value_estimate_cost_noise
    new_agent.state_value_estimate_cost = min(max(new_agent.discount_factor, 0.1), 0.7)
    if random.random() < 0.5:
        new_agent.entropy_cost = agent2.entropy_cost
    new_agent.entropy_cost += np.random.normal(size = 1).item() * entropy_cost_noise
    new_agent.entropy_cost = min(max(new_agent.discount_factor, 0.0001), 0.05)

    return new_agent

def update_population(env, optimizer_func, population, parents_population, best_score, new_generation_new_individuals = 5):
    """
    Updates population with new individuals which are the result of crossing over and mutation of two parents agents.
    Removes the same amount of random agents from the population.

    Inputs:
    - env (HarlowExperimentEnv): environment.
    - optimizer_func (torch.Optim): optimizer to use for training.
    - population (list): current population which consists of tuples (agent, score).
    - parents_population (list) : parents individuals (part of current population) for creating new individuals.
    - best_score (int): the best score for the individual in the population registered so far.
    - new_generation_new_individuals (int, default = 5): the number of individuals to create (and the old ones to remove).
    """
    ###################################################################
    ## Fill out the following then remove
    raise NotImplementedError("Student exercise: complete update of the population logic.")
    ###################################################################

    # Create new individuals with progress bar
    new_individuals = []
    for _ in tqdm(range(new_generation_new_individuals), desc="Creating New Individuals"):
        agent1, agent2 = random.choices(..., k=2)
        new_agent = create_new_agent(agent1[0], agent2[0])
        score, _ = evaluate_individual(env, ..., optimizer_func)
        # Evaluate whether best score has increased
        best_score = max(score, best_score)
        new_individuals.append((new_agent, score))

    # Remove random old individuals with progress bar
    for _ in tqdm(range(new_generation_new_individuals), desc="Removing Old Individuals"):
        population.pop(random.randint(0, len(population) - 1))

    return population + new_individuals, best_score

# Selection - random
random.seed(42)

# Define environment
env = HarlowExperimentEnv()

# Define agent and optimizer
agent = ActorCritic(hidden_size=20)
optimizer_func = optim.RMSprop

# GA constants
num_generations = 1
tournament_size = 20
parents_size = 4
new_generation_new_individuals = 5

mean_population_scores = []

# Measure the time for creating the initial population
population, best_score = create_initial_population(env, optimizer_func)

for generation in range(1, num_generations + 1):

    # Selection
    sorted_population = sorted(population, key=lambda x: x[1], reverse=True)
    tournament_population = sorted_population[:tournament_size]

    # Random choice of parents from tournament population
    parents_population = random.choices(tournament_population, k=parents_size)

    # Update population
    population, best_score = update_population(env, optimizer_func, population, parents_population, best_score)
    mean_population_scores.append(np.mean([agent_score[1] for agent_score in population]))