"""Sort and group svd banks and time slices by sample rates."""

# Copyright (C) 2024 Yun-Jing Huang

import itertools
import sys
from collections import Counter, OrderedDict, defaultdict

import numpy as np
import torch
from lal.utils import CacheEntry
from sgnts.base import Offset
from sgnts.transforms.resampler import DOWN_HALF_LENGTH, UP_HALF_LENGTH

from sgnl.svd_bank import Bank, parse_bank_files


def group_and_read_banks(
    svd_bank: list[str],
    source_ifos: list[str] | None = None,
    nsubbank_pretend: int = 0,
    nslice=-1,
    verbose=False,
):
    """Read a list of svd banks file names into bank objects, and group by ifo and
    bankid

    Args:
        svd_bank:
            list[str], a list of svd bank file paths
        source_ifos:
            list[str], a list of ifo the source is processing. Used for checking
            consistency of ifos between the svd banks and data sources
        nsubbank_pretend:
            int, default 0. If not 0, pretend there are this many subbanks by copying
            the first subbank read in this many times. Used for benchmarking.
        nslice:
            int, default -1. If nslice > 0, only read in this number of time slices
        verbose:
            bool, be verbose
    """
    # Read SVD banks
    svd_bank_cache = [CacheEntry.from_T050017(path) for path in svd_bank]
    svd_bank_cache.sort(key=lambda cache_entry: cache_entry.description)
    svd_banks = []

    for _key, seq in itertools.groupby(
        svd_bank_cache, key=lambda cache_entry: cache_entry.description
    ):
        svd_group = dict(
            (cache_entry.observatory, cache_entry.url) for cache_entry in seq
        )
        svd_banks.append(svd_group)

    ifoset = set([tuple(s.keys()) for s in svd_banks])
    if len(ifoset) != 1:
        raise ValueError("The ifos have different sets of svd bank files provided.")

    ifos = sorted(list(ifoset)[0])
    if source_ifos and ifos != sorted(source_ifos):
        raise ValueError(
            f"Data source ifos: {source_ifos} must be the same as svd bank ifos: {ifos}"
        )

    banks: dict[str, list] = {ifo: [] for ifo in ifos}
    if nsubbank_pretend:
        # Pretend we are filtering multiple banks by copying the first bank
        #   many times
        svd_bank_url_dict = svd_banks[0]
        banks_per_svd = parse_bank_files(svd_bank_url_dict, verbose=verbose, fast=True)
        for _ in range(nsubbank_pretend):
            for ifo in ifos:
                banks[ifo].extend([banks_per_svd[ifo][0]])
    else:
        for svd_bank_url_dict in svd_banks:
            banks_per_svd = parse_bank_files(
                svd_bank_url_dict, verbose=verbose, fast=True
            )
            for ifo in ifos:
                banks[ifo].extend(banks_per_svd[ifo])

    # TODO: fix nslice reading
    if nslice > 0:
        for ifo in ifos:
            bifo = banks[ifo]
            for sub in bifo:
                if len(sub.bank_fragments) < nslice:
                    raise ValueError(
                        f"nslice: {nslice} greater than number of bank \
                            fragments: {len(sub.bank_fragments)}"
                    )
                sub.bank_fragments = sub.bank_fragments[:nslice]

    if verbose:
        print("Using nsubbanks", len(banks[ifo]), flush=True, file=sys.stderr)

    return OrderedDict(sorted(banks.items()))


