from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl import app
from absl import flags
from absl import logging

import math 
from math import ceil 
from math import floor 
from math import exp 
from math import log

import numpy as np 
from numpy import dot
from numpy import arange 
from numpy import zeros 
from numpy import ones 
from numpy import multiply 
from numpy import diagflat as diag 
from numpy.random import rand 
from numpy.linalg import svd 
from numpy.linalg import cholesky as chol 

import pandas as pd 
from pandas import concat

import scipy 
from scipy.special import gamma 
from scipy.special import gammaincc


def getallobj(p,n,R,opt_type):

  # calculate all objective values for different shifts
  allobj=zeros(n)
  l=len(p)
  if opt_type == 1:
    for k in range(n):
      o1=np.sum(multiply((p[1:l-k-1] - p[1+k:l-1]), np.log(p[1:l-k-1] / p[1+k:l-1])))           
      q_plus= p[l-1]*np.power(R, arange(0,k+1,1))
      q_minus=p[0]*np.power(R, arange(k,-1,- 1))
      o2=sum(multiply((p[l-k-1:l] - q_plus), np.log(p[l-k-1:l] / q_plus))) + sum(multiply((p[1:2+k] - q_minus), np.log(p[1:2+k] / q_minus)))
      o3=(p[0] + p[l-1])*k*(1 - R ** (k+1)) / (1 - R)*(- log(R))

      allobj[k]=o1 + o2 + o3
  else:
    for k in arange(n):
      allobj[k]=sum(multiply((p[k:l] - p[:l-k]),np.log(p[:l-k] / p[:l-k])))
  
  return allobj

def getgrad(p,k,R,opt_type):
  # Calculate gradient for a shift of k
  l=len(p)
  if opt_type == 1:
    rat1=np.concatenate((p[2 + k:],p[l-1]*np.power(R, arange(1,k+1)))) / p[1:l-1]
    rat2=np.concatenate((p[0]*np.power(R,arange(k,0,-1)),p[0:l-k-2])) / p[1:l-1]
    gradmid=(- np.log(rat1) + 1 - rat1) + (- np.log(rat2) + 1 - rat2)

    Nfactor=(1+k)*(1 - R ** (1+k)) / (1 - R)*(- log(R))
    gradN=sum(multiply(np.power(R,arange(0,k+1)), np.log(rat1[-k-1:]) + 1 - 1.0 / rat1[-k-1:])) + Nfactor
    gradmN=sum(multiply(np.power(R,arange(k,-1,-1)),(np.log(rat2[:k+1]) + 1 - 1.0 / rat2[:k+1]))) + Nfactor

    grad=np.concatenate(([gradmN],gradmid,[gradN]))
  
  else:
    rat=p[-l+k:] / p[:l-k]
    rati=p[:l-k] / p[-l+k:]
    logr=np.log(rat)
    grad=np.concatenate((zeros((k+1,1)),[logr + 1 - rati], [- logr + 1 - rat],zeros((k+1,1))))
    
  return grad

