"""
Infinite horizon principal-agent problem with hidden effort 
This calibration uses the Phelan Townsend Calibration
Uses the PuLP linear programming solver http://code.google.com/p/pulp-or/

Created by Ryan Banerjee, 30 May 2009"""

from __future__ import division
import pulp as pp
import scipy as sp
import scipy.optimize
import time
import matplotlib.pyplot as plt
import pytrix.tseries as ts
import bkfilter

# Helper functions
def makeVarNames(i, j, k, l, zfillfactor=2):
    """Creates as string for the variable names with padding:
    Pa[0i]q[0j]c[0k]w[0l]"""
    p = zfillfactor
    return 'Pa' + str(i).zfill(p) + 'q' + str(j).zfill(p)\
            + 'c' + str(k).zfill(p) + 'w' + str(l).zfill(p)

def buildProb(Action, Output, Cons, W, Cprob):
    """Build strings for linear program which only need to be done once
    Returns
    pmf - list of strings of probabilities
    sumToOne - all probabilities must sum to one
    uprincipalPeriod - one period utility funtion 
    """
    
    # Create a list of strings to store names probabilities
    # They are strings 'Pa[i]q[j]c[k]w[l]'
    # when put into a dictionary, they can be indexed easily with
    # appropriately named strings
    pmf=[]
    for i in xrange(0,Action):
        for j in xrange(0,Output):
            for k in xrange(0,Cons) :
                for l in xrange(0,W):
                    pmf.append(makeVarNames(i,j,k,l))

    # A dictionary called 'Vars' is created to reference the variables
    # The ( , ,0 = min, 1 = max) value each variable can take
    vars = pp.LpVariable.dicts("Pr", pmf, 0, 1)

    # Constraints on probabilities (from definitions of probabilities)
    # probabilities sum to one
    sumToOne=[]
    for i in xrange(0,Action):
        for j in xrange(0,Output):
            for k in xrange(0,Cons) :
                for l in xrange(0,W):
                    pmfTemp = makeVarNames(i,j,k,l)
                    sumToOne.append(vars[pmfTemp])

    # Period utility for policy iteration step
    uprincipalPeriod = sp.zeros((Action*Output*Cons*W))
    iCount = 0
    for i in xrange(0,Action):
        for j in xrange(0,Output):
            for k in xrange(0,Cons) :
                for l in xrange(0,W):
                    uprincipalPeriod[iCount] = output[j] - cons[k]
                    iCount += 1

    # Constraints on conditional probabilities
    condProb=[]
    for i in xrange(0,Action):
        for j in xrange(0,Output):
            condProbij=[]
            for k in xrange(0,Cons) :
                for l in xrange(0,W):
                    pmfTemp = makeVarNames(i,j,k,l)
                    condProbij.append(vars[pmfTemp])
                    for q in xrange(0,Output):
                        qpmfTemp = makeVarNames(i,q,k,l)
                        condProbij.append(-vars[qpmfTemp] * Cprob[i,j])

            # condProb is a list of lists
            condProb.append(condProbij)

    return [pmf, sumToOne, condProb, vars, uprincipalPeriod]

