import collections
from functools import reduce
from pathlib import Path

import numpy as np
import pandas as pd
from tqdm import tqdm

from mepylome.dtypes.arrays import ArrayType
from mepylome.dtypes.cache import memoize
from mepylome.dtypes.idat import IdatParser
from mepylome.dtypes.manifests import Manifest
from mepylome.dtypes.probes import Channel, ProbeType
from mepylome.utils.varia import normexp_get_xs

ENDING_GRN = "_Grn.idat"
ENDING_RED = "_Red.idat"
ENDING_GZ = ".gz"

NEUTRAL_BETA = 0.49


def is_valid_idat_basepath(basepath):
    if not isinstance(basepath, list):
        basepath = [basepath]
    return all(
        (
            Path(str(x) + ENDING_GRN).exists()
            or Path(str(x) + ENDING_GRN + ENDING_GZ).exists()
        )
        and (
            Path(str(x) + ENDING_RED).exists()
            or Path(str(x) + ENDING_RED + ENDING_GZ).exists()
        )
        for x in basepath
    )


def idat_basepaths(files):
    # If basenames is dir take all files in it
    if isinstance(files, list):
        _files = files
    elif Path(files).expanduser().is_dir():
        _files = Path(files).expanduser().rglob("*_[GR][re][nd].idat*")
    else:
        _files = [files]
    # Remove file endings
    _files = [
        Path(
            str(name)
            .replace(ENDING_GZ, "")
            .replace(ENDING_GRN, "")
            .replace(ENDING_RED, "")
        ).expanduser()
        for name in _files
    ]
    # Remove duplicates, keep ordering
    _files = list(dict.fromkeys(_files))
    return _files


def idat_paths_from_basenames(basenames):
    grn_idat_files = np.array(
        [Path(str(name) + ENDING_GRN) for name in basenames]
    )
    red_idat_files = np.array(
        [Path(str(name) + ENDING_RED) for name in basenames]
    )

    def check_and_fix(files):
        not_existing = [i for i, path in enumerate(files) if not path.exists()]
        files[not_existing] = [
            x.parent / (x.name + ENDING_GZ) for x in files[not_existing]
        ]
        not_found = next(
            (x for x in files[not_existing] if not x.exists()), None
        )
        return not_found

    not_found = check_and_fix(grn_idat_files)
    not_found = (
        check_and_fix(red_idat_files) if not_found is None else not_found
    )
    if not_found is not None:
        idat_file = str(not_found).replace(ENDING_GZ, "")
        raise FileNotFoundError(f"IDAT file not found: {idat_file}.")
    return grn_idat_files, red_idat_files


