import tensorflow as tf
import tensorflow_probability as tfp
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")
tfd = tfp.distributions
import random
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import pandas as pd
import os

# tf.config.experimental_run_functions_eagerly(True)
# # tf.enable_eager_execution()
# tf.compat.v1.enable_eager_execution()
# # tf.compat.v1.disable_eager_execution()

# os.environ['CUDA_VISIBLE_DEVICES'] = ''    

# if tf.test.gpu_device_name():
#     print('GPU found')
# else:
#     print("No GPU found")

# # GPU found
# print("TF version:", tf.__version__)

############################################################
#                        IAC
############################################################

#https://rstudio-pubs-static.s3.amazonaws.com/258436_5c7f6f9a84bd47aeaa33ee763e57a531.html

def tau(M, kappa = 1000):
#   autocorr is the UNintegrated autocorrelation curve
    autocorr = auto_corr_fast(M, kappa)    
#   tau = 1 + 2*sum(G)
    return 1 + 2*np.sum(autocorr), autocorr

def auto_corr_fast(M, kappa = 1000):  
#   The autocorrelation has to be truncated at some point so there are enough
#   data points constructing each lag. Let kappa be the cutoff
    M = M - np.mean(M)
    N = len(M)
    fvi = np.fft.fft(M, n=2*N)
#   G is the autocorrelation curve
    G = np.real( np.fft.ifft( fvi * np.conjugate(fvi) )[:N] )
    G /= N - np.arange(N); G /= G[0]
    G = G[:kappa]
    return G

def compare_IAC(Ms,labels):
#   loop through each magnetization chain
    for ind,M in enumerate(Ms):
#       get IAC and autocorrelation curve
        IAC,G = tau(M)
        plt.plot(np.arange(len(G)),G,label="{}: IAC = {:.2f}".\
                                            format(labels[ind],IAC))
    plt.legend(loc='best',fontsize=14)
    plt.tight_layout()
    plt.show()
#     plt.savefig('project_plot3.png')
#     plt.close()

def calculate_IAC(chain, kappa = 1000):
    
    chain.columns = range(chain.shape[1])
    total_IAC_chain = 0
    
    for i in range(chain.shape[1]):
        IAC_chain, G_chain = tau(chain[i], kappa = 1000)

        total_IAC_chain += IAC_chain
    print('Mean IAC for chain:', total_IAC_chain/chain.shape[1])
    
    return total_IAC_chain/chain.shape[1]

############################################################
#                         ESS
############################################################

def calculate_ESS(chain):
    
    chain.columns = range(chain.shape[1])

    ess_chain = 0

    for i in range(chain.shape[1]):
        ess_chain += tfp.mcmc.effective_sample_size(chain[i], filter_beyond_positive_pairs=True).numpy()/chain.shape[1]

    print("Mean ESS for chain:", ess_chain)
    
    return ess_chain










