from daredare import *

set_variables("k c")
set_exovariables("z")
set_shocks("x")
set_parameters("alpha beta gamma delta aa")

equations = [
    Equation(c + k , (1 - delta)*k(-1) + k(-1)**alpha*(1 + z)*aa),
    Equation(0 , -1/(1 + beta)*c(1)**(-gamma)*(1 - delta + k**(-1 + alpha)*(1 + z(1))*aa*alpha) + c**(-gamma)),
    Equation(z(+1),x(+1))
]

parameters_values = {
    alpha : 0.5,
    gamma : 0.5,
    delta : 0.02,
    beta : 0.05,
    aa : 0.5
}

covariances = matrix([[0.01]])

init_values = {
    k:(1/aa/alpha*(beta + delta))**(1/(-1 + alpha)),
    c:k**alpha*aa - delta*k
}

############################


model = Model("pyramst",lookup=True)


#model.check_all(print_info=True,print_eq_info=True)
# Try to solve the model using Uhligs toolkit
model = model
model.check_all('uhlig')

#umodel = model.to_uhlig_model()
#print(umodel.variables)
#print(umodel.exovariables)
#print(umodel.equations)
#umodel.check_all(print_info=True,print_eq_info=True)

#print("Compatible with Uhlig formulation : " + str(umodel.is_uhlig_compatible()))

#solver = Solver(model,parameters_values,init_values,covariances)
solver = Solver(model,lookup=True)

ss = solver.find_steady_state()

res = solver.solve_with_uhligs_toolkit()

print('PP',res.PP)
print('QQ',res.QQ)