#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Policy gradient stimators
@author: Matteo Papini
"""

import torch
import potion.common.torch_utils as tu
from potion.common.misc_utils import unpack, discount
from potion.common.torch_utils import tensormat, jacobian
from potion.estimation.moments import incr_mean, incr_var

def off_gpomdp_estimator(batch, disc, policy, target_params, 
                         baselinekind='avg', 
                         result='mean',
                         shallow=False):
    """off-policy G(PO)MDP policy gradient estimator
       
    batch: list of N trajectories generated by behavioral policy. Each trajectory is a tuple 
        (states, actions, rewards, mask). Each element of the tuple is a 
        tensor where the first dimension is time.
    disc: discount factor
    policy: behavioral policy, the one used to collect the data
    target_params: parameters of the policy to evaluate
    baselinekind: kind of baseline to employ in the estimator. 
        Either 'avg' (average reward, default), 'peters' 
        (variance-minimizing),  or 'zero' (no baseline)
    result: whether to return the final estimate ('mean', default), or the 
        single per-trajectory estimates ('samples')
    shallow: whether to use precomputed score functions (only available
        for shallow policies)
    """
    if shallow:
        return _shallow_off_gpomdp_estimator(batch, disc, policy, target_params,
                                             baselinekind, result)
    else:
        #Data
        N = len(batch)
        states, actions, rewards, mask, _ = unpack(batch) #NxHxd_s, NxHxd_a, NxH, NxH
        H = rewards.shape[1]
        m = policy.num_params()
        disc_rewards = discount(rewards, disc) #NxH
        
        #Behavioral
        behavioral_params = policy.get_flat()
        behavioral_logps = policy.log_pdf(states, actions) #NxH
        
        #Target
        policy.set_from_flat(target_params)
        target_logps = policy.log_pdf(states, actions) * mask #NxH
        cm_logps = torch.cumsum(target_logps, 1) #NxH
        
        log_iws = torch.cumsum((target_logps - behavioral_logps) * mask, 1) #NxH
        stabilizers, _ = torch.max(log_iws, dim=0, keepdim=True) #NxH
        
        if baselinekind == 'peters':
            jac = jacobian(policy, cm_logps.view(-1)).reshape((N,H,m)) #NxHxm   
            b_num = torch.sum(tensormat((jac * torch.exp(log_iws - stabilizers).unsqueeze(1))**2, 
                                        disc_rewards), 0) #Hxm
            b_den = torch.sum((jac * torch.exp(log_iws - stabilizers).unsqueeze(1))**2, 0) #Hxm
            baseline = b_num / b_den #Hxm
            baseline[baseline != baseline] = 0
            values = disc_rewards.unsqueeze(2) - baseline.unsqueeze(0) #NxHxm
            _samples = torch.sum(tensormat(values * torch.exp(log_iws).unsqueeze(1) * jac, mask), 1) #Nxm
        else:
            if baselinekind == 'avg':
                baseline = torch.mean(disc_rewards, 0) #H
            else:
                baseline = torch.zeros(1) #1
            values = (disc_rewards - baseline) * mask #NxH
            
            _samples = torch.stack([tu.flat_gradients(policy, cm_logps[i,:], 
                                                  values[i,:] * torch.exp(log_iws)[i,:])
                                           for i in range(N)], 0) #Nxm
        
        #restore (behavioral) policy
        policy.set_from_flat(behavioral_params)
        if result == 'samples':
            return _samples #Nxm
        else:
            return torch.mean(_samples, 0) #m

def _shallow_off_gpomdp_estimator(batch, disc, policy, target_params, 
                         baselinekind, result):
    with torch.no_grad():
        states, actions, rewards, mask, _ = unpack(batch) # NxHxm, NxHxd, NxH, NxH
        disc_rewards = discount(rewards, disc) #NxH
        
        #Behavioral
        behavioral_params = policy.get_flat()
        behavioral_logps = policy.log_pdf(states, actions) #NxH
        
        #Target
        policy.set_from_flat(target_params)
        target_logps = policy.log_pdf(states, actions) #NxH
        scores = policy.score(states, actions) #NxHxM
        
        #restore (behavioral) policy
        policy.set_from_flat(behavioral_params)

        G = torch.cumsum(tensormat(scores, mask), 1) #NxHxm
        n_k = torch.sum(mask, dim=0) #H
        n_k[n_k==0.] = 1.
              
        log_iws = torch.cumsum((target_logps - behavioral_logps) * mask, 1) #NxH
        stabilizers, _ = torch.max(log_iws, dim=0, keepdim=True) #NxH
        
        if baselinekind == 'peters': #Mastrangelo's variant
            baseline = torch.sum(tensormat(G ** 2, disc_rewards * torch.exp(2*(log_iws - stabilizers))), 0) / \
                            torch.sum(tensormat(G ** 2, torch.exp(2*(log_iws - stabilizers))), 0) #Hxm
            baseline[baseline != baseline] = 0 #removes non-real values
            values = disc_rewards.unsqueeze(2) - baseline.unsqueeze(0) #NxHxm
        else:
            if baselinekind == 'zero':
                baseline = torch.zeros_like(disc_rewards[0]) #H
            elif baselinekind == 'avg':
                baseline = torch.mean(disc_rewards, 0) #H
            else:
                raise NotImplementedError
            baseline[baseline != baseline] = 0 #removes non-real values
            values = (disc_rewards - baseline.unsqueeze(0)).unsqueeze(2) #NxHxm     
        
        _samples = torch.sum(tensormat(G * values, mask * torch.exp(log_iws)), 1) #Nxm
        if result == 'samples':
            return _samples #Nxm
        else:
            return torch.mean(_samples, 0) #m

def off_reinforce_estimator(batch, disc, policy, target_params, 
                         baselinekind='avg', 
                         result='mean',
                         shallow=False):
    """off-policy reinforce policy gradient estimator
       
    batch: list of N trajectories generated by behavioral policy. Each trajectory is a tuple 
        (states, actions, rewards, mask). Each element of the tuple is a 
        tensor where the first dimension is time.
    disc: discount factor
    policy: behavioral policy, the one used to collect the data
    target_params: parameters of the policy to evaluate
    baselinekind: kind of baseline to employ in the estimator. 
        Either 'avg' (average reward, default), 'peters' 
        (variance-minimizing),  or 'zero' (no baseline)
    result: whether to return the final estimate ('mean', default), or the 
        single per-trajectory estimates ('samples')
    shallow: whether to use precomputed score functions (only available
        for shallow policies)
    """
    if shallow:
        return _shallow_off_reinforce_estimator(batch, disc, policy, target_params, 
                                 baselinekind, result)
    else:
        #data
        N = len(batch)
        states, actions, rewards, mask, _ = unpack(batch) #NxHxm, NxHxd, NxH, NxH
        
        disc_rewards = discount(rewards, disc) #NxH
        rets = torch.sum(disc_rewards, 1) #N
        
        #Behavioral
        behavioral_params = policy.get_flat()
        behavioral_logps = policy.log_pdf(states, actions) #NxH
        
        #Target
        policy.set_from_flat(target_params)
        target_logps = policy.log_pdf(states, actions) * mask #NxH
    
        log_iws = torch.sum((target_logps - behavioral_logps) * mask, 1) #N
        stabilizers, _ = torch.max(log_iws, dim=0, keepdim=True) #N
        
        if baselinekind == 'peters':
            logp_sums = torch.sum(target_logps, 1) #N
            jac = jacobian(policy, logp_sums) #Nxm   
            b_num = torch.sum((jac * torch.exp(log_iws - stabilizers).unsqueeze(1))**2 * rets.unsqueeze(1), 0) #m
            b_den = torch.sum((jac * torch.exp(log_iws - stabilizers).unsqueeze(1))**2, 0) #m
            baseline = b_num / b_den #m
            baseline[baseline != baseline] = 0
            values = rets.unsqueeze(1) - baseline.unsqueeze(0) #Nxm
            _samples = jac * values * torch.exp(log_iws).unsqueeze(1)
        else:
            if baselinekind == 'avg':
                baseline = torch.mean(rets, 0) #1
            else:
                baseline = torch.zeros(1) #1
            baseline[baseline != baseline] = 0
            values = rets - baseline #N
            
            if result == 'mean':
                logp_sums = torch.sum(target_logps, 1)
                grad = tu.flat_gradients(policy, logp_sums, values * torch.exp(log_iws)) / N
                #restore (behavioral) policy
                policy.set_from_flat(behavioral_params)
                return grad
            
            _samples = torch.stack([tu.flat_gradients(policy, target_logps[i,:]) * 
                                                         values[i] * torch.exp(log_iws)[i]
                                           for i in range(N)], 0) #Nxm
        
        #restore (behavioral) policy
        policy.set_from_flat(behavioral_params)
        if result == 'samples':
            return _samples #Nxm
        else:
            return torch.mean(_samples, 0) #m

def _shallow_off_reinforce_estimator(batch, disc, policy, target_params, 
                         baselinekind, result):
    with torch.no_grad():
        states, actions, rewards, mask, _ = unpack(batch) # NxHxm, NxHxd, NxH, NxH
        disc_rewards = discount(rewards, disc) #NxH 
        returns = torch.sum(disc_rewards, 1) #N
        
        #Behavioral
        behavioral_params = policy.get_flat()
        behavioral_logps = policy.log_pdf(states, actions) #NxH
        
        #Target
        policy.set_from_flat(target_params)
        target_logps = policy.log_pdf(states, actions) #NxH
        scores = policy.score(states, actions) #NxHxM
        
        #restore (behavioral) policy
        policy.set_from_flat(behavioral_params)        

        G = torch.sum(tensormat(scores, mask), 1) #Nxm
        
        log_iws = torch.sum((target_logps - behavioral_logps) * mask, 1) #N
        stabilizers, _ = torch.max(log_iws, dim=0, keepdim=True) #N
        
        if baselinekind == 'peters': #Mastrangelo's variant
            baseline = torch.sum(returns.unsqueeze(1) * (torch.exp(log_iws - stabilizers).unsqueeze(1) * G)**2, 0) / \
                            torch.sum((torch.exp(log_iws - stabilizers).unsqueeze(1) * G)**2, 0) #m
            baseline[baseline != baseline] = 0 #removes non-real values
            values = returns.unsqueeze(1) - baseline.unsqueeze(0) #Nxm
        else:
            if baselinekind == 'zero':
                baseline = torch.zeros(1) #1
            elif baselinekind == 'avg':
                baseline = torch.mean(returns, 0) #1
            else:
                raise NotImplementedError
            baseline[baseline != baseline] = 0 #removes non-real values
            values = (returns - baseline.unsqueeze(0)).unsqueeze(1) #Nxm     
        
        _samples = torch.exp(log_iws).unsqueeze(1) * values * G #Nxm
        if result == 'samples':
            return _samples #Nxm
        else:
            return torch.mean(_samples, 0) #m        
        
"""Testing"""
if __name__ == '__main__':
    from potion.actors.continuous_policies import ShallowGaussianPolicy as Gauss
    from potion.simulation.trajectory_generators import generate_batch
    from potion.common.misc_utils import seed_all_agent
    import potion.envs
    import gym.spaces
    from potion.estimation.gradients import gpomdp_estimator
    
    env = gym.make('ContCartPole-v0')
    env.seed(0)
    seed_all_agent(0)
    N = 100
    H = 100
    disc = 0.99
    pol = Gauss(4,1, mu_init=[0.,0.,0.,0.], learn_std=True)
    
    batch = generate_batch(env, pol, H, N)
    
    on = gpomdp_estimator(batch, disc, pol, baselinekind='peters', 
                         shallow=True)
    
    off = off_gpomdp_estimator(batch, disc, pol, pol.get_flat(), 
                               baselinekind='peters',
                               shallow=True)
    print(on, off)
    