__author__ = 'Tom Schaul, tom@idsia.ch'

""" 
Doing RL when an environment model (transition matrices and rewards) are available,
and the states are observed by a feature vector for each state:

 - a feature map (fMap) is a 2D array of features, one row per state.

(otherwise same representation as for in policyiteration.py)

Here we provide some algorithms when the values are estimated linearly from the features.

"""

# TODO: optimize LSTDQ by vectorization

import random
from scipy import ravel, zeros, outer, dot, tile, transpose, tensordot
from scipy.linalg import lstsq
from numpy.matlib import repmat


from pybrain.rl.learners.modelbased.policyiteration import randomPolicy, greedyQPolicy, collapsedTransitions, policyIteration


def trueFeatureStats(T, R, fMap, discountFactor, stateProp=1, MAT_LIMIT=1e8):
    """ Gather the statistics needed for LSTD,
    assuming infinite data (true probabilities).
    Option: if stateProp is  < 1, then only a proportion of all 
    states will be seen as starting state for transitions """
    dim = len(fMap)
    numStates = len(T)
    statMatrix = zeros((dim, dim))
    statResidual = zeros(dim)
    ss = list(range(numStates))
    repVersion = False
    
    if stateProp < 1:
        ss = random.sample(ss, int(numStates * stateProp))
    elif dim * numStates**2 < MAT_LIMIT:
        repVersion = True
    
    # two variants, depending on how large we can afford our matrices to become.        
    if repVersion:    
        tmp1 = tile(fMap, (numStates,1,1))
        tmp2 = transpose(tmp1, (2,1,0))
        tmp3 = tmp2 - discountFactor * tmp1            
        tmp4 = tile(T, (dim,1,1))
        tmp4 *= transpose(tmp1, (1,2,0))
        statMatrix = tensordot(tmp3, tmp4, axes=[[0,2], [1,2]]).T
        statResidual = dot(R, dot(fMap, T).T)
    else:
        for sto in ss:
            tmp = fMap - discountFactor * repmat(fMap[:, sto], numStates, 1).T
            tmp2 = fMap * repmat(T[:, sto], dim, 1)
            statMatrix += dot(tmp2, tmp.T)             
            statResidual += R[sto] * dot(fMap, T[:, sto])
    return statMatrix, statResidual


def LSTD_values(T, R, fMap, discountFactor, **kwargs):
    """ Least-squares temporal difference algorithm. """
    statMatrix, statResidual = trueFeatureStats(T, R, fMap, discountFactor,**kwargs)
    weights = lstsq(statMatrix, statResidual)[0]
    return dot(weights, fMap)


def LSTD_Qvalues(Ts, policy, R, fMap, discountFactor):
    """ LSTDQ is like LSTD, but with features replicated 
    once for each possible action.
    
    Returns Q-values in a 2D array. """
    numA = len(Ts)
    dim = len(Ts[0])
    numF = len(fMap)
    fMapRep = zeros((numF * numA, dim * numA))
    for a in range(numA):
        fMapRep[numF * a:numF * (a + 1), dim * a:dim * (a + 1)] = fMap

    statMatrix = zeros((numF * numA, numF * numA))
    statResidual = zeros(numF * numA)
    for sto in range(dim):
        r = R[sto]
        fto = zeros(numF * numA)
        for nextA in range(numA):
            fto += fMapRep[:, sto + nextA * dim] * policy[sto][nextA]
        for sfrom in range(dim):
            for a in range(numA):
                ffrom = fMapRep[:, sfrom + a * dim]
                prob = Ts[a][sfrom, sto]
                statMatrix += outer(ffrom, ffrom - discountFactor * fto) * prob
                statResidual += ffrom * r * prob

    Qs = zeros((dim, numA))
    w = lstsq(statMatrix, statResidual)[0]
    for a in range(numA):
        Qs[:,a] = dot(w[numF*a:numF*(a+1)], fMap)
    return Qs


def LSPI_policy(fMap, Ts, R, discountFactor, initpolicy=None, maxIters=20):
    """ LSPI is like policy iteration, but Q-values are estimated based 
    on the feature map. 
    Returns the best policy found. """
    if initpolicy is None:
        policy, _ = randomPolicy(Ts) 
    else:
        policy = initpolicy
    
    while maxIters > 0:
        Qs = LSTD_Qvalues(Ts, policy, R, fMap, discountFactor)
        newpolicy = greedyQPolicy(Qs)
        if sum(ravel(abs(newpolicy - policy))) < 1e-3:
            return policy, collapsedTransitions(Ts, policy)
        policy = newpolicy
        maxIters -= 1
    return policy, collapsedTransitions(Ts, policy)


def LSTD_PI_policy(fMap, Ts, R, discountFactor, initpolicy=None, maxIters=20):
    """ Alternative version of LSPI using value functions
    instead of state-action values as intermediate.
    """
    def veval(T):
        return LSTD_values(T, R, fMap, discountFactor)
    return policyIteration(Ts, R, discountFactor, VEvaluator=veval, 
                           initpolicy=initpolicy, maxIters=maxIters)
