"""
Slightly edited verions of some functions from postprocess from DNest4/python/classic.py
"""

import os
import copy
import numpy as np
from matplotlib import pyplot as plt

def logsumexp(values):
  biggest = np.max(values)
  x = values - biggest
  result = np.log(np.sum(np.exp(x))) + biggest
  return result


def logdiffexp(x1, x2):
  biggest = x1
  xx1 = x1 - biggest
  xx2 = x2 - biggest
  result = np.log(np.exp(xx1) - np.exp(xx2)) + biggest
  return result


def loadtxt_rows(filename, rows, single_precision=False):
  """
  Load only certain rows
  """
  # Open the file
  f = open(filename, "r")

  # Storage
  results = {}

  # Row number
  i = 0

  # Number of columns
  ncol = None

  while(True):
    # Read the line and split by whitespace
    line = f.readline()
    cells = line.split()

    # Quit when you see a different number of columns
    if ncol is not None and len(cells) != ncol:
      break

    # Non-comment lines
    if cells[0] != "#":
      # If it's the first one, get the number of columns
      if ncol is None:
        ncol = len(cells)

      # Otherwise, include in results
      if i in rows:
        if single_precision:
          results[i] = np.array([float(cell) for cell in cells], dtype="float32")
        else:
          results[i] = np.array([float(cell) for cell in cells])
      i += 1

  results["ncol"] = ncol
  return results