def getH(p,n,evals,R,opt_type):
 

  # Calculate the Hessian matrix
  l=len(p)

  if opt_type == 1:
  
    d=zeros(l)
    for k in range(n):
      dN = sum(np.power(R, arange(0,k+1)))/p[l-1] + sum(p[-k-1:]) / p[l-1] ** 2
      dmN=sum(np.power(R,arange(k-1,-1,-1))) / p[0] + sum(p[1:1+k]) / p[0] ** 2
      dmid=2.0/p[1:l-1] + (np.concatenate((p[2+k:], np.power(R, arange(1,k+1))*p[l-1])) + np.concatenate((np.power(R, arange(k,0,-1))*p[0], p[:l-2-k] )))/ p[1:l-1]**2
      d=d + evals[k]*np.concatenate(([dmN],dmid,[dN]))
    
    H=np.diagflat(d)

    q=zeros((n,1))
    for k in range(n):
      j = 1 + k
      q[k]=sum(multiply(evals[k:n],(- np.power(R,arange(0,n-k))) / p[j] - 1 / p[0]))
    H[0,1:n+1]=q.reshape(-1)
    H[1:n+1,0]=q.reshape(-1)


    # deal with i=l
    q=zeros((n,1))
    for k in range(n):
      i = l-k-1
      q[k]=sum(multiply(evals[k:n], -np.power(R,arange(0,n-k))/p[i] - 1.0/p[l-1]))
    H[l-1,range(l-2,l-n-2,-1) ]=q.reshape(-1)
    H[range(l-2,l-n-2,-1),l-1]=q.reshape(-1)  

    for k in range(1,l-2):
      kmax = min(n,l-k-1)
      q=multiply(evals[:kmax], (-1/p[k]-1.0/p[k+1:k+1+kmax]))
      H[k, k+1:k+1+kmax]=q.reshape(-1)
      H[k+1:k+1+kmax,k]=q.reshape(-1)        
  else:
    d=zeros(l,1)
    for a in arange(1,n).reshape(-1):
      d=d + dot(evals(a),(concat([[1.0 / p(arange(1,l - a)) + p(arange(1 + a,l)) / p(arange(1,l - a)) ** 2],[zeros(a,1)]]) + concat([[zeros(a,1)],[1.0 / p(arange(1 + a,l)) + p(arange(1,l - a)) / p(arange(1 + a,l)) ** 2]])))
    H=diag(d)
    for i in arange(1,l - 1).reshape(-1):
      amax=min(n,l - i)
      q=multiply(evals(arange(1,amax)),(- 1 / p(i) - 1.0 / p(arange(i + 1,i + amax))))
      H[i,arange(i + 1,i + amax)]=q
      H[arange(i + 1,i + amax),i]=q
  
  return H

def sample(x,CDF,r,n):

  # produce one sample from the distribution
  
  l=len(x)
  i=find(rand(1) < CDF,1)
  s=x(i)
  if i == 1 or i == l:
    # sample from a geometric with parameter R
    j=floor(log(rand(1)) / log(r))
    if i == 1:
      s=s - j / n
    else:
      s=s + j / n
  
  s=s + (rand(1) - 0.5) / n
  return s