# Put below into a linear programming function
def solvePoint(v, wprime, wi, fullInfo, pmf, sumToOne, condProb):
    """Function to solve for optimal policy and the value to the 
    principal given the state variable promised utility w[i] 
    Solve v(w[i]) = max u( ) + beta * vprime
    Inputs:
    v = v0: Previous iteration of being in state for principal
    wprime = w: promised amount tomorrow
    wi = w[i]: Promised amount
    fullInfo = False or True"""

    
    # Create the pulp.prob variable to contain the problem data
    prob = pp.LpProblem("One period principal agent problem", pp.LpMaximize)

    # A dictionary called 'Vars' is created to reference the variables
    # The ( , ,0 = min, 1 = max) value each variable can take
    #vars = pp.LpVariable.dicts("Pr", pmf, 0, 1)
    
    # Generate utilities in each state
    uagent = sp.zeros((Action,Output,Cons,W))
    uprincipal = sp.zeros((Action,Output,Cons,W))
    for i in xrange(0,Action):
        for j in xrange(0,Output):
            for k in xrange(0,Cons) :
                for l in xrange(0,W):
                    uprincipal[i,j,k,l] = output[j] - cons[k] + betaP * v[l] 
                    uagent[i,j,k,l] = cons[k]**(s) / (s)\
                            + (amax - action[i])**(s) / (s)\
                            + betaA * wprime[l]

    # The objective function is added to 'prob' first
    objectiveP = []
    objectiveA = []
    for i in xrange(0,Action):
        for j in xrange(0,Output):
            for k in xrange(0,Cons) :
                for l in xrange(0,W):
                    pmfTemp = makeVarNames(i,j,k,l) 
                    objectiveP.append(uprincipal[i,j,k,l] * vars[pmfTemp])
                    objectiveA.append(uagent[i,j,k,l] * vars[pmfTemp])

    # Add the objective function to the prob object first
    prob += pp.lpSum(objectiveP), "Utility of principal"

    # IR constaint (1 equation)
    prob += pp.lpSum(objectiveA) == wi, "IR constraint"

    # IC constraints A x A  of them
    if fullInfo == False:
        for i in xrange(0,Action):  # loop for A's
            for ahati in xrange(0,Action):      # loop for Ahat - the IC constraint starts here
                IC = []
                for j in xrange(0,Output):
                    for k in xrange(0,Cons) :
                        for l in xrange(0,W):
                            pmfTemp = makeVarNames(i,j,k,l)
                            if i != ahati:
                                IC.append((uagent[i,j,k,l]\
                                        - uagent[ahati,j,k,l]\
                                        * Cprob[ahati,j] / Cprob[i,j])\
                                        * vars[pmfTemp])

                # Add IC constraint to prob
                if i != ahati:
                    prob += pp.lpSum(IC) >= 0, 'IC a ' + str(i) + ' ahat ' + str(ahati)

    prob += pp.lpSum(sumToOne) == 1, "Probs sum to one"

    # Add to constraints
    j = 0
    for i in condProb:
        prob += pp.lpSum(i) == 0, 'CondProb a' + str(j)
        j += 1

    # The problem data is written to an .lp file
    prob.writeLP("PrincipalAgent.lp")

    # The problem is solved using PuLP's choice of Solver
    prob.solve()

    # The status of the solution is printed to the screen
    print "Status:", pp.LpStatus[prob.status]

    # For each variable printed with it's optimum value
    pvalues={}
    for v in prob.variables():
        pvalues[v.name] = v.varValue

    # Put into a policy list
    policyw =[]
    for i in xrange(0,Action):
        for j in xrange(0,Output):
            for k in xrange(0,Cons):
                for l in xrange(0,W):
                    pmfTemp = 'Pr_' + makeVarNames(i,j,k,l)
                    policyw.append(pvalues[pmfTemp])

    # Print the optimised objective function
    print "Utility of principal is = ", pp.value(prob.objective)
    #print "Utility of agent is = ", utilityAgent

    return [pp.value(prob.objective), policyw]

def policyStep(v, policy):
    """Policy iteration step, inputs
    v1 - starting value for fixed point iteration
    policy - w x (A*Q*C*W) matrix of probabilites
    uprincipal one period utility of principal at this point
    """
    out = v - betaP * sp.dot(policy,
                            sp.kron(sp.ones(Action*Output*Cons), v))\
            - sp.dot(policy, uprincipalPeriod)
            
    return out
    

###########################################
# Start of program
###########################################
# The order for all indexing is [a, q, c, w]

start = time.time()

# Fullinformation problem
fullInfo = False    # True of False

# States - could put this in a dictionary?
# Action
actionMax = 0.6
actionMin = 0
Action = 4
action = sp.linspace(actionMin, actionMax, Action)

# Output (or q)
outputMin = 1
outputMax = 2
Output = 2
output = sp.linspace(outputMin, outputMax, Output)

# Consumption - transfer from principal to agent
consMin = 0
consMax = 2.25 #2.25
Cons = 55 #55
cons = sp.linspace(consMin, consMax, Cons)


# Conditional probabilities of output realisation conditional on effort P(q|a)
# Must sum to one across the rows
Cprob = sp.zeros((Action, Output))
Cprob[0,0] = 0.9
Cprob[0,1] = 0.1
Cprob[1,0] = 0.6
Cprob[1,1] = 0.4
Cprob[2,0] = 0.4
Cprob[2,1] = 0.6
Cprob[3,0] = 0.25
Cprob[3,1] = 0.75


# Define preferences and utilities over the state grid
# Principal
# Risk neutral
betaP  = 0.8 

# Agent
betaA = betaP
# u(a,c) = c**(s)/(s) + (amax-a)**(s)/(s)
s = 0.5
amax = 1

# Promised utility
wMin = 1 / (1 - betaA) * ((consMin**s / s) + (amax - actionMax)**s / s)
wMax = 1 / (1 - betaA) * ((consMax**s / s) + (amax - actionMin)**s / s) 
W = 25 #25
w = sp.linspace(wMin, wMax, W)
 

# Get some variables we need repeatedly in the LP stage
pmf, sumToOne, condProb, vars, uprincipalPeriod = buildProb(Action, Output, Cons, W, Cprob)

# Initial guess for future value for principal
v0 = sp.ones(W)
v1 = sp.ones(W)
tol = 1e-5
diff = 1

# Value function iteration
iter = 0
while diff > tol:
    policy = sp.zeros((W, Action * Output * Cons * W))
    for i in xrange(0,W):
        v1[i], policy[i] = solvePoint(v0, w, w[i], fullInfo,\
                pmf, sumToOne, condProb)

    diff = max(abs(v1-v0))
    print 'diff =' , diff
    # Policy function iteration step
    vpol = sp.optimize.fsolve(policyStep, v1, args = (policy))
    v0 = vpol.copy()
    #v0 = v1.copy()  # Note assignment (=) creates a pointer - careful
    iter += 1