class RawData:
    def __init__(self, basenames):
        _basenames = idat_basepaths(basenames)

        self.probes = [path.name.replace(ENDING_GZ, "") for path in _basenames]

        grn_idat_files, red_idat_files = idat_paths_from_basenames(_basenames)

        grn_idat = [
            IdatParser(str(filepath), intensity_only=True)
            for filepath in grn_idat_files
        ]
        red_idat = [
            IdatParser(str(filepath), intensity_only=True)
            for filepath in red_idat_files
        ]

        array_types = [
            ArrayType.from_probe_count(len(idat.illumina_ids))
            for idat in grn_idat + red_idat
        ]

        if len(set(array_types)) != 1:
            raise ValueError("Array types must all be the same.")

        self.array_type = array_types[0]

        all_illumina_ids = [idat.illumina_ids for idat in grn_idat + red_idat]

        if all(
            [
                np.array_equal(all_illumina_ids[0], arr)
                for arr in all_illumina_ids
            ]
        ):
            self.ids = all_illumina_ids[0]
            self._grn = np.array([idat.probe_means for idat in grn_idat])
            self._red = np.array([idat.probe_means for idat in red_idat])
        else:
            self.ids = reduce(
                np.intersect1d, [idat.illumina_ids for idat in grn_idat]
            )
            self._grn = np.array(
                [
                    idat.probe_means[
                        np.isin(
                            idat.illumina_ids, self.ids, assume_unique=True
                        )
                    ]
                    for idat in grn_idat
                ]
            )
            self._red = np.array(
                [
                    idat.probe_means[
                        np.isin(
                            idat.illumina_ids, self.ids, assume_unique=True
                        )
                    ]
                    for idat in red_idat
                ]
            )
        self.manifest = Manifest(self.array_type)

        self._grn_df = None
        self._red_df = None
        self.methylated = None
        self.unmethylated = None

    @property
    def grn(self):
        if self._grn_df is None:
            self._grn_df = pd.DataFrame(
                self._grn.T, index=self.ids, columns=self.probes, dtype="int32"
            )
        return self._grn_df

    @property
    def red(self):
        if self._red_df is None:
            self._red_df = pd.DataFrame(
                self._red.T, index=self.ids, columns=self.probes, dtype="int32"
            )
        return self._red_df

    def __repr__(self):
        title = "RawData():"
        lines = [
            title + "\n" + "*" * len(title),
            f"array_type: {self.array_type}",
            f"manifest: {self.manifest.array_type}",
            f"probes:\n{self.probes}",
            f"ids:\n{self.ids}",
            f"_grn:\n{self._grn}",
            f"_red:\n{self._red}",
            f"grn:\n{self.grn}",
            f"red:\n{self.red}",
        ]
        return "\n\n".join(lines)


@memoize
def overlap_indices(left_arr, right_arr):
    """Compute the indices of overlapping elements between two arrays.

    This function finds the common elements (indices) between two input arrays
    and returns their positions in both arrays. It uses pandas Index objects
    and memoization for performance improvement.

    Example:
        >>> left_arr = ['a', 'b', 'c', 'd']
        >>> right_arr = ['b', 'c', 'e']
        >>> left_idx, right_idx = overlap_indices(left_arr, right_arr)
        >>> print(left_idx)
        [1 2]
        >>> print(right_idx)
        [0 1]
    """
    if not isinstance(left_arr, pd.Index):
        left_arr = pd.Index(left_arr)
    if not isinstance(right_arr, pd.Index):
        right_arr = pd.Index(right_arr)
    common_indices = left_arr.intersection(right_arr)
    left_index = left_arr.get_indexer(common_indices)
    right_index = right_arr.get_indexer(common_indices)
    return left_index, right_index


