# Copyright (C) 2024 Andrea Raffo <andrea.raffo@ibv.uio.no>
#
# SPDX-License-Identifier: MIT

import datetime
import functools
import json
import pathlib
import shutil
from importlib.metadata import version
from typing import Any, Dict, List, Optional, Sequence, Union

import h5py
import hictkpy
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
import pandas as pd
import seaborn as sns
from matplotlib.ticker import EngFormatter, ScalarFormatter

from stripepy.utils.stripe import Stripe

no_frills_in_images = False  # TODO safe removal of no_frills_in_images

fruit_punch = sns.blend_palette(["white", "red"], as_cmap=True)


class Result(object):
    """
    A class used to represent the results generated by stripepy call.

    Attributes
    ----------
    chrom: str
        name of the chromosomes to which the Result instance belongs to
    empty: bool
        check whether any stripe has been registered with the Result instance
    roi: Optional[Dict[str, List[int]]]
        the region of interest associated with the Result instance
    min_persistence: float
        the minimum persistence used during computation
    """

    def __init__(self, chrom: str):
        """
        Parameters
        ----------

        chrom: str
            chromosome name
        """
        self._chrom = chrom
        self._roi = None
        self._min_persistence = None

        self._ut_all_minimum_points = None
        self._ut_all_maximum_points = None
        self._ut_persistence_of_all_minimum_points = None
        self._ut_persistence_of_all_maximum_points = None

        self._lt_all_minimum_points = None
        self._lt_all_maximum_points = None
        self._lt_persistence_of_all_minimum_points = None
        self._lt_persistence_of_all_maximum_points = None

        self._ut_persistent_minimum_points = None
        self._ut_persistent_maximum_points = None
        self._ut_persistence_of_minimum_points = None
        self._ut_persistence_of_maximum_points = None
        self._ut_pseudodistribution = None

        self._lt_persistent_minimum_points = None
        self._lt_persistent_maximum_points = None
        self._lt_persistence_of_minimum_points = None
        self._lt_persistence_of_maximum_points = None
        self._lt_pseudodistribution = None

        self._ut_stripes = None
        self._lt_stripes = None

    @property
    def _valid_attributes(self) -> List[str]:
        """
        Get the list of valid attributes
        """
        return [a.removeprefix("_lt_") for a in dir(self) if a.startswith("_lt_")]

    @property
    def empty(self) -> bool:
        return self._lt_stripes is None and self._ut_stripes is None

    @property
    def chrom(self) -> str:
        return self._chrom

    @property
    def roi(self) -> Optional[Dict[str, List[int]]]:
        return self._roi

    @property
    def min_persistence(self) -> float:
        if self._min_persistence is None:
            raise RuntimeError('Attribute "min_persistence" is not set')

        return self._min_persistence

    def get(self, name: str, location: str) -> Union[List[Stripe], npt.NDArray[int], npt.NDArray[float]]:
        """
        Get the value associated with the given attribute name and location.

        Parameters
        ----------
        name: str
            name of the attribute to be fetched
        location: str
            location of the attribute to be fetched. Should be "LT" or "UT"

        Returns
        -------
        attribute
            the value associated with the given name and location.
        """
        if location not in {"LT", "UT"}:
            raise ValueError("Location should be UT or LT")

        attr_name = f"_{location.lower()}_{name}"
        if not hasattr(self, attr_name):
            raise AttributeError(
                f"No attribute named \"{name}\". Valid attributes are: {', '.join(self._valid_attributes)}"
            )

        attr = getattr(self, attr_name)
        if name == "stripes" and attr is None:
            return []

        if attr is None:
            raise RuntimeError(f'Attribute "{name}" for "{location}" is not set')

        return attr

    def get_stripes_descriptor(self, descriptor: str, location: str) -> Union[npt.NDArray[float], npt.NDArray[int]]:
        """
        Get the stripe descriptor for the given location.

        Parameters
        ----------
        descriptor: str
            name of the descriptor to be fetched
        location: str
            location of the attribute to be fetched. Should be "LT" or "UT"

        Returns
        -------
        descriptor
            the value associated with the given descriptor and location.
        """
        if location not in {"LT", "UT"}:
            raise ValueError("Location should be UT or LT")

        if not hasattr(Stripe, descriptor):
            raise AttributeError(f'Stripe instance does not have an attribute named "{descriptor}"')

        stripes = self.get("stripes", location)

        if descriptor in {"seed", "left_bound", "right_bound", "top_bound", "bottom_bound"}:
            dtype = int
        else:
            dtype = float

        return np.array([getattr(stripe, descriptor) for stripe in stripes], dtype=dtype)

    def get_stripe_geo_descriptors(self, location: str) -> pd.DataFrame:
        """
        Fetch all geometric descriptors at once.

        Parameters
        ----------
        location: str
            location of the attribute to be fetched. Should be "LT" or "UT"

        Returns
        -------
        descriptors
            the table with the geometric descriptors associated with the Result instance
        """
        descriptors = [
            "seed",
            "top_persistence",
            "left_bound",
            "right_bound",
            "top_bound",
            "bottom_bound",
        ]

        return pd.DataFrame(
            {descriptor: self.get_stripes_descriptor(descriptor, location) for descriptor in descriptors}
        )

    def get_stripe_bio_descriptors(self, location: str) -> pd.DataFrame:
        """
        Fetch all biological descriptors at once.

        Parameters
        ----------
        location: str
            location of the attribute to be fetched. Should be "LT" or "UT"

        Returns
        -------
        descriptors
            the table with the biological descriptors associated with the Result instance
        """
        descriptors = [
            "inner_mean",
            "outer_mean",
            "rel_change",
            "inner_std",
        ]

        return pd.DataFrame(
            {descriptor: self.get_stripes_descriptor(descriptor, location) for descriptor in descriptors}
        )

    def set_roi(self, coords: Dict[str, List[int]]):
        """
        Set the region of interest (RoI) for the current Result instance.

        Parameters
        ----------
        coords: Dict[str, List[int]]
            a dictionary with the coordinates of the region of interest.
            The dictionary should contain two keys: "genomic" and "matrix".
            The value associated with the "genomic" key should be a list of 4 integers
            representing the genomic coordinates of the region of interest.
            The value associated with the "matrix" key should be a list of 4 integers
            representing the matrix coordinates of the region of interest.
        """
        if self._roi is not None:
            raise RuntimeError("roi has already been set")

        self._roi = coords

    def set_min_persistence(self, min_persistence: float):
        """
        Set the minimum persistence used during computation.

        Parameters
        ----------
        min_persistence: float
        """
        if self._min_persistence is not None:
            raise RuntimeError("min_persistence has already been set")

        self._min_persistence = min_persistence

    def set(self, name: str, data: Union[Sequence[int], Sequence[float], Sequence[Stripe]], location: str):
        """
        Set the attribute corresponding to the given attribute name and location.

        Parameters
        ----------
        name: str
           name of the attribute to be set.
           Supported attributes are:
              * all_minimum_points
              * all_maximum_points
              * persistence_of_all_minimum_points
              * persistence_of_all_maximum_points
              * persistent_minimum_points
              * persistent_maximum_points
              * persistence_of_minimum_points
              * persistence_of_maximum_points
              * pseudodistribution
              * stripes

        data:
            data to be registered with the Result instance
        location: str
            location of the attribute to be registered. Should be "LT" or "UT"
        """
        if location not in {"LT", "UT"}:
            raise ValueError("Location should be UT or LT")

        attr_name = f"_{location.lower()}_{name}"
        if not hasattr(self, attr_name):
            raise AttributeError(
                f"No attribute named \"{name}\". Valid attributes are: {', '.join(self._valid_attributes)}"
            )

        if getattr(self, attr_name) is not None:
            raise RuntimeError(f'Attribute "{name}" for {location} has already been set')

        setattr(self, attr_name, np.array(data))


