import numpy as np
from copy import deepcopy
import pdb, cProfile, pstats
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from scipy.sparse import csr_matrix, lil_matrix, isspmatrix_csr, isspmatrix_bsr
from scipy.io import mmwrite

from pyamg.gallery import poisson, linear_elasticity
from pyamg.gallery.diffusion import diffusion_stencil_2d
from pyamg.gallery.stencil import stencil_grid
from pyamg.util.utils import symmetric_rescaling
from pyamg.aggregation.aggregation import smoothed_aggregation_solver
from adaptive_pairwise import adaptive_pairwise_solver

# pyAMG diffusion stencil is rotated from what it looks like, i.e. each row of 3
# elements corresponds to a fixed x-value, and 3 y-values. This means that 
# the nodes are organized something like 
#   ind = iy + ix*Nx
# N gives interior DOF, so hx = 1/(Nx+1), hy = 1/(Ny+1). The inverse transform
# should look like
#   x = (floor(ind/Ny) + 1)hx
#   y = (rem(ind/Ny) + 1)hy
def ind_to_coordinates(i, Nx, Ny):

    ix = (1.0 + np.floor(i/Ny)) * 1.0/(Nx+1)
    iy = (1.0 + np.remainder(i,Ny)) * 1.0/(Ny+1)
    return [ix,iy] 

# Weight defined as follows:
#   - Get unit vector between i,j 
#   - Unit vector of direction theta given by [acos(theta), asin(theta)] 
#   - Take absolute value of dot product between two vectors = |cos(phi)|,
#     where phi is angle between two vectors. Maximum at parallel, 
#     minimum at orthogonal. 
# Takes in 
#   - Nodes - i,j
#   - Number of unknowns in x and y direction - Nx,Ny
#   - Angle of anisotropy - theta 
#   - power to raise weight to - 2
def get_weight(i, j, Nx, Ny, theta, dim=2, power=2):

    if dim > 2:
        raise ValueError("Only implemented for 2-dimensions.")

    if i == j:
        return 0.0
    else:
        [ix,iy] = ind_to_coordinates(i, Nx, Ny)
        [jx,jy] = ind_to_coordinates(j, Nx, Ny)
        if ix == jx:
            phi = 0.0
        else:
            phi = np.arctan( float(jy-iy)/(jx-ix) )
        
        w = np.abs( np.cos(phi)*np.cos(theta) + np.sin(phi)*np.sin(theta) )
        return w**power


def get_geometric_weights(A, theta, Nx, Ny):

    n = A.shape[0]
    if isspmatrix_csr(A):
        W = A.copy()
        # Get angle between nodes for (i,j)th entry in A
        for i in range(0,n):
            lower_ind = W.indptr[i]
            upper_ind = W.indptr[i+1]
            for j_ind in range(lower_ind,upper_ind):
                j = W.indices[j_ind]
                W.data[j_ind] = get_weight(i,j,Nx,Ny,theta)

    elif isspmatrix_bsr(A): 
        raise ValueError("Not implemented for BSR matrices yet.")

    return W


# Get geometric weights for angle between nodes in a neighborhood of degree 2. 
def get_geometric_weights_d2(A, theta, Nx, Ny):

    n = A.shape[0]
    # Get maximum possible number of paths of length 2 in A
    paths_d1 = np.max( A.indptr[1:(n+1)] - A.indptr[0:n] )
    max_paths = n*paths_d1*paths_d1
    # Preallocate empty sparse array
    rowptr = np.zeros((n+1,))
    data = np.zeros((max_paths,))
    colinds = np.zeros((max_paths,))
    next_ind = 0

    if isspmatrix_csr(A):
        # Loop over each node in A
        for i in range(0,n):
            lower_ind1 = A.indptr[i]
            upper_ind1 = A.indptr[i+1]
            neighbors = []

            # Loop over d1 and d2 neighbors for ith node, save neighbors in set.
            for d1_ind in range(lower_ind1,upper_ind1):
                d1_node = A.indices[d1_ind]
                lower_ind2 = A.indptr[d1_node]
                upper_ind2 = A.indptr[d1_node+1]
                for d2_ind in range(lower_ind2,upper_ind2):
                    d2_node = A.indices[d2_ind]
                    neighbors.append(d2_node)

            # Find unique set of neighbors, get weight w(i,j)
            neighbors = np.unique(neighbors)
            for node in neighbors:
                colinds[next_ind] = node
                data[next_ind] = get_weight(i,node,Nx,Ny,theta)
                next_ind += 1

            # Set row-pointer for ith row
            rowptr[i+1] = next_ind

    elif isspmatrix_bsr(A): 
        raise ValueError("Not implemented for BSR matrices yet.")

    W = csr_matrix((data[0:next_ind],colinds[0:next_ind],rowptr),shape=A.shape)
    return W


