# This script checks that the new and old version of the simple 
# model produce consistent results
# it makes draws from the prior and computes the value of hc 
# using the old and new versions of the code. 

import numpy as np
import matplotlib.pyplot as plt
import math

import sys
# new version of the model
sys.path.append('../agnostic_model/')
import hmodel
# old version of the model 
sys.path.append('../../../PTAInference/code/agnostic_model/sampling/')
import hModelCircular


def priorBounds():
    '''
    prior ranges on the parameters
    '''	
    logno = [-20.0,3.0]
    beta  = [-2.0,7.0]
    gamma = [0.2,5.0]
    alpha = [-3.0,3.0]
    delta = [-1.0,2.0]

    priors = [logno,beta,gamma,alpha,delta]

    return priors


def makeDraw(priors):
    '''
    make random draws from the priors
    '''
    draw = np.zeros(len(priors))
    for i,p in enumerate(priors): 
        paramDraw = p[0]+np.random.rand()*(p[1]-p[0])
        draw[i] = paramDraw 
    return tuple(draw)


def main():
        
    # integration limits
    MlowIntLimit  = 10.0**6.0	
    MhighIntLimit = 10.0**11.0
    zlowIntLimit  = 0.0 
    zhighIntLimit = 5.0

    # get the prior
    bounds = priorBounds()

    # how many draws should we do?
    nDrawsToDo = 10000
    
    # set up for results
    hnew = np.zeros(nDrawsToDo)
    hold = np.zeros(nDrawsToDo)
    absDiff = np.zeros(nDrawsToDo)
    percentDiff = np.zeros(nDrawsToDo)


    # looping and drawing from prior each time
    for j in range(nDrawsToDo):

        # get the draw from the prior
        theta = makeDraw(bounds)

        # get the new value of h
        hnew[j] = hmodel.hmodel(theta,\
                             MlowIntLimit,MhighIntLimit,\
                             zlowIntLimit,zhighIntLimit)

        # get the old value of h
        hold[j] = hModelCircular.hmodel(theta,\
                                     MlowIntLimit,MhighIntLimit,\
                                     zlowIntLimit,zhighIntLimit)
                                     
        # difference between old and new (absolute and percentage)
        absDiff[j] = abs(hnew[j]-hold[j])
        percentDiff[j] = 100*(abs(hnew[j]-hold[j])/hold[j])
        

    # for the binning
    minimumh = math.floor(min(min(np.log10(hnew)),min(np.log10(hold))))
    maximumh = math.ceil(max(max(np.log10(hnew)),max(np.log10(hold))))
    bins = np.arange(minimumh, maximumh, 0.5)

    fig, (ax1, ax2, ax3) = plt.subplots(1,3)
    
    # plotting hold and hnew
    ax1.hist(np.log10(hnew),histtype='step',label='new',bins=bins,lw=2,color='#E69F00')
    ax1.hist(np.log10(hold),histtype='step',label='old',bins=bins,lw=2,color='#0072B2',ls=':')
    ax1.set_xlabel(r'$\log_{10} h_{\rm c}$',fontsize=14)
    ax1.set_ylabel('count',fontsize=14)
    ax1.legend()
    
    # plotting difference
    ax2.hist(np.log10(absDiff),histtype='step',label='abs(hnew-hold)',lw=2,color='#009E73')
    ax2.set_xlabel(r'$\log_{10} |h_{\rm c, new} - h_{\rm c, old}|$',fontsize=14)
    ax2.set_ylabel('count',fontsize=14)
    ax2.legend()
    
    # plotting percentage difference
    ax3.hist(percentDiff,histtype='step',label='%-difference',lw=2,color='#009E73')
    ax3.set_xlabel(r'percent difference $\left(100*\frac{|h_{\rm c, new}-h_{\rm c, old}|}{h_{\rm c, old}}\right)$',fontsize=14)
    ax3.set_ylabel('count',fontsize=14)
    ax3.legend()
    
    ax2.set_title('{} draws from prior'.format(nDrawsToDo),fontsize=14)
    
    fig.set_figwidth(20)
    plt.tight_layout()

    plt.savefig('model_check.png')
    plt.savefig('model_check.pdf')
    plt.show()
    
    
   

if __name__ == "__main__":
    main()
