import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

import numpy as np
import pandas as pd

from scipy.stats import spearmanr
from scipy.spatial.distance import jensenshannon

"""
the probability of a learner's belief that the underlying utility function is U, 
given observed demo d is:
P_L(U | d) \propto P(d | U) * P(U)

How do we get P(d | U)?
Intuitively, P(d | U) should capture how "good" demo d is at conveying utility function U.
A demo consists of a set of objects X, and the target object selected t. It follows that
P(d | U) = P(t | X, U) \propto exp(U(X)^alpha)
"""

def compute_demo_score(all_possible_games,
                        is_first_agent, 
                        reward_matrices,
                        game,
                        teacher_alpha,
                        alpha=1):
    """
    Given the set of all possible demos that could be shown the student, compute the 
    score associated with each demo

    Arguments:
        all_possible_games (torch.Tensor): 
            size (batch_size, 560, num_choices_in_listener_context, obj_encoding_len)
        is_first_agent (Bool):
            True if first agent in the chain, false otherwise
        reward_matrices (torch.Tensor):
            size (batch_size, num_colors*num_shapes, num_colors + num_shapes + 1)
        game (SignalingBanditsGame):
            need this to call the game.compute_reward function
        alpha (int):
            temperature parameter for the softmax

    Return:
        demo_scores (torch.Tensor):
            size (batch_size, 560)
    """
    batch_size = all_possible_games.shape[0]
    num_games = all_possible_games.shape[1]

    rewards_only = game.compute_rewards(all_possible_games, reward_matrices)    # (batch_size, num_games, 3)

    max_vals, _ = torch.max(rewards_only, dim=-1)
    numerator = torch.exp(alpha * max_vals)
    denominator = torch.sum(torch.exp(alpha * rewards_only), dim=-1)

    demo_scores = numerator / denominator
    #breakpoint()
    #demo_scores = torch.exp(100 * demo_scores.log())
    demo_scores = F.softmax(teacher_alpha * demo_scores, dim=-1)
    
    return demo_scores


def compute_correlation(neural_game_scores, bayesian_demo_scores):
    """
    Compute the strength of correlation between the scores over all possible games generated 
    by the neural model, and the scores over all possible demos generated by the Bayesian model
    """
    #breakpoint()
    jsd = jensenshannon(neural_game_scores.detach().cpu() + 0.00000001, bayesian_demo_scores.detach().cpu() + 0.00000001, axis=-1)
    mean_jsd = np.mean(jsd)

    if np.isinf(mean_jsd): breakpoint()
    
    return mean_jsd


def compute_mean_bayesian_score(neural_game_scores, bayesian_demo_scores):
    """
    Return the Bayesian model score associated with the highest scoring demo
    as indicated by the neural model. The goal here is to determine whether the neural model's
    prediction for the best demo aligns at all with the Bayesian model
    """
    batch_size = neural_game_scores.shape[0]
    idx = torch.argmax(neural_game_scores, dim=-1)
    bayesian_demo_scores_indexed = bayesian_demo_scores[range(batch_size), idx]
    mean = torch.mean(bayesian_demo_scores_indexed)
    #breakpoint()
    return mean.item()



def bayesian_teacher(all_possible_games,
                    is_first_agent, 
                    reward_matrices,
                    game,
                    teacher_alpha,
                    games_for_eval):

    """
    Return:
    scores_i (torch.Tensor): These are the scores over the eval games. shape (batch_size, num_games_for_eval, num_choices_in_game)
    demo_i (torch.Tensor): shape (batch_size, num_demos, num_choices_in_game, obj_encoding_len+1)
    """

    # scores_i are just the raw ground truth values (so we will assume the teacher is 100% accurate)
    scores_i = game.compute_rewards(games_for_eval, reward_matrices) 

    # demo_i is the concatenation of the best demo and the teacher's selection (assuming the teacher
    # always makes the optimal choice)
    bayesian_demo_scores = compute_demo_score(all_possible_games,
                                                is_first_agent, 
                                                reward_matrices,
                                                game,
                                                teacher_alpha)
    indices = bayesian_demo_scores.argmax(dim=-1)

    demo_only = all_possible_games[torch.arange(all_possible_games.shape[0]), indices].to(reward_matrices.device)

    # now compute the rewards associated with each demo
    rewards_associated_with_demo = game.compute_rewards(demo_only.unsqueeze(1), reward_matrices)

    # assume the teacher always picks the highest reward
    argmax = rewards_associated_with_demo.argmax(-1).squeeze()
    num_classes = demo_only.shape[1]
    argmax_as_onehot = F.one_hot(argmax, num_classes=num_classes).unsqueeze(-1).float()
    
    demo_i = torch.cat([demo_only, argmax_as_onehot], dim=-1)                             
    
    return demo_i, scores_i

