import scipy
import scipy.optimize
import scipy.linalg
import numpy as np
import sys
import warnings

import simulation
import utils


# Different optimization modes.
ABSOLUTE_AT_LEAST = 0    # Use absolute error (error only when there are traits missing).
ABSOLUTE_EXACT = 1       # Use absolute error (1-norm).
QUADRATIC_AT_LEAST = 2   # Use quadratic error (error only when there are traits missing).
QUADRATIC_EXACT = 3      # Use quadratic error (Frobenius norm).


def CheckGradient(f, x, h=1e-4, max_reldiff=1e-4):
  fx, grad = f(x)  # Evaluate function value at original point
  assert x.shape == grad.shape, 'Variable and gradient must have the same shape'
  passed = True
  numerical_grad = np.empty_like(x)
  # Iterate over all indexes in x
  it = np.nditer(x, flags=['multi_index'], op_flags=['readwrite'])
  while not it.finished:
    ix = it.multi_index
    x[ix] -= h
    y1 = f(x)[0]
    x[ix] += 2 * h
    y2 = f(x)[0]
    numgrad = (y2 - y1) / (2 * h)
    x[ix] -= h
    numerical_grad[ix] = numgrad
    # Compare gradients
    reldiff = abs(numgrad - grad[ix]) / max(1, abs(numgrad), abs(grad[ix]))
    if reldiff > max_reldiff and passed:
      # Only print the first error.
      with warnings.catch_warnings():
        warnings.simplefilter('ignore', np.ComplexWarning)  # Ignore complex warning.
        print utils.Highlight('Gradient check failed.', utils.RED, bold=True)
        print 'First gradient error found at index %s' % str(ix)
        print 'Your gradient: %f \t Numerical gradient: %f' % (grad[ix], numgrad)
      passed = False
    it.iternext()  # Step to next dimension.
  if passed:
    print utils.Highlight('Gradient check passed!', utils.GREEN, bold=True)
  return numerical_grad


# To give feedback during the optimization.
current_iteration = 0
def Print(x, f, accepted):
  global current_iteration
  current_iteration += 1
  output = '\r%i) Cost: %.4f, Accepted %s' % (current_iteration, f, str(accepted))
  sys.stdout.write(output)
  sys.stdout.write(' ' * (utils.GetTerminalWidth() - len(output)))
  sys.stdout.flush()


# Computes V * exp_wt * U.
# By construction the exponential of our matrices are always real-valued.
def Expm(V, exp_wt, U):
    return np.real(V.dot(np.diag(exp_wt)).dot(U))


# Computes the mixture of costs.
#
# || Y(t) - Y_desired ||_2^2 + \alpha * t^2 + \beta * || X(t) - X(t + \nu) ||_2^2
#
# When setting \alpha to zero, the time t needs to be specified (through specified_time).
# When \beta is greater than zero, \nu needs to be specified.
#
# This function outputs both the cost and its Jacobian w.r.t. to the flattened rates (parameters)
# and time (when \alpha > 0).
#
# When the mode is *_AT_LEAST, the margin needs to be specified. The cost will penalize only missing traits (up to a margin).
# When the mode is QUADRATIC_*, the distance to Y_desired is measured quadratically.
# When the mode is ABSOLUTE_*, the distance to Y_desired is measured with the absolute value.
def Cost(parameters,
         Y_desired,
         A, X_init, Q,
         alpha=1.0, specified_time=None,
         beta=5.0, nu=1.0,
         mode=QUADRATIC_EXACT, margin=None):
  # Sanity checks.
  assert alpha >= 0.0 and beta >= 0.0
  assert (alpha == 0.0) == (specified_time is not None)
  assert (beta == 0.0) == (nu is None)
  assert (mode in (QUADRATIC_EXACT, ABSOLUTE_EXACT)) == (margin is None)

  # Prepare variable depending on whether t part of the parameters.
  num_nodes = A.shape[0]
  num_species = X_init.shape[1]
  num_traits = Q.shape[1]
  if specified_time is None:
    t = parameters[-1]
    num_parameters_i = (np.size(parameters) - 1) / num_species
    grad_all = np.zeros(np.size(parameters))
  else:
    t = specified_time
    num_parameters_i = np.size(parameters) / num_species
    grad_all = np.zeros(np.size(parameters))

  # Reshape adjacency matrix to make sure.
  Adj = A.astype(float).reshape((num_nodes, num_nodes))
  Adj_flatten = Adj.flatten().astype(bool)  # Flatten boolean version.

  # Loop through the species to compute the cost value.
  # At the same time, prepare the different matrices.
  Ks = []                     # K_s
  eigenvalues = []            # w
  eigenvectors = []           # V.T
  eigenvectors_inverse = []   # U.T
  exponential_wt = []         # exp(eigenvalues * t).
  x_matrix = []               # Pre-computed X matrices.
  x0s = []                    # Avoids reshaping.
  qs = []                     # Avoids reshaping.
  xts = []                    # Keeps x_s(t).
  inside_norm = np.zeros((num_nodes, num_traits))  # Will hold the value prior to using the norm.
  for s in range(num_species):
    x0 = X_init[:, s].reshape((num_nodes, 1))
    q = Q[s, :].reshape((1, num_traits))
    x0s.append(x0)
    qs.append(q)
    k_ij = parameters[s * num_parameters_i:(s + 1) * num_parameters_i]
    # Create K from individual k_{ij}.
    K = np.zeros(Adj_flatten.shape)
    K[Adj_flatten] = k_ij
    K = K.reshape((num_nodes, num_nodes))
    np.fill_diagonal(K, -np.sum(K, axis=0))
    # Store K.
    Ks.append(K)
    # Perform eigen-decomposition to compute matrix exponential.
    w, V = scipy.linalg.eig(K, right=True)
    U = scipy.linalg.inv(V)
    wt = w * t
    exp_wt = np.exp(wt)
    xt = Expm(V, exp_wt, U).dot(x0)
    inside_norm += xt.dot(q)
    # Store the transpose of these matrices for later use.
    eigenvalues.append(w)
    eigenvectors.append(V.T)
    eigenvectors_inverse.append(U.T)
    exponential_wt.append(exp_wt)
    xts.append(xt)
    # Pre-build X matrix.
    with warnings.catch_warnings():
      warnings.simplefilter('ignore', RuntimeWarning)  # We don't care about 0/0 on the diagonal.
      X = np.subtract.outer(exp_wt, exp_wt) / (np.subtract.outer(wt, wt) + 1e-10)
    np.fill_diagonal(X, exp_wt)