import numpy as np
from numpy.fft import fft, ifft, rfft, irfft
from numpy.linalg import norm
import scipy.ndimage

from matplotlib import pyplot as plt
import matplotlib as mpl
mpl.rcParams["figure.figsize"] = [8,6] # or 7, 4 or 10,8
mpl.rcParams["lines.linewidth"] = 2
mpl.rcParams["lines.markersize"] = 4
mpl.rcParams['mathtext.fontset'] = 'cm'
mpl.rcParams.update({'font.size': 20})

# Helper routine to print out matrices nicely
def matprint(mat, fmt="g",roundToDecimal=2):
  # from https://gist.github.com/braingineer/d801735dac07ff3ac4d746e1f218ab75
  # Modified to round
  if roundToDecimal is not None:
    mat = np.round(mat,decimals=roundToDecimal)
  #col_maxes = [max([len(("{:"+fmt+"}").format(x)) for x in col]) for col in mat.T]
  if np.min( mat.flat ) < 0:
    col_maxes = [8 for col in mat.T] # quick hack to deal with fmt='.1e'
  else:
    col_maxes = [7 for col in mat.T] # quick hack to deal with fmt='.1e'
  for x in mat:
    for i, y in enumerate(x):
      if abs(y) >= 1000:
        fmt = '.1e'
      else:
        fmt = 'g'
      print(("{:"+str(col_maxes[i])+fmt+"}").format(y), end="  ")
    print("")

"""
Explanation: <a href="https://colab.research.google.com/github/stephenbeckr/convex-optimization-class/blob/master/Homeworks/HW04/APPM5630_HW4_helperFunctions.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
Homework 4
University of Colorado Boulder, Spring 2021
APPM 5630 Advanced Convex Optimization, by Stephen Becker
1D blurring problem
This notebook contains helper functions (utilities) that may be useful in future homeworks
1. implicit2explict
2. Adjoint check
3. Power method to find spectral norms
First, setup some packages, and make a helper routine for printing matrices, etc.
End of explanation
"""


def implicit2explicit( linearFcn, domainSize ):
    """ domainSize can be an integer, for a pure vector,
    or can be a tuple of integers """
    x = np.zeros( domainSize )
    # use ravel not flatten, as we don't want to make a copy

    # Determine sizes
    n = x.size
    x.ravel()[0] = 1  # first unit vector
    y = np.asarray( linearFcn( x ) ) # output may be scalar or list (not array) so convert it
    m = y.size
    x.ravel()[0] = 0 # undo it

    # Allocate array
    A = np.zeros( (m,n), dtype=y.dtype ) # allow for complex values
    A[:,0] = y.ravel()

    # Loop over all inputs (we already did the first one)
    for col in range(1,n):
      x.ravel()[col] = 1  # first unit vector
      y = linearFcn( x )
      A[:,col] = np.asarray(y).ravel()
      x.ravel()[col] = 0 # undo it

    return A

"""
Explanation: implicit2explicit
We'll get a little bit fancy and allow functions that have matrix inputs (theoretically, this is no big deal, since matrices in $\mathbb{R}^{m\times n}$ are isomorphic to vectors in $\mathbb{R}^{mn}$ when we use the inner product $\langle X, Y\rangle = \text{trace}(X^Y)$ where $X^=\overline{X}^T$ is the adjoint.  This was not necessary for the homework problem.
Note that you can specify domainSize which could be something like n for a 1D vector of length $n$, or sometimes (n,1) depending on how the function expects it. For matrices, set domainSize to (m,n).
End of explanation
"""


def implicit2explicit_square_buggy( linearFcn, n ):
  I   = np.eye(n)
  y   = linearFcn( I[0] )
  A   = np.zeros( (n,n), dtype=y.dtype )
  for col in range(0,n):
    A[col] = linearFcn( I[col] )
  return A

# linearFcn = fft
A   = np.arange(16).reshape(4,4) # 4 x 4 matrix
linearFcn = lambda x : A@x
A_buggy = implicit2explicit_square_buggy( linearFcn, 4 )
print("We think it is:")
matprint(A_buggy)
print("It should be:")
matprint(A)

"""
Explanation: Be careful! A few students made what seemed like a valid implementation of implicit2explicit (specialized to the square operator case) using the identity matrix, but then ran into trouble due to Python's row-based indexing. See the issue below:
End of explanation
"""


def implicit2explicit_square_fixed( linearFcn, n ):
  I   = np.eye(n)
  y   = linearFcn( I[0] )
  A   = np.zeros( (n,n), dtype=y.dtype )
  for col in range(0,n):
    A[:,col] = linearFcn( I[:,col] )
  return A