class ResultFile(object):
    """
    A class used to read and write StripePy results to a HDF5 file.

    Attributes
    ----------
    path: pathlib.Path
        the path to the opened file
    assembly: str
        the name of the reference genome assembly used to generate the file
    resolution: int
        the resolution of the Hi-C matrix used to generate the file
    creation_date: datetime.datetime
        the file creation date
    format: str
        the file format string
    format_url: str
        the URL where the file format is documented
    format_version: int
        the format version of the file currently opened
    generated_by: str
        the name of the tool used to generate the opened file
    normalization: Optional[str]
        the name of the normalization used to generate the data stored in the given file
    chromosomes: Dict[str, int]
        the chromosomes associated with the opened file

    """

    def __init__(self, path: pathlib.Path, mode: str = "r"):
        if mode not in ["r", "w"]:
            raise ValueError('mode should be "r" or "w"')

        self._path = path
        self._mode = mode
        self._chroms = None

        self._h5 = h5py.File(self._path, self._mode)

        if self._mode == "r":
            self._validate(self._h5)
            self._chroms = {
                chrom.decode("utf-8"): size for chrom, size in zip(self._h5["/chroms/name"], self._h5["/chroms/length"])
            }

        self._version = self._h5.attrs.get("format-version", 1)
        self._attrs = dict(self._h5.attrs)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._h5.close()

    @property
    def path(self) -> pathlib.Path:
        return self._path

    @functools.cached_property
    def assembly(self) -> str:
        return self._h5.attrs["assembly"]

    @functools.cached_property
    def resolution(self) -> int:
        return int(self._h5.attrs["bin-size"])

    @functools.cached_property
    def creation_date(self) -> datetime.datetime:
        return datetime.datetime.fromisoformat(self._h5.attrs["creation-date"])

    @functools.cached_property
    def format(self) -> str:
        return self._h5.attrs["format"]

    @functools.cached_property
    def format_url(self) -> str:
        return self._h5.attrs["format-url"]

    @functools.cached_property
    def format_version(self) -> int:
        return self._h5.attrs["format-version"]

    @functools.cached_property
    def generated_by(self) -> str:
        return self._h5.attrs["generated-by"]

    @functools.cached_property
    def metadata(self) -> Dict[str, Any]:
        return json.loads(self._h5.attrs["metadata"])

    @functools.cached_property
    def normalization(self) -> Optional[str]:
        norm = self._h5.attrs["normalization"]
        if norm == "NONE":
            return None
        return norm

    @property
    def chromosomes(self) -> Dict[str, int]:
        return self._chroms

    @staticmethod
    def _validate(h5: h5py.File):
        """
        Perform a basic sanity check on the metadata of the current file
        """
        format = h5.attrs.get("format")  # noqa
        format_version = h5.attrs.get("format-version")
        try:
            if format is None:
                raise RuntimeError('attribute "format" is missing')

            if format_version is None:
                raise RuntimeError('attribute "format-version" is missing')

            if format != "HDF5::StripePy":
                raise RuntimeError(f'unrecognized file format: expected "HDF5::StripePy", found "{format}"')

            if format_version != 1:
                raise RuntimeError(
                    f'unsupported file format version "{format_version}". At present only version 1 is supported'
                )
        except RuntimeError as e:
            raise RuntimeError(
                f'failed to validate input file "{h5.filename}": {e}: file is corrupt or was not generated by StripePy.'
            )

    def get_min_persistence(self, chrom: str) -> float:
        """
        Get the minimum persistence associated with the given chromosome.

        Parameters
        ----------
        chrom: str
            chromosome name

        Returns
        -------
            the minimum persistence
        """
        if chrom not in self._chroms:
            raise KeyError(f'File "{self.path}" does not have data for chromosome "{chrom}"')

        return float(self._h5[f"/{chrom}/global-pseudo-distributions"].attrs["min_persistence_used"])

    def get(self, chrom: str, field: str, location: str) -> pd.DataFrame:
        """
        Get the data associated with the given chromosome, field, and location.

        Parameters
        ----------
        chrom: str
            chromosome name
        field: str
            name of the field to be fetched.
            Supported names:
                * pseudodistribution
                * all_minimum_points
                * persistence_of_all_minimum_points
                * all_maximum_points
                * persistence_of_all_maximum_points
                * geo_descriptors
                * bio_descriptors
        location: str
            location of the attribute to be registered. Should be "LT" or "UT"

        Returns
        -------
            the data associated with the given chromosome, field, and location
        """
        if chrom not in self._chroms:
            raise KeyError(f'File "{self.path}" does not have data for chromosome "{chrom}"')

        if location not in {"LT", "UT"}:
            raise ValueError("Location should be UT or LT")

        mappings = {
            "pseudodistribution": f"/{chrom}/global-pseudo-distributions/{location}/pseudo-distribution",
            "all_minimum_points": f"/{chrom}/global-pseudo-distributions/{location}/minima_pts_and_persistence",
            "persistence_of_all_minimum_points": f"/{chrom}/global-pseudo-distributions/{location}/minima_pts_and_persistence",
            "all_maximum_points": f"/{chrom}/global-pseudo-distributions/{location}/maxima_pts_and_persistence",
            "persistence_of_all_maximum_points": f"/{chrom}/global-pseudo-distributions/{location}/maxima_pts_and_persistence",
            "geo_descriptors": f"/{chrom}/stripes/{location}/geo-descriptors",
            "bio_descriptors": f"/{chrom}/stripes/{location}/bio-descriptors",
        }

        if field not in mappings:
            raise KeyError(f"Unknown field \"{field}\". Valid fields are {', '.join(mappings.keys())}")

        path = mappings[field]

        if field not in {"geo_descriptors", "bio_descriptors"}:
            data = self._h5[path][:]
            if field.startswith("persistence"):
                data = data[1, :]
            elif field.endswith("points"):
                data = data[0, :]
            return pd.DataFrame({field: data})

        df = pd.DataFrame(data=self._h5[path], columns=self._h5[path].attrs["col_names"])

        if field == "geo_descriptors":
            df = df.rename(
                columns={
                    "seed persistence": "top_persistence",
                    "L-boundary": "left_bound",
                    "R_boundary": "right_bound",
                    "U-boundary": "top_bound",
                    "D-boundary": "bottom_bound",
                }
            )
            for col in ("seed", "left_bound", "right_bound", "top_bound", "bottom_bound"):
                df[col] = df[col].astype(int)
            return df

        return df.rename(
            columns={
                "inner mean": "inner_mean",
                "outer mean": "outer_mean",
                "relative change": "rel_change",
                "standard deviation": "inner_std",
            }
        )

    def init_file(self, matrix_file: hictkpy.File, normalization: Optional[str], metadata: Dict[str, Any]):
        """
        Initialize the current file.

        This method must be called when opening a ResultFile for writing before adding data to the file.

        Parameters
        ----------
        matrix_file: hictkpy.File
            handle of the file that is used to compute the results stored in this file
        normalization: Optional[str]
            name of the normalization used to compute the results stored in this file
        metadata: Dict[str, Any]
            dictionary with the metadata to be associated with this file.
            The dictionary should contain values that can be encoded as a JSON string
        """
        if normalization is None:
            normalization = "NONE"

        self._chroms = matrix_file.chromosomes(include_ALL=False)

        self._h5.attrs["assembly"] = matrix_file.attributes().get("assembly", "unknown")
        self._h5.attrs["bin-size"] = matrix_file.resolution()
        self._h5.attrs["creation-date"] = datetime.datetime.now().isoformat()
        self._h5.attrs["format"] = "HDF5::StripePy"
        self._h5.attrs["format-url"] = "https://github.com/paulsengroup/StripePy"
        self._h5.attrs["format-version"] = 1
        self._h5.attrs["generated-by"] = f"StripePy v{version('stripepy-hic')}"
        self._h5.attrs["metadata"] = json.dumps(metadata, indent=2)
        self._h5.attrs["normalization"] = normalization

        chroms = matrix_file.chromosomes(include_ALL=False)
        self._h5.create_group("/chroms")
        self._h5.create_dataset("/chroms/name", data=list(chroms.keys()))
        self._h5.create_dataset("/chroms/length", data=list(chroms.values()))

    def write_descriptors(self, result: Result):
        """
        Read the descriptors from the given Result object and write them to the opened file.

        Parameters
        ----------
        result: Result
            results to be added to the opened file
        """
        grp = self._h5.create_group(f"/{result.chrom}/global-pseudo-distributions/")

        grp.attrs.create("min_persistence_used", result.min_persistence)

        for location in ["UT", "LT"]:
            grp = self._h5.create_group(f"/{result.chrom}/global-pseudo-distributions/{location}")
            grp.create_dataset(
                "pseudo-distribution",
                data=result.get("pseudodistribution", location),
                compression="gzip",
                compression_opts=4,
                shuffle=True,
            )
            grp.create_dataset(
                "minima_pts_and_persistence",
                data=np.array(
                    [
                        result.get("all_minimum_points", location),
                        result.get("persistence_of_all_minimum_points", location),
                    ]
                ),
                compression="gzip",
                compression_opts=4,
                shuffle=True,
            )
            grp.create_dataset(
                "maxima_pts_and_persistence",
                data=np.array(
                    [
                        result.get("all_maximum_points", location),
                        result.get("persistence_of_all_maximum_points", location),
                    ]
                ),
                compression="gzip",
                compression_opts=4,
                shuffle=True,
            )

            grp = self._h5.create_group(f"/{result.chrom}/stripes/{location}")
            descriptors = result.get_stripe_geo_descriptors(location)
            dset = grp.create_dataset(
                "geo-descriptors",
                data=descriptors.to_numpy(),
                compression="gzip",
                compression_opts=4,
                shuffle=True,
            )
            dset.attrs.create("col_names", data=descriptors.columns.tolist())

            descriptors = result.get_stripe_bio_descriptors(location)
            dset = grp.create_dataset(
                "bio-descriptors",
                data=descriptors.to_numpy(),
                compression="gzip",
                compression_opts=4,
                shuffle=True,
            )
            dset.attrs.create("col_names", data=descriptors.columns.tolist())