# ----------------------------------------------------------------------------- #
# ----------------------------------------------------------------------------- #

N 			= 200
problem_dim = 2
epsilon 	= 0.01
theta 		= 3.0*np.pi/16
zero_rhs    = True
rand_guess  = True

# 1d Poisson 
if problem_dim == 1:
	grid_dims = [N,1]
	A = poisson((N,), format='csr')
# 2d Poisson
elif problem_dim == 2:
	grid_dims = [N,N]
	stencil = diffusion_stencil_2d(epsilon,theta)
	A = stencil_grid(stencil, grid_dims, format='csr')

# Scale out diagonal
[d,d,A]   = symmetric_rescaling(A)
vec_size  = A.shape[0]
# W = get_geometric_weights(A, theta, N, N)
# W.eliminate_zeros()

pdb.set_trace()

# Zero right hand side or sin(pi x)
if zero_rhs:
    b = np.zeros((vec_size,1))
    # If zero rhs and zero initial guess, throw error
    if not rand_guess:
        print "Zero rhs and zero initial guess converges trivially."
# Note, this vector probably doesn't make sense in 2d... 
else: 
    b = np.sin(math.pi*np.arange(0,vec_size)/(vec_size-1.0))
    # b = np.array([np.sin(math.pi*np.arange(0,vec_size)/(vec_size-1.0)),np.ones(vec_size)]).T

# Random vs. zero initial guess
if rand_guess:
    x0 = np.random.rand(vec_size,1)
else:
    x0 = np.zeros(vec_size,1)

asa_residuals = []

# ----------------------------------------------------------------------------- #
# ----------------------------------------------------------------------------- #

# General multilevel parameters
# -----------------------------
max_levels         = 20         # Max levels in hierarchy
max_coarse         = 30         # Max points allowed on coarse grid
tol                = 1e-8       # Residual convergence tolerance
is_pdef            = True       # Assume matrix positive definite (only for aSA)
keep_levels        = False      # Also store SOC, aggregation, and tentative P operators
diagonal_dominance = False      # Avoid coarsening diagonally dominant rows 
coarse_solver = 'pinv'
accel = 'cg'

aggregate = ('pairwise', {'matchings': 2, 'algorithm': 'drake_C'})
# interp_smooth = ('jacobi', {'omega': 4.0/3.0, 'degree':1 } )
interp_smooth = None
relaxation = ('gauss_seidel', {'sweep': 'forward', 'iterations': 1} )
# relaxation = ('jacobi', {'omega': 4.0/3.0} )
improve_candidates = [('gauss_seidel', {'sweep': 'symmetric', 'iterations': 5})]
# improve_candidates = None
test_iterations = 15
desired_convergence = 0.7
additive = False
reconstruct = False
use_ritz = False

# ----------------------------------------------------------------------------- #
# ----------------------------------------------------------------------------- #

ml_asa = adaptive_pairwise_solver(A, initial_targets = None, symmetry = 'symmetric',
                                  desired_convergence = desired_convergence,
                                  test_iterations = test_iterations, 
                                  test_accel=accel,
                                  strength = None,
                                  aggregate = aggregate,
                                  presmoother = relaxation,
                                  postsmoother = relaxation,
                                  max_levels = max_levels, max_coarse = max_coarse,
                                  coarse_solver=coarse_solver, additive=additive,
                                  reconstruct=reconstruct, use_ritz=use_ritz,
                                  improve_candidates=improve_candidates)
grid = ml_asa.operator_complexity()
cycle = ml_asa.cycle_complexity()

sol = ml_asa.solve(b, x0, tol=tol, residuals=asa_residuals, accel=accel, additive=additive)
asa_conv_factors = np.zeros((len(asa_residuals)-1,1))
for i in range(0,len(asa_residuals)-1):
  asa_conv_factors[i] = asa_residuals[i]/asa_residuals[i-1]

print "Adaptive SA/AMG - ", np.mean(asa_conv_factors[1:])
print " operator complexity - ", grid
print " cycle complexity - ", cycle, "\n"

# pdb.set_trace()