A_bugFixed = implicit2explicit_square_fixed( linearFcn, 4 )
print("Now it works:")
matprint(A_bugFixed)

"""
Explanation: to fix this bug, we don't use A[col] and so on, we use A[:,col] to make sure we're getting a column, not a row:
End of explanation
"""


A   = np.arange(12).reshape(3,4) # 3 x 4 matrix
def sampleLinearFunction(x):
  return A@x

def trickierLinearFunction(X):
  """ 'Tricky' since a matrix domain """
  return np.trace(X)

def veryTrickierLinearFunction(X):
  """ 'Trickier' since complex valued """
  return 1j*np.trace(X)

Amat = implicit2explicit( sampleLinearFunction, 4)
Amat = implicit2explicit( sampleLinearFunction, (4,1) ) # also works
print("The matrix we found is")
matprint(Amat)
print("The original matrix was")
matprint(A)

print("The trace of a 3x3 matrix, as a matrix operator:")
Amat = implicit2explicit( trickierLinearFunction, (3,3) )
matprint(Amat)

print("The trace of a 3x3 matrix, as a matrix operator, with imaginary numbers:")
Amat = implicit2explicit( veryTrickierLinearFunction, (3,3) )
matprint(Amat)

"""
Explanation: ... so if you had this bug, you may have been swapping the transpose with the forward operator. This often shows up in your results by an incorrect shift
Now test our implicit2explicit function
End of explanation
"""


def test_adjoint( A, At, domainSize, nRep = 10):
  rng   = np.random.default_rng()
  for rep in range(nRep):
    x  = rng.normal( size=domainSize )
    Ax = A(x)
    y  = rng.normal( size=Ax.shape)
    Aty= At(y)

    er  = np.vdot(Ax,y) - np.vdot(x,Aty)
    er  = np.abs(er)/np.sqrt( norm(x)*norm(y) ) # nice scaling
    print(f"Relative error in adjoint: {er:.2e}")


A   = np.arange(12).reshape(3,4) # test it on something rectangular (not square) to help expose more bugs

linop     = lambda x : A@x
linop_adj = lambda y : A.T@y

test_adjoint( linop, linop_adj, A.shape[1] )

"""
Explanation: Adjoint test
Let's make a different test of the adjoint
This one will scale easily to huge values of $N$
The adjoint of a linear operator $A$ is the operator (let's call it $T$, but later we'll denote it $A^*$ or $A^\dagger$ or $A^H$ in physics) such that for any inputs $x$ and $y$ of the right size,
$$ \langle Ax, y \rangle = \langle x, Ty \rangle. $$
This inner product we use np.vdot for, as it does the right thing (adds in a complex conjugat, and automatically vectorizes it, so it's the right thing to do for a matrix).  np.dot does the right thing only if our input is 1D and real-valued.
So, we'll just check this equation for several random choices of $x$ and $y$.  This is no way to prove that the adjoint is implemented correctly, but if we pass a lot of these checks, it's good evidence we've done it correctly.  If we fail a check, then we know for sure that there is a problem.
Note that for a matrix, there's also the formula 
$$ A^ = \overline{A}^T$$
meaning the adjoint is the transpose of the complex conjugate. So for a real-valued matrix, the adjoint is just the transpose*.
Btw, do not confuse adjoint with "adjugate", which is a classical term used when finding the inverse of a matrix by hand
End of explanation
"""


A   = np.arange(12).reshape(3,4) # test it on something rectangular (not square) to help expose more bugs

normA = norm(A,ord=2)
print("Spectral norm of A is", normA )

def powerMethod(A, At, domainSize=None, x=None, iters=100, tol=1e-6):
  if x is None:
    if domainSize is None:
      raise ValueError("need domain size or x0 to be specified")
    rng   = np.random.default_rng()
    x = rng.normal( size=domainSize )
  normalization = norm( x.ravel() ) # Euclidean/Frobenius norm
  for k in range(iters):
    x = np.real_if_close( At(A(x)) )
    oldNormalization = normalization
    normalization = norm( x.ravel() )
    x /= normalization  # do this in-place
    if abs(oldNormalization - normalization)/np.max( [1e-12,normalization] ) < tol:
      print("Reached tolerance after",k,"iterations.")
      break
  return np.sqrt(normalization)

normA_v2 = powerMethod(linop,linop_adj, A.shape[1], iters=1000, tol=1e-12)
print("Spectral norm of A is", normA_v2, "using power method" )
print("  so the error in the power method is", abs(normA-normA_v2) )

"""
Explanation: Power method to estimate spectral norms
End of explanation
"""


