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

# Load the data
with open('data_voxel_reconR7.5_cut30.0.txt', 'rb') as json_file:
  data = json.load(json_file)

# List all the keys
keys = data[0].keys()
print("We have " + str(len(keys)) + " keys in the dictionary")
for key in sorted(keys):
  print(key)

for d in data:
  # Halo density profile around the voids found in realspace
  r1    = d["profile_Halo_REAL_Void_REAL_radius"]
  delta = d["profile_Halo_REAL_Void_REAL_delta"]
  
  # Halo-Void cross correlation function in realspace
  r2  = d["CCF_radial_Halo_REAL_Void_REAL_radius"]
  xir = d["CCF_radial_Halo_REAL_Void_REAL_xir"]

  # Plot it (these two cruves should match)
  plt.plot(r1,delta,color="red")
  plt.plot(r2,xir,color="purple")
plt.show()

for d in data:
  # Fetch the multipoles of the cross-correlation of Halos in RSD and Voids in realspace
  r1  = d["CCF_multipole_Halo_RSD_Void_REAL_radius"]
  xi0 = d["CCF_multipole_Halo_RSD_Void_REAL_xi0"]
  xi2 = d["CCF_multipole_Halo_RSD_Void_REAL_xi2"]
  xi4 = d["CCF_multipole_Halo_RSD_Void_REAL_xi4"]
  
  # Fetch the multipoles of the cross-correlation of Halos in RSD and Voids in reconstructed space
  r2      = d["CCF_multipole_Halo_RSD_Void_RECON_radius"]
  xi0_rec = d["CCF_multipole_Halo_RSD_Void_RECON_xi0"]
  xi2_rec = d["CCF_multipole_Halo_RSD_Void_RECON_xi2"]
  xi4_rec = d["CCF_multipole_Halo_RSD_Void_RECON_xi4"]

  # Plot it (these two curves should match if reconstruction works well)
  plt.plot(r1,xi0,color="red")
  plt.plot(r1,xi2,color="blue")
  plt.plot(r1,xi4,color="purple")
  
  plt.plot(r2,xi0_rec,color="red", linestyle='dashed')
  plt.plot(r2,xi2_rec,color="blue", linestyle='dashed')
  plt.plot(r2,xi4_rec,color="purple", linestyle='dashed')
plt.show()
  
# Compare the quadrupole we get for the case of finding voids in realspace,
# reconstructed space and redshiftspace
for d in data:
  # Realspace
  r1        = d["CCF_multipole_Halo_RSD_Void_REAL_radius"]
  xi2_real  = d["CCF_multipole_Halo_RSD_Void_REAL_xi2"]
  
  # Reconstruction space
  r2        = d["CCF_multipole_Halo_RSD_Void_RECON_radius"]
  xi2_recon = d["CCF_multipole_Halo_RSD_Void_RECON_xi2"]
  
  # Redshiftspace
  r3        = d["CCF_multipole_Halo_RSD_Void_RSD_radius"]
  xi2_rsd   = d["CCF_multipole_Halo_RSD_Void_RSD_xi2"]

  # Plot it
  plt.plot(r1,xi2_real,color="red")
  plt.plot(r2,xi2_recon,color="blue")
  plt.plot(r3,xi2_rsd,color="purple")
  plt.legend()
plt.show()