import os
import pdb
import sys
import h5py
import glob
import cycler
import pickle
import numpy as np
from scipy import stats
import matplotlib as mpl
from matplotlib import rc
import astropy.units as u
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.offsetbox import OffsetImage, AnnotationBbox

from astropy.cosmology import FlatLambdaCDM
cosmo = FlatLambdaCDM(H0=71, Om0=0.220, Ob0=0.02258*(0.71**2), name='OuterRim')

sys.path.insert(0, "/Users/joe/repos/shearfit/shearfit")
sys.path.insert(0, "/home/hollowed/repos/shearfit/shearfit")
from analytic_profiles import NFW
from astropy.cosmology import WMAP7
from lensing_system import obs_lens_system
from mass_concentration import child2018 as cm
from fit_profile import fit_nfw_profile_lstq as fit
from fit_profile import fit_nfw_profile_gridscan as fit_gs
    
def imscatter(x, y, image, ax=None, zoom=1):
    if ax is None:
        ax = plt.gca()
    try:
        image = plt.imread(image)
    except TypeError:
        # Likely already an array...
        pass
    im = OffsetImage(image, zoom=zoom)
    x, y = np.atleast_1d(x, y)
    artists = []
    for x0, y0 in zip(x, y):
        ab = AnnotationBbox(im, (x0, y0), xycoords='data', frameon=False)
        artists.append(ax.add_artist(ab))
    ax.update_datalim(np.column_stack([x, y]))
    ax.autoscale()
    return artists


# =================================================================================================
# =================================================================================================