class MethylData:
    def __init__(self, data=None, file=None, prep="illumina"):
        if data is None and file is None:
            raise ValueError("'data' or 'file' must be given.")
        if data is None:
            data = RawData(file)
        elif not isinstance(data, RawData):
            raise ValueError("'data' is not of type 'RawData'.")
        self._grn = data._grn
        self._red = data._red
        self.array_type = data.array_type
        self.ids = data.ids
        self.manifest = data.manifest
        self.probes = data.probes
        self.data = data
        self._grn_df = None
        self._red_df = None
        self._methylated_df = None
        self._unmethylated_df = None
        if prep == "illumina":
            self.preprocess_illumina()
        elif prep == "swan":
            self.preprocess_swan()
        elif prep == "noob":
            self.preprocess_noob()
        elif prep == "raw":
            self.preprocess_raw_cached()
        else:
            raise ValueError(f"invalid 'prep' value {prep}")

    @property
    def grn(self):
        if self._grn_df is None:
            self._grn_df = pd.DataFrame(
                self._grn.T,
                index=self.ids,
                columns=self.probes,
                dtype="float32",
            )
        return self._grn_df

    @property
    def red(self):
        if self._red_df is None:
            self._red_df = pd.DataFrame(
                self._red.T,
                index=self.ids,
                columns=self.probes,
                dtype="float32",
            )
        return self._red_df

    @property
    def methylated(self):
        if self._methylated_df is None:
            self._methylated_df = pd.DataFrame(
                self.methyl.T,
                index=self.methyl_ilmnid,
                columns=self.probes,
                dtype="float32",
            )
            self._methylated_df.index.name = "IlmnID"
        return self._methylated_df

    @property
    def unmethylated(self):
        if self._unmethylated_df is None:
            self._unmethylated_df = pd.DataFrame(
                self.unmethyl.T,
                index=self.methyl_ilmnid,
                columns=self.probes,
                dtype="float32",
            )
            self._unmethylated_df.index.name = "IlmnID"
        return self._unmethylated_df

    def preprocess_illumina(self):
        self._methylated_df = None
        self._unmethylated_df = None
        self.illumina_control_normalization()
        self.illumina_bg_correction()
        self.preprocess_raw_cached()

    def illumina_control_normalization(self, reference=0):
        at_controls = self.manifest.control_address(["NORM_A", "NORM_T"])
        cg_controls = self.manifest.control_address(["NORM_C", "NORM_G"])

        grn_average = np.mean(
            self._grn[:, np.isin(self.ids, cg_controls, assume_unique=True)],
            axis=1,
        )
        red_average = np.mean(
            self._red[:, np.isin(self.ids, at_controls, assume_unique=True)],
            axis=1,
        )

        ref = (grn_average + red_average)[reference] / 2
        grn_factor = ref / grn_average
        red_factor = ref / red_average

        self._grn = grn_factor[:, np.newaxis] * self._grn
        self._red = red_factor[:, np.newaxis] * self._red

    def illumina_bg_correction(self):
        neg_controls = self.manifest.control_address("NEGATIVE")

        grn_bg = np.sort(
            self._grn[:, np.isin(self.ids, neg_controls, assume_unique=True)]
        )[:, 30]
        red_bg = np.sort(
            self._red[:, np.isin(self.ids, neg_controls, assume_unique=True)]
        )[:, 30]

        self._grn = np.maximum(self._grn - grn_bg[:, np.newaxis], 0)
        self._red = np.maximum(self._red - red_bg[:, np.newaxis], 0)

    def preprocess_raw(self):
        type_1 = self.manifest.probe_info(ProbeType.ONE)
        type_2 = self.manifest.probe_info(ProbeType.TWO)
        type_1_red = type_1[type_1.Color_Channel.values == Channel.RED.value]
        type_1_grn = type_1[type_1.Color_Channel.values == Channel.GRN.value]
        man_idx_np = np.sort(
            np.concatenate(
                [
                    type_1.IlmnID.index,
                    type_2.IlmnID.index,
                ]
            )
        )
        self.preprocess_raw_methylated(
            man_idx_np, type_1_grn, type_1_red, type_2
        )
        self.preprocess_raw_unmethylated(
            man_idx_np, type_1_grn, type_1_red, type_2
        )

    def preprocess_raw_methylated(self, man_idx_np, t1_grn, t1_red, t2):
        red = self.red
        grn = self.grn
        df = pd.DataFrame(
            np.nan,
            index=man_idx_np,
            columns=red.columns,
            dtype="float32",
        )
        df.loc[t1_red.index] = red.loc[t1_red["AddressB_ID"].values].values
        df.loc[t1_grn.index] = grn.loc[t1_grn["AddressB_ID"].values].values
        df.loc[t2.index] = grn.loc[t2["AddressA_ID"].values].values
        df["IlmnID"] = self.manifest.data_frame.IlmnID.values[man_idx_np]
        self._methylated_df = df.set_index("IlmnID")

    def preprocess_raw_unmethylated(self, man_idx_np, t1_grn, t1_red, t2):
        red = self.red
        grn = self.grn
        df = pd.DataFrame(
            np.nan,
            index=man_idx_np,
            columns=red.columns,
            dtype="float32",
        )
        df.loc[t1_red.index] = red.loc[t1_red["AddressA_ID"].values].values
        df.loc[t1_grn.index] = grn.loc[t1_grn["AddressA_ID"].values].values
        df.loc[t2.index] = red.loc[t2["AddressA_ID"].values].values
        df["IlmnID"] = self.manifest.data_frame.IlmnID.values[man_idx_np]
        self._unmethylated_df = df.set_index("IlmnID")

    @memoize
    def cached_indices(manifest, illumina_ids):
        type_1 = manifest.probe_info(ProbeType.ONE)
        type_2 = manifest.probe_info(ProbeType.TWO)
        type_1_red = type_1[type_1.Color_Channel.values == Channel.RED.value]
        type_1_grn = type_1[type_1.Color_Channel.values == Channel.GRN.value]
        idx = pd.Index(
            np.sort(
                np.concatenate(
                    [
                        type_1.IlmnID.index,
                        type_2.IlmnID.index,
                    ]
                )
            )
        )
        ids = pd.Index(illumina_ids)
        ids_1_red_a = ids.get_indexer(type_1_red["AddressA_ID"])
        ids_1_red_b = ids.get_indexer(type_1_red["AddressB_ID"])
        ids_1_grn_a = ids.get_indexer(type_1_grn["AddressA_ID"])
        ids_1_grn_b = ids.get_indexer(type_1_grn["AddressB_ID"])
        ids_2_____a = ids.get_indexer(type_2["AddressA_ID"])
        idx_1_red__ = idx.get_indexer(type_1_red.index)
        idx_1_grn__ = idx.get_indexer(type_1_grn.index)
        idx_2______ = idx.get_indexer(type_2.index)
        return {
            "idx": idx.values,
            "ids_1_red_a": ids_1_red_a,
            "ids_1_red_b": ids_1_red_b,
            "ids_1_grn_a": ids_1_grn_a,
            "ids_1_grn_b": ids_1_grn_b,
            "ids_2_____a": ids_2_____a,
            "idx_1_red__": idx_1_red__,
            "idx_1_grn__": idx_1_grn__,
            "idx_2______": idx_2______,
        }

    @memoize
    def cached_control_indices(manifest, illumina_ids):
        ids = pd.Index(illumina_ids)
        control_probes = manifest.control_data_frame
        control_probes = control_probes[
            control_probes.Address_ID.isin(ids)
        ].reset_index(drop=True)
        idx_cg = control_probes[
            control_probes.Control_Type.isin(["NORM_C", "NORM_G"])
        ].index.values
        idx_at = control_probes[
            control_probes.Control_Type.isin(["NORM_A", "NORM_T"])
        ].index.values
        ids_ctr = ids.get_indexer(control_probes["Address_ID"])
        return {
            "ids_ctr": ids_ctr,
            "idx_at": idx_at,
            "idx_cg": idx_cg,
        }

    def preprocess_raw_cached(self):
        ci = MethylData.cached_indices(self.manifest, self.ids)
        self.methyl = np.full((len(self.probes), len(ci["idx"])), np.nan)
        self.methyl[:, ci["idx_1_red__"]] = self._red[:, ci["ids_1_red_b"]]
        self.methyl[:, ci["idx_1_grn__"]] = self._grn[:, ci["ids_1_grn_b"]]
        self.methyl[:, ci["idx_2______"]] = self._grn[:, ci["ids_2_____a"]]
        self.unmethyl = np.full((len(self.probes), len(ci["idx"])), np.nan)
        self.unmethyl[:, ci["idx_1_red__"]] = self._red[:, ci["ids_1_red_a"]]
        self.unmethyl[:, ci["idx_1_grn__"]] = self._grn[:, ci["ids_1_grn_a"]]
        self.unmethyl[:, ci["idx_2______"]] = self._red[:, ci["ids_2_____a"]]
        self.methyl_index = ci["idx"]
        self.methyl_ilmnid = self.manifest.data_frame.IlmnID.values[ci["idx"]]

    def swan_bg_intensity(self):
        neg_controls = self.manifest.control_address("NEGATIVE")
        grn_med = np.median(
            self._grn[:, np.isin(self.ids, neg_controls, assume_unique=True)],
            axis=1,
        )
        red_med = np.median(
            self._red[:, np.isin(self.ids, neg_controls, assume_unique=True)],
            axis=1,
        )
        bg_intensity = np.mean([grn_med, red_med], axis=0)
        return bg_intensity

    @staticmethod
    def swan_indices(manifest, methyl_index):
        all_ncpgs = (
            manifest.data_frame[["Probe_Type", "N_CpG"]]
            .loc[methyl_index]
            .reset_index(drop=True)
        )
        subset_sizes = all_ncpgs.groupby(
            ["Probe_Type", "N_CpG"], dropna=False
        ).size()
        subset_size = min(subset_sizes.loc[(slice(None), slice(1, 3))])
        all_indices = {}
        random_subset_indices = {}
        for probe_type in [ProbeType.ONE, ProbeType.TWO]:
            all_ncpts_type = all_ncpgs[all_ncpgs.Probe_Type == probe_type]
            all_indices[probe_type] = all_ncpts_type.index.values
            all_ncpts_type.reset_index(drop=True, inplace=True)
            indices = []
            for ncpgs in range(1, 4):
                ids = all_ncpts_type.index[all_ncpts_type.N_CpG == ncpgs]
                ids_subset = np.random.permutation(ids)[:subset_size]
                indices.append(ids_subset)
            random_subset_indices[probe_type] = np.sort(
                np.concatenate(indices)
            )
        return all_indices, random_subset_indices

    def preprocess_swan(self):
        self._methylated_df = None
        self._unmethylated_df = None
        self.preprocess_raw_cached()
        bg_intensity = self.swan_bg_intensity()
        all_indices, random_subset_indices = MethylData.swan_indices(
            self.manifest, self.methyl_index
        )
        self.methyl = MethylData._preprocess_swan(
            self.methyl, bg_intensity, all_indices, random_subset_indices
        )
        self.unmethyl = MethylData._preprocess_swan(
            self.unmethyl, bg_intensity, all_indices, random_subset_indices
        )

    @staticmethod
    def _preprocess_swan(
        intensity, bg_intensity, all_indices, random_subset_indices
    ):
        from scipy.stats import rankdata

        random_subset_one = all_indices[ProbeType.ONE][
            random_subset_indices[ProbeType.ONE]
        ]
        random_subset_two = all_indices[ProbeType.TWO][
            random_subset_indices[ProbeType.TWO]
        ]
        sorted_subset_intensity = (
            np.sort(intensity[:, random_subset_one], axis=1)
            + np.sort(intensity[:, random_subset_two], axis=1)
        ) / 2
        swan = np.full(intensity.shape, np.nan)
        for i in range(len(intensity)):
            for probe_type in [ProbeType.ONE, ProbeType.TWO]:
                curr_intensity = intensity[i, all_indices[probe_type]]
                x = rankdata(curr_intensity) / len(curr_intensity)
                xp = np.sort(x[random_subset_indices[probe_type]])
                fp = sorted_subset_intensity[i, :]
                interp = np.interp(x=x, xp=xp, fp=fp)
                intensity_min = np.min(
                    curr_intensity[random_subset_indices[probe_type]]
                )
                intensity_max = np.max(
                    curr_intensity[random_subset_indices[probe_type]]
                )
                interp[x > np.max(xp)] += (
                    curr_intensity[x > np.max(xp)] - intensity_max
                )
                interp[x < np.min(xp)] += (
                    curr_intensity[x < np.min(xp)] - intensity_min
                )
                interp[interp <= 0] = bg_intensity[i]
                swan[i, all_indices[probe_type]] = interp
        return swan

    def preprocess_noob(self, offset=15, dye_method="single"):
        self._methylated_df = None
        self._unmethylated_df = None
        self.preprocess_raw_cached()

        ci = MethylData.cached_indices(self.manifest, self.ids)

        grn_oob = np.concatenate(
            [self._grn[:, ci["ids_1_red_a"]], self._grn[:, ci["ids_1_red_b"]]],
            axis=1,
        )
        red_oob = np.concatenate(
            [self._red[:, ci["ids_1_grn_a"]], self._red[:, ci["ids_1_grn_b"]]],
            axis=1,
        )

        methyl = self.methyl
        unmethyl = self.unmethyl
        methyl[methyl <= 0] = 1
        unmethyl[unmethyl <= 0] = 1

        grn_m = methyl[:, ci["idx_1_grn__"]]
        grn_u = unmethyl[:, ci["idx_1_grn__"]]
        grn_2 = methyl[:, ci["idx_2______"]]

        xf_grn = np.concatenate([grn_m, grn_u, grn_2], axis=1)
        xs_grn = normexp_get_xs(xf_grn, controls=grn_oob, offset=offset)

        cumsum = np.cumsum([0, grn_m.shape[1], grn_u.shape[1], grn_2.shape[1]])
        slice_grn_m = slice(cumsum[0], cumsum[1])
        slice_grn_u = slice(cumsum[1], cumsum[2])
        slice_grn_2 = slice(cumsum[2], cumsum[3])

        red_m = methyl[:, ci["idx_1_red__"]]
        red_u = unmethyl[:, ci["idx_1_red__"]]
        red_2 = unmethyl[:, ci["idx_2______"]]

        xf_red = np.concatenate([red_m, red_u, red_2], axis=1)
        xs_red = normexp_get_xs(xf_red, controls=red_oob, offset=offset)

        cumsum = np.cumsum([0, red_m.shape[1], red_u.shape[1], red_2.shape[1]])
        slice_red_m = slice(cumsum[0], cumsum[1])
        slice_red_u = slice(cumsum[1], cumsum[2])
        slice_red_2 = slice(cumsum[2], cumsum[3])

        methyl[:, ci["idx_1_grn__"]] = xs_grn["xs"][:, slice_grn_m]
        unmethyl[:, ci["idx_1_grn__"]] = xs_grn["xs"][:, slice_grn_u]

        methyl[:, ci["idx_1_red__"]] = xs_red["xs"][:, slice_red_m]
        unmethyl[:, ci["idx_1_red__"]] = xs_red["xs"][:, slice_red_u]

        methyl[:, ci["idx_2______"]] = xs_grn["xs"][:, slice_grn_2]
        unmethyl[:, ci["idx_2______"]] = xs_red["xs"][:, slice_red_2]

        # Dye correction

        cci = MethylData.cached_control_indices(self.manifest, self.ids)

        grn_control = self._grn[:, cci["ids_ctr"]]
        red_control = self._red[:, cci["ids_ctr"]]

        xcs_grn = normexp_get_xs(
            grn_control, param=xs_grn["param"], offset=offset
        )
        xcs_red = normexp_get_xs(
            red_control, param=xs_red["param"], offset=offset
        )

        grn_avg = np.mean(xcs_grn["xs"][:, cci["idx_cg"]], axis=1)
        red_avg = np.mean(xcs_red["xs"][:, cci["idx_at"]], axis=1)

        red_grn_ratio = red_avg / grn_avg

        if dye_method == "single":
            red_factor = 1 / red_grn_ratio
            grn_factor = np.array([1, 1])
        elif dye_method == "reference":
            ref_idx = np.argmin(np.abs(red_grn_ratio - 1))
            ref = (grn_avg + red_avg)[ref_idx] / 2
            if np.isnan(ref):
                raise ValueError(
                    "'ref_idx' refers to an array that is not present"
                )
            grn_factor = ref / grn_avg
            red_factor = ref / red_avg
        else:
            raise ValueError("dye_method must be 'single' or 'reference'")

        red_factor = red_factor.reshape(-1, 1)
        methyl[:, ci["idx_1_red__"]] *= red_factor
        unmethyl[:, ci["idx_1_red__"]] *= red_factor
        unmethyl[:, ci["idx_2______"]] *= red_factor

        if dye_method == "reference":
            grn_factor = grn_factor.reshape(-1, 1)
            methyl[:, ci["idx_1_grn__"]] *= grn_factor
            unmethyl[:, ci["idx_1_grn__"]] *= grn_factor
            methyl[:, ci["idx_2______"]] *= grn_factor

        self.methyl = methyl
        self.unmethyl = unmethyl

    @property
    def beta(self):
        beta = self.get_beta(self.methyl, self.unmethyl)
        return pd.DataFrame(
            beta.T, columns=self.probes, index=self.methyl_ilmnid
        )

    def converted_beta(self, cpgs=None, fill=NEUTRAL_BETA):
        if cpgs is None:
            cpgs = self.manifest.methylation_probes
        beta = self.get_beta(self.methyl, self.unmethyl)
        converted = np.full((len(self.probes), len(cpgs)), fill)
        left_idx, right_idx = overlap_indices(cpgs, self.methyl_ilmnid)
        converted[:, left_idx] = beta[:, right_idx]
        converted[np.isnan(converted)] = fill
        return pd.DataFrame(converted.T, columns=self.probes, index=cpgs)

    @staticmethod
    def get_beta(
        methylated, unmethylated, offset=0, beta_threshold=0, min_zero=True
    ):
        assert offset >= 0, "offset must be non-negative"
        assert (
            0 <= beta_threshold <= 0.5
        ), "beta_threshold must be between 0 and 0.5"

        if min_zero:
            methylated = np.maximum(methylated, 0)
            unmethylated = np.maximum(unmethylated, 0)
        # Ignore division by zero
        with np.errstate(divide="ignore", invalid="ignore"):
            beta = methylated / (methylated + unmethylated + offset)

        if beta_threshold > 0:
            beta = np.minimum(
                np.maximum(beta, beta_threshold), 1 - beta_threshold
            )

        return beta

    def __repr__(self):
        title = "MethylData():"
        lines = [
            title + "\n" + "*" * len(title),
            f"array_type: {self.array_type}",
            f"manifest: {self.manifest.array_type}",
            f"probes:\n{self.probes}",
            f"_grn:\n{self._grn}",
            f"_red:\n{self._red}",
            f"grn:\n{self.grn}",
            f"red:\n{self.red}",
            f"methylated:\n{self.methylated}",
            f"unmethylated:\n{self.unmethylated}",
        ]
        if hasattr(self, "intensity"):
            lines.append(f"intensity:\n{self.intensity}")
        return "\n\n".join(lines)


@memoize
class ReferenceMethylData:
    def __init__(self, files=None, prep="illumina"):
        idat_files = idat_basepaths(files)
        reference_files = collections.defaultdict(list)
        self._methyl = {}
        for idat_file in tqdm(
            idat_files, desc="Categorizing reference IDAT files"
        ):
            raw_data = RawData(idat_file)
            array_type = raw_data.array_type
            reference_files[array_type].append(idat_file)
        for array_type, file_list in tqdm(
            reference_files.items(), desc="Processing reference IDAT files"
        ):
            raw_data = RawData(file_list)
            self._methyl[array_type] = MethylData(raw_data, prep=prep)

    def __getitem__(self, array_type):
        if array_type not in self._methyl:
            raise ValueError(
                f"No reference files found for array type {array_type.value}"
            )
        return self._methyl[array_type]

    # def __reduce__(self):
    # return (self.__class__, (self.files, self.prep))
