import numpy as np
import scipy.linalg as la
import scipy.sparse as sp


# Only for periodic BC because we have advection only in x direction
def getUpwindMatrix(N, dx):
     
  #stencil    = [-1.0, 1.0]
  #zero_pos = 2
  #coeff      = 1.0
  
  #stencil    = [1.0, -4.0, 3.0]
  #coeff      = 1.0/2.0
  #zero_pos   = 3
  
  #stencil    = [1.0, -6.0, 3.0, 2.0]
  #coeff      = 1.0/6.0
  #zero_pos   = 3
  
  #stencil  = [-5.0, 30.0, -90.0, 50.0, 15.0]
  #coeff    = 1.0/60.0
  #zero_pos = 4
  
  stencil = [3.0, -20.0, 60.0, -120.0, 65.0, 12.0]
  coeff   = 1.0/60.0
  zero_pos = 5
  
  first_col = np.zeros(N)
  
  # Because we need to specific first column (not row) in circulant, flip stencil array
  first_col[0:np.size(stencil)] = np.flipud(stencil)

  # Circulant shift of coefficient column so that entry number zero_pos becomes first entry
  first_col = np.roll(first_col, -np.size(stencil)+zero_pos, axis=0)

  return sp.csc_matrix( coeff*(1.0/dx)*la.circulant(first_col) )

def getMatrix(N, dx, bc_left, bc_right):
  stencil = [1.0, -8.0, 0.0, 8.0, -1.0]
  range   = [ -2,   -1,   0,   1,    2]
  A       = sp.diags(stencil, range, shape=(N,N))
  A       = sp.lil_matrix(A)

  assert bc_left in ['periodic','neumann','dirichlet'], "Unknown type of BC"
  if bc_left in ['periodic']:
    assert bc_right in ['periodic'], "Periodic BC can only be selected for both sides simultaneously"
  
  if bc_left in ['periodic']:
    A[0,N-2] = stencil[0]
    A[0,N-1] = stencil[1]
    A[1,N-1] = stencil[0]

  if bc_right in ['periodic']:
    A[N-2,0] = stencil[4]
    A[N-1,0] = stencil[3]
    A[N-1,1] = stencil[4]

  if bc_left in ['neumann']:
    A[0,:] = np.zeros(N)
    A[0,0] = -8.0
    A[0,1] =  8.0
    A[1,0] = -8.0 + 4.0/3.0
    A[1,1] = -1.0/3.0

  if bc_right in ['neumann']:
    A[N-1,:]   = np.zeros(N)
    A[N-2,N-1] = 8.0 - 4.0/3.0
    A[N-2,N-2] = 1.0/3.0
    A[N-1,N-1] =  8.0
    A[N-1,N-2] = -8.0

  if bc_left in ['dirichlet']:
    A[0,:] = np.zeros(N)
    A[0,1] = 6.0

  if bc_right in ['dirichlet']:
    A[N-1,:]   = np.zeros(N)
    A[N-1,N-2] = -6.0

  A = 1.0/(12.0*dx)*A

  return sp.csc_matrix(A)

def getBCLeft(value, N, dx, type):

  assert type in ['periodic','neumann','dirichlet'], "Unknown type of BC"

  b = np.zeros(N)
  if type in ['dirichlet']:
    b[0] = -6.0*value
    b[1] =  1.0*value
  
  if type in ['neumann']:
    b[0] = 4.0*dx*value
    b[1] = -(2.0/3.0)*dx*value

  return (1.0/(12.0*dx))*b

def getBCRight(value, N, dx, type):

  assert type in ['periodic','neumann','dirichlet'], "Unknown type of BC"

  b = np.zeros(N)
  if type in ['dirichlet']:
    b[N-2] = -1.0*value
    b[N-1] =  6.0*value
  
  if type in ['neumann']:
    b[N-2] = -(2.0/3.0)*dx*value
    b[N-1] =  4.0*dx*value

  return (1.0/(12.0*dx))*b
