from __future__ import absolute_import, print_function
import numpy as np
import warnings
from numpy.polynomial.hermite import hermvander
from six.moves import xrange
from .common import Baseline

try:
  from cvxopt import matrix as cvx_matrix, solvers
except ImportError:
  from scipy.optimize import linprog
  HAS_CVXOPT = False
else:
  HAS_CVXOPT = True

np.set_printoptions(precision=4, suppress=True)
_callback_state = {'last_nit':0, 'last_phase':0}


def mario_baseline(bands, intensities, poly_order=10, max_iters=None,
                   verbose=False, tol=1e-2):
  '''Solves a linear program: min_u f'u s.t. -P'u <= -s
  Where u are coefficients of a Hermite polynomial.'''
  bands = bands.astype(float)
  intensities = intensities.astype(float)
  if max_iters is None:
    max_iters = len(bands) * 10
  opts = dict(maxiter=max_iters, disp=verbose, tol=tol)
  callback = _linprog_callback if verbose else None
  # Flip intensities upside down.
  maxval = intensities.max() + 500
  s = maxval - intensities
  # Keep trying to solve until we succeed.
  for order in xrange(poly_order, 0, -1):
    result, P = _mario_helper(bands, s, order, opts, callback)
    if verbose:
      print('With order %d:' % order, result['status'])
    if result['x'] is not None:
      break
  else:
    warnings.warn('mario_baseline didnt find a fit at any order')
    return np.zeros_like(s)
  baseline = P.dot(np.array(result['x']).ravel())
  # Flip it back over.
  return maxval - baseline


def _mario_helper(bands, s, poly_order, opts, callback):
  # Build the polynomial basis over the bands.
  P = hermvander(bands, poly_order-1)
  f = P.sum(axis=0)

  if HAS_CVXOPT:
    solvers.options['show_progress'] = opts['disp']
    solvers.options['maxiters'] = opts['maxiter']
    solvers.options['abstol'] = opts['tol']
    solvers.options['reltol'] = opts['tol']
    solvers.options['feastol'] = 1e-100  # For some reason this helps.
    try:
      res = solvers.lp(cvx_matrix(f), cvx_matrix(-P), cvx_matrix(-s))
    except ValueError as e:
      # This can be thrown when poly_order is too large for the data size.
      res = {'status': e.message, 'x': None}
    return res, P

  res = linprog(f, A_ub=-P, b_ub=-s, bounds=(-np.inf,np.inf), options=opts,
                callback=callback)
  res = {'status': res.message, 'x': res.x if res.success else None}
  return res, P


def _linprog_callback(xk, nit=0, phase=0, tableau=None, **kwargs):
  obj = -tableau[-1, -1]
  new_state = False
  if _callback_state['last_phase'] != phase:
    new_state = True
    _callback_state['last_phase'] = phase
  if phase == 1:
    if new_state:
      print('--- Phase 1: Find a feasible point. ---')
      print('Iter\tObjective')
    _callback_state['last_nit'] = nit
    print('%d\t%g' % (nit, obj))
  else:
    if new_state:
      print('--- Phase 2: Minimize using simplex. ---')
      print('Iter\tObjective')
    print('%d\t%g' % (nit - _callback_state['last_nit'], obj))


class Mario(Baseline):
  def __init__(self, poly_order=10, max_iters=None, verbose=False, tol=1e-2):
    self.poly_order_ = poly_order
    self.max_iters_ = max_iters
    self.verbose_ = verbose
    self.tol_ = tol

  def _fit_one(self, bands, intensities):
    return mario_baseline(bands, intensities, self.poly_order_,
                          self.max_iters_, self.verbose_, self.tol_)

  def param_ranges(self):
    return {'poly_order_': (1, 12, 'integer')}