class nfw_profile_fitter:

    def __init__(self, NFW_dir):
        """
        This class contains routines to perform an example run of the shearfit package, fitting an NFW profile 
        to background source data as obtained from ray-tracing through a pure NFW particle realization. The 
        process is as follows:
        1) sources positions, redshifts, and shears are read in from ray-tacing outputs
        4) an NFW profile is fit to the scaled tangential shear ΔΣ = γ_t * Σ_c, using three methods:
            - allow both the concentration and halo radius to float
            - only fit the halo radius, inferring the concentration from a c-M relation
            - do a grid scan over the parameter pair including radius and concentration
        
        Parameters
        ----------
        NFW_dir : string
            Path to an output directory generated by the modules in mpwl_raytrace/NFW_test_cases.
            This directory is assumed to contain an HDF5 file giving ray-traced maps, as well as a properties.csv
            file, containing the intrinsic halo properties from the simulation.
        """
        
        print('\n\n======= Initializing fitter object for dataset at {} ========'.format(NFW_dir.split('/')[-1]))
        zlpl_dirs = np.array(glob.glob('{}/halo*'.format(NFW_dir)))
        zl = np.array([float(d.split('/')[-1].split('zl')[-1].split('_')[0]) for d in zlpl_dirs])
        zs = np.array([float(d.split('/')[-1].split('zs')[-1].split('_')[0]) for d in zlpl_dirs])
        N = np.array([int(d.split('/')[-1].split('N')[-1].split('_')[0]) for d in zlpl_dirs])
        rfrac = np.array([float(d.split('/')[-1].split('r200c')[0].split('_')[-1]) for d in zlpl_dirs])
        rfrac_los = np.array([float(d.split('/')[-1].split('r200clos')[0].split('_')[-1]) for d in zlpl_dirs])
        nsrcs = np.array([int(d.split('/')[-1].split('nsrcs')[-1].split('_')[0]) for d in zlpl_dirs])
        lenspix = np.array([int(d.split('/')[-1].split('lenspix')[-1].split('_')[0]) for d in zlpl_dirs])
        print('Found {} lenses'.format(len(zl)))
        
        sort = np.argsort(zl)
        self.zlpl_dirs, self.zl, self.zs, self.N, self.rfrac, self.rfrac_los, self.nsrcs, self.lenspix = \
            zlpl_dirs[sort], zl[sort], zs[sort], N[sort], rfrac[sort], rfrac_los[sort], nsrcs[sort], lenspix[sort]
        self.out_dir = '{}/profile_fits'.format(NFW_dir)
        
        self.var_strs = {'zl':r'$z_l$', 'zs':r'$z_s$', 'N':r'$\mathrm{particle\>count}$', 
                         'fov':r'$\mathrm{fov\>size/extent\>of\>particle generation}$', 
                         'rmin':r'$r_{\mathrm{fit,min}}$', 'rmax':r'$r_{\mathrm{fit,max}}$',
                         'losClip':r'$\Delta z_\mathrm{clip}$', 'nsrcs':r'$N\>\mathrm{sources}$',
                         'lenspix':r'$\sqrt{N_\mathrm{lensmap\>pixels}}$'}

    
    # ------------------------------------------------------------------------------------------------------------


    def fit_halos(self, rmax = None, rmin = 0, bootstrap = False, bin_data = False, rbins = 25, 
                  overwrite = False, inputs = 'mocks', single_halo_plots = False, grid_scan = True, grid_N = 100):
        """
       Use the shearfit package to fit the profiles of all halos present in NFW_dir, saving the resultinng 
       differential surface mass profile and mass estimation to self.out_dir

        Parameters
        ----------
        rmax : float array, optional
            The maximum radial distance of sources to include in the fit in units of r200c (e.g. rmin = 2.0 will
            trim the source popiulation beyond 2*r200c. Defaults to None, in which case rmax will automatically be 
            set to the full radial extent of the FOV as given in the NFW input. This default will also be imposed if
            the input rmax is greater than the size of the FOV. If an array is provided, the entire fitting procedure
            will be done per array entry, varying the rmax "hyperparameter" 
        rmin : float, optional
            The minimum radial distance of sources to include in the fit in units of r200c (e.g. rmin = 0.3 will
            remove the inner 300kpc/h of source information). Defaults to 0.
        bin_data : bool, optional
            whether or not to fit to shears averaged in radial bins, rather than to each individual source.
            Defaults to False.
        rbins : int, optional
            Number of bins distribute over the radial range of the data, if bin_data == True. Defaults to 25.
        bootstrap : bool, optional
            Whether or not to run bootstrap error estimation on the fitted profiles. Defaults to False, which is much
            faster.
        overwrite : bool, optional
            whether or not to force running of profile fits, even if already exist. If False, then this function does 
            nothing in the case that profile fitting outputs are already present in self.out_dir. This is simply to 
            allow calling code the convenience of not having to comment out this function call if plotting etc. needs 
            to be redone. Defaults to False.
        inputs : string, optional
            The kind of lensing inputs to use; 'mocks' to use the interpolated mock catalogs, and 'grid' to
            use raytracing results on the full set of map gridpoints
        single_halo_plots : boolean, optional
            Whether or not to render a profile fitting plot per-halo. If True, perform a  grid scan in addition to 
            the two profile fits to plot residuals over parameter space, saving plots to the class-wide output 
            direcotry set in the constructor. If False, return after fit and skip parameter sweep, plot nothing. 
            Defaults to False. 
        grid_scan : boolean, optional
            If single_halo_plot is True, this option toggles the ability to perform a parameter sweep in the 
            radius-concentration plane, plotting the resulting cost function as a colormap on a subplot. 
            Defaults to True
        grid_N : boolean
            The number of sample points in each dimension of the parameter grid, if grid_scan=True. Default is n=100
        """

        # loop over all halos in input dir
        for i in range(len(self.zlpl_dirs)):
            print('\nreading lensing mock data for lens zl={:.2f}, zs={:.2f}, N={}, rfrac={:.2f}, '\
                  'rfrac_los={:.2f}, nsrcs={}, lenspix={}'.format( self.zl[i], self.zs[i], self.N[i], 
                  self.rfrac[i], self.rfrac_los[i], self.nsrcs[i], self.lenspix[i]))
            
            # read lensing maps and get FOV size
            [sim_lens, true_profile] = self._read_nfw_shears(self.zlpl_dirs[i], inputs)
            
            pfs = glob.glob('{}/properties.csv'.format(self.zlpl_dirs[i]))
            assert len(pfs) == 1, "Exactly one properties file is expected in {}".format(self.zlpl_dirs[i])
            props_file = pfs[0]
            props = np.genfromtxt(props_file, delimiter=',', names=True)
            fov_edge = props['boxRadius_Mpc']
            
            # loop over all chocies of rmax, record scaled ranges to self.rmax, self.rmin attributes
            self.rmin = true_profile.r200c * rmin
            self.rmax = np.zeros(len(rmax))
            
            for k in range(len(rmax)): 
                # scale rmin, rmax to Mpc
                if(rmax[k] is None or rmax[k] > fov_edge): this_rmax = fov_edge
                else: this_rmax = true_profile.r200c * rmax[k]
                assert self.rmin < this_rmax, "inner radial cut rmin ({}) cannot exceed outer radial cut rmax ({})!".format(self.rmin, this_rmax)
                
                output_suffix = 'zl{:.2f}_zs{:.2f}_N{}_fov{:.2f}_losClip{}_rmin{:.2f}_rmax{:.2f}_'\
                                'nsrcs{}_lenspix{}_{}'.format(self.zl[i], self.zs[i], self.N[i], 
                                self.rfrac[i], self.rfrac_los[i], self.rmin, this_rmax, self.nsrcs[i], 
                                self.lenspix[i], inputs)
               
                if(len(glob.glob('{}/*{}.obj'.format(self.out_dir, output_suffix))) == 0 or overwrite):
                   
                    if not os.path.exists(self.out_dir): os.makedirs(self.out_dir)
                    else: [os.remove(obj) for obj in glob.glob('{}/*{}.obj'.format(self.out_dir, output_suffix))]

                    print('------------ fitting in radial range [{:.2f}, {:.2f}] Mpc ------------'.format(
                                                                                      self.rmin, this_rmax))
                    self._fit_nfw(sim_lens, true_profile, rmin=self.rmin, rmax=this_rmax,
                                  make_plot=single_halo_plots, bin_data=bin_data, rbins=rbins, 
                                  bootstrap=bootstrap, sfx=output_suffix, grid_scan=grid_scan, grid_N=grid_N)    
                else:
                    print('fitting files already exist for {}; skipping'.format(output_suffix))
        
    
    # ------------------------------------------------------------------------------------------------------------


    def _read_nfw_shears(self, halo_dir, inputs='mocks'):
        """
        Reads the lensing data for a single halo as output by mpwl_raytrace. THe halo property information
        is loaded into a analytic_profiles.NFW object from shearfit, and the backgrouns source data (shears, etc)
        are loaded into a lensing_)system.obs_lens_system object from shearfit. Both of these objects are returned

        Parameters
        ----------
        halo_dir : string
            Top-level directory containing halo particle data and lensing maps

        Returns
        -------
        sim_lens : obs_lens_system object instance
            A shearfit obs_lens_system object, containing source positions, redshfits, and lensing quantities
        true_profile : NFW object instance
            A shearfit NFW object, containing true input halo properties (to compare fit result)
        inputs : string, optional
            The kind of lensing inputs to use; 'mocks' to use the interpolated mock catalogs, and 'grid' to
            use raytracing results on the full set of map gridpoints
        """

        # get ray-trace hdf5 and properties csv
        pfs = glob.glob('{}/properties.csv'.format(halo_dir))
        assert len(pfs) == 1, "Exactly one properties file is expected in {}".format(halo_dir)
        
        if(inputs == 'mocks'):
            rtfs = glob.glob('{}/lensing_maps_zs*/*mock*.hdf5'.format(halo_dir))
        elif(inputs == 'grid'):
            rtfs = glob.glob('{}/lensing_maps_zs*/*raytrace*.hdf5'.format(halo_dir))
        else:
            raise RuntimeError('Unknown input type {}! Must be \'mocks\' or \'grid\'.'.format(inptus))
        
        # read lens properties from csv, source data from hdf5
        props_file = pfs[0]
        props = np.genfromtxt(props_file, delimiter=',', names=True)
        zl = props['halo_redshift']
        r200c = props['sod_halo_radius']
        c = props['sod_halo_cdelta']
        #c_err = props['sod_halo_cdelta_error']
        c_err = 0
        m200c = props['sod_halo_mass']
        true_profile = NFW(r200c, c, zl, c_err = c_err, cosmo=cosmo)
 
        raytrace_file = h5py.File(rtfs[0], 'r')
        nplanes = len(list(raytrace_file.keys()))
        t1, t2, y1, y2, k, zs = [], [], [], [], [], []

        # stack data from each source plane
        for i in range(nplanes):
            plane_key = list(raytrace_file.keys())[i]
            plane = raytrace_file[plane_key]
            
            #ignore this plane if infront of the halo
            plane_z = plane['zs'][0] 
            if(plane_z < zl): continue
            
            # get dimensions of shear inputs 
            # if 2d grid--> raytracing outputs, 
            # if 1d vector--> interpolated mocks)
            dim = len(np.shape(plane['shear1']))
            
            if(dim == 1):
                t1 = np.hstack([t1, np.ravel(plane['x1'][:])])
                t2 = np.hstack([t2, np.ravel(plane['x2'][:])])
            
            if(dim == 2): 
                try:
                    # if output_positions was turned on in lensing run 
                    x1 = np.hstack([t1, np.ravel(plane['x1'][:])])
                    x2 = np.hstack([t2, np.ravel(plane['x2'][:])])
                except:
                    # get grid point angular posiitons
                    nnn = len(plane['shear1'][0]) # fov width in gridpoints
                    bsz_arc = float(props['boxRadius_arcsec']*2) # fov width in arcsec
                    dsx_arc = bsz_arc/nnn
                    x1 = np.linspace(0, bsz_arc-dsx_arc, nnn) - bsz_arc/2.0 + dsx_arc/2.0
                    x2 = np.linspace(0, bsz_arc-dsx_arc, nnn) - bsz_arc/2.0 + dsx_arc/2.0
                t1, t2 = np.meshgrid(x1, x2)
                t1 = np.ravel(t1)
                t2 = np.ravel(t2)
        
            y1 = np.hstack([y1, np.ravel(plane['shear1'][:])])
            y2 = np.hstack([y2, np.ravel(plane['shear2'][:])])
            k = np.hstack([k, np.ravel(plane['kappa0'][:])])
            zs = np.hstack([zs, np.ones(len(t1)-len(zs)) * plane_z])
       
        # trim the fov borders by 10% to be safe
        mask = np.logical_and(np.abs(t1)<props['boxRadius_arcsec']*0.9, 
                              np.abs(t2)<props['boxRadius_arcsec']*0.9)
        t1 = t1[mask]
        t2 = t2[mask]
        zs = zs[mask]
        y1 = y1[mask]
        y2 = y2[mask]
        k = k[mask]
        
        raytrace_file.close()
        sim_lens = obs_lens_system(zl)
        sim_lens.set_background(t1, t2, zs, y1=y1, y2=y2, k=k)
        
        # temporary for Lindsey
        writeout=False
        if(writeout):
            print('writing sources out')
            from astropy.io import fits
            c1 = fits.Column(name='theta1', array=t1, format='D')
            c2 = fits.Column(name='theta2', array=t2, format='D')
            c3 = fits.Column(name='shear1', array=y1, format='D')
            c4 = fits.Column(name='shaer2', array=y2, format='D')
            c5 = fits.Column(name='kappa', array=k, format='D')
            c6 = fits.Column(name='redshift', array=zs, format='D')
            
            p1 = fits.Column(name='zl', array=[zl], format='D')
            p2 = fits.Column(name='r200c', array=[r200c], format='D')
            p3 = fits.Column(name='conc', array=[c], format='D')
            p4 = fits.Column(name='conc_err', array=[c_err], format='D')
            p5 = fits.Column(name='m200c', array=[m200c], format='D')
           
            h = fits.PrimaryHDU()
            t = fits.BinTableHDU.from_columns([c1, c2, c3, c4, c5, c6])
            p = fits.BinTableHDU.from_columns([p1, p2, p3, p5])
            hl = fits.HDUList([h, t, p])
            hl.writeto('lensing_example.fits', overwrite=True)
            sys.exit()

        return [sim_lens, true_profile]

    
    # ------------------------------------------------------------------------------------------------------------
    

    def _fit_nfw(self, lens, true_profile, rmin, rmax, bin_data, rbins, bootstrap, make_plot, sfx, grid_scan, grid_N):
        """
        Calls the fitting routines in shearfit on the loaded data, saving the result to both numpy files 
        and png figures

        Parameters
        ----------
        lens : obs_lens_system object instance
            A shearfit obs_lens_system object, containing source positions, redshfits, and lensing quantities
        true_profile : NFW object instance
            A shearfit NFW object, containing true input halo properties (to compare fit result)
        rmin : float
            The minimum radial distance of sources to include in the fit in Mpc (e.g. rmin = 0.3 will
            remove the inner 300kpc/h of source information). Defaults to 0.
        rmax : float
            The maximum radial distance of sources to include in the fit in Mpc (e.g. rmin = 2.0 will
            trim the source popiulation beyond 2*r200c. Defaults to None, in which case rmax will automatically be 
            set to the full radial extent of the FOV as given in the NFW input. This default will also be imposed if
            the input rmax is greater than the size of the FOV.
        bin_data : bool
            whether or not to fit to shears averaged in radial bins, rather than to each individual source.
            Defaults to False.
        rbins : int
            Number of bins distribute over the radial range of the data, if bin_data == True. Defaults to 25.
        bootstrap : bool
            Whether or not to run bootstrap error estimation on the fitted profiles. Defaults to False, which is much
            faster.
        make_plot : boolean
            Whether or not to render a profile fitting plot per-halo. If True, perform a  grid scan in addition to 
            the two profile fits to plot residuals over parameter space, saving plots to the class-wide output 
            direcotry set in the constructor. If False, return after fit and skip parameter sweep, plot nothing. 
            Defaults to False.
        sfx : string
           Suffix to append to end of output filenames (for marking halo uniqueness)
        grid_scan : boolean
            Whether or not to perform a grid scan in the radius and concentration parameters, and plot the resulting 
            cost function as a colormap on a second subplot
        grid_N : boolean
            The number of sample points in each dimension of the parameter grid, if grid_scan=True
        """

        zl = lens.zl
        r200c = true_profile.r200c
        m200c = true_profile.radius_to_mass()
        c = true_profile.c
        
        # fit the concentration and radius
        print('fitting with floating concentration')
        fitted_profile = NFW(0.75, 3.0, zl)
        fit(lens, fitted_profile, r200_bounds = [0.1, 15], conc_bounds = [1, 10], rmin=rmin, rmax=rmax, 
            bootstrap=bootstrap, bin_data=bin_data, bins=rbins)

        # and now do it again, iteratively using a c-M relation instead of fitting for c
        print('fitting with inferred c-M concentration')
        fitted_cm_profile = NFW(0.75, 3.0, zl)
        fit(lens, fitted_cm_profile, r200_bounds = [0.1, 15], cM_relation='child2018', rmin=rmin, rmax=rmax, 
            bootstrap=bootstrap, bin_data=bin_data, bins=rbins)
        
        # write out fitting result
        print('r200c_fit = {:.4f}; m200c_fit = {:.4e}; c_fit = {:.2f}'.format(
               fitted_profile.r200c, fitted_profile.radius_to_mass(), fitted_profile.c))
        print('r200c_cm = {:.4f}; m200c_cm = {:.4e}; c_cm = {:.2f}'.format(
               fitted_cm_profile.r200c, fitted_cm_profile.radius_to_mass(), fitted_cm_profile.c))
        print('r200c_true = {:.4f}; m200c_true = {:.4e}; c_true = {:.2f}'.format(r200c, m200c, c))
        
        # save fitted profiles
        fits = [fitted_profile, fitted_cm_profile, true_profile, lens]
        fit_dirs = ['fitted_profile', 'fitted_cm_profile', 'true_profile', 'source_data']
        for i in range(len(fits)):
            fit_file = open('{}/{}_{}.obj'.format(self.out_dir, fit_dirs[i], sfx), 'wb')
            pickle.dump(fits[i], fit_file)

        if(not make_plot): return
        else: self._plot_halo_profile(lens, rmin, rmax, true_profile, fitted_profile, fitted_cm_profile, 
                                      bootstrap, bin_data, rbins, sfx, grid_scan, grid_N)


    # ------------------------------------------------------------------------------------------------------------


    def _plot_halo_profile(self, lens, rmin, rmax, true_profile, fitted_profile, fitted_cm_profile,
                           bootstrap, bin_data, rbins, sfx, grid_scan, grid_N):
        """
        Plots a single halo profile, with points per-source, the true proile, and both fit variants. Includes a 
        second panel with a parameter sweep colormap if grid_scan is True.

        Parameters
        ----------
        lens : obs_lens_system object instance
            A shearfit obs_lens_system object, containing source positions, redshfits, and lensing quantities
        rmin : float
            The minimum radial distance of sources to include in the plot in Mpc (e.g. rmin = 0.3 will
            remove the inner 300kpc/h of source information). Defaults to 0.
        rmax : float
            The maximum radial distance of sources to include in the plot in Mpc (e.g. rmin = 2.0 will
            trim the source popiulation beyond 2*r200c. Defaults to None, in which case rmax will automatically be 
            set to the full radial extent of the FOV as given in the NFW input. This default will also be imposed if
            the input rmax is greater than the size of the FOV.
        true_profile : NFW object instance
            A shearfit NFW object, containing true input halo properties (to compare fit result)
        fitted_profile : shearfit NFW object
            NFW object containing the final fitted form of the halo
        fitted_cm_profile : shearfit NFW object
            NFW object containing the final fitted form of the halo using the c-M relation during minimization
        bootstrap : bool
            Whether or not to run bootstrap error estimation on the fitted profiles. Defaults to False, which is much
            faster.
        bin_data : bool
            whether or not to fits were done with shears averaged in radial bins, rather than to each individual 
            source.Defaults to False.
        rbins : int
            Number of bins used to perform passed fits over the radial range of the data, if bin_data == True. 
            Defaults to 25.
        sfx : string
           Suffix to append to end of output filenames (for marking halo uniqueness) 
        grid_scan : boolean
            Whether or not to perform a grid scan in the radius and concentration parameters, and plot the 
            resulting cost function as a colormap on a second subplot
        grid_N : boolean
            The number of sample points in each dimension of the parameter grid, if grid_scan=True
        """
        
        r200c = true_profile.r200c
        m200c = true_profile.radius_to_mass()
        c = true_profile.c
        
        zl = lens.zl
        lens.set_radial_cuts(rmin, rmax)
        bg = lens.get_background()
        r = bg['r']
        yt = bg['yt']
        k = bg['k']
        zs = bg['zs']
        if(bin_data):
            binned_data = lens.calc_delta_sigma_binned(nbins=rbins, return_edges=True, 
                                                 return_std=True, return_gradients=True)
            delsig = binned_data['delta_sigma_mean']
            r = binned_data['r_mean']
        else:
            delsig = lens.calc_delta_sigma()
        
        # get profiles from fit results
        rsamp = np.linspace(min(r), max(r), 1000)
        dSigma_true = true_profile.delta_sigma(rsamp)
        if(bootstrap):
            [dSigma_fitted, dSigma_fitted_err] = fitted_profile.delta_sigma(rsamp, bootstrap=True)
            [dSigma_fitted_cm, dSigma_fitted_cm_err] = fitted_cm_profile.delta_sigma(rsamp, bootstrap=True)
        else:
            dSigma_fitted = fitted_profile.delta_sigma(rsamp, bootstrap=False)
            dSigma_fitted_err = np.zeros(len(dSigma_fitted))
            dSigma_fitted_cm = fitted_cm_profile.delta_sigma(rsamp, bootstrap=False)
            dSigma_fitted_cm_err = np.zeros(len(dSigma_fitted_cm))
         
        # visualize fits
        print('plotting profile fits')
        rc('text', usetex=True)
        color = plt.cm.plasma(np.linspace(0.2, 0.8, 3))
        mpl.rcParams['axes.prop_cycle'] = cycler.cycler('color', color)

        if(grid_scan): 
            f = plt.figure(figsize=(12,6))
            ax = f.add_subplot(121)
        else: 
            f = plt.figure(figsize=(6,6))
            ax = f.add_subplot(111)
        rR = rsamp/r200c
      
        # ----- plot data -----
        if(bin_data):
            delsig_std = binned_data['delta_sigma_std']
            delsig_err = binned_data['delta_sigma_se_mean']
            r_std = binned_data['r_std']
            r_err = binned_data['r_se_mean']
            bin_widths = np.diff(binned_data['bin_edges'])
            
            ax.errorbar(r/r200c, delsig, xerr=bin_widths/2/r200c, yerr=delsig_std, fmt='none', capsize=2, color='k',
                        label=r'$\gamma_{\mathrm{NFW}}\Sigma_c,\>\>y:\mathrm{std},\>x:\mathrm{bin\>width},$')
            ax.errorbar(r/r200c, delsig, xerr=r_err/r200c, yerr=delsig_err, fmt='none', ecolor='k', 
                        elinewidth=6, alpha=0.2, label=r'$\gamma_{\mathrm{NFW}}\Sigma_c,\>\>\mathrm{bin\>SEM}$')
            ax.plot(r/r200c, delsig, '--k', ms=10, label=r'$\gamma_{\mathrm{NFW}}\Sigma_c$')
            
        else:
            ax.loglog(r/r200c, delsig, 'xk', label=r'$\gamma_{\mathrm{NFW}}\Sigma_c$', alpha=0.33)
        
        # ----- plot profile fits -----
        ax.loglog(rR, dSigma_true, '--', label=r'$\Delta\Sigma_\mathrm{{NFW}},\>\>r_{{200c}}={:.3f}; c={:.3f}$'\
                                                .format(r200c, c), color=color[0], lw=2)
        ax.loglog(rR, dSigma_fitted, label=r'$\Delta\Sigma_\mathrm{{fit}},\>\>r_{{200c}}={:.3f}; c={:.3f}$'\
                                            .format(fitted_profile.r200c, fitted_profile.c), color=color[1], lw=2)
        ax.loglog(rR, dSigma_fitted_cm,label=r'$\Delta\Sigma_{{\mathrm{{fit}},c-M}},\>\>r_{{200c}}={:.3f}; c={:.3f}$'\
                                               .format(fitted_cm_profile.r200c, fitted_cm_profile.c), 
                                               color=color[2], lw=2)
        ax.fill_between(rR, dSigma_fitted - dSigma_fitted_err.T[0], 
                               dSigma_fitted + dSigma_fitted_err.T[1], 
                               color=color[1], alpha=0.2, lw=0)
        ax.fill_between(rR, dSigma_fitted_cm - dSigma_fitted_cm_err.T[0], 
                               dSigma_fitted_cm + dSigma_fitted_cm_err.T[1], 
                               color=color[2], alpha=0.33, lw=0)

        # ----- format -----
        ax.legend(fontsize=12, loc='best')
        ax.set_xlabel(r'$r/R_{200c}$', fontsize=14)
        ax.set_ylabel(r'$\Delta\Sigma\>\>\lbrack\mathrm{M}_\odot\mathrm{pc}^{-2}\rbrack$', fontsize=14)
        ax.set_title(r'$z_\mathrm{{lens}} = {}$'.format(zl), fontsize=14)
        ax.set_xlim([0.2, 6])
        ax.set_ylim([2, 300])
        
        if(grid_scan):
            # now do a grid scan
            print('doing grid scan')
            gridscan_profile = NFW(0.75, 3.0, zl)
            grid_r_bounds = [0.01, 1.2]
            grid_c_bounds = [0.5, 10]
            [grid_pos, grid_res] = fit_gs(lens, gridscan_profile, r200_bounds = grid_r_bounds, 
                                          conc_bounds = grid_c_bounds, n=grid_N, bin_data=bin_data, bins=rbins)

            # plot fit cost in the radius-concentration plane
            print('plotting grid scan')
            ax2 = f.add_subplot(122)
            chi2 = ax2.pcolormesh(grid_pos[0], grid_pos[1], (1/grid_res)/(np.max(1/grid_res)), cmap='plasma')
            ax2.plot([r200c], [c], 'xw', ms=10, mew=1)
            ax2.plot([r200c], [c], 'xk', ms=10, label=r'$\mathrm{{truth}}$')
            ax2.errorbar(fitted_profile.r200c, fitted_profile.c, 
                         xerr=fitted_profile.r200c_err, yerr=fitted_profile.c_err, 
                         ms=10, marker='.', c=color[1], label=r'$\mathrm{{fit}}$')
            ax2.errorbar(fitted_cm_profile.r200c, fitted_cm_profile.c, 
                         xerr=fitted_cm_profile.r200c_err, yerr=fitted_cm_profile.c_err, 
                         ms=10, marker='.', c=color[2], label=r'${\mathrm{{fit\>w/}c\mathrm{-}M}}$')

            # include c-M relation curve
            tmp_profile = NFW(1,1,zl)
            tmp_m200c = np.zeros(len(grid_pos[0][0]))
            for i in range(len(tmp_m200c)):
                tmp_profile.r200c = grid_pos[0][0][i]
                tmp_m200c[i] = tmp_profile.radius_to_mass()
            tmp_c, tmp_dc = cm(tmp_m200c, zl, cosmo)
            ax2.plot(grid_pos[0][0], tmp_c, '--k', lw=2, label=r'$c\mathrm{-}M\mathrm{\>relation\>(Child+2018)}$')
            ax2.fill_between(grid_pos[0][0], tmp_c - tmp_dc, tmp_c + tmp_dc, color='k', alpha=0.1, lw=0)

            # format
            ax2.set_xlim(grid_r_bounds)
            ax2.set_ylim(grid_c_bounds)
            ax2.legend(fontsize=12, loc='upper right')
            cbar = f.colorbar(chi2, ax=ax2)
            cbar.set_label(r'$\mathrm{min}(\chi^2)/\chi^2$', fontsize=14)
            ax2.set_xlabel(r'$R_{200c}\>\>\left[\mathrm{Mpc}\right]$', fontsize=14)
            ax2.set_ylabel(r'$c_{200c}$', fontsize=14)
        
        plt.tight_layout() 
        if(bin_data): binstr = '_{}bins_'.format(rbins)
        else: binstr = '_'
        f.savefig('{}/shearprof_fit{}_{}.png'.format(
                   self.out_dir, binstr, sfx), dpi=300)
    
        
    # ------------------------------------------------------------------------------------------------------------


    def plot_mass_convergence(self, vary_var, sfx=None):
        '''
        Plot the mass estimated from many runs of fit_halos(), results of which are found in self.out_dir

        Parameters
        ---------
        vary_var : string
            which variable to plot convergence with respect to. This should be one of the quantities given in the 
            output directory suffix as defined in fit_halos(). Choices are: zl, zs, N, fov, rmin, rmax. It is 
            assumed that user has ensured that all output subdirs in self.out_dir have only one of these variables
            changing, to make for a meaningul convergence study.
        sfx : string, optional
            Suffix to add to output figure filename (to prevent overwrites if testing convergence on different ranges 
            of one variable). Default is None.
        '''
      
        print('\n--------------- reading fits and plotting mass convergence ---------------')

        fit_out = glob.glob('{}/fitted_profile*.obj'.format(self.out_dir))
        m = np.array([pickle.load(open(f, 'rb')).radius_to_mass() for f in fit_out])
        cmfit_out = glob.glob('{}/fitted_cm_profile*.obj'.format(self.out_dir))
        m_cm = np.array([pickle.load(open(f, 'rb')).radius_to_mass() for f in cmfit_out])
        truth_out = glob.glob('{}/true*.obj'.format(self.out_dir))
        m_true = np.array([pickle.load(open(f, 'rb')).radius_to_mass() for f in truth_out])
      
        vary_val_m = np.array([float(f.split('/')[-1].split(vary_var)[-1].split('_')[0].strip('.obj')) 
                               for f in fit_out])
        vary_val_cm = np.array([float(f.split('/')[-1].split(vary_var)[-1].split('_')[0].strip('.obj')) 
                                for f in cmfit_out])
        vary_val_truth = np.array([float(f.split('/')[-1].split(vary_var)[-1].split('_')[0].strip('.obj')) 
                                   for f in truth_out])

        sm = np.argsort(vary_val_m)
        scm = np.argsort(vary_val_cm)
        st = np.argsort(vary_val_truth)
 
        f = plt.figure()
        ax = f.add_subplot(111)
        ax.plot(vary_val_truth[st], m_true[st]/m_true[st], '--k', label="truth")
        ax.plot(vary_val_m[sm], m[sm]/m_true[st], '-dr', label="unconstrained fit")
        ax.plot(vary_val_cm[scm], m_cm[scm]/m_true[st], '-db', label="c-M relation fit")

        ax.set_xlabel(self.var_strs[vary_var], fontsize=14)
        ax.set_ylabel(r'$m_{200c,\mathrm{fit}} / m_{200c,\mathrm{truth}}$', fontsize=14)
        ax.legend()
        
        if(sfx is None):
            plt.savefig('{}/{}_mass_convergence.png'.format(self.out_dir, vary_var), dpi=300)
        else:
            plt.savefig('{}/{}_mass_convergence_{}.png'.format(self.out_dir, vary_var, sfx), dpi=300)

    
    # ------------------------------------------------------------------------------------------------------------


    def plot_profile_convergence(self, vary_var, sfx=None):
        '''
        Plot the halo profiles from many runs of fit_halos(), results of which are found in self.out_dir

        Parameters
        ---------
        vary_var : string
            which variable to plot convergence with respect to. This should be one of the quantities given in the 
            output directory suffix as defined in fit_halos(). Choices are: zl, zs, N, fov, rmin, rmax. It is 
            assumed that user has ensured that all output subdirs in self.out_dir have only one of these variables
            changing, to make for a meaningul convergence study.
        sfx : string, optional
            Suffix to add to output figure filename (to prevent overwrites if testing convergence on different 
            ranges of one variable). Default is None.
        '''
        
        print('\n--------------- reading fits and plotting profile convergence ---------------')
       
        # get true halo profile
        truth_out = glob.glob('{}/true*.obj'.format(self.out_dir))
        true_profiles = np.array([pickle.load(open(f, 'rb')) for f in truth_out])
        vary_val = np.array([float(f.split('/')[-1].split(vary_var)[-1].split('_')[0].strip('.obj')) 
                               for f in truth_out])
        sorter = np.argsort(vary_val)
        true_profiles = true_profiles[sorter]
        
        r200c = [t.r200c for t in true_profiles]
        c = [t.c for t in true_profiles]
        zl = [t.zl for t in true_profiles]

        # get profiles from fit results -- no bootstrap errors here
        fit_out = glob.glob('{}/fitted_profile*.obj'.format(self.out_dir))
        fitted_profiles = np.array([pickle.load(open(f, 'rb')) for f in fit_out])
        
        rmin = [float(p.split('rmin')[-1].split('_')[0]) for p in fit_out]
        rmax = [float(p.split('rmax')[-1].split('_')[0].strip('.obj')) for p in fit_out]
        rsamp = [np.linspace(rmin[i], rmax[i], 1000) for i in range(len(rmin))]
        rsamp_full = np.linspace(min(rmin), max(rmax), 1000)
        dSigma_fitted = [fitted_profiles[i].delta_sigma(rsamp_full, bootstrap=False) 
                         for i in range(len(fitted_profiles))]
       
        # sort profiles by increasing vary_val
        vary_val = np.array([float(f.split('/')[-1].split(vary_var)[-1].split('_')[0].strip('.obj')) 
                               for f in fit_out])
        sorter = np.argsort(vary_val)
        vary_val = vary_val[sorter]
        fitted_profiles = fitted_profiles[sorter]
        dSigma_fitted = np.array(dSigma_fitted)[sorter]
        rsamp = np.array(rsamp)[sorter]
        
        # get fitted halo props
        all_r200c = [p.r200c for p in fitted_profiles]
        all_r200c_err = [p.r200c_err for p in fitted_profiles]
        all_c = [p.c for p in fitted_profiles]
        all_c_err = [p.c_err for p in fitted_profiles] 
       
        # plot
        f = plt.figure(figsize=(12,6))
        ax = f.add_subplot(121) 
        colors = plt.cm.viridis(np.linspace(0, 1, len(dSigma_fitted)))
        
        for i in range(len(dSigma_fitted)):
            ax.loglog(rsamp[i]/all_r200c[i], dSigma_fitted[i], color=colors[i], lw=2)
        # dummy plot for legend
        ax.plot([-100, -101], [-100, -101], color='k', label=r'$\Delta\Sigma_\mathrm{{fits}}$')
        
        # plot true profile (include radius, concentration in legend if they are constant for all halos)
        if(len(np.unique(r200c)) == 1 and len(np.unique(c)) == 1):
            label=r'$\Delta\Sigma_\mathrm{{true}},\>\>r_{{200c}}={:.3f}; c={:.3f}$'.format(r200c[0], c[0])
        else:
            label=r'$\Delta\Sigma_\mathrm{{true}}$'.format(r200c[0], c[0])
        # dummy for legend
        ax.plot([-100, -101], [-100, -101], '--', color='k', label=label, lw=1.75) 
        
        for i in range(len(true_profiles)):
            dSigma_true = true_profiles[i].delta_sigma(rsamp[i], bootstrap=False)
            ax.loglog(rsamp[i]/r200c[i], dSigma_true, '--', color=colors[i])  

        # format
        ax.legend(fontsize=12, loc='best')
        ax.set_xlabel(r'$r/R_{200c}$', fontsize=14)
        ax.set_ylabel(r'$\Delta\Sigma\>\>\lbrack\mathrm{M}_\odot\mathrm{pc}^{-2}\rbrack$', fontsize=14)
        ax.set_xlim([0.2, 6])
        ax.set_ylim([2, 300])

        # plot convergence in the radius-concentration plane
        ax2 = f.add_subplot(122)
        grid_r = np.linspace(0.6, 1.2, 100)
        grid_c = np.linspace(3, 6, 100)
       
        ax2.plot(all_r200c, all_c, '-k', lw=0.5)
        for i in range(len(all_c)):
            ax2.errorbar(all_r200c[i], all_c[i], xerr=all_r200c_err[i], yerr=all_c_err[i], 
                         ms=10, marker='.', c=colors[i])
        # dummy for legend
        ax2.plot([-100], [-100], '.k', ms=10, label=r'$\mathrm{{fit}}$')
       
        if(len(np.unique(r200c)) != 1 and len(np.unique(c)) != 1):
            for i in range(len(r200c)):
                ax2.plot([r200c[i]], [c[i]], 'x', c=colors[i], ms=10)
            # dummy for legend
            ax2.plot([-100], [-100], 'xk', ms=10, label=r'$\mathrm{{truth}}$')
        else:
            ax2.plot([r200c[0]], [c[0]], 'xk', ms=10, label=r'$\mathrm{{truth}}$')

        # include c-M relation curve (if redshift of lens is fixed)
        if(len(np.unique(zl)) == 0):
            tmp_profile = NFW(1,1,zl[0])
            tmp_m200c = np.zeros(len(grid_r))
            for i in range(len(tmp_m200c)):
                tmp_profile.r200c = grid_r[i]
                tmp_m200c[i] = tmp_profile.radius_to_mass()
            tmp_c, tmp_dc = cm(tmp_m200c, zl, cosmo)
            ax2.plot(grid_r, tmp_c, '--k', lw=2, label=r'$c\mathrm{-}M\mathrm{\>relation\>(Child+2018)}$')
            ax2.fill_between(grid_r, tmp_c - tmp_dc, tmp_c + tmp_dc, color='k', alpha=0.1, lw=0)

        # format
        ax2.set_xlim([min(grid_r), max(grid_r)])
        ax2.set_ylim([min(grid_c), max(grid_c)])
        ax2.legend(fontsize=12, loc='upper right')
        ax2.set_xlabel(r'$R_{200c}\>\>\left[\mathrm{Mpc}\right]$', fontsize=14)
        ax2.set_ylabel(r'$c_{200c}$', fontsize=14)
        ax.set_xlim([0.2, 6])
        ax.set_ylim([2, 300])
        
        cmap = mpl.cm.viridis
        sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=min(vary_val), vmax=max(vary_val)))
        sm._A = []
        cbar = plt.colorbar(sm)
        
        cbar.set_label(self.var_strs[vary_var], fontsize=14)
        plt.tight_layout() 

        if(sfx is None):
            f.savefig('{}/{}_profile_convergence.png'.format(self.out_dir, vary_var), dpi=300)
        else:
            f.savefig('{}/{}_profile_convergence_{}.png'.format(self.out_dir, vary_var, sfx), dpi=300)
    
    
    # ------------------------------------------------------------------------------------------------------------


    def plot_data_convergence(self, vary_var, rmin=0, rmax=None, bin_data=False, rbins=25, sfx=None, 
                              plot_gradient=False, plot_zdepx=True):
        '''
        Plot the scaled shear data from many runs of fit_halos(), results of which are found in self.out_dir

        Parameters
        ---------
        vary_var : string
            which variable to plot convergence with respect to. This should be one of the quantities given in the 
            output directory suffix as defined in fit_halos(). Choices are: zl, zs, N, fov, rmin, rmax. It is 
            assumed that user has ensured that all output subdirs in self.out_dir have only one of these variables
            changing, to make for a meaningul convergence study.
        rmin : float
            The minimum radial distance of sources to include in the plot in Mpc (e.g. rmin = 0.3 will
            remove the inner 300kpc/h of source information). Defaults to 0.
        rmax : float
            The maximum radial distance of sources to include in the plot in Mpc (e.g. rmin = 2.0 will
            trim the source popiulation beyond 2*r200c. Defaults to None, in which case rmax will automatically be 
            set to the full radial extent of the FOV as given in the NFW input.
        bin_data : bool
            whether or not to fits were done with shears averaged in radial bins, rather than to each individual 
            source.Defaults to False.
        rbins : int
            Number of bins used to perform passed fits over the radial range of the data, if bin_data == True. 
            Defaults to 25.
        sfx : string, optional
            Suffix to add to output figure filename (to prevent overwrites if testing convergence on different 
            ranges of one variable). Default is None.
        plot_gradient : bool, optional
            Whether or not to render a second subplot which displays the gradient of each bin. Defaults to False, 
            in which case the fractional bias of the data wrt to the true input profile is plotted.
        plot_zdepx : bool, optional
            Whether or not to render a subplot inserted within the main figure, which shows the redshift dependence
            of the differential surface density at the dimensionless radius x = r/r200c = 1. Over plotted on the 
            data is the theoretical expectation of ΔΣ at x=1, which is (H_0 / H(z))**2 for the NFW profile. This only
            makes sense to turn on if the lens redshift is varying, and if the halo radii and concentrations are all
            the same (since ΔΣ at the scale radius also depends on those quantities). This assumption is made explicit
            in the plotting of this panel.
        '''
        
        print('\n--------------- reading sources and plotting data convergence ---------------')
        
        # get true halo profile
        truth_out = glob.glob('{}/true*.obj'.format(self.out_dir))
        true_profiles = np.array([pickle.load(open(f, 'rb')) for f in truth_out])
        vary_val = np.array([float(f.split('/')[-1].split(vary_var)[-1].split('_')[0].strip('.obj')) 
                               for f in truth_out])
        sorter = np.argsort(vary_val)
        true_profiles = true_profiles[sorter]
        r200c = np.array([float(t.r200c) for t in true_profiles])
        rs = np.array([float(t.rs) for t in true_profiles])
        c = np.array([float(t.c) for t in true_profiles])
        dc = (200/3) * c**3 / (np.log(1+c)-c/(1+c))
        zl = np.array([float(t.zl) for t in true_profiles])

        # find and sort lenses by increasing vary_val
        sources_out = glob.glob('{}/source_data*.obj'.format(self.out_dir))
        lenses = np.array([pickle.load(open(f, 'rb')) for f in sources_out])        
        vary_val = np.array([float(f.split('/')[-1].split(vary_var)[-1].split('_')[0].strip('.obj')) 
                               for f in sources_out])
        sorter = np.argsort(vary_val)
        vary_val = vary_val[sorter]
        lenses = lenses[sorter]
        
        # prepare figure
        colors = plt.cm.viridis(np.linspace(0, 1, len(lenses)))
        f = plt.figure(figsize=(6,6))
        data_lw = 0.85
        data_ms = 3
       
        if(not plot_zdepx): 
            ax = f.add_subplot(211)
            ax2 = f.add_subplot(212)
        else:
            ax = f.add_subplot(311)
            ax2 = f.add_subplot(312)
            ax3 = f.add_subplot(313)
            dsig_exp_rs = np.array([t.delta_sigma([t.rs]) for t in true_profiles])
            ax3.plot(np.sort(zl), dsig_exp_rs/dsig_exp_rs, '--k')
         
        # read data and plot for each lens
        for i in range(len(lenses)):
            
            # ----- read -----
            true_profile = true_profiles[i]
            
            lens = lenses[i]
            lens.set_radial_cuts(rmin, rmax)
            bg = lens.get_background()
            all_r = bg['r']
            yt = bg['yt']
            k = bg['k']
            zs = bg['zs']
            print('binning data for lens {}/{}'.format(i+1, len(lenses)))
            binned_data = lens.calc_delta_sigma_binned(nbins=rbins, return_edges=True, 
                                                 return_std=True, return_gradients=True)
            
            if(bin_data):
                delsig = binned_data['delta_sigma_mean']
                r = binned_data['r_mean'] 
            else:
                delsig = lens.calc_delta_sigma()
                r = all_r
             
            # ----- plot data -----
            print('plotting')
            if(bin_data):
                delsig_std = binned_data['delta_sigma_std']
                delsig_err = binned_data['delta_sigma_se_mean']
                r_std = binned_data['r_std']
                r_err = binned_data['r_se_mean']
                bin_widths = np.diff(binned_data['bin_edges'])
                ax.errorbar(r/r200c[i], delsig, xerr=bin_widths/2/r200c[i], yerr=delsig_std, fmt='none', capsize=2, 
                            ecolor=colors[i], elinewidth=data_lw, capthick=data_lw)
                ax.errorbar(r/r200c[i], delsig, xerr=r_err/r200c[i], yerr=delsig_err, fmt='none', ecolor=colors[i],
                            elinewidth=6, alpha=0.2)
                ax.plot(r/r200c[i], delsig, '-o', c=colors[i], lw=data_lw, ms=data_ms) 
            else:
                ax.plot(r/r200c[i], delsig, 'x', c=colors[i], alpha=0.33, ms=data_ms)
            
            dSigma_true = true_profile.delta_sigma(r, bootstrap=False)
            ax.loglog(r/r200c[i], dSigma_true, '--', color=colors[i]) 
 
            # plot bin gradients or biases
            if(plot_gradient):
                grad = binned_data['bin_grad']
                ax2.plot(r/r200c[i], grad, '-x', c=colors[i], lw=data_lw, ms=data_ms)
            else:
                delsig_true = true_profile.delta_sigma(r, bootstrap=False)
                ax2.errorbar(r/r200c[i], delsig/delsig_true, 
                            xerr=bin_widths/2/r200c[i], yerr=delsig_std/delsig_true, fmt='none', capsize=2, 
                            ecolor=colors[i], elinewidth=data_lw, capthick=data_lw)
                ax2.plot(r/r200c[i], delsig/delsig_true, '-o', c=colors[i], lw=data_lw, ms=data_ms)
                ax2.plot(ax.get_xlim(), [1,1], '--k', lw=1)
                ax2.set_xscale('log')
                ax2.set_xlim(ax.get_xlim())
           
            # plot ΔΣ at scale radius as a fucntion of redshift
            if(plot_zdepx):
                # get all points within 1kpc of x=1
                bin_width = np.diff(binned_data['r_mean'])[0]
                rs_mask = np.abs(all_r - rs[i]) <= 0.001
                dimensionless_r = all_r[rs_mask]
                delsig_rs = lens.calc_delta_sigma()[rs_mask]
                delsig_rs_err = np.std(delsig_rs)/np.sqrt(len(delsig_rs))
                pdb.set_trace()
                ax3.errorbar([zl[i]], [np.mean(delsig_rs) / dsig_exp_rs[i]], 
                              xerr=[0], yerr=[delsig_rs_err / dsig_exp_rs[i]], fmt='.', 
                              capsize=2, color=colors[i], ecolor=colors[i], elinewidth=data_lw, capthick=data_lw)
         
        # dummy plot to get error bars and true profiles in legend exaxctly once
        ax.errorbar(-100, -100, xerr=1, y_std=1, fmt='none', capsize=2, ecolor='k', elinewidth=data_lw,
                    label=r'$\gamma_{\mathrm{NFW}}\Sigma_c,\>\>y:\mathrm{std},\>x:\mathrm{bin\>width},$')
        ax.errorbar(-100, -100, xerr=10, y_err=10, fmt='none', ecolor='k', elinewidth=6, alpha=0.2, 
                    label=r'$\gamma_{\mathrm{NFW}}\Sigma_c,\>\>\mathrm{bin\>sem}$')
        ax.plot(-100, -100, '-ok', label=r'$\gamma_{\mathrm{NFW}}\Sigma_c$', lw=data_lw, ms=data_ms)         
        ax.plot([-100, -101], [-100, -101], '--', color='k', label=r'$\Delta\Sigma_\mathrm{{true}}$')
        
        # ----- format -----
        ax.set_xscale('log')
        ax.set_yscale('log')
        ax.set_ylabel(r'$\Delta\Sigma\>\>\lbrack\mathrm{M}_\odot\mathrm{pc}^{-2}\rbrack$', fontsize=14)
        ax.set_title(r'$z_\mathrm{{lens}} = {}$'.format(zl), fontsize=14)
        ax.legend(fontsize=8, loc='best')
        ax2.set_xlabel(r'$r/R_{200c}$', fontsize=14)
        if(plot_gradient): ax2.set_ylabel(r'$\mathrm{best\>fit\>slope}$', fontsize=14)
        else: ax2.set_ylabel(r'$\Delta\Sigma/\Delta\Sigma_\mathrm{true}$', fontsize=14)
        if(plot_zdepx):
            ax3.set_xlabel(r'$z_l$', fontsize=12)
            ax3.set_ylabel(r'$\Delta\Sigma(r_s)/\Delta\Sigma_{\mathrm{true}}$', fontsize=12)
        
        # colorbar 
        cmap = mpl.cm.viridis
        sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=min(vary_val), vmax=max(vary_val)))
        sm._A = []
        cbar = plt.colorbar(sm, ax=[ax, ax2]) 
        cbar.set_label(self.var_strs[vary_var], fontsize=14)

        # ----- save -----
        if(sfx is None):
            f.savefig('{}/{}_data_convergence.png'.format(self.out_dir, vary_var), dpi=300)
        else:
            f.savefig('{}/{}_data_convergence_{}.png'.format(self.out_dir, vary_var, sfx), dpi=300)