class ANSI:
    RED = "\033[31m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    BLUE = "\033[34m"
    MAGENTA = "\033[35m"
    CYAN = "\033[36m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"
    ENDC = "\033[0m"


def remove_and_create_folder(path: pathlib.Path, force: bool):
    path = pathlib.Path(path)

    # Deleting folders:
    if path.exists():
        if not force:
            raise RuntimeError(f"output folder {path} already exists. Pass --force to overwrite it.")
        else:
            shutil.rmtree(path)

    # Create the folder:
    path.mkdir(parents=True)


def create_folders_for_plots(path: pathlib.Path):
    path = pathlib.Path(path)

    folders4plots = [
        path,
        path / "1_preprocessing",
        path / "2_TDA",
        path / "3_shape_analysis",
        path / "4_biological_analysis",
        path / "3_shape_analysis" / "local_pseudodistributions",
    ]

    # Creating folders:
    for folder2create in folders4plots:
        folder2create.mkdir(parents=True)

    return folders4plots


def format_ticks(ax, x=True, y=True, rotate=True):
    """
    Function taken from https://cooltools.readthedocs.io/en/latest/notebooks/viz.html
    :param ax:      an Axes object.
    :param x:       if True, it formats labels in engineering notation for the x-axis
    :param y:       if True, it formats labels in engineering notation for the y-axis
    :param y:       if True, it formats labels in engineering notation for the y-axis
    :param rotate:  if True, it rotates labels in the x-axis by 45 degrees.
    :return:        -
    """

    if x:
        ax.xaxis.set_major_formatter(EngFormatter("b"))
        ax.xaxis.tick_bottom()
    else:
        ax.xaxis.set_major_formatter(ScalarFormatter())
        ax.xaxis.tick_bottom()

    if y:
        ax.yaxis.set_major_formatter(EngFormatter("b"))
    else:
        ax.yaxis.set_major_formatter(ScalarFormatter())

    if rotate:
        ax.tick_params(axis="x", rotation=45)


def pseudodistrib(
    pseudo_distrib,
    IoI,
    resolution,
    coords2scatter=None,
    colors=None,
    output_folder=None,
    file_name=None,
    title=None,
    display=False,
):
    # TODO remove
    """
    :param pseudo_distrib:          1D ndarray representing a scalar function sampled over a uniform mesh
    :param  IoI:                    refers to the Interval of Interest [IoI[0], IoI[1]] (e.g., in genomic coordinates)
                                    where the scalar function was sampled on; see also plot_in_bp
    :param resolution:              resolution of the Hi-C matrix
    :param coords2scatter:          list of lists of genomic coordinates; each list of genomic coordinates determines
                                    a point cloud as follows: for genomic coordinate x, we sample the value of
                                    pseudo_distrib at x; each point cloud is scatterplotted with a specific color,
                                    potentially given in input, read below;
                                    if set to None, nothing happens
    :param colors:                  list of colors, one color per list of genomic coordinates (see coords2scatter);
                                    if set to None, use red
    :param output_folder:           path to folder where to save the image
    :param file_name:               name of the file to be created
    :param title:                   title to give to the image
    :param display:                 if False, it does not display the plot
    :return:                        -
    """

    fig, ax = plt.subplots(1, 1)
    ax.plot(
        range(IoI[0], IoI[1], resolution),
        pseudo_distrib[int(IoI[0] / resolution) : int(IoI[1] / resolution)],
        color="red",
        linewidth=0.5,
        linestyle="solid",
    )
    if coords2scatter is not None:
        for n, cur_coords2scatter in enumerate(coords2scatter):
            if colors is not None:
                color = colors[n]
            else:
                color = "red"
            ax.plot(
                [cur_coord2scatter * resolution for cur_coord2scatter in cur_coords2scatter],
                pseudo_distrib[cur_coords2scatter],
                marker=".",
                linestyle="",
                markersize=6 * 1.5,
                color=color,
            )
    ax.xaxis.set_major_formatter(EngFormatter("b"))

    if no_frills_in_images is False:
        if title is not None:
            fig.suptitle(title)

    plt.ylim((0, 1))
    ax.set_xlabel("genomic coordinates (bp)")
    ax.set_ylabel("pseudo-distribution")
    fig.tight_layout()
    ax.grid(True)
    # plt.axis('scaled')

    if output_folder is not None and file_name is not None:
        plt.savefig(f"{output_folder}/{file_name}", bbox_inches="tight")

    if display:
        plt.show()
    plt.close()


def pseudodistrib_and_HIoIs(
    pseudo_distrib, IoIs, resolution, colors=None, output_folder=None, file_name=None, title=None, display=False
):
    # TODO remove
    """
    :param pseudo_distrib:          1D ndarray representing a scalar function sampled over a uniform mesh
    :param  IoIs:                   list of lists, where the innermost lists are pairs of coordinates; the first pair
                                    refers to the Interval of Interest [IoI[0], IoI[1]] (e.g., in genomic coordinates)
                                    where pseudo-distribution is plotted; the remaining pairs define sub-regions to be
                                    plotted in (potentially) different colors
                                    where the scalar function was sampled on; see also plot_in_bp
    :param resolution:              resolution of the Hi-C matrix
    :param colors:                  list of colors, one color per pair of genomic coordinates (see IoIs);
                                    if set to None, use red
    :param output_folder:           path to folder where to save the image
    :param file_name:               name of the file to be created
    :param title:                   title to give to the image
    :param display:                 if False, it does not display the plot
    :return:                        -
    """
    fig, ax = plt.subplots(1, 1)
    for IoI, color in zip(IoIs, colors):
        ax.plot(
            range(IoI[0], IoI[1], resolution),
            pseudo_distrib[int(IoI[0] / resolution) : int(IoI[1] / resolution)],
            color=color,
            linewidth=0.5,
            linestyle="solid",
        )
    ax.xaxis.set_major_formatter(EngFormatter("b"))
    if no_frills_in_images is False:
        if title is not None:
            fig.suptitle(title)
    ax.set_xlabel("genomic coordinates (bp)")
    ax.set_ylabel("pseudo-distribution")
    fig.tight_layout()
    ax.grid(True)
    # plt.axis('scaled')
    if output_folder is not None and file_name is not None:
        plt.savefig(f"{output_folder}/{file_name}", bbox_inches="tight")
    if display:
        plt.show()
    plt.close()


def HiC_and_sites(
    I,
    sites,
    RoI,
    resolution,
    where=None,
    plot_in_bp=False,
    output_folder=None,
    file_name=None,
    title=None,
    display=False,
):  # TODO remove
    """
    :param I:                  Hi-C matrix to be plotted as image and saved
    :param sites:              list of locations of interest
    :param  RoI:               refers to the region of interest [RoI[0], RoI[1]]x[RoI[2], RoI[3]]
                               (e.g., in genomic coordinates); see also plot_in_bp
    :param where:              if "lower" (resp. "upper"), then it plots sites only on the lower (resp. upper) part of
                               the Hi-C matrix; otherwise, it plots sites spanning the whole Hi-C matrix
    :param plot_in_bp:         if True, labels are set in base pairs
    :param output_folder:      path to folder where to save the image
    :param file_name:          name of the file to be created
    :param title:              title to give to the image
    :param display:            if False, it does not display the plot
    :return:                   -
    """

    fig, ax = plt.subplots(1, 1)
    img = ax.matshow(I, vmax=np.amax(I), extent=(RoI[0], RoI[1], RoI[3], RoI[2]), cmap=fruit_punch)
    if plot_in_bp:
        format_ticks(ax)

    for site in [site * resolution for site in sites]:
        if where == "lower":
            ax.plot(
                [site, site], [site, RoI[1] - 1 * resolution], color=(0.0, 0.0, 1.0), linestyle="dashed", linewidth=1
            )
        elif where == "upper":
            ax.plot(
                [site, site], [RoI[0] + 1 * resolution, site], color=(0.0, 0.0, 1.0), linestyle="dashed", linewidth=1
            )
        else:
            ax.plot(
                [site, site],
                [RoI[0] + 1 * resolution, RoI[2] - 1 * resolution],
                color=(0.0, 0.0, 1.0),
                linestyle="dashed",
                linewidth=1,
            )

    if no_frills_in_images is False:
        plt.colorbar(img)
        if title is not None:
            fig.suptitle(title)
    else:
        plt.axis("off")

    fig.set_dpi(256)
    plt.axis("scaled")
    fig.tight_layout()

    if output_folder is not None and file_name is not None:
        plt.savefig(f"{output_folder}/{file_name}", bbox_inches="tight")

    if display is True:
        plt.show()
    plt.close()


def HiC_and_HIoIs(
    I,
    HIoIs,
    RoI,
    resolution,
    where=None,
    plot_in_bp=False,
    output_folder=None,
    file_name=None,
    title=None,
    display=False,
):
    # TODO: remove
    """
    :param I:                   Hi-C matrix to be plotted as image and saved
    :param HIoIs:               list of lists, where the innermost lists are pairs of elements
    :param  RoI:                refers to the region of interest [RoI[0], RoI[1]]x[RoI[2], RoI[3]]
                                (e.g., in genomic coordinates); see also plot_in_bp
    :param resolution:              resolution of the Hi-C matrix
    :param where:               if "lower" (resp. "upper"), then it plots sites only on the lower (resp. upper) part of
                                the Hi-C matrix; otherwise, it plots sites spanning the whole Hi-C matrix
    :param plot_in_bp:          if True, labels are set in base pairs
    :param output_folder:       path to folder where to save the image
    :param file_name:           name of the file to be created
    :param title:               title to give to the image
    :param display:             if False, it does not display the plot
    :return:                    -
    """

    fig, ax = plt.subplots(1, 1)
    img = ax.matshow(I, vmax=np.amax(I), extent=(RoI[0], RoI[1], RoI[3], RoI[2]), cmap=fruit_punch)
    if plot_in_bp:
        format_ticks(ax)

    for HIoI in HIoIs:
        if where == "lower":
            ax.plot(
                [RoI[0] + HIoI[0], RoI[0] + HIoI[0]],
                [RoI[2] + HIoI[0], RoI[2] + (I.shape[0] - 1) * resolution],
                color=(0.0, 0.0, 1.0),
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + HIoI[1] + resolution, RoI[0] + HIoI[1] + resolution],
                [RoI[2] + HIoI[1], RoI[2] + (I.shape[0] - 1) * resolution],
                color=(0.0, 0.0, 1.0),
                linestyle="dashed",
                linewidth=1,
            )
        elif where == "upper":
            ax.plot(
                [RoI[0] + HIoI[0], RoI[0] + HIoI[0]],
                [RoI[2] + 1 * resolution, RoI[2] + HIoI[0]],
                color=(0.0, 0.0, 1.0),
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + HIoI[1] + resolution, RoI[0] + HIoI[1] + resolution],
                [RoI[2] + 1 * resolution, RoI[2] + HIoI[1] + resolution],
                color=(0.0, 0.0, 1.0),
                linestyle="dashed",
                linewidth=1,
            )
        else:
            ax.plot(
                [RoI[0] + HIoI[0], RoI[0] + HIoI[0]],
                [RoI[2] + 1 * resolution, RoI[2] + (I.shape[0] - 1) * resolution],
                color=(0.0, 0.0, 1.0),
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + HIoI[1] + resolution, RoI[0] + HIoI[1] + resolution],
                [RoI[2] + 1 * resolution, RoI[2] + (I.shape[0] - 1) * resolution],
                color=(0.0, 0.0, 1.0),
                linestyle="dashed",
                linewidth=1,
            )

    if no_frills_in_images is False:
        plt.colorbar(img)
        if title is not None:
            fig.suptitle(title)
    else:
        plt.axis("off")

    fig.set_dpi(256)
    plt.axis("scaled")
    fig.tight_layout()

    if output_folder is not None and file_name is not None:
        plt.savefig(f"{output_folder}/{file_name}", bbox_inches="tight")

    if display:
        plt.show()
    plt.close()