def main(n=200, xmax=8, C=0.25, c_exp=2, opt_type=1, r=0.9, tol=1e-08, verbose=2):

  N = math.ceil(n*xmax)
  xmax = N / n

  x=arange(-N,N+1) / n
  l=len(x)

  # The two constraints (probability normalization, and the variance constraint) are handled by A*p=b as follows. 
  # Note that even though the variance constraint is an inequality, it will always be tight at optimal, so we treat it as an equality constraint.

  if opt_type == 1:
    c=zeros(l)
    for i in range(l):
      if x[i]>0:
        c[i] = dot(((x[i] + 1 / 2 / n) ** (c_exp + 1) - (x[i] - 1 / 2 / n) ** (c_exp + 1)),(n / (c_exp + 1)))
      elif x[i]==0:
        c[i] = (1 / 2 / n) ** c_exp / (1 + c_exp)
      else:
        c[i]=dot(((abs(x[i]) + 1 / 2 / n) ** (c_exp + 1) - (abs(x[i]) - 1 / 2 / n) ** (c_exp + 1)),(n / (c_exp + 1)))

    # slight over-estimate, but getting the exact value requires a nasty infinite sum. This upper bounds the sum by an integral. 
    # This still gives a true bound, without much loss
    cN = n*r**(-N-1/2)*(-n*np.log(r))**(-c_exp-1)*gamma(c_exp+1)*gammaincc(c_exp+1,-np.log(r)*(N-1/2))
    c[0]=cN
    c[l-1]=cN
    d = np.ones(l)
    d[0] = d[l-1] = 1/(1-r)  
    A = np.stack((c,d),axis = 0)      
    # A = [c,d]
  else:
    # If we're doing the under-estimating model, let c be the minimal value of the cost function within each quantization bin.
    c=max(0,abs(x)-1/(2*n)) ** c_exp
    A=np.stack((c,ones(l)),axis = 0)

  b = np.stack(([C],[1]),axis = 0)
  # We start with a guess for the distribution p.
  # This is a distribution designed to have farily heavy tails (which seems to make the solver work better), and to roughly satisfy the cost constraint.
  q = C ** (- 1 - 2 / c_exp)
  p = 1/(1+q*abs(x)**(c_exp+2))
  # make sure that p is normalized (it actually doesn't need to satisfy the cost constraint exactly initially -- the algorithm with ensure that)
  p = p/dot(A[1],p) 
  iter=0

  t=1

  isfeasible = False

  allobj = getallobj(p,n,r,opt_type)
  primobj=np.amax(allobj)
  fval=np.log(sum(np.exp(t*(allobj - primobj)))) / t + primobj



  while 1:

    iter = iter + 1
    eta = np.exp(t*(allobj - primobj))
    eta=eta / sum(eta)
    allgrads=zeros((l,n))

    for k in range(n):
      allgrads[:,k]=getgrad(p,k,r,opt_type)    
    grad=np.matmul(allgrads,eta.reshape(200,1))

    rpri=np.matmul(A,p.reshape(len(p),1))-b

    # The next few lines are typically the majority of the run-time.
    # Compute Hessian matrix
    H=getH(p,n,eta,r,opt_type)
    # strictly positive definite. This doesn't much change overall performance
    fullH = H + np.matmul(np.matmul(allgrads,np.diagflat(t*eta)),allgrads.T) - (1 - 1e-05)*t*np.matmul(grad,grad.T)
    
    try:
      R=scipy.linalg.cholesky(fullH)

      gtilde=np.linalg.solve(R.T,grad)

      AR=np.matmul(A, np.linalg.inv(R))
      U,S,V = svd(AR,full_matrices=False)
      #U,S,V=svd(AR,'econ',nargout=3)
      s=S
      vtilde = - gtilde - np.matmul(V.T,multiply(1.0/s.reshape(2,1), np.matmul(U.T,rpri))) + np.matmul(V.T,(np.matmul(V,gtilde)))
      v=np.linalg.solve(R,vtilde)

    except:
      if verbose >=1:
          print('Chlesky failed')
      vw = np.linalg.solve(np.stack((np.concatenate((fullH, A.T)), np.concatenate((A, zeros(2)))),axis = 0), np.stack((-grad, -rpri),axis = 0))
      v = vw[:]

    nwt_dec=np.matmul(- grad.T,v) / 2

    # implementing a line search. We move to p+dst*v
    dst=1
    while 1:
      pnew = p + dst*v.reshape(-1)
      if min(pnew) > 0:

        allobj=getallobj(pnew,n,r,opt_type)
        primobj=max(allobj)
        newfval=np.log(sum(np.exp(t*(allobj - primobj)))) / t + primobj
        if not isfeasible:
          break
        if newfval < fval + 0.1*dst*(np.matmul(grad.T,v)):
          break
        if dst < 1e-08:
          break
      dst=dst / 2

    if  not isfeasible and dst == 1:
      isfeasible=True
      if verbose >= 1:
          print('Feasible!')
    p=pnew
    fval=newfval
    if verbose >= 1:
      print('iter=%d  dst=%1.1e  primobj=%1.4f  nwt_dec=%1.2e  t=%1.1e \n' %(iter,dst,primobj,nwt_dec,t));
    # if verbose >=2:
    #     drawnow(draw_fig)
    if isfeasible:
      # n/e/t is an estimate for the duality gap, if we are at exact optimality for the soft-max problem. 
      # Thus nwt_dec+n/e/1 is an estimate for the total duality gap
      if nwt_dec + n / exp(1) / t < tol:
          break
      # if we take a full Newton step or get extremely close to optimal, then increase t.
      if dst == 1 or nwt_dec < tol / 2:
          # This is a fairly modest increase in t.
          # This seems to make the solver the most robust, if not necessarily the most efficient.
          t=t*1.25
    
  cdf = np.cumsum(multiply(p,A[1].T))
    
  filename=('test_data/cactus_cdf_v%.2f.csv' %C)
  np.savetxt(filename, cdf, delimiter=",")

  filename=('test_data/cactus_x_v%.2f.csv' %C)
  np.savetxt(filename, x, delimiter=",")
      
  filename=('test_data/cactus_pdf_v%.2f.csv' %C)
  np.savetxt(filename, p, delimiter=",")
      
#   return x, p 

if __name__ == '__main__':
  app.run(main(n=200, xmax=8, C=0.25, c_exp=2, opt_type=1, r=0.9, tol=1e-05, verbose=2))