# ==================================================================================================================
# ==================================================================================================================



if(__name__ == "__main__"):
   
    # turn to True to force overwrite pre-existing fits
    overwrite = False
    
    fitter = nfw_profile_fitter(NFW_dir = '/Users/joe/repos/repo_user/nfw_lensing_runs/output_newtest')
    fitter.fit_halos(rmax=[None], rmin=0.2, single_halo_plots=True, grid_scan=False, overwrite=overwrite, 
                     bin_data=True, rbins=30, inputs='grid')
    fitter.plot_mass_convergence(vary_var='zl')
    fitter.plot_profile_convergence(vary_var='zl')
    fitter.plot_data_convergence(vary_var='zl', rmin=0.2, rmax=None, bin_data=True, 
                                 rbins=30, plot_gradient=False, plot_zdepx=False)
    sys.exit()
    
    fitter = nfw_profile_fitter(NFW_dir = '/Users/joe/repos/repo_user/nfw_lensing_runs/output_prop/vary_zl')
    fitter.fit_halos(rmax=[None], rmin=0.2, single_halo_plots=True, grid_scan=False, overwrite=overwrite, 
                     bin_data=True, rbins=30, inputs='grid')
    fitter.plot_mass_convergence(vary_var='zl')
    fitter.plot_profile_convergence(vary_var='zl')
    fitter.plot_data_convergence(vary_var='zl', rmin=0.2, rmax=None, bin_data=True, 
                                 rbins=30, plot_gradient=False, plot_zdepx=False)
    
    
    fitter = nfw_profile_fitter(NFW_dir = '/Users/joe/repos/repo_user/nfw_lensing_runs/output/vary_rfrac_0.5')
    fitter.fit_halos(rmax=[None], rmin=0.2, single_halo_plots=True, grid_scan=False, overwrite=overwrite, 
                     bin_data=True, rbins=30, inputs='grid')
    fitter.plot_mass_convergence(vary_var='fov')
    fitter.plot_profile_convergence(vary_var='fov')
    fitter.plot_data_convergence(vary_var='fov', rmin=0.2, rmax=None, bin_data=True, 
                                 rbins=30, plot_gradient=False, plot_zdepx=False)
     
    fitter = nfw_profile_fitter(NFW_dir = '/Users/joe/repos/repo_user/nfw_lensing_runs/output/vary_rfrac_0.2')
    fitter.fit_halos(rmax=[None], rmin=0.2, single_halo_plots=True, grid_scan=False, overwrite=overwrite, 
                     bin_data=True, rbins=30, inputs='grid')
    fitter.plot_mass_convergence(vary_var='fov')
    fitter.plot_profile_convergence(vary_var='fov')
    fitter.plot_data_convergence(vary_var='fov', rmin=0.2, rmax=None, bin_data=True, 
                                 rbins=30, plot_gradient=False)
    
    fitter = nfw_profile_fitter(NFW_dir = '/Users/joe/repos/repo_user/nfw_lensing_runs/output/vary_zl')
    fitter.fit_halos(rmax=[None], rmin=0.2, single_halo_plots=True, grid_scan=False, overwrite=overwrite, 
                     bin_data=True, rbins=30, inputs='grid')
    fitter.plot_mass_convergence(vary_var='zl')
    fitter.plot_profile_convergence(vary_var='zl')
    fitter.plot_data_convergence(vary_var='zl', rmin=0.2, rmax=None, bin_data=True, 
                                 rbins=30, plot_gradient=False)
    
    