def plot_stripes(
    I,
    LT_HIoIs,
    LT_VIoIs,
    UT_HIoIs,
    UT_VIoIs,
    RoI,
    resolution,
    plot_in_bp=False,
    output_folder=None,
    file_name=None,
    title=None,
    display=False,
):
    # TODO remove
    """
    :param I:                  Hi-C matrix to be plotted as image and saved
    :param LT_HIoIs:           Horizontal Intervals of Interest (lower-triangular part)
    :param LT_VIoIs:           Vertical Intervals of Interest (lower-triangular part)
    :param UT_HIoIs:           Horizontal Intervals of Interest (upper-triangular part)
    :param UT_VIoIs:           Vertical Intervals of Interest (upper-triangular part)
    :param RoI:                refers to the region of interest [RoI[0], RoI[1]]x[RoI[2], RoI[3]]
                               (e.g., in genomic coordinates); see also plot_in_bp
    :param plot_in_bp:         if True, labels are set in base pairs
    :param output_folder:      path to folder where to save the image
    :param file_name:          name of the file to be created
    :param title:              title to give to the image
    :param display:            if False, it does not display the plot
    :return:                   -
    """

    fig, ax = plt.subplots(1, 1)
    ax.matshow(I, vmax=np.amax(I), extent=(RoI[0], RoI[1], RoI[3], RoI[2]), cmap=fruit_punch)
    if plot_in_bp:
        format_ticks(ax)

    # Upper-triangular candidates:
    for num_window, (HIoI, VIoI) in enumerate(zip(UT_HIoIs, UT_VIoIs)):
        low_idx_up_opt = max(1, VIoI[0]) * resolution
        low_idx_dw_opt = min(VIoI[1], I.shape[0] - 1) * resolution
        low_idx_sx_opt = max(1, HIoI[0]) * resolution
        low_idx_dx_opt = min(HIoI[1], I.shape[1] - 1) * resolution
        if low_idx_up_opt != low_idx_dw_opt and low_idx_sx_opt != low_idx_dx_opt:
            ax.plot(
                [RoI[0] + low_idx_sx_opt, RoI[2] + low_idx_sx_opt],
                [RoI[0] + low_idx_dw_opt, RoI[2] + low_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + low_idx_dx_opt, RoI[2] + low_idx_dx_opt],
                [RoI[0] + low_idx_dw_opt, RoI[2] + low_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + low_idx_sx_opt, RoI[2] + low_idx_dx_opt],
                [RoI[0] + low_idx_dw_opt, RoI[2] + low_idx_dw_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + low_idx_sx_opt, RoI[2] + low_idx_dx_opt],
                [RoI[0] + low_idx_up_opt, RoI[2] + low_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            # plt.text(RoI[0] + low_idx_sx_opt, RoI[2] + low_idx_dw_opt + 15, str(num_window))

    for num_window, (HIoI, VIoI) in enumerate(zip(LT_HIoIs, LT_VIoIs)):
        upp_idx_up_opt = max(1, VIoI[0]) * resolution
        upp_idx_dw_opt = min(VIoI[1], I.shape[0] - 1) * resolution
        upp_idx_sx_opt = max(1, HIoI[0]) * resolution
        upp_idx_dx_opt = min(HIoI[1], I.shape[0] - 1) * resolution
        if upp_idx_up_opt != upp_idx_dw_opt and upp_idx_sx_opt != upp_idx_dx_opt:
            ax.plot(
                [RoI[0] + upp_idx_sx_opt, RoI[2] + upp_idx_sx_opt],
                [RoI[0] + upp_idx_dw_opt, RoI[2] + upp_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + upp_idx_dx_opt, RoI[2] + upp_idx_dx_opt],
                [RoI[0] + upp_idx_dw_opt, RoI[2] + upp_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + upp_idx_sx_opt, RoI[2] + upp_idx_dx_opt],
                [RoI[0] + upp_idx_dw_opt, RoI[2] + upp_idx_dw_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + upp_idx_sx_opt, RoI[2] + upp_idx_dx_opt],
                [RoI[0] + upp_idx_up_opt, RoI[2] + upp_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            # plt.text(RoI[0] + upp_idx_sx_opt, RoI[2] + upp_idx_up_opt - 15, str(num_window))

    if no_frills_in_images is False:
        if title is not None:
            fig.suptitle(title)
    else:
        plt.axis("off")

    fig.set_dpi(256)
    fig.tight_layout()

    if output_folder is not None and file_name is not None:
        plt.savefig(f"{output_folder}/{file_name}", bbox_inches="tight")

    if display:
        plt.show()
    plt.close()


def plot_stripes_and_peaks(
    I,
    LT_HIoIs,
    LT_VIoIs,
    UT_HIoIs,
    UT_VIoIs,
    LT_peaks_ids,
    UT_peaks_ids,
    RoI,
    resolution,
    plot_in_bp=False,
    output_folder=None,
    file_name=None,
    title=None,
    display=False,
):
    # TODO CORREGGERE LA TRASLAZIONE DI ROI
    """
    :param I:                  Hi-C matrix to be plotted as image and saved
    :param LT_HIoIs:           Horizontal Intervals of Interest (lower-triangular part)
    :param LT_VIoIs:           Vertical Intervals of Interest (lower-triangular part)
    :param UT_HIoIs:           Horizontal Intervals of Interest (upper-triangular part)
    :param UT_VIoIs:           Vertical Intervals of Interest (upper-triangular part)
    :param RoI:                refers to the region of interest [RoI[0], RoI[1]]x[RoI[2], RoI[3]]
                               (e.g., in genomic coordinates); see also plot_in_bp
    :param plot_in_bp:         if True, labels are set in base pairs
    :param output_folder:      path to folder where to save the image
    :param file_name:          name of the file to be created
    :param title:              title to give to the image
    :param display:            if False, it does not display the plot
    :return:                   -
    """

    fig, ax = plt.subplots(1, 1)
    ax.matshow(I, vmax=np.amax(I), extent=(RoI[0], RoI[1], RoI[3], RoI[2]), cmap=fruit_punch)
    if plot_in_bp:
        format_ticks(ax)

    # Upper-triangular candidates:
    for num_window, (HIoI, VIoI, peaks_ids) in enumerate(zip(UT_HIoIs, UT_VIoIs, UT_peaks_ids)):
        low_idx_up_opt = max(1, VIoI[0]) * resolution
        low_idx_dw_opt = min(VIoI[1], I.shape[0] - 1) * resolution
        low_idx_sx_opt = max(1, HIoI[0]) * resolution
        low_idx_dx_opt = min(HIoI[1], I.shape[1] - 1) * resolution
        if low_idx_up_opt != low_idx_dw_opt and low_idx_sx_opt != low_idx_dx_opt:
            ax.plot(
                [RoI[0] + low_idx_sx_opt, RoI[2] + low_idx_sx_opt],
                [RoI[0] + low_idx_dw_opt, RoI[2] + low_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + low_idx_dx_opt, RoI[2] + low_idx_dx_opt],
                [RoI[0] + low_idx_dw_opt, RoI[2] + low_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + low_idx_sx_opt, RoI[2] + low_idx_dx_opt],
                [RoI[0] + low_idx_dw_opt, RoI[2] + low_idx_dw_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + low_idx_sx_opt, RoI[2] + low_idx_dx_opt],
                [RoI[0] + low_idx_up_opt, RoI[2] + low_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )

            transl_peak_roots_ids = [RoI[0] + (low_idx_sx_opt + low_idx_dx_opt) / 2 for _ in peaks_ids]
            transl_peaks_ids = [peak_idx * resolution + RoI[2] for peak_idx in peaks_ids]
            ax.scatter(transl_peak_roots_ids, transl_peaks_ids, c="blue")

    for num_window, (HIoI, VIoI, peaks_ids) in enumerate(zip(LT_HIoIs, LT_VIoIs, LT_peaks_ids)):
        upp_idx_up_opt = max(1, VIoI[0]) * resolution
        upp_idx_dw_opt = min(VIoI[1], I.shape[0] - 1) * resolution
        upp_idx_sx_opt = max(1, HIoI[0]) * resolution
        upp_idx_dx_opt = min(HIoI[1], I.shape[0] - 1) * resolution
        if upp_idx_up_opt != upp_idx_dw_opt and upp_idx_sx_opt != upp_idx_dx_opt:
            ax.plot(
                [RoI[0] + upp_idx_sx_opt, RoI[2] + upp_idx_sx_opt],
                [RoI[0] + upp_idx_dw_opt, RoI[2] + upp_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + upp_idx_dx_opt, RoI[2] + upp_idx_dx_opt],
                [RoI[0] + upp_idx_dw_opt, RoI[2] + upp_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + upp_idx_sx_opt, RoI[2] + upp_idx_dx_opt],
                [RoI[0] + upp_idx_dw_opt, RoI[2] + upp_idx_dw_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )
            ax.plot(
                [RoI[0] + upp_idx_sx_opt, RoI[2] + upp_idx_dx_opt],
                [RoI[0] + upp_idx_up_opt, RoI[2] + upp_idx_up_opt],
                "blue",
                linestyle="dashed",
                linewidth=1,
            )

            transl_peak_roots_ids = [RoI[0] + (upp_idx_sx_opt + upp_idx_dx_opt) / 2 for _ in peaks_ids]
            transl_peaks_ids = [peak_idx * resolution + RoI[2] for peak_idx in peaks_ids]
            ax.scatter(transl_peak_roots_ids, transl_peaks_ids, c="blue")
            # plt.text(RoI[0] + upp_idx_sx_opt, RoI[2] + upp_idx_up_opt - 15, str(num_window))

    if no_frills_in_images is False:
        if title is not None:
            fig.suptitle(title)
    else:
        plt.axis("off")

    fig.set_dpi(256)
    fig.tight_layout()

    if output_folder is not None and file_name is not None:
        plt.savefig(f"{output_folder}/{file_name}", bbox_inches="tight")

    if display:
        plt.show()
    plt.close()