# postprocess from DNest4/python/classic.py, with minor edits to data loading
def postprocess(temperature=1., numResampleLogX=1, plot=True, rundir=".", \
                cut=0., save=True, zoom_in=True, compression_bias_min=1., verbose=True,\
                compression_scatter=0., moreSamples=1., compression_assert=None, single_precision=False):
  try:
    levels_orig = np.atleast_2d(np.loadtxt(os.path.join(rundir, "levels.txt"), comments="#"))
  except:
    return None # some error during file reading
  try:
    sample_info = np.atleast_2d(np.loadtxt(os.path.join(rundir, "sample_info.txt"), comments="#"))
  except:
    return None # some error during file reading

  # Remove regularisation from levels_orig if we asked for it
  if compression_assert is not None:
    levels_orig[1:,0] = -np.cumsum(compression_assert*np.ones(levels_orig.shape[0] - 1))

  cut = int(cut*sample_info.shape[0])
  sample_info = sample_info[cut:, :]

  if plot:
    plt.figure(1)
    plt.plot(sample_info[:,0], "k")
    plt.xlabel("Iteration")
    plt.ylabel("Level")

    plt.figure(2)
    plt.subplot(2,1,1)
    plt.plot(np.diff(levels_orig[:,0]), "k")
    plt.ylabel("Compression")
    plt.xlabel("Level")
    xlim = plt.gca().get_xlim()
    plt.axhline(-1., color='g')
    plt.axhline(-np.log(10.), color='g', linestyle="--")
    plt.ylim(ymax=0.05)

    plt.subplot(2,1,2)
    good = np.nonzero(levels_orig[:,4] > 0)[0]
    plt.plot(levels_orig[good,3]/levels_orig[good,4], "ko-")
    plt.xlim(xlim)
    plt.ylim([0., 1.])
    plt.xlabel("Level")
    plt.ylabel("MH Acceptance")

  # Convert to lists of tuples
  logl_levels = [(levels_orig[i,1], levels_orig[i, 2]) for i in range(0, levels_orig.shape[0])] # logl, tiebreaker
  logl_samples = [(sample_info[i, 1], sample_info[i, 2], i) for i in range(0, sample_info.shape[0])] # logl, tiebreaker, id
  logx_samples = np.zeros((sample_info.shape[0], numResampleLogX))
  logp_samples = np.zeros((sample_info.shape[0], numResampleLogX))
  logP_samples = np.zeros((sample_info.shape[0], numResampleLogX))
  P_samples = np.zeros((sample_info.shape[0], numResampleLogX))
  logz_estimates = np.zeros((numResampleLogX, 1))
  H_estimates = np.zeros((numResampleLogX, 1))

  # Find sandwiching level for each sample
  sandwich = sample_info[:,0].copy().astype('int')
  for i in range(0, sample_info.shape[0]):
    while sandwich[i] < levels_orig.shape[0]-1 and logl_samples[i] > logl_levels[sandwich[i] + 1]:
      sandwich[i] += 1

  for z in range(0, numResampleLogX):
    # Make a monte carlo perturbation of the level compressions
    levels = levels_orig.copy()
    compressions = -np.diff(levels[:,0])
    compressions *= compression_bias_min + (1. - compression_bias_min)*np.random.rand()
    compressions *= np.exp(compression_scatter*np.random.randn(compressions.size))
    levels[1:, 0] = -compressions
    levels[:, 0] = np.cumsum(levels[:,0])

    # For each level
    for i in range(0, levels.shape[0]):
      # Find the samples sandwiched by this level
      which = np.nonzero(sandwich == i)[0]
      logl_samples_thisLevel = [] # (logl, tieBreaker, ID)
      for j in range(0, len(which)):
        logl_samples_thisLevel.append(copy.deepcopy(logl_samples[which[j]]))
      logl_samples_thisLevel = sorted(logl_samples_thisLevel)
      N = len(logl_samples_thisLevel)

      # Generate intermediate logx values
      logx_max = levels[i, 0]
      if i == levels.shape[0]-1:
        logx_min = -1E300
      else:
        logx_min = levels[i+1, 0]
      Umin = np.exp(logx_min - logx_max)

      if N == 0 or numResampleLogX > 1:
        U = Umin + (1. - Umin)*np.random.rand(len(which))
      else:
        U = Umin + (1. - Umin)*np.linspace(1./(N+1), 1. - 1./(N+1), N)
      logx_samples_thisLevel = np.sort(logx_max + np.log(U))[::-1]

      for j in range(0, which.size):
        logx_samples[logl_samples_thisLevel[j][2]][z] = logx_samples_thisLevel[j]

        if j != which.size - 1:
          left = logx_samples_thisLevel[j+1]
        elif i == levels.shape[0]-1:
          left = -1E300
        else:
          left = levels[i+1][0]

        if j != 0:
          right = logx_samples_thisLevel[j-1]
        else:
          right = levels[i][0]

        logp_samples[logl_samples_thisLevel[j][2]][z] = np.log(0.5) + logdiffexp(right, left)

    logl = sample_info[:,1]/temperature

    logp_samples[:,z] = logp_samples[:,z] - logsumexp(logp_samples[:,z])
    logP_samples[:,z] = logp_samples[:,z] + logl
    logz_estimates[z] = logsumexp(logP_samples[:,z])
    logP_samples[:,z] -= logz_estimates[z]
    P_samples[:,z] = np.exp(logP_samples[:,z])
    H_estimates[z] = -logz_estimates[z] + np.sum(P_samples[:,z]*logl)

    if plot:
      plt.figure(3)

      plt.subplot(2,1,1)
      plt.plot(logx_samples[:,z], sample_info[:,1], 'k.', label='Samples')
      plt.plot(levels[1:,0], levels[1:,1], 'g.', label='Levels')
      plt.legend(numpoints=1, loc='lower left')
      plt.ylabel('log(L)')
      plt.title(str(z+1) + "/" + str(numResampleLogX) + ", log(Z) = " + str(logz_estimates[z][0]))
      # Use all plotted logl values to set ylim
      combined_logl = np.hstack([sample_info[:,1], levels[1:, 1]])
      combined_logl = np.sort(combined_logl)
      lower = combined_logl[int(0.1*combined_logl.size)]
      upper = combined_logl[-1]
      diff = upper - lower
      lower -= 0.05*diff
      upper += 0.05*diff
      if zoom_in:
        plt.ylim([lower, upper])
      xlim = plt.gca().get_xlim()

    if plot:
      plt.subplot(2,1,2)
      plt.plot(logx_samples[:,z], P_samples[:,z], 'k.')
      plt.ylabel('Posterior Weights')
      plt.xlabel('log(X)')
      plt.xlim(xlim)

  P_samples = np.mean(P_samples, 1)
  P_samples = P_samples/np.sum(P_samples)
  logz_estimate = np.mean(logz_estimates)
  logz_error = np.std(logz_estimates)
  H_estimate = np.mean(H_estimates)
  H_error = np.std(H_estimates)
  ESS = np.exp(-np.sum(P_samples*np.log(P_samples+1E-300)))

  errorbar1 = ""
  errorbar2 = ""
  if numResampleLogX > 1:
    errorbar1 += " +- " + str(logz_error)
    errorbar2 += " +- " + str(H_error)

  if verbose:
    print("log(Z) = " + str(logz_estimate) + errorbar1)
    print("Information = " + str(H_estimate) + errorbar2 + " nats.")
    print("Effective sample size = " + str(ESS))

  # Resample to uniform weight
  N = int(moreSamples*ESS)
  w = P_samples
  w = w/np.max(w)
  rows = np.empty(N, dtype="int64")
  for i in range(0, N):
    while True:
      which = np.random.randint(sample_info.shape[0])
      if np.random.rand() <= w[which]:
        break
    rows[i] = which + cut

  # Get header row
  f = open(os.path.join(rundir, "sample.txt"), "r")
  line = f.readline()
  if line[0] == "#":
    header = line[1:]
  else:
    header = ""
  f.close()

  try:
    sample = loadtxt_rows(os.path.join(rundir, "sample.txt"), set(rows), single_precision)
  except:
    return None # some error during file reading
  posterior_sample = None
  if single_precision:
    posterior_sample = np.empty((N, sample["ncol"]), dtype="float32")
  else:
    posterior_sample = np.empty((N, sample["ncol"]))

  for i in range(0, N):
    posterior_sample[i, :] = sample[rows[i]]

  if save:
    np.savetxt(os.path.join(rundir, 'weights.txt'), w)
    if single_precision:
      np.savetxt(os.path.join(rundir, "posterior_sample.txt"), posterior_sample, fmt="%.7e", header=header)
    else:
      np.savetxt(os.path.join(rundir, "posterior_sample.txt"), posterior_sample, header=header)

  if plot:
    plt.show()

  return [logz_estimate, H_estimate, logx_samples, posterior_sample]

