from __future__ import division
import scipy.linalg

import autograd.numpy as anp
from autograd.numpy.numpy_wrapper import wrap_namespace
from autograd.numpy.linalg import atleast_2d_col as al2d

anp.set_printoptions(precision=3)

wrap_namespace(scipy.linalg.__dict__, globals())  # populates module namespace


def _flip(a, trans):
    if anp.iscomplexobj(a):
        return 'H' if trans in ('N', 0) else 'N'
    else:
        return 'T' if trans in ('N', 0) else 'N'

def make_grad_solve_triangular(ans, a, b, trans=0, lower=False, **kwargs):
    tri = anp.tril if (lower ^ (_flip(a, trans) == 'N')) else anp.triu
    transpose = lambda x: x if _flip(a, trans) != 'N' else x.T

    def solve_triangular_grad(g):
        v = al2d(solve_triangular(a, g, trans=_flip(a, trans), lower=lower))
        return -transpose(tri(anp.dot(v, al2d(ans).T)))
    return solve_triangular_grad
solve_triangular.defgrad(make_grad_solve_triangular)
solve_triangular.defgrad(lambda ans, a, b, trans=0, lower=False, **kwargs: lambda g:
    solve_triangular(a, g, trans=_flip(a, trans), lower=lower), argnum=1)


def make_grad_sqrtm(ans, A, **kwargs):
    def sqrtm_grad(g):
        return solve_lyapunov(ans, g)
    return sqrtm_grad
sqrtm.defgrad(make_grad_sqrtm)

from __future__ import division
import scipy.linalg

import autograd.numpy as anp
from autograd.numpy.numpy_wrapper import wrap_namespace
from autograd.numpy.linalg import atleast_2d_col as al2d

wrap_namespace(scipy.linalg.__dict__, globals())  # populates module namespace


def _flip(a, trans):
    if anp.iscomplexobj(a):
        return 'H' if trans in ('N', 0) else 'N'
    else:
        return 'T' if trans in ('N', 0) else 'N'

def make_grad_solve_triangular(ans, a, b, trans=0, lower=False, **kwargs):
    tri = anp.tril if (lower ^ (_flip(a, trans) == 'N')) else anp.triu
    transpose = lambda x: x if _flip(a, trans) != 'N' else x.T

    def solve_triangular_grad(g):
        v = al2d(solve_triangular(a, g, trans=_flip(a, trans), lower=lower))
        return -transpose(tri(anp.dot(v, al2d(ans).T)))
    return solve_triangular_grad
solve_triangular.defgrad(make_grad_solve_triangular)
solve_triangular.defgrad(lambda ans, a, b, trans=0, lower=False, **kwargs: lambda g:
    solve_triangular(a, g, trans=_flip(a, trans), lower=lower), argnum=1)


def make_grad_sqrtm(ans, A, **kwargs):
    def sqrtm_grad(g):
        return solve_lyapunov(ans, g)
    return sqrtm_grad
sqrtm.defgrad(make_grad_sqrtm)

remove set_printoptions call from debugging