# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
"""
A script to calculate overlap between two IC realisations of the same
simulation. The matching is performed for haloes whose total particles mass is
    - CSiBORG: > 1e13 Msun/h,
    - Quijote: > 1e14 Msun/h,
since Quijote has much lower resolution than CSiBORG.
"""
from argparse import ArgumentParser
from copy import deepcopy
from datetime import datetime
from distutils.util import strtobool

import numpy
from scipy.ndimage import gaussian_filter

try:
    import csiborgtools
except ModuleNotFoundError:
    import sys

    sys.path.append("../")
    import csiborgtools


def pair_match(nsim0, nsimx, simname, sigma, verbose):
    """
    Calculate overlaps between two simulations.

    Parameters
    ----------
    nsim0 : int
        The reference simulation IC index.
    nsimx : int
        The cross simulation IC index.
    simname : str
        Simulation name.
    sigma : float
        Smoothing scale in number of grid cells.
    verbose : bool
        Verbosity flag.

    Returns
    -------
    None
    """
    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
    smooth_kwargs = {"sigma": sigma, "mode": "wrap"}

    if simname == "csiborg":
        overlapper_kwargs = {"box_size": 2048, "bckg_halfsize": 475}
        mass_kind = "fof_totpartmass"
        bounds = {mass_kind: (1e13, None)}
        cat0 = csiborgtools.read.CSiBORGHaloCatalogue(
            nsim0, paths, bounds=bounds, load_fitted=False,
            with_lagpatch=True)
        catx = csiborgtools.read.CSiBORGHaloCatalogue(
            nsimx, paths, bounds=bounds, load_fitted=False,
            with_lagpatch=True)
    elif simname == "quijote":
        overlapper_kwargs = {"box_size": 512, "bckg_halfsize": 256}
        mass_kind = "group_mass"
        bounds = {mass_kind: (1e14, None)}
        cat0 = csiborgtools.read.QuijoteHaloCatalogue(
            nsim0, paths, 4, load_fitted=False, with_lagpatch=True)
        catx = csiborgtools.read.QuijoteHaloCatalogue(
            nsimx, paths, 4, load_fitted=False, with_lagpatch=True)
    else:
        raise ValueError(f"Unknown simulation name: `{simname}`.")

    halomap0 = csiborgtools.read.read_h5(
        paths.particles(nsim0, simname))["halomap"]
    parts0 = csiborgtools.read.read_h5(
        paths.initmatch(nsim0, simname, "particles"))["particles"]
    hid2map0 = {hid: i for i, hid in enumerate(halomap0[:, 0])}

    halomapx = csiborgtools.read.read_h5(
        paths.particles(nsimx, simname))["halomap"]
    partsx = csiborgtools.read.read_h5(
        paths.initmatch(nsimx, simname, "particles"))["particles"]
    hid2mapx = {hid: i for i, hid in enumerate(halomapx[:, 0])}

    if verbose:
        print(f"{datetime.now()}: calculating the background density fields.",
              flush=True)
    overlapper = csiborgtools.match.ParticleOverlap(**overlapper_kwargs)
    delta_bckg = overlapper.make_bckg_delta(parts0, halomap0, hid2map0, cat0,
                                            verbose=verbose)
    delta_bckg = overlapper.make_bckg_delta(partsx, halomapx, hid2mapx, catx,
                                            delta=delta_bckg, verbose=verbose)

    if verbose:
        print(f"{datetime.now()}: NGP crossing the simulations.", flush=True)
    matcher = csiborgtools.match.RealisationsMatcher(
        mass_kind=mass_kind, **overlapper_kwargs)
    match_indxs, ngp_overlap = matcher.cross(cat0, catx, parts0, partsx,
                                             halomap0, halomapx, delta_bckg,
                                             verbose=verbose)

    # We want to store the halo IDs of the matches, not their array positions
    # in the catalogues.
    match_hids = deepcopy(match_indxs)
    for i, matches in enumerate(match_indxs):
        for j, match in enumerate(matches):
            match_hids[i][j] = catx["index"][match]

    fout = paths.overlap(nsim0, nsimx, smoothed=False)
    if verbose:
        print(f"{datetime.now()}: saving to ... `{fout}`.", flush=True)
    numpy.savez(fout, ref_hids=cat0["index"], match_hids=match_hids,
                ngp_overlap=ngp_overlap)

    if not sigma > 0:
        return

    if verbose:
        print(f"{datetime.now()}: smoothing the background field.", flush=True)
    gaussian_filter(delta_bckg, output=delta_bckg, **smooth_kwargs)

    # We calculate the smoothed overlap for the pairs whose NGP overlap is > 0.
    smoothed_overlap = matcher.smoothed_cross(cat0, catx, parts0, partsx,
                                              halomap0, halomapx, delta_bckg,
                                              match_indxs, smooth_kwargs,
                                              verbose=verbose)

    fout = paths.overlap(nsim0, nsimx, smoothed=True)
    if verbose:
        print(f"{datetime.now()}: saving to ... `{fout}`.", flush=True)
    numpy.savez(fout, smoothed_overlap=smoothed_overlap, sigma=sigma)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--nsim0", type=int,
                        help="Reference simulation IC index.")
    parser.add_argument("--nsimx", type=int,
                        help="Cross simulation IC index.")
    parser.add_argument("--simname", type=str, help="Simulation name.")
    parser.add_argument("--sigma", type=float, default=0,
                        help="Smoothing scale in number of grid cells.")
    parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
                        default=False, help="Verbosity flag.")
    args = parser.parse_args()

    pair_match(args.nsim0, args.nsimx, args.simname, args.sigma, args.verbose)