import numpy as np
import matplotlib.pyplot as plt

from tfr_inference import estimate_data_parameters, get_fields, create_mock, myprint, generateMBData
import borg_velocity.utils as utils

def main():
    
    myprint('Beginning')
    
    # Get some parameters from the data
    sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma, hyper_m_sigma = estimate_data_parameters()
    print('sigma_m', sigma_m)
    print('sigma_eta', sigma_eta)
    print('hyper_eta_mu', hyper_eta_mu)
    print('hyper_eta_sigma', hyper_eta_sigma)
    
    # Other parameters to use
    L = 500.0
    N = 64
    xmin = -L/2
    R_lim = L / 2
    Rmax = 100
    Nt = 5_000
    alpha = 1.4
    mthresh = 11.25
    a_TFR = -23
    b_TFR = -8.2
    sigma_TFR = 0.3
    sigma_v = 150
    Nint_points = 201 
    Nsig = 10
    frac_sigma_r = 0.07  # WANT A BETTER WAY OF DOING THIS - ESTIMATE THROUGH SIGMAS FROM TFR
    interp_order = 1
    bias_epsilon = 1.e-7
    
    # Make mock
    np.random.seed(123)
    cpar, dens, vel = get_fields(L, N, xmin)
    RA, Dec, czCMB, m_true, eta_true, m_obs, eta_obs, xtrue, vbulk = create_mock(
            Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh,
            a_TFR, b_TFR, sigma_TFR, sigma_m, sigma_eta, 
            hyper_eta_mu, hyper_eta_sigma, sigma_v, 
            interp_order=interp_order, bias_epsilon=bias_epsilon)
    MB_pos = generateMBData(RA, Dec, czCMB, L, N, R_lim, Nsig, Nint_points, sigma_v, frac_sigma_r)
    
    r = np.sqrt(np.sum(MB_pos**2, axis=0))    
    zcosmo = utils.z_cos(r, cpar.omega_m)
    mu_true = 5 * np.log10((1 + zcosmo) * r / cpar.h) + 25 

    gamma_inv2 = 1 + sigma_m ** 2 * (hyper_eta_sigma**2 + sigma_eta**2) / (b_TFR**2 * hyper_eta_sigma**2 * sigma_eta**2 + sigma_TFR**2 * (hyper_eta_sigma**2 + sigma_eta**2))
    print(gamma_inv2)
    
    delta0 = (mthresh - m_obs) / sigma_m
    den = (b_TFR**2 * hyper_eta_sigma**2 * sigma_eta**2 + (sigma_TFR**2 + sigma_m**2) * (hyper_eta_sigma**2 + sigma_eta**2))
    delta1 = ((m_obs - a_TFR - b_TFR * eta_obs)[:,None] - mu_true) * sigma_m * hyper_eta_sigma**2 / den
    delta2 = ((m_obs - a_TFR - b_TFR * hyper_eta_sigma)[:,None] - mu_true) * sigma_m * hyper_eta_sigma**2 / den
    
    delta = delta0[:,None] + delta1 + delta2
    
    fig, axs = plt.subplots(1, 3, figsize=(15,4))
    hist_kwargs = dict(density=True, histtype='step', bins=30)
    x = delta.flatten()
    print('delta', x.min(), x.max())
    axs[0].hist(x, **hist_kwargs)
    x = gamma_inv2
    print('gamma_inv2', x.min(), x.max())
    axs[1].hist(x, **hist_kwargs)
    x = (delta * np.sqrt(gamma_inv2)).flatten()
    print('delta/gamma', x.min(), x.max())
    axs[2].hist(x, **hist_kwargs)
    axs[0].set_title(r'$\delta$')
    axs[1].set_title(r'$\gamma^{-2}$')
    axs[2].set_title(r'$\delta / \gamma$')
    fig.tight_layout()
    fig.savefig('tfr_integral_params.png')
    
    
    return

if __name__ == "__main__":
    main()