end = time.time()
print "Time elapsed = ", end - start, "seconds"

# Get graphs of policy function - as in Phelan Townsend
pmf=[]
for i in xrange(0,Action):
    for j in xrange(0,Output):
        for k in xrange(0,Cons) :
            for l in xrange(0,W):
                pmf.append(makeVarNames(i,j,k,l))

eAction=[]
eCons = sp.zeros((W, Action, Output))
eWprime = sp.zeros((W, Action, Output))
for wi in xrange(0,W):
    condProbAQ = sp.zeros((Action, Output))
    actionTemp = 0
    for x in xrange(0, Output * Action * Cons * W):
        i = int(pmf[x][2:4])    # unpack action state
        j = int(pmf[x][5:7])    # unpack output
        k = int(pmf[x][8:10])    # unpack consumption
        l = int(pmf[x][11:13])  # unpack wprime
        actionTemp += action[i] * policy[wi,x]
        condProbAQ[i,j] += policy[wi,x]
        eCons[wi,i,j] += cons[k] * policy[wi,x]
        eWprime[wi,i,j] += w[l] * policy[wi,x]

    # Compute expected consumption and Wprim conditional on (a,q) pair and w
    eCons[wi] = eCons[wi] / condProbAQ
    eWprime[wi] = eWprime[wi] / condProbAQ
    # Expected action conditional on w
    eAction.append(actionTemp)
    eAction

plt.figure()
plt.plot(w, v0)
plt.figure()
plt.plot(w, eAction)
plt.figure()
waxis = sp.transpose(sp.append([w],[w], axis=0))
for i in xrange(0,Action):
    plt.plot(waxis, eCons[:,i,:])
plt.figure()
for i in xrange(0,Action):
    plt.plot(waxis, eWprime[:,i,:])


# Simulate model - put this in a function (of initial w and noOfSims)

# Initial promise
promise = sp.ceil(W/2)      # Location of lowest w so it does not bind

# Generate array of RV's
noOfSims = 150 
rv = sp.random.random(noOfSims) # Draw from u[0,1]

simQ = sp.zeros(noOfSims)
simA = sp.zeros(noOfSims)
simC = sp.zeros(noOfSims)
simWprime = sp.zeros(noOfSims)

# Create a index of states 
indexA=[]
indexQ=[]
indexC=[]
indexW=[]
for i in xrange(0,Action):
    for j in xrange(0,Output):
        for k in xrange(0,Cons) :
            for l in xrange(0,W):
                indexA.append(i)
                indexQ.append(j)
                indexC.append(k)
                indexW.append(l)

# Begin iteration
for i in xrange(0,noOfSims):
    pdf = policy[promise]
    cdf = pdf.cumsum()
    b = cdf - rv[i] < 0
    index = b.argmin()  # Find index of first false
    # Find q, a, c, wprime indices
    simQ[i] = output[indexQ[index]]
    simA[i] = action[indexA[index]]
    simC[i] = cons[indexC[index]]
    promise = indexW[index]
    simWprime[i] = w[promise]

# Filter the data as I do with the actual data
fSimQ = ts.hpfilter(simQ, penalty=10, percentage_dev=True)[1]
fSimA = ts.hpfilter(simA, penalty=10, percentage_dev=True)[1]
fSimC = ts.hpfilter(simC, penalty=10, percentage_dev=True)[1]
fSimWprime = ts.hpfilter(simWprime, penalty=10, percentage_dev=True)[1]

# Can also detrend with the bkfilter
bkSimQ = bkfilter.bkfilter(simQ, 3, 8, 3)
bkSimA = bkfilter.bkfilter(simA, 3, 8, 3)
bkSimC = bkfilter.bkfilter(simC, 3, 8, 3)
bkSimWprime = bkfilter.bkfilter(simWprime, 3, 8, 3)


# P and T use initial data - since there are somewhat absorbing regions
hpcorrQC = sp.corrcoef(fSimQ[:100], fSimC[:100])[0,1]
hpcorrQW = sp.corrcoef(fSimQ[:100], fSimWprime[:100])[0,1]
hpcorrQA = sp.corrcoef(fSimQ[:100], fSimA[:100])[0,1]

bkcorrQC = sp.corrcoef(bkSimQ[5:100], bkSimC[5:100])[0,1]
bkcorrQW = sp.corrcoef(bkSimQ[5:100], bkSimWprime[5:100])[0,1]
bkcorrQA = sp.corrcoef(bkSimQ[5:100], bkSimA[5:100])[0,1]
print "HP filtered correlations"
print "Corr(Q,C) = ", hpcorrQC
print "Corr(Q,W') = ", hpcorrQW
print "Corr(Q,A) = ", hpcorrQA
print "\nBand-Pass filter correlations"
print "Corr(Q,C) = ", bkcorrQC
print "Corr(Q,W') = ", bkcorrQW
print "Corr(Q,A) = ", bkcorrQA

plt.show()