class SortedBank:
    """Sort and group svd banks and time slices by sample rate

    Args:
        banks:
            dict[str, list[Bank]], a dictionary keyed by ifos, with values of lists
            of subbanks as gstlal bank objects
        device:
            str, the torch device
        dtype:
            torch.dtype, the torch data type
        memory_format:
            torch.memory_format, the memory format for the svd tensors
        nsubbank_pretend:
            int, default 0. If not 0, pretend there are this many subbanks by copying
            the first subbank read in this many times. Used for benchmarking.
        nslice:
            int, default -1. If nslice > 0, only read in this number of time slices
        verbose:
            bool, be verbose
    """

    def __init__(
        self,
        banks: dict[str, list[Bank]],
        device: str = "cpu",
        dtype: torch.dtype = torch.float32,
        memory_format: torch.memory_format = torch.contiguous_format,
        nsubbank_pretend: int = 0,
        nslice: int = -1,
        verbose: bool = False,
    ):
        self.device = device
        self.dtype = dtype
        self.memory_format = memory_format
        self.verbose = verbose

        temp = torch.empty(1, dtype=dtype)
        self.cdtype = torch.complex(temp, temp).dtype
        self.nsubbank_pretend = nsubbank_pretend

        self.bank_metadata, self.reordered_bank = self.prepare_metadata(banks)
        # Prepare tensors for LLOID methods
        (
            self.coeff_sv_cat,
            self.bases_cat,
            self.template_ids,
            self.end_time_delta,
            self.subbankids,
            self.bankids_map,
            self.sngls,
            self.autocorrelation_banks,
            self.autocorrelation_length_mask,
            self.autocorrelation_lengths,
            self.horizon_distance_funcs,
            self.template_durations,
        ) = self.prepare_tensors(self.bank_metadata, self.reordered_bank)

        del self.reordered_bank
        del banks

    def prepare_metadata(self, bank):
        """Determine rates and template properties across subbanks

        bank_metadata.keys() = ['ifos', 'nifo', 'nsubbank', 'maxrate', 'unique_rates',
        'nrates', 'nfilter_samples', 'ntempmax', 'delay_per_rate', 'sorted_rates']

        Arguments:
        ----------
        bank:
            gstlal Bank class
        """
        bank_metadata = {}
        ifos = sorted(bank.keys())
        bank_metadata["ifos"] = ifos
        bank_metadata["nifo"] = len(ifos)

        # assume all ifos have same banks
        bank0 = bank[ifos[0]]

        nsubbank = len(bank0)  # number of subbanks in each ifo
        bank_metadata["nsubbank"] = nsubbank

        # determine some properties across banks
        # determine rates
        maxrate = bank0[0].sample_rate_max
        bank_metadata["maxrate"] = maxrate  # assume same for all banks
        unique_rates = dict(
            sorted(
                Counter([bi.rate for b in bank0 for bi in b.bank_fragments]).items(),
                reverse=True,
            )
        )
        bank_metadata["unique_rates"] = unique_rates
        bank_metadata["nrates"] = len(unique_rates)

        # number of samples in each filter, assume it's the same for all templates
        bank_metadata["nfilter_samples"] = (
            bank0[0].bank_fragments[0].orthogonal_template_bank.shape[1]
        )

        # maximum number of templates across all subbanks
        bank_metadata["ntempmax"] = max(
            sub.bank_fragments[0].mix_matrix.shape[1] for sub in bank0
        )
        delay_per_rate = dict.fromkeys(unique_rates.keys())
        for rate in unique_rates:
            delay_per_rate[rate] = max(
                int(bf.start * rate)
                for banki in bank.values()
                for sub in banki
                for bf in sub.bank_fragments
                if rate == bf.rate
            )
        bank_metadata["delay_per_rate"] = delay_per_rate

        sorted_rates, reordered_bank = self.sort_bank_by_rates(bank, bank_metadata)
        bank_metadata["sorted_rates"] = sorted_rates

        return bank_metadata, reordered_bank

    def sort_bank_by_rates(self, bank, bank_metadata):
        """
        Determine the ids to upsample from rate1 to rate2
        This places templates that need to be upsampled from and to
            the same rate next to each other, which allows the templates
            to be upsampled together and allows a simpler slicing method
        The following example shows id placement
        slice  | 0  | 1   |  2 |  3 |
        -----------------------------
        bank0: [2048, 1024, 512, 512]
        bank1: [2048, 1024, 512, 256]
        bank2: [2048,  512, 256]

        bases[sample_rate][upsample_rate]:

        bases[2048][()]
            = [bank0slice0, bank1slice0, bank2slice0]
        bases[1024][(2048,)]
            = [bank0slice1, bank1slice1]
        bases[512][(2048, 1024)]
            = [bank0slice2, bank0slice3, bank1slice2]
            # templates in the same bank will be placed next to each other,
                to allow more convenient summing in sum_same_rates()

        bases[512][(2048,)]
            = [bank2slice1]
        bases[256][(2048, 1024, 512)]
            = [bank1slice3]
        bases[256][(2048, 512)]
            = [bank2slice2]

        * The reason that bases[512][(2048, 1024)] is in a different group than
            bases[512][(2048,)] is because they will go into different upsamplers
        * The reason that bases[256][(2048, 1024, 512)] is in a different group
            than bases[256][(2048, 512)], even though they are upsampling to the
            same rate (512), is because due to the different upsampling path
            256->512->1024->2048 vs. 256->512->2048, the time stamps of the buffers
            will be different (to account for the upsampling padding)

        ----------

        sorted_rates = {
            from_rate: {
                to_rate: {
                    '_ids_counter':,
                    '_addto_counter':,
                    '_addids':,
                    'counts':,
                    'nbmax':,
                    'ntempmax':,
                    'segments':,
                    'sum_same_rate_slices':,
                    'uppad':,
                    'addslice':,
                    'metadata': {
                         bank_order: {
                             'ids':,'nslice':,
                         }
                    },
                }
            }
        }

        """

        ifos = bank_metadata["ifos"]
        bank0 = bank[ifos[0]]
        nsubbank = len(bank0)  # number of subbanks in each ifo
        unique_rates = bank_metadata["unique_rates"]
        maxrate = bank_metadata["maxrate"]

        #
        # Sort the unique rates, this might alter the order of the template banks
        #
        sorted_unique_rates0 = []
        unique_rates_by_bank = []
        for j in range(nsubbank):
            bk = bank0[j]
            if self.nsubbank_pretend:
                bankid = bk.bank_id + "_" + str(j)
            else:
                bankid = bk.bank_id
            unique_rates0 = dict(
                sorted(
                    Counter([bi.rate for bi in bk.bank_fragments]).items(), reverse=True
                )
            )
            unique_rates_by_bank.append(tuple(unique_rates0.keys()))
            # for key in unique_rates:
            #    if key not in unique_rates0:
            #        unique_rates0[key] = 0
            unique_rates0["bankid"] = bankid
            sorted_unique_rates0.append(unique_rates0)

        # sorted_unique_rates = sorted(
        #    # sorted_unique_rates, key=lambda x: x.keys(), reverse=True
        #    sorted_unique_rates,
        #    key=operator.itemgetter(*list(unique_rates.keys())),
        #    reverse=True,
        # )

        unique_rates_by_bank = sorted(set(unique_rates_by_bank), reverse=True)
        sorted_unique_rates = []

        for urb in unique_rates_by_bank:
            for s in sorted_unique_rates0:
                s2 = list(s.keys())
                s2.remove("bankid")
                if list(urb) == s2 and s not in sorted_unique_rates:
                    sorted_unique_rates.append(s)

        # if self.nsubbank_pretend:
        #    sorted_unique_rates *= self.nsubbank_pretend

        if self.verbose:
            print("sorted_unique_rates", flush=True, file=sys.stderr)
            for a in sorted_unique_rates:
                print(a, flush=True, file=sys.stderr)

        # check if nsubbank_pretend is true
        # bankids = [a["bankid"] for a in sorted_unique_rates]
        # nsubbank_pretend = False
        # if len(bankids) > 1 and len(set(bankids)) == 1:
        #    if self.verbose:
        #        print("nsubbank_pretend", flush=True, file=sys.stderr)
        #    nsubbank_pretend = True

        # reorder bankid
        bankid_order = {}
        for i, a in enumerate(sorted_unique_rates):
            bankid = a["bankid"]
            bankid_order[bankid] = i

        reorder = {}
        # reordered bank
        if self.nsubbank_pretend == 0:
            for ifo in ifos:
                reorder[ifo] = {}
                for j in range(nsubbank):
                    bk = bank[ifo][j]
                    bankid = bk.bank_id
                    order = bankid_order[bankid]
                    reorder[ifo][order] = bk
        else:
            if self.verbose:
                print("nsubbank_pretend", flush=True, file=sys.stderr)
            for ifo in ifos:
                reorder[ifo] = {i: b for i, b in enumerate(bank[ifo])}

        #
        # Construct sorted_rates, determine id placement of timeslices
        #
        sorted_rates = defaultdict(dict)
        addto_ids = {}
        for a in sorted_unique_rates:
            temp = {}
            for from_rate in unique_rates:
                from_bank = sorted_rates[from_rate]
                bankid = a["bankid"]
                k = list(a.keys())
                if from_rate in k and a[from_rate] > 0:
                    to_rate = tuple(
                        ki
                        for ki in k
                        if type(ki) is int and ki > from_rate and a[ki] > 0
                    )
                    if to_rate not in from_bank:
                        from_bank[to_rate] = {
                            "_ids_counter": 0,
                            "_addto_counter": 0,
                            "_addids": [],
                            "counts": 0,
                            "nbmax": [],
                            "ntempmax": [],
                            "segments": [],
                            "sum_same_rate_slices": None,
                            "metadata": {},
                        }
                    to_bank = from_bank[to_rate]
                    if from_rate != maxrate:
                        to_bank["counts"] += a[from_rate]
                    temp[from_rate] = to_bank["_addto_counter"]
                    to_bank["_addto_counter"] += 1
            addto_ids[bankid] = temp

        # find the ids each bank should addto after upsampling
        # FIXME: is there a more elegant way?
        for _bankid, a in addto_ids.items():
            for from_rate in unique_rates:
                from_bank = sorted_rates[from_rate]
                k = list(a.keys())
                if from_rate in k:
                    to_rate = tuple(
                        ki for ki in k if type(ki) is int and ki > from_rate
                    )
                    if from_rate != maxrate:
                        from_bank[to_rate]["_addids"].append(a[to_rate[-1]])

        sorted_rates[maxrate][()]["counts"] = nsubbank

        urates = list(unique_rates.keys())
        downpads = {r: None for r in urates}
        downpad = 0
        downpads[urates[0]] = 0
        for urate in urates[1:]:
            downpad += Offset.fromsamples(DOWN_HALF_LENGTH, urate)
            downpads[urate] = downpad

        # loop over all subbanks and update metadata in each rate group
        for j in range(nsubbank):
            bk = reorder[ifos[0]][j]
            rates = [bf.rate for bf in bk.bank_fragments]
            # urates = np.array(sorted(set(rates), reverse=True))
            urates = list(sorted(set(rates), reverse=True))
            urs = urates[1:]
            # uppad = {
            # #r0: sum(Offset.fromsamples(UP_HALF_LENGTH, ri) for ri in urs[urs >= r0])
            # r0: sum([Offset.fromsamples(UP_HALF_LENGTH, ri) for ri in urs[urs >= r0]])
            # for r0 in urs
            # }
            uppads = {r: None for r in urates}
            uppad = 0
            uppads[maxrate] = 0
            for ri in urs:
                uppad += Offset.fromsamples(UP_HALF_LENGTH, ri)
                uppads[ri] = uppad

            for bi, bf in enumerate(bk.bank_fragments):
                rate = bf.rate
                to_rate = tuple(r for r in urates if r > rate)
                rate_group = sorted_rates[rate][to_rate]
                if j not in rate_group["metadata"]:
                    rate_group["metadata"][j] = {
                        "ids": [],
                        "nslice": [],
                    }
                mdata = rate_group["metadata"][j]
                mdata["bankid"] = bk.bank_id
                mdata["nslice"].append(bi)
                rate_group["uppad"] = uppads[rate]
                rate_group["downpad"] = downpads[rate]
                rate_group["shift"] = uppads[rate] + downpads[rate]
                rate_group["segments"].append((bf.start, bf.end))
                if rate != maxrate:
                    if rate == rates[bi - 1]:
                        mdata["ids"].append(rate_group["_ids_counter"])
                        rate_group["_ids_counter"] += 1
                        if len(mdata["ids"]) > 1:
                            rate_group["sum_same_rate_slices"] = []
                    else:
                        mdata["ids"].append(rate_group["_ids_counter"])
                        rate_group["_ids_counter"] += 1
                else:
                    mdata["ids"] = [j]

        for ifo in ifos:
            for j in range(nsubbank):
                bk = reorder[ifo][j]
                rates = [bf.rate for bf in bk.bank_fragments]
                for bf in bk.bank_fragments:
                    rate = bf.rate
                    urates = np.array(sorted(set(rates), reverse=True))
                    to_rate = tuple(r for r in urates if r > rate)
                    rate_group = sorted_rates[rate][to_rate]
                    rate_group["nbmax"].append(bf.mix_matrix.shape[0])
                    rate_group["ntempmax"].append(bf.mix_matrix.shape[1])

        for from_rate, v in sorted_rates.items():
            for _to_rate, rate_group in v.items():
                if from_rate != maxrate:
                    addids = rate_group["_addids"]
                    rate_group["addslice"] = slice(addids[0], addids[-1] + 1)
                rate_group["nbmax"] = max(rate_group["nbmax"])
                rate_group["ntempmax"] = max(rate_group["ntempmax"])
                mdata = rate_group["metadata"]
                if rate_group["sum_same_rate_slices"] is not None:
                    for md in mdata.values():
                        ids = md["ids"]
                        rate_group["sum_same_rate_slices"].append(
                            slice(ids[0], ids[-1] + 1)
                        )

        return sorted_rates, reorder

    def prepare_tensors(self, bank_metadata, reordered_bank):
        """
        Prepare large tensors to store input and output of methods in LLOID
        coeff_sv:      all the coeff_sv from different banks
        bases:         all the bases from different banks
        """
        print(
            "Preparing tensors for LLOID methods...",
            end="",
            flush=True,
            file=sys.stderr,
        )

        dtype = self.dtype
        device = self.device

        ifos = bank_metadata["ifos"]
        nfilter_samples = bank_metadata["nfilter_samples"]
        nsubbank = bank_metadata["nsubbank"]
        sorted_rates = bank_metadata["sorted_rates"]

        # outputs
        bases_by_rate = {
            r1: {r2: {} for r2 in rb.keys()} for r1, rb in sorted_rates.items()
        }
        coeff_sv_by_rate = {
            r1: {r2: {} for r2 in rb.keys()} for r1, rb in sorted_rates.items()
        }

        # construct big tensors of data, bases, and coeff_sv, grouped by sample rates
        for from_rate, rbr in sorted_rates.items():
            for to_rate, rate_group in rbr.items():
                count = rate_group["counts"]
                mdata = rate_group["metadata"]
                nbm = rate_group["nbmax"]
                ntempmax = rate_group["ntempmax"]

                for ifo in ifos:

                    # group the bases by sample rate
                    bases_by_rate[from_rate][to_rate][ifo] = torch.zeros(
                        size=(count, nbm, nfilter_samples),
                        device=device,
                        dtype=dtype,
                    )

                    # group the coeff by sample rate
                    coeff_sv_by_rate[from_rate][to_rate][ifo] = torch.zeros(
                        size=(count, ntempmax, nbm),
                        device=device,
                        dtype=dtype,
                    )

                    # fill in the bases and coeff tensors!
                    # for k, ifo in enumerate(ifos):
                    bifo = reordered_bank[ifo]
                    for bank_order, md in mdata.items():
                        bifo_order = bifo[bank_order]
                        ids = md["ids"]
                        nslices = md["nslice"]
                        for id0, nslice in zip(ids, nslices):
                            this_slice = bifo_order.bank_fragments[nslice]
                            assert this_slice.rate == from_rate
                            b = this_slice.orthogonal_template_bank
                            c = this_slice.mix_matrix.T
                            bases_by_rate[from_rate][to_rate][ifo][
                                id0, : b.shape[0], :
                            ] = torch.tensor(b, device=device, dtype=dtype)
                            coeff_sv_by_rate[from_rate][to_rate][ifo][
                                id0, : c.shape[0], : c.shape[1]
                            ] = torch.tensor(c, device=device, dtype=dtype)
                bases = bases_by_rate[from_rate][to_rate][ifo].view(
                    -1,
                    nfilter_samples,
                )
                bases = (
                    bases.unsqueeze(1).unsqueeze(0).to(memory_format=self.memory_format)
                )
                bases_by_rate[from_rate][to_rate][ifo] = bases.view(
                    count, nbm, nfilter_samples
                )
                # benchmark conv methods
                rate_group["conv_group"] = True
                """
                if self.device == "cpu":
                    rb["conv_group"] = True
                else:
                    if self.verbose:
                        print("Benchmarking conv...", flush=True, file=sys.stderr)
                    tgroup = benchmark.Timer(
                        stmt="SNRSlices.conv_group(mask, data, nbm, basesr)",
                        setup="from greg.filtering.snr_slices import SNRSlices",
                        globals={
                            "mask": [True] * nifo,
                            "data": data_by_rate[from_rate][to_rate],
                            "nbm": nbm,
                            "basesr": bases_by_rate[from_rate][to_rate],
                        },
                    )

                    tgroupm = tgroup.timeit(500).median

                    tloop = benchmark.Timer(
                        stmt="SNRSlices.conv_loop(mask, data, nbm, basesr)",
                        setup="from greg.filtering.snr_slices import SNRSlices",
                        globals={
                            "mask": [True] * nifo,
                            "data": data_by_rate[from_rate][to_rate],
                            "nbm": nbm,
                            "basesr": bases_by_rate[from_rate][to_rate],
                        },
                    )

                    tloopm = tloop.timeit(500).median

                    if self.verbose:
                        print(
                            f"{from_rate=} {to_rate=}",
                            "conv_group",
                            tgroupm,
                            "conv_loop",
                            tloopm,
                            flush=True,
                            file=sys.stderr,
                        )
                    if tgroupm < tloopm:
                        rb["conv_group"] = True
                    else:
                        rb["conv_group"] = False
                """

        # Get template ids
        #   Assume same ids for all ifos for the same bank
        #   Init template ids as -1 for banks with ntemp < ntempmax,
        #   the template id for empty entries will be -1
        ntempmax = bank_metadata["ntempmax"]
        template_ids = (
            torch.ones(size=(nsubbank, ntempmax // 2), dtype=torch.int32) * -1
        )
        subbankids = []
        sngls = []
        end_time_delta = torch.zeros(size=(nsubbank,), dtype=torch.long)
        # in seconds
        template_durations = np.zeros((nsubbank, ntempmax // 2))
        bankids_map = defaultdict(list)
        horizon_distance_funcs = {}
        autocorrelation_lengths = {}
        for j in range(nsubbank):
            sngl = reordered_bank[ifos[0]][j].sngl_inspiral_table
            template_ids0 = torch.tensor([row.template_id for row in sngl])
            template_ids[j, : template_ids0.shape[0]] = template_ids0

            ends0 = [row.end.ns() for row in sngl]
            assert len(set(ends0)) == 1, "there are different end times in a subbank"
            end_time_delta[j] = list(set(ends0))[0]

            subbank_id = reordered_bank[ifos[0]][j].bank_id
            if self.nsubbank_pretend:
                bank_id = "%04d" % int(subbank_id.split("_")[0]) + "_" + str(j // 2)
            else:
                bank_id = "%04d" % int(subbank_id.split("_")[0])
            horizon_distance_funcs[bank_id] = reordered_bank[ifos[0]][
                j
            ].horizon_distance_func
            subbankids.append(subbank_id)
            bankids_map[bank_id].append(j)
            sngl0 = {row.template_id: row for row in sngl}
            sngls.append(sngl0)
            for k, row in enumerate(sngl):
                template_durations[j, k] = row.template_duration
            autocorrelation_lengths[bank_id] = reordered_bank[ifos[0]][
                j
            ].autocorrelation_bank.shape[1]

        bankids_map = OrderedDict(sorted(bankids_map.items()))

        # Write out single inspiral table
        # sngl0 = reordered_bank[ifos[0]][0].sngl_inspiral_table
        # row = sngl0[0]
        # keys = [a for a in dir(row) if not a.startswith('__') and not
        #           callable(getattr(row, a))]
        # import h5py
        # with h5py.File('sngl_inspiral_table.h5', "a") as f:
        #    for j in range(nsubbank):
        #        sngl = reordered_bank[ifos[0]][j].sngl_inspiral_table
        #        for row in sngl:
        #            group = str(row.template_id)
        #            f.create_group(group)
        #            for k in keys:
        #                v = getattr(row, k)
        #                if k == 'end':
        #                    v = float(v)
        #                f[group][k] = v

        # Trigger generator
        # Get the autocorrelation_bank
        acl = [
            reordered_bank[ifo][j].autocorrelation_bank.shape[1]
            for i, ifo in enumerate(ifos)
            for j in range(nsubbank)
        ]
        max_acl = max(acl)
        if len(set(acl)) > 1:
            print("Warning: different autocorrelation lengths among banks")
            autocorrelation_length_mask = {}
        else:
            autocorrelation_length_mask = None

        autocorrelation_banks = {}
        for ifo in ifos:
            autocorrelation_banks[ifo] = torch.zeros(
                size=(nsubbank, ntempmax // 2, max_acl),
                device=device,
                dtype=self.cdtype,
            )
            if autocorrelation_length_mask is not None:
                autocorrelation_length_mask[ifo] = torch.ones(
                    size=(nsubbank, ntempmax // 2, max_acl),
                    device=device,
                    dtype=bool,
                )

        for ifo in ifos:
            for j in range(nsubbank):
                acorr = reordered_bank[ifo][j].autocorrelation_bank
                # this is for adjusting to the bank used for impulse test
                # if acorr.shape[0] > ntempmax // 2:
                #    acorr = acorr[: ntempmax // 2]
                acorr_len = acorr.shape[1]
                if acorr_len < max_acl:
                    pad = (max_acl - acorr_len) // 2
                    autocorrelation_banks[ifo][j, : acorr.shape[0], pad:-pad] = (
                        torch.tensor(acorr)
                    )
                    autocorrelation_length_mask[ifo][j, :, :pad] = False
                    autocorrelation_length_mask[ifo][j, :, -pad:] = False
                else:
                    autocorrelation_banks[ifo][j, : acorr.shape[0], :] = torch.tensor(
                        acorr
                    )
                # autocorrelation_banks[ifo][j, : acorr.shape[0], : acorr.shape[1]] = (
                #     torch.tensor(acorr)
                # )
        print(" Done.", flush=True, file=sys.stderr)

        return (
            coeff_sv_by_rate,
            bases_by_rate,
            template_ids,
            end_time_delta,
            subbankids,
            bankids_map,
            sngls,
            autocorrelation_banks,
            autocorrelation_length_mask,
            autocorrelation_lengths,
            horizon_distance_funcs,
            template_durations,
        )
