"""

A class to carry localization data.

"""
import copy
import logging
import time
import warnings
from itertools import accumulate

import numpy as np
import pandas as pd
from google.protobuf import json_format, text_format

try:
    from scipy.spatial import QhullError
except ImportError:
    from scipy.spatial.qhull import QhullError  # needed for Python 3.7

import locan.data.hulls
from locan import locdata_id  # is required to use locdata_id as global variable
from locan.constants import PROPERTY_KEYS, PropertyKey
from locan.data import metadata_pb2
from locan.data.metadata_utils import _modify_meta, metadata_to_formatted_string
from locan.data.region import Region, RoiRegion

__all__ = ["LocData"]

logger = logging.getLogger(__name__)

class LocData:
    """
    This class carries localization data, aggregated properties and meta data.

    Data consist of individual elements being either localizations or other `LocData` objects. Both, localizations and
    `Locdata` objects have properties. Properties come from the original data or are added by analysis procedures.

    Parameters
    ----------
    references : LocData, list(LocData), None
        A `LocData` reference or an array with references to `LocData` objects referring to the selected localizations
        in dataset.
    dataframe : pandas.DataFrame, None
        Dataframe with localization data.
    indices : slice object, list(int), None
        Indices for dataframe in references that makes up the data. `indices` refers to index label, not position.
    meta : locan.data.metadata_pb2.Metadata, dictionary
        Metadata about the current dataset and its history.

    Attributes
    ----------
    references : LocData, list(LocData), None
        A LocData reference or an array with references to LocData objects referring to the selected localizations
        in dataframe.
    dataframe : pandas.DataFrame, None
        Dataframe with localization data.
    indices : slice object, list(int), None
        Indices for dataframe in references that makes up the data.
    meta : locan.data.metadata_pb2.Metadata
        Metadata about the current dataset and its history.
    properties : pandas.DataFrame
        List of properties generated from data.
    coordinate_labels : list of str
        The available coordinate properties.
    dimension : int
        Number of coordinates available for each localization (i.e. size of `coordinate_labels`).
    """

    count = 0
    """int: A counter for counting LocData instantiations (class attribute)."""

    def __init__(
        self, references=None, dataframe=pd.DataFrame(), indices=None, meta=None
    ):
        self.__class__.count += 1

        self.references = references
        self.dataframe = dataframe
        self.indices = indices
        self.meta = metadata_pb2.Metadata()
        self.properties = {}

        # regions and hulls
        self._region = None
        self._bounding_box = None
        self._oriented_bounding_box = None
        self._convex_hull = None
        self._alpha_shape = None
        self._inertia_moments = None

        self.coordinate_labels = sorted(
            list(
                set(self.data.columns).intersection(
                    {"position_x", "position_y", "position_z"}
                )
            )
        )

        self.dimension = len(self.coordinate_labels)

        self._update_properties()

        # meta
        global locdata_id
        locdata_id += 1
        self.meta.identifier = str(locdata_id)

        self.meta.creation_time.GetCurrentTime()
        self.meta.source = metadata_pb2.DESIGN
        self.meta.state = metadata_pb2.RAW

        self.meta.element_count = len(self.data.index)
        if "frame" in self.data.columns:
            self.meta.frame_count = len(self.data["frame"].unique())

        if meta is None:
            pass
        elif isinstance(meta, dict):
            for key, value in meta.items():
                setattr(self.meta, key, value)
        else:
            self.meta.MergeFrom(meta)

    def _update_properties(self):
        self.properties["localization_count"] = len(self.data.index)

        # property for mean spatial coordinates (centroids)
        self.properties.update(dict(self.data[self.coordinate_labels].mean()))

        self.bounding_box  # update self._bounding_box

    def __del__(self):
        """Updating the counter upon deletion of class instance."""
        self.__class__.count -= 1

    def __len__(self):
        """Return the length of data, i.e. the number of elements (localizations or collection elements)."""
        return len(self.data.index)

    def __getstate__(self):
        """Modify pickling behavior."""
        # Copy the object's state from self.__dict__ to avoid modifying the original state.
        state = self.__dict__.copy()
        # Serialize the unpicklable protobuf entries.
        json_string = json_format.MessageToJson(
            self.meta, including_default_value_fields=False
        )
        state["meta"] = json_string
        return state

    def __setstate__(self, state):
        """Modify pickling behavior."""
        # Restore instance attributes.
        self.__dict__.update(state)
        # Restore protobuf class for meta attribute
        self.meta = metadata_pb2.Metadata()
        self.meta = json_format.Parse(state["meta"], self.meta)

    def __copy__(self):
        """
        Create a shallow copy of locdata (keeping all references) with the following exceptions:
        (i) The class variable `count` is increased for the copied LocData object.
        (ii) Metadata keeps the original meta.creation_time while meta.modification_time and meta.history is updated.
        """
        new_locdata = LocData(self.references, self.dataframe, self.indices, meta=None)
        new_locdata._region = self._region
        # meta
        meta_ = _modify_meta(
            self, new_locdata, function_name="LocData.copy", parameter=None, meta=None
        )
        new_locdata.meta = meta_
        return new_locdata

    def __deepcopy__(self, memodict=None):
        """
        Create a deep copy of locdata (including all references) with the following exceptions:
        (i) The class variable `count` is increased for all deepcopied LocData objects.
        (ii) Metadata keeps the original meta.creation_time while meta.modification_time and meta.history is updated.
        """
        if memodict is None:
            memodict = {}
        new_locdata = LocData(
            copy.deepcopy(self.references, memodict),
            copy.deepcopy(self.dataframe, memodict),
            copy.deepcopy(self.indices, memodict),
            meta=None,
        )
        new_locdata._region = self._region
        # meta
        meta_ = _modify_meta(
            self,
            new_locdata,
            function_name="LocData.deepcopy",
            parameter=None,
            meta=None,
        )
        new_locdata.meta = meta_
        return new_locdata

    @property
    def bounding_box(self):
        """Hull object: Return an object representing the axis-aligned minimal bounding box."""
        if self._bounding_box is None:
            try:
                self._bounding_box = locan.data.hulls.BoundingBox(self.coordinates)
                self.properties["region_measure_bb"] = self._bounding_box.region_measure
                if self._bounding_box.region_measure:
                    self.properties["localization_density_bb"] = (
                        self.properties["localization_count"]
                        / self._bounding_box.region_measure
                    )
                if self._bounding_box.subregion_measure:
                    self.properties[
                        "subregion_measure_bb"
                    ] = self._bounding_box.subregion_measure
            except ValueError:
                warnings.warn(
                    "Properties related to bounding box could not be computed.",
                    UserWarning,
                )
        return self._bounding_box

    @property
    def convex_hull(self):
        """Hull object: Return an object representing the convex hull of all localizations."""
        if self._convex_hull is None:
            try:
                self._convex_hull = locan.data.hulls.ConvexHull(self.coordinates)
                self.properties["region_measure_ch"] = self._convex_hull.region_measure
                if self._convex_hull.region_measure:
                    self.properties["localization_density_ch"] = (
                        self.properties["localization_count"]
                        / self._convex_hull.region_measure
                    )
            except (TypeError, QhullError):
                warnings.warn(
                    "Properties related to convex hull could not be computed.",
                    UserWarning,
                )
        return self._convex_hull

    @property
    def oriented_bounding_box(self):
        """Hull object: Return an object representing the oriented minimal bounding box."""
        if self._oriented_bounding_box is None:
            try:
                self._oriented_bounding_box = locan.data.hulls.OrientedBoundingBox(
                    self.coordinates
                )
                self.properties[
                    "region_measure_obb"
                ] = self._oriented_bounding_box.region_measure
                if self._oriented_bounding_box.region_measure:
                    self.properties["localization_density_obb"] = (
                        self.properties["localization_count"]
                        / self._oriented_bounding_box.region_measure
                    )
                self.properties["orientation_obb"] = self._oriented_bounding_box.angle
                self.properties[
                    "circularity_obb"
                ] = self._oriented_bounding_box.elongation
            except TypeError:
                warnings.warn(
                    "Properties related to oriented bounding box could not be computed.",
                    UserWarning,
                )
        return self._oriented_bounding_box

    @property
    def alpha_shape(self):
        """Hull object: Return an object representing the alpha-shape of all localizations."""
        return self._alpha_shape

    def update_alpha_shape(self, alpha):
        """Compute the alpha shape for specific `alpha` and update `self.alpha_shape`.

        Parameters
        ----------
        alpha : float
            Alpha parameter specifying a unique alpha complex.

        Returns
        -------
        LocData
            The modified object
        """
        try:
            if self._alpha_shape is None:
                self._alpha_shape = locan.data.hulls.AlphaShape(
                    points=self.coordinates, alpha=alpha
                )
            else:
                self._alpha_shape.alpha = alpha

            self.properties["region_measure_as"] = self._alpha_shape.region_measure
            try:
                self.properties["localization_density_as"] = (
                    self._alpha_shape.n_points_alpha_shape
                    / self._alpha_shape.region_measure
                )
            except ZeroDivisionError:
                self.properties["localization_density_as"] = float("nan")

        except TypeError:
            warnings.warn(
                "Properties related to alpha shape could not be computed.", UserWarning
            )
        return self

    def update_alpha_shape_in_references(self, alpha):
        """
        Compute the alpha shape for each element in `locdata.references` and update `locdata.dataframe`.

        Returns
        -------
        LocData
            The modified object
        """
        if isinstance(self.references, list):
            for reference in self.references:
                reference.update_alpha_shape(alpha=alpha)
            new_df = pd.DataFrame(
                [reference.properties for reference in self.references]
            )
            new_df.index = self.data.index
            self.dataframe.update(new_df)
            new_columns = [
                column for column in new_df.columns if column in self.dataframe.columns
            ]
            new_df.drop(columns=new_columns, inplace=True, errors="ignore")
            self.dataframe = pd.concat([self.dataframe, new_df], axis=1)
        return self

    @property
    def inertia_moments(self):
        """Inertia moments are returned as computed by :func:`locan.data.properties.inertia_moments`."""
        if self._inertia_moments is None:
            try:
                self._inertia_moments = locan.data.properties.inertia_moments(
                    self.coordinates
                )
                self.properties["orientation_im"] = self._inertia_moments.orientation
                self.properties["circularity_im"] = self._inertia_moments.eccentricity
            except TypeError:
                warnings.warn(
                    "Properties related to inertia_moments could not be computed.",
                    UserWarning,
                )
        return self._inertia_moments

    def update_inertia_moments_in_references(self):
        """
        Compute inertia_moments for each element in locdata.references and update locdata.dataframe.

        Returns
        -------
        LocData
            The modified object
        """
        if isinstance(self.references, list):
            for reference in self.references:
                reference.inertia_moments  # request property to update
            new_df = pd.DataFrame(
                [reference.properties for reference in self.references]
            )
            new_df.index = self.data.index
            self.dataframe.update(new_df)
            new_columns = [
                column for column in new_df.columns if column in self.dataframe.columns
            ]
            new_df.drop(columns=new_columns, inplace=True, errors="ignore")
            self.dataframe = pd.concat([self.dataframe, new_df], axis=1)
        return self

    @property
    def region(self):
        """RoiRegion object: Return the region that supports all localizations."""
        return self._region

    @region.setter
    def region(self, region):
        if region is not None:
            if region.dimension != self.dimension:
                raise TypeError(
                    "Region dimension and coordinates dimension must be identical."
                )
            elif len(self) != len(region.contains(self.coordinates)):
                logger.warning("Not all coordinates are within region.")

        if isinstance(region, (Region, RoiRegion)) or region is None:
            self._region = region

        elif isinstance(
            region, dict
        ):  # legacy code to deal with deprecated RoiLegacy_0
            region_ = RoiRegion(**region)
            if region_ is not None:
                if region_.dimension != self.dimension:
                    raise TypeError(
                        "Region dimension and coordinates dimension must be identical."
                    )
                elif len(self) != len(region_.contains(self.coordinates)):
                    logger.warning("Not all coordinates are within region.")
            self._region = region_

        else:
            raise TypeError

        # property for region measures
        if self._region is not None:
            if self._region.region_measure:
                self.properties["region_measure"] = self._region.region_measure
                self.properties["localization_density"] = (
                    self.meta.element_count / self._region.region_measure
                )
            if self._region.subregion_measure:
                self.properties["subregion_measure"] = self._region.subregion_measure

    @property
    def data(self):
        """pandas.DataFrame: Return all elements either copied from the reference or referencing the current
        dataframe. """
        if isinstance(self.references, LocData):
            # we refer to the localization data by its index label, not position
            # in other words we decided not to use iloc but loc
            # df = self.references.data.loc[self.indices]  ... but this does not work in pandas.
            # also see:
            # https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#deprecate-loc-reindex-listlike
            try:
                df = self.references.data.loc[self.indices]
            except KeyError:
                df = self.references.data.loc[
                    self.references.data.index.intersection(self.indices)
                ]
            df = pd.merge(
                df, self.dataframe, left_index=True, right_index=True, how="outer"
            )
            return df
        else:
            return self.dataframe

    @property
    def coordinates(self):
        """ndarray: Return all coordinate values. """
        return self.data[self.coordinate_labels].values

    @property
    def centroid(self):
        """ndarray: Return coordinate values of the centroid
        (being the property values for all coordinate labels)."""
        return np.array(
            [
                self.properties[coordinate_label]
                for coordinate_label in self.coordinate_labels
            ]
        )

    @classmethod
    def from_dataframe(cls, dataframe=pd.DataFrame(), meta=None):
        """
        Create new LocData object from pandas.DataFrame with localization data.

        Parameters
        ----------
        dataframe : pandas.DataFrame
            Localization data.
        meta : locan.data.metadata_pb2.Metadata
            Metadata about the current dataset and its history.

        Returns
        -------
        LocData
            A new LocData instance with dataframe representing the concatenated data.
        """
        dataframe = dataframe
        meta_ = metadata_pb2.Metadata()

        meta_.source = metadata_pb2.DESIGN
        meta_.state = metadata_pb2.RAW
        meta_.history.add(name="LocData.from_dataframe")

        if meta is None:
            pass
        elif isinstance(meta, dict):
            for key, value in meta.items():
                setattr(meta_, key, value)
        else:
            meta_.MergeFrom(meta)

        return cls(dataframe=dataframe, meta=meta_)

    @classmethod
    def from_coordinates(cls, coordinates=(), coordinate_labels=None, meta=None):
        """
        Create new LocData object from a sequence of localization coordinates.

        Parameters
        ----------
        coordinates : sequence of tuples with shape (n_loclizations, dimension)
            Sequence of tuples with localization coordinates
        coordinate_labels : sequence of str
            The available coordinate properties.
        meta : locan.data.metadata_pb2.Metadata
            Metadata about the current dataset and its history.

        Returns
        -------
        LocData
            A new LocData instance with dataframe representing the concatenated data.
        """
        if np.size(coordinates):
            dimension = len(coordinates[0])

            if coordinate_labels is None:
                coordinate_labels = ["position_x", "position_y", "position_z"][
                    0:dimension
                ]
            else:
                if all(cl in PROPERTY_KEYS for cl in coordinate_labels):
                    coordinate_labels = coordinate_labels
                else:
                    raise ValueError(
                        "The given coordinate_labels are not standard property keys."
                    )

            dataframe = pd.DataFrame.from_records(
                data=coordinates, columns=coordinate_labels
            )

        else:
            dataframe = pd.DataFrame()

        meta_ = metadata_pb2.Metadata()
        meta_.source = metadata_pb2.DESIGN
        meta_.state = metadata_pb2.RAW
        meta_.history.add(name="LocData.from_coordinates")

        if meta is None:
            pass
        elif isinstance(meta, dict):
            for key, value in meta.items():
                setattr(meta_, key, value)
        else:
            meta_.MergeFrom(meta)

        return cls(dataframe=dataframe, meta=meta_)

    @classmethod
    def from_selection(cls, locdata, indices=slice(0, None), meta=None):
        """
        Create new LocData object from selected elements in another `LocData`.

        Parameters
        ----------
        locdata : LocData
            Locdata object from which to select elements.
        indices : slice object, list(int), None
            Index labels for elements in locdata that make up the new data.
            Note that contrary to usual python slices, both the start and the stop are included
            (see pandas documentation). `Indices` refer to index value not position in list.
        meta : locan.data.metadata_pb2.Metadata
            Metadata about the current dataset and its history.

        Returns
        -------
        LocData
            A new LocData instance with dataframe representing the selected data.

        Note
        ----
        No error is raised if indices do not exist in locdata.
        """
        references = locdata
        indices = indices
        meta_ = metadata_pb2.Metadata()
        meta_.CopyFrom(locdata.meta)
        try:
            meta_.ClearField("identifier")
        except ValueError:
            pass

        try:
            meta_.ClearField("element_count")
        except ValueError:
            pass

        try:
            meta_.ClearField("frame_count")
        except ValueError:
            pass

        meta_.modification_time.GetCurrentTime()
        meta_.state = metadata_pb2.MODIFIED
        meta_.ancestor_identifiers.append(locdata.meta.identifier)
        meta_.history.add(name="LocData.from_selection")

        if meta is None:
            pass
        elif isinstance(meta, dict):
            for key, value in meta.items():
                setattr(meta_, key, value)
        else:
            meta_.MergeFrom(meta)

        new_locdata = cls(references=references, indices=indices, meta=meta_)
        new_locdata.region = references.region
        return new_locdata

    @classmethod
    def from_collection(cls, locdatas, meta=None):
        """
        Create new LocData object by collecting LocData objects.

        Parameters
        ----------
        locdatas : list of LocData
            Locdata objects to collect.
        meta : locan.data.metadata_pb2.Metadata
            Metadata about the current dataset and its history.

        Returns
        -------
        LocData
            A new LocData instance with dataframe representing the concatenated data.
        """
        references = locdatas
        dataframe = pd.DataFrame([ref.properties for ref in references])

        meta_ = metadata_pb2.Metadata()

        meta_.source = metadata_pb2.DESIGN
        meta_.state = metadata_pb2.RAW
        meta_.ancestor_identifiers[:] = [ref.meta.identifier for ref in references]
        meta_.history.add(name="LocData.from_collection")

        if meta is None:
            pass
        elif isinstance(meta, dict):
            for key, value in meta.items():
                setattr(meta_, key, value)
        else:
            meta_.MergeFrom(meta)

        return cls(references=references, dataframe=dataframe, meta=meta_)

    @classmethod
    def concat(cls, locdatas, meta=None):
        """
        Concatenate LocData objects.

        Parameters
        ----------
        locdatas : list of LocData
            Locdata objects to concatenate.
        meta : locan.data.metadata_pb2.Metadata
            Metadata about the current dataset and its history.

        Returns
        -------
        LocData
            A new LocData instance with dataframe representing the concatenated data.
        """

        dataframe = pd.concat([i.data for i in locdatas], ignore_index=True, sort=False)

        # concatenate references also if None
        references = []
        for locdata in locdatas:
            try:
                references.extend(locdata.references)
            except TypeError:
                references.append(locdata.references)

        # check if all elements are None
        if not any(references):
            references = None

        meta_ = metadata_pb2.Metadata()

        meta_.creation_time.GetCurrentTime()
        meta_.source = metadata_pb2.DESIGN
        meta_.state = metadata_pb2.MODIFIED
        meta_.ancestor_identifiers[:] = [dat.meta.identifier for dat in locdatas]
        meta_.history.add(name="concat")

        if meta is None:
            pass
        elif isinstance(meta, dict):
            for key, value in meta.items():
                setattr(meta_, key, value)
        else:
            meta_.MergeFrom(meta)

        return cls(references=references, dataframe=dataframe, meta=meta_)

    @classmethod
    def from_chunks(
        cls,
        locdata,
        chunks=None,
        chunk_size=None,
        n_chunks=None,
        order="successive",
        drop=False,
        meta=None,
    ):
        """
        Divide locdata in chunks of localization elements.

        Parameters
        ----------
        locdatas : list of LocData
            Locdata objects to concatenate.
        chunks : list[tuples]
            Localization chunks as defined by a list of index-tuples
        chunk_size : int, None
            Number of localizations per chunk. One of `chunk_size` or `n_chunks` must be different from None.
        n_chunks : int, None
            Number of chunks. One of `chunk_size` or `n_chunks` must be different from None.
        order : str
            The order in which to select localizations. One of 'successive' or 'alternating'.
        drop : bool
            If True the last chunk will be eliminated if it has fewer localizations than the other chunks.
        meta : locan.data.metadata_pb2.Metadata
            Metadata about the current dataset and its history.

        Returns
        -------
        LocData
            A new LocData instance with references and dataframe elements representing the individual chunks.
        """
        n_nones = sum(element is None for element in [chunks, chunk_size, n_chunks])

        if n_nones != 2:
            raise ValueError(
                "One and only one of `chunks`, `chunk_size` or `n_chunks` must be different from None."
            )
        elif chunks is not None:
            index_lists = list(chunks)
        else:
            if chunk_size is not None:
                if (len(locdata) % chunk_size) == 0:
                    n_chunks = len(locdata) // chunk_size
                else:
                    n_chunks = len(locdata) // chunk_size + 1
            else:  # if n_chunks is not None
                if (len(locdata) % n_chunks) == 0:
                    chunk_size = len(locdata) // n_chunks
                else:
                    chunk_size = len(locdata) // (n_chunks - 1)

            if order == "successive":
                if (len(locdata) % chunk_size) == 0:
                    chunk_sizes = [chunk_size] * n_chunks
                else:
                    chunk_sizes = [chunk_size] * (n_chunks - 1) + [
                        (len(locdata) % chunk_size)
                    ]
                cum_chunk_sizes = list(accumulate(chunk_sizes))
                cum_chunk_sizes.insert(0, 0)
                index_lists = [
                    locdata.data.index[slice(lower, upper)]
                    for lower, upper in zip(cum_chunk_sizes[:-1], cum_chunk_sizes[1:])
                ]

            elif order == "alternating":
                index_lists = [
                    locdata.data.index[slice(i_chunk, None, n_chunks)]
                    for i_chunk in range(n_chunks)
                ]

            else:
                raise ValueError(f"The order {order} is not implemented.")

        if drop and len(index_lists) > 1 and len(index_lists[-1]) < len(index_lists[0]):
            index_lists = index_lists[:-1]

        references = [
            LocData.from_selection(locdata=locdata, indices=index_list)
            for index_list in index_lists
        ]
        dataframe = pd.DataFrame([ref.properties for ref in references])

        meta_ = metadata_pb2.Metadata()

        meta_.creation_time.GetCurrentTime()
        meta_.source = metadata_pb2.DESIGN
        meta_.state = metadata_pb2.RAW
        meta_.ancestor_identifiers[:] = [ref.meta.identifier for ref in references]
        meta_.history.add(name="LocData.chunks")

        if meta is None:
            pass
        elif isinstance(meta, dict):
            for key, value in meta.items():
                setattr(meta_, key, value)
        else:
            meta_.MergeFrom(meta)

        return cls(references=references, dataframe=dataframe, meta=meta_)

    def reset(self, reset_index=False):
        """
        Reset hulls and properties. This is needed after the dataframe attribute has been modified in place.

        Note
        ----
        Should be used with care because metadata is not updated accordingly.
        The region property is not changed.
        Better to just re-instantiate with `LocData.from_dataframe()` or use `locdata.update()`.

        Parameters
        ----------
        reset_index : Bool
            Flag indicating if the index is reset to integer values. If True the previous index values are discarded.

        Returns
        -------
        LocData
            The modified object
        """
        if reset_index is True:
            self.dataframe.reset_index(drop=True, inplace=True)

        self.properties = {}
        self._bounding_box = None
        self._oriented_bounding_box = None
        self._convex_hull = None
        self._alpha_shape = None

        self._update_properties()

        return self

    def update(self, dataframe, reset_index=False, meta=None):
        """
        Update the dataframe attribute in place.

        Use this function rather than setting locdata.dataframe directly in order to automatically update
        the attributes for dimension, coordinate_labels, hulls, properties, and metadata.

        Parameters
        ----------
        dataframe : pandas.DataFrame, None
            Dataframe with localization data.
        reset_index : Bool
            Flag indicating if the index is reset to integer values. If True the previous index values are discarded.
        meta : locan.data.metadata_pb2.Metadata
            Metadata about the current dataset and its history.

        Returns
        -------
        LocData
            The modified object
        """
        local_parameter = locals()
        del local_parameter[
            "dataframe"
        ]  # dataframe is obvious and possibly large and should not be repeated in meta.

        if self.references is not None:
            self.reduce(reset_index=reset_index)
            logger.warning(
                "LocData.reduce() was applied since self.references was not None."
            )

        self.dataframe = dataframe
        self.coordinate_labels = sorted(
            list(
                set(self.data.columns).intersection(
                    {"position_x", "position_y", "position_z"}
                )
            )
        )
        self.dimension = len(self.coordinate_labels)
        self.reset(reset_index=reset_index)  # update hulls and properties

        # update meta
        self.meta.modification_time.GetCurrentTime()
        self.meta.state = metadata_pb2.MODIFIED
        self.meta.history.add(name="LocData.update", parameter=str(local_parameter))

        self.meta.element_count = len(self.data.index)
        if "frame" in self.data.columns:
            self.meta.frame_count = len(self.data["frame"].unique())

        if meta is None:
            pass
        elif isinstance(meta, dict):
            for key, value in meta.items():
                setattr(self.meta, key, value)
        else:
            self.meta.MergeFrom(meta)

        return self

    def reduce(self, reset_index=False):
        """
        Clean up references.

        This includes to update `Locdata.dataframe` and set `LocData.references` and `LocData.indices` to None.

        Parameters
        ----------
        reset_index : Bool
            Flag indicating if the index is reset to integer values. If True the previous index values are discarded.

        Returns
        -------
        LocData
            The modified object
        """
        if self.references is None:
            pass
        elif isinstance(self.references, (LocData, list)):
            self.dataframe = self.data
            self.indices = None
            self.references = None
        else:
            raise ValueError("references has undefined value.")

        if reset_index is True:
            self.dataframe.reset_index(drop=True, inplace=True)

        return self

    def update_convex_hulls_in_references(self):
        """
        Compute the convex hull for each element in locdata.references and update locdata.dataframe.

        Returns
        -------
        LocData
            The modified object
        """
        if isinstance(self.references, list):
            for reference in self.references:
                reference.convex_hull  # request property to update reference._convex_hull

            new_df = pd.DataFrame(
                [reference.properties for reference in self.references]
            )
            new_df.index = self.data.index
            self.dataframe.update(new_df)
            new_columns = [
                column for column in new_df.columns if column in self.dataframe.columns
            ]
            new_df.drop(columns=new_columns, inplace=True, errors="ignore")
            self.dataframe = pd.concat([self.dataframe, new_df], axis=1)
        return self

    def update_oriented_bounding_box_in_references(self):
        """
        Compute the oriented bounding box for each element in locdata.references and update locdata.dataframe.

        Returns
        -------
        LocData
            The modified object
        """
        if isinstance(self.references, list):
            for reference in self.references:
                reference.oriented_bounding_box  # request property to update reference._convex_hull
            new_df = pd.DataFrame(
                [reference.properties for reference in self.references]
            )
            new_df.index = self.data.index
            self.dataframe.update(new_df)
            new_columns = [
                column for column in new_df.columns if column in self.dataframe.columns
            ]
            new_df.drop(columns=new_columns, inplace=True, errors="ignore")
            self.dataframe = pd.concat([self.dataframe, new_df], axis=1)
        return self

    def projection(self, coordinate_labels):
        """
        Reduce dimensions by projecting all localization coordinates onto selected coordinates.

        Parameters
        ----------
        coordinate_labels : str, list
            The coordinate labels to project onto.

        Returns
        -------
            LocData
        """
        local_parameter = locals()

        if isinstance(coordinate_labels, str):
            coordinate_labels = [coordinate_labels]

        new_locdata = copy.deepcopy(self)

        # reduce coordinate dimensions
        coordinate_labels_to_drop = [
            label for label in self.coordinate_labels if label not in coordinate_labels
        ]
        columns = self.data.columns
        new_columns = [
            column for column in columns if column not in coordinate_labels_to_drop
        ]
        dataframe = new_locdata.data[new_columns]

        # update
        _meta = metadata_pb2.Metadata()
        _meta.history.add(name="LocData.projection", parameter=str(local_parameter))
        # other updates are done in the coming update call.

        new_locdata = new_locdata.update(dataframe=dataframe, meta=_meta)

        return new_locdata

    def print_meta(self):
        """
        Print Locdata.metadata.

        See Also
        --------
        :func:`locan.data.metadata_utils.metadata_to_formatted_string`
        """
        print(metadata_to_formatted_string(self.meta))

    def print_summary(self):
        """
        Print a summary containing the most common metadata keys.
        """
        meta_ = metadata_pb2.Metadata()
        if self.meta.HasField("file"):
            meta_.file.CopyFrom(self.meta.file)
        meta_.identifier = self.meta.identifier
        meta_.comment = self.meta.comment
        meta_.creation_time.CopyFrom(self.meta.creation_time)
        if self.meta.HasField("modification_time"):
            meta_.modification_time.CopyFrom(self.meta.modification_time)
        meta_.source = self.meta.source
        meta_.state = self.meta.state
        meta_.element_count = self.meta.element_count
        meta_.frame_count = self.meta.frame_count

        print(metadata_to_formatted_string(meta_))


<gh_stars>1-10
#!/usr/bin/python

import sys
import time
import threading
import numpy
import string
from math import sqrt,exp,log,pi,atan,sin
from scipy.integrate import quad

def Calc_Lifetime(Pmed,param,I0,twiss,ex,ey,sigP,sigS):
# float Pmed = residual gas pressure (nTorr)
# Dicionary param = basic machine parameters
# array I0,twiss,acc = current distribution(A), twiss parameters and momentum acceptance along the ring (%)
# array ex,ey,sigP,sigS = Calculated values for emittances, energy spread and bunch length for each bunch

	#Twiss parameters
	s=numpy.zeros(len(twiss))
	betax=numpy.zeros(len(twiss))
	alphax=numpy.zeros(len(twiss))
	betay=numpy.zeros(len(twiss))
	alphay=numpy.zeros(len(twiss))
	Dx=numpy.zeros(len(twiss))
	Dpx=numpy.zeros(len(twiss))
	Dy=numpy.zeros(len(twiss))
	Dpy=numpy.zeros(len(twiss))
	accp=numpy.zeros(len(twiss))
	accn=numpy.zeros(len(twiss))

	s=twiss[:,0]
	betax=twiss[:,2]
	alphax=twiss[:,3]
	betay=twiss[:,6]
	alphay=twiss[:,7]
	Dx=twiss[:,4]
	Dpx=twiss[:,5]
	Dy=twiss[:,8]
	Dpy=twiss[:,9]
	accp=twiss[:,10]
	accn=twiss[:,11]
	Ds=numpy.zeros(len(twiss))
	acc=numpy.zeros(len(twiss))

	Ds=s-numpy.roll(s,1)
	Ds[0]=0
	acc=numpy.minimum(accp,accn)

	#Calculate average beta functions
	betax_avg=numpy.average(betax,weights=Ds)
	betay_avg=numpy.average(betay,weights=Ds)
	#print "<betax> = ",betax_avg
	#print "<betay> = ",betay_avg
	#print "<acc> = ",numpy.average(acc,weights=Ds)#, "<accp> = ",numpy.average(accp,weights=Ds), "<accn> = ",numpy.average(accn,weights=Ds)

	#Machine parameters
	C = param['C'] #Circunference (m)
	frev=param['C']/param['cluz'] #Rev. freq (Hz)
	theta_x = sqrt(param['Ax']/betax_avg)
	theta_y = sqrt(param['Ay']/betay_avg)
	R=theta_y/theta_x
	FR=pi+(R**2+1.0)*sin(2*atan(R))+2.0*(R**2.0-1.0)*atan(R)

	#Elastic Scattering Lifetime
	Telas=10.25*2.0*pi/FR*(param['En']/1.0e+09)**2*param['Ay']/(betay_avg*Pmed)

	#Inelastic Scattering Lifetime
	Tine=1/(0.0065*Pmed*log(1./(numpy.average(acc,weights=Ds))))

	#Touschek Lifetime Calculation
	sigx=numpy.zeros(len(twiss))
	sigy=numpy.zeros(len(twiss))
	epsilon=numpy.zeros(len(twiss))
	Tv=numpy.zeros(len(twiss))
	dsdT=numpy.zeros(len(twiss)-1)

	sigx=numpy.sqrt(betax*ex+(Dx*sigP)**2)
	sigy=numpy.sqrt(betay*ey+(Dy*sigP)**2)
	epsilonp=(accp)**2*(betax/ex)/(1957.0*param['En']/1.0e+09)**2
	epsilonn=(accn)**2*(betax/ex)/(1957.0*param['En']/1.0e+09)**2
	Tvp=(5.39e17*(param['En']/1.0e+09)**2*(accp)**3*sigx*sigy*sigS/(De(epsilonp)*C))**(-1.0)
	Tvn=(5.39e17*(param['En']/1.0e+09)**2*(accn)**3*sigx*sigy*sigS/(De(epsilonn)*C))**(-1.0)
 	aux=0.5*(numpy.average(Tvp,weights=Ds)+numpy.average(Tvn,weights=Ds))
 	Ttous=1.0/(aux*I0)

	return (Ttous,Tine,Telas)

def Calc_Lifetime_Matlab(Pmed,param,twiss,ex,ey,sigP,sigS):
# float Pmed = residual gas pressure (nTorr)
# Dicionary param = basic machine parameters
# array I0,twiss,acc = current distribution(A), twiss parameters and momentum acceptance along the ring (%)
# array ex,ey,sigP,sigS = Calculated values for emittances, energy spread and bunch length for each bunch

	#Twiss parameters
	s=numpy.zeros(len(twiss))
	betax=numpy.zeros(len(twiss))
	alphax=numpy.zeros(len(twiss))
	betay=numpy.zeros(len(twiss))
	alphay=numpy.zeros(len(twiss))
	Dx=numpy.zeros(len(twiss))
	Dpx=numpy.zeros(len(twiss))
	Dy=numpy.zeros(len(twiss))
	Dpy=numpy.zeros(len(twiss))
	accp=numpy.zeros(len(twiss))
	accn=numpy.zeros(len(twiss))

	s=twiss[:,0]
	betax=twiss[:,2]
	alphax=twiss[:,3]
	betay=twiss[:,6]
	alphay=twiss[:,7]
	Dx=twiss[:,4]
	Dpx=twiss[:,5]
	Dy=twiss[:,8]
	Dpy=twiss[:,9]
	accp=twiss[:,10]
	accn=twiss[:,11]
	Ds=numpy.zeros(len(twiss))
	acc=numpy.zeros(len(twiss))

	Ds=s-numpy.roll(s,1)
	Ds[0]=0
	acc=numpy.minimum(accp,accn)

	#Calculate average beta functions
	betax_avg=numpy.average(betax,weights=Ds)
	betay_avg=numpy.average(betay,weights=Ds)
	#print "<betax> = ",betax_avg
	#print "<betay> = ",betay_avg
	#print "<acc> = ",numpy.average(acc,weights=Ds)#, "<accp> = ",numpy.average(accp,weights=Ds), "<accn> = ",numpy.average(accn,weights=Ds)

	#Machine parameters
	C = param['C'] #Circunference (m)
	frev=param['C']/param['cluz'] #Rev. freq (Hz)
	gamma=param['gamma']
	Np=param['Np']
	theta_x = sqrt(param['Ax']/betax_avg)
	theta_y = sqrt(param['Ay']/betay_avg)
	R=theta_y/theta_x
	FR=pi+(R**2+1.0)*sin(2*atan(R))+2.0*(R**2.0-1.0)*atan(R)

	#Elastic Scattering Lifetime
	Telas=10.25*2.0*pi/FR*(param['En']/1.0e+09)**2*param['Ay']/(betay_avg*Pmed)

	#Inelastic Scattering Lifetime
	Tine=1/(0.0065*Pmed*log(1./(numpy.average(acc,weights=Ds))))

	#Touschek Lifetime Calculation
	sigx=numpy.zeros(len(twiss))
	sigy=numpy.zeros(len(twiss))
	epsilon=numpy.zeros(len(twiss))
	Tv=numpy.zeros(len(twiss))
	dsdT=numpy.zeros(len(twiss)-1)

	#bunch size and volume
	sigx=numpy.sqrt(betax*ex+(Dx*sigP)**2)
	sigy=numpy.sqrt(betay*ey+(Dy*sigP)**2)
	V=sigS*sigx*sigy

	# parameters
	Sx2=ex*betax
	factor=betax*Dpx+alphax*Dx
	A1=1.0/(4.0*sigP**2)+(Dx**2+factor**2)/(4.0*Sx2)
	B1=betax*factor/(2.0*Sx2)
	C1=betax**2/(4.0*Sx2)-B1**2/(4.0*A1)

	#Epsilon factors
	#epsilonp=(2.0*numpy.sqrt(C1)/gamma*accp)**2
	#epsilonn=(2.0*numpy.sqrt(C1)/gamma*accn)**2
	epsilonp=(2.0*numpy.sqrt(C1)*accp)**2/(1957.0*param['En']/1.0e+09)**2
	epsilonn=(2.0*numpy.sqrt(C1)*accn)**2/(1957.0*param['En']/1.0e+09)**2

	#Lifetime
	Tvp=9.4718e-23*Np/(gamma**2)*1/(accp**3)*De(epsilonp)/V
	Tvn=9.4718e-23*Np/(gamma**2)*1/(accn**3)*De(epsilonn)/V
 	aux=0.5*(numpy.average(Tvp,weights=Ds)+numpy.average(Tvn,weights=Ds))*3600
 	Ttous=1.0/aux

	return (Ttous,Tine,Telas)

def De(e):

	out=numpy.zeros(len(e))
	for j in range(len(e)):
		(int1,err1)=quad(integrand1,e[j],numpy.inf)
		(int2,err2)=quad(integrand2,e[j],numpy.inf)
		out[j]=0.5*sqrt(e[j])*(-3.0*exp(-e[j])+e[j]*int1+int2*(3.0*e[j]-e[j]*log(e[j])+2.0))
	return out

def integrand1(x):
	return log(x)/x*exp(-x)

def integrand2(x):
	return exp(-x)/x


import numpy as np
import libs.configs.config_v1 as cfg
from PIL import Image, ImageFont, ImageDraw, ImageEnhance
from scipy.misc import imresize

FLAGS = cfg.FLAGS
_DEBUG = False

def draw_img(step, image, name='', image_height=1, image_width=1, rois=None):
    img = np.uint8(image/0.1*127.0 + 127.0)
    img = Image.fromarray(img)
    return img.save(FLAGS.train_dir + 'test_' + name + '_' +  str(step) +'.jpg', 'JPEG')

def draw_bbox(step, image, name='', image_height=1, image_width=1, bbox=None, label=None, gt_label=None, mask=None, prob=None, iou=None, vis_th=0.5, vis_all=False, ignore_bg=True):
    source_img = Image.fromarray(image)
    b, g, r = source_img.split()
    source_img = Image.merge("RGB", (r, g, b))
    draw = ImageDraw.Draw(source_img)
    color = '#0000ff'
    if mask is not None:
        m = np.array(mask*255.0)
        m = np.transpose(m,(0,3,1,2))
    if bbox is not None:
        for i, box in enumerate(bbox):
            if label is not None and not np.all(box==0):
                if prob is not None:
                    if ((prob[i,label[i]] > vis_th) or (vis_all is True)) and ((ignore_bg is True) and (label[i] > 0)) :
                        if gt_label is not None:
                            if gt_label is not None and len(iou) > 1:
                                text  = cat_id_to_cls_name(label[i]) + ' : ' + cat_id_to_cls_name(gt_label[i]) + ' : ' + str(iou[i])[:3]
                            else:
                                text  = cat_id_to_cls_name(label[i]) + ' : ' + cat_id_to_cls_name(gt_label[i]) + ' : ' + str(prob[i][label[i]])[:4]

                            if label[i] != gt_label[i]:
                                color = '#ff0000'#draw.text((2+bbox[i,0], 2+bbox[i,1]), cat_id_to_cls_name(label[i]) + ' : ' + cat_id_to_cls_name(gt_label[i]), fill='#ff0000')
                            else:
                                color = '#0000ff'
                        else:
                            text = cat_id_to_cls_name(label[i]) + ' : ' +  "{:.3f}".format(prob[i][label[i]])  #str(i)#+
                        draw.text((2+bbox[i,0], 2+bbox[i,1]), text, fill=color)

                        if _DEBUG is True:
                            print("plot",label[i], prob[i,label[i]])
                        draw.rectangle(box,fill=None,outline=color)

                        if mask is not None:
                            # print("mask number: ",i)
                            box = np.floor(box).astype('uint16')
                            bbox_w = box[2]-box[0]
                            bbox_h = box[3]-box[1]
                            mask_color_id = np.random.randint(35)
                            color_img = color_id_to_color_code(mask_color_id)* np.ones((bbox_h,bbox_w,1)) * 255
                            color_img = Image.fromarray(color_img.astype('uint8')).convert('RGBA')
                            #color_img = Image.new("RGBA", (bbox_w,bbox_h), np.random.rand(1,3) * 255 )
                            resized_m = imresize(m[i][label[i]], [bbox_h, bbox_w], interp='bilinear') #label[i]
                            resized_m[resized_m >= 128] = 128
                            resized_m[resized_m < 128] = 0
                            resized_m = Image.fromarray(resized_m.astype('uint8'), 'L')
                            #print(box)
                            #print(resized_m)

                            source_img.paste(color_img , (box[0], box[1]), mask=resized_m)

                        #return source_img.save(FLAGS.train_dir + 'est_imgs/' + name + '_' +  str(step) +'.jpg', 'JPEG')

                    else:
                        if _DEBUG is True:
                            print("skip",label[i], prob[i,label[i]])
                else:
                    text = cat_id_to_cls_name(label[i])
                    draw.text((2+bbox[i,0], 2+bbox[i,1]), text, fill=color)
                    draw.rectangle(box,fill=None,outline=color)

    return source_img.save(FLAGS.train_dir + 'est_imgs/' + name + '_' +  str(step) +'.jpg', 'JPEG')

def cat_id_to_cls_name(catId):
    cls_name = np.array([  'background', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
                       'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
                       'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
                       'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
                       'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
                       'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
                       'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
                       'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
                       'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
                       'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
                       'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
                       'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
                       'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
                       'scissors', 'teddy bear', 'hair drier', 'toothbrush'])
    return cls_name[catId]

def color_id_to_color_code(colorId):
    color_code = np.array([[178, 31, 53],
                           [216, 39, 53],
                           [255, 116, 53],
                           [255, 161, 53],
                           [255, 203, 53],
                           [255, 255, 53],
                           [0, 117, 58],
                           [0, 158, 71],
                           [22, 221, 53],
                           [0, 82, 165],
                           [0, 121, 231],
                           [0, 169, 252],
                           [104, 30, 126],
                           [125, 60, 181],
                           [189, 122, 246],
                           [234, 62, 112],
                           [198, 44, 58],
                           [243, 114, 82],
                           [255, 130, 1],
                           [255, 211, 92],
                           [138, 151, 71],
                           [2, 181, 160],
                           [75, 196, 213],
                           [149, 69, 103],
                           [125, 9, 150],
                           [169, 27, 176],
                           [198, 30, 153],
                           [207, 0, 99],
                           [230, 21, 119],
                           [243, 77, 154],
                           [144, 33, 71],
                           [223, 40, 35],
                           [247, 106, 4],
                           [206, 156, 72],
                           [250, 194, 0],
                           [254, 221, 39],
                           ])
    return color_code[colorId]


from __future__ import unicode_literals, division, print_function, absolute_import
from builtins import object
from collections import defaultdict
from copy import deepcopy
import numpy as np
from scipy.sparse import csr_matrix, lil_matrix

class ContextModel(object):

    def __init__(self, sentences, min_count=5, window=5, wordlist=[], progress=1000, forward=True, backward=True):
        """
        sentences: list/generator of lists of words
        in case this is based on a pretrained word2vec model, give the index2word attribute as wordlist

        Attributes:
            - min_count: how often a word has to occur at least
            - window: how many words in a word's context should be considered
            - word2index: {word:idx}
            - index2word: [word1, word2, ...]
            - wcounts: {word: frequency}
            - featmat: n_voc x n_voc sparse array with weighted context word counts for every word
            - progress: after how many sentences a progress printout should occur (default 1000)
        """
        self.progress = progress
        self.min_count = min_count
        self.window = window
        self.forward = forward
        self.backward = backward
        self.build_windex(sentences, wordlist)
        self._get_raw_context_matrix(sentences)

    def build_windex(self, sentences, wordlist=[]):
        """
        go through all the sentences and get an overview of all used words and their frequencies
        """
        # get an overview of the vocabulary
        vocab = defaultdict(int)
        for sentence_no, sentence in enumerate(sentences):
            if not sentence_no % self.progress:
                print("PROGRESS: at sentence #%i, processed %i words and %i unique words" % (sentence_no, sum(vocab.values()), len(vocab)))
            for word in sentence:
                vocab[word] += 1
        print("collected %i unique words from a corpus of %i words and %i sentences" % (len(vocab), sum(vocab.values()), sentence_no + 1))
        # assign a unique index to each word and remove all words with freq < min_count
        self.wcounts, self.word2index, self.index2word = {}, {}, []
        if not wordlist:
            wordlist = [word for word, c in vocab.items() if c >= self.min_count]
        for word in wordlist:
            self.word2index[word] = len(self.word2index)
            self.index2word.append(word)
            self.wcounts[word] = vocab[word]

    def _get_raw_context_matrix(self, sentences):
        """
        compute the raw context matrix with weighted counts
        it has an entry for every word in the vocabulary
        """
        # make the feature matrix
        featmat = lil_matrix((len(self.index2word), len(self.index2word)), dtype=float)
        for sentence_no, sentence in enumerate(sentences):
            if not sentence_no % self.progress:
                print("PROGRESS: at sentence #%i" % sentence_no)
            sentence = [word if word in self.word2index else None for word in sentence]
            # forward pass
            if self.forward:
                for i, word in enumerate(sentence[:-1]):
                    if word:
                        # get all words in the forward window
                        wwords = sentence[i + 1:min(i + 1 + self.window, len(sentence))]
                        for j, w in enumerate(wwords, 1):
                            if w:
                                featmat[self.word2index[word], self.word2index[w]] += 1.  # /j
            # backwards pass
            if self.backward:
                sentence_back = sentence[::-1]
                for i, word in enumerate(sentence_back[:-1]):
                    if word:
                        # get all words in the forward window of the backwards sentence
                        wwords = sentence_back[i + 1:min(i + 1 + self.window, len(sentence_back))]
                        for j, w in enumerate(wwords, 1):
                            if w:
                                featmat[self.word2index[word], self.word2index[w]] += 1.  # /j
        print("PROGRESS: through with all the sentences")
        self.featmat = csr_matrix(featmat)

    def get_context_matrix(self, fill_diag=True, norm='count'):
        """
        for every word in the sentences, create a vector that contains the counts of its context words
        (weighted by the distance to it with a max distance of window)
        Inputs:
            - norm: if the feature matrix should be normalized to contain ones on the diagonal
                    (--> average context vectors)
            - fill_diag: if diagonal of featmat should be filled with word counts
        Returns:
            - featmat: n_voc x n_voc sparse array with weighted context word counts for every word
        """
        featmat = deepcopy(self.featmat)
        # fill up the diagonals with the total counts of each word --> similarity matrix
        if fill_diag:
            featmat = lil_matrix(featmat)
            for i, word in enumerate(self.index2word):
                featmat[i, i] = self.wcounts[word]
            featmat = csr_matrix(featmat)
        assert ((featmat - featmat.transpose()).data**2).sum() < 2.220446049250313e-16, "featmat not symmetric"
        # possibly normalize by the max counts
        if norm in ("count", "max"):
            normmat = lil_matrix(featmat.shape, dtype=float)
            if norm == "count":
                print("normalizing feature matrix by word count")
                normmat.setdiag([1. / self.wcounts[word] for word in self.index2word])
            elif norm == "max":
                print("normalizing feature matrix by max counts")
                normmat.setdiag([1. / v[0] if v[0] else 1. for v in featmat.max(axis=1).toarray()])
            featmat = csr_matrix(normmat) * featmat  # row in featmat multiplied by entry on diagonal
        return featmat

    def get_local_context_matrix(self, tokens):
        """
        compute a local context matrix. it has an entry for every token, even if it is not present in the vocabulary
        Inputs:
            - tokens: list of words
        Returns:
            - local_featmat: size len(set(tokens)) x n_vocab
            - tok_idx: {word: index} to map the words from the tokens list to an index of the featmat
        """
        # for every token we still only need one representation per document
        tok_idx = {word: i for i, word in enumerate(set(tokens))}
        featmat = lil_matrix((len(tok_idx), len(self.index2word)), dtype=float)
        # clean out context words we don't know
        known_tokens = [word if word in self.word2index else None for word in tokens]
        # forward pass
        if self.forward:
            for i, word in enumerate(tokens[:-1]):
                # get all words in the forward window
                wwords = known_tokens[i + 1:min(i + 1 + self.window, len(known_tokens))]
                for j, w in enumerate(wwords, 1):
                    if w:
                        featmat[tok_idx[word], self.word2index[w]] += 1. / j
        # backwards pass
        if self.backward:
            tokens_back = tokens[::-1]
            known_tokens_back = known_tokens[::-1]
            for i, word in enumerate(tokens_back[:-1]):
                # get all words in the forward window of the backwards sentence, incl. word itself
                wwords = known_tokens_back[i + 1:min(i + 1 + self.window, len(known_tokens_back))]
                for j, w in enumerate(wwords, 1):
                    if w:
                        featmat[tok_idx[word], self.word2index[w]] += 1. / j
        featmat = csr_matrix(featmat)
        # normalize matrix
        normmat = lil_matrix((featmat.shape[0], featmat.shape[0]), dtype=float)
        normmat.setdiag([1. / v[0] if v[0] else 1. for v in featmat.max(axis=1).toarray()])
        featmat = csr_matrix(normmat) * featmat
        return featmat, tok_idx


from bs4 import BeautifulSoup as bs
from label_navi import calLinkChaotic as clc
from label_navi import calLinkRatio as clr
import scipy as sp
import numpy as np
import sys
import matplotlib.pyplot as plt

fileName = sys.argv[1]
dom = bs(open('../data/en-original/' + fileName))

nodeList = []
for des in dom.descendants:
    try:
        if des.name != None and des.find_all('a'):
            if ~np.isnan(clc(des)) and clc(des) >= 0:
                nodeList.append([clc(des), clr(des), len(des.find_all('a')), des.name])
    except:
        print(des.name)
        break
dataList = []
for node in nodeList:
    if node[2] > 1:
        dataList.append( [node[0], node[1] ])
dataList = np.array(dataList)

print(len(dataList))

k = len(dataList)
x = dataList[:,0]
y = dataList[:,1]
#n, bins, patches = P.hist(x, 10, histtype='step', stacked=True, fill=True)

plt.scatter( range( 1, (k+1)), x )

plt.grid()
plt.show()


<reponame>davisidarta/dynamo-release<filename>dynamo/preprocessing/utils.py
import numpy as np
import pandas as pd
import statsmodels.api as sm
from scipy.sparse import issparse, csr_matrix
import warnings
# from functools import reduce
from sklearn.decomposition import PCA, TruncatedSVD

# ---------------------------------------------------------------------------------------------------
# symbol conversion related
def convert2gene_symbol(input_names, scopes='ensembl.gene'):
    """Convert ensemble gene id to official gene names using mygene package.

    Parameters
    ----------
        input_names: list-like
            The ensemble gene id names that you want to convert to official gene names. All names should come from the same
            species.
        scopes: `list-like` or `None` (default: `None`)
            Scopes are needed when you use non-official gene name as your gene indices (or adata.var_name). This
            arugument corresponds to type of types of identifiers, either a list or a comma-separated fields to specify
            type of input qterms, e.g. “entrezgene”, “entrezgene,symbol”, [“ensemblgene”, “symbol”]. Refer to official
            MyGene.info docs (https://docs.mygene.info/en/latest/doc/query_service.html#available_fields) for full list
            of fields.

    Returns
    -------
        var_pd: `pd.Dataframe`
            A pandas dataframe that includes the following columns:
            query: the input ensmble ids
            _id: identified id from mygene
            _score: confidence of the retrieved official gene name.
            symbol: retrieved official gene name
    """

    try:
        import mygene
    except ImportError:
        raise ImportError("You need to install the package `mygene` (pip install mygene --user) "
                          "See https://pypi.org/project/mygene/ for more details.")

    import mygene
    mg = mygene.MyGeneInfo()

    ensemble_names = [i.split('.')[0] for i in input_names]
    var_pd = mg.querymany(ensemble_names, scopes=scopes, fields='symbol', as_dataframe=True, df_index=True)
    #var_pd.drop_duplicates(subset='query', inplace=True) # use when df_index is not True
    var_pd = var_pd.loc[~var_pd.index.duplicated(keep='first')]

    return var_pd

def convert2symbol(adata, scopes=None, subset=True):
    if np.all(adata.var_names.str.startswith('ENS')) or scopes is not None:
        prefix = adata.var_names[0]
        if scopes is None:
            if prefix[:4] == 'ENSG' or prefix[:7] == 'ENSMUSG':
                scopes = 'ensembl.gene'
            elif prefix[:4] == 'ENST' or prefix[:7] == 'ENSMUST':
                scopes = 'ensembl.transcript'
            else:
                raise Exception('Your adata object uses non-official gene names as gene index. \n'
                                'Dynamo finds those IDs are neither from ensembl.gene or ensembl.transcript and thus cannot '
                                'convert them automatically. \n'
                                'Please pass the correct scopes or first convert the ensemble ID to gene short name '
                                '(for example, using mygene package). \n'
                                'See also dyn.pp.convert2gene_symbol')

        adata.var['query'] = [i.split('.')[0] for i in adata.var.index]
        if scopes is str:
            adata.var[scopes] = adata.var.index
        else:
            adata.var['scopes'] = adata.var.index

        warnings.warn('Your adata object uses non-official gene names as gene index. \n'
                      'Dynamo is converting those names to official gene names.')
        official_gene_df = convert2gene_symbol(adata.var_names, scopes)
        merge_df = adata.var.merge(official_gene_df, left_on='query', right_on='query', how='left').set_index(
            adata.var.index)
        adata.var = merge_df
        valid_ind = np.where(merge_df['notfound'] != True)[0]

        if subset:
            adata._inplace_subset_var(valid_ind)
            adata.var.index = adata.var['symbol'].values.copy()
        else:
            indices = np.array(adata.var.index)
            indices[valid_ind] = adata.var.loc[valid_ind, 'symbol'].values.copy()
            adata.var.index = indices

    return adata

# ---------------------------------------------------------------------------------------------------
# implmentation of Cooks' distance (but this is for Poisson distribution fitting)

# https://stackoverflow.com/questions/47686227/poisson-regression-in-statsmodels-and-r

# from __future__ import division, print_function

# https://stats.stackexchange.com/questions/356053/the-identity-link-function-does-not-respect-the-domain-of-the-gamma-family
def _weight_matrix(fitted_model):
    """Calculates weight matrix in Poisson regression

    Parameters
    ----------
    fitted_model : statsmodel object
        Fitted Poisson model

    Returns
    -------
    W : 2d array-like
        Diagonal weight matrix in Poisson regression
    """
    return np.diag(fitted_model.fittedvalues)

def _hessian(X, W):
    """Hessian matrix calculated as -X'*W*X

    Parameters
    ----------
    X : 2d array-like
        Matrix of covariates

    W : 2d array-like
        Weight matrix

    Returns
    -------
    hessian : 2d array-like
        Hessian matrix
    """
    return -np.dot(X.T, np.dot(W, X))

def _hat_matrix(X, W):
    """Calculate hat matrix = W^(1/2) * X * (X'*W*X)^(-1) * X'*W^(1/2)

    Parameters
    ----------
    X : 2d array-like
        Matrix of covariates

    W : 2d array-like
        Diagonal weight matrix

    Returns
    -------
    hat : 2d array-like
        Hat matrix
    """
    # W^(1/2)
    Wsqrt = W ** (0.5)

    # (X'*W*X)^(-1)
    XtWX = -_hessian(X=X, W=W)
    XtWX_inv = np.linalg.inv(XtWX)

    # W^(1/2)*X
    WsqrtX = np.dot(Wsqrt, X)

    # X'*W^(1/2)
    XtWsqrt = np.dot(X.T, Wsqrt)

    return np.dot(WsqrtX, np.dot(XtWX_inv, XtWsqrt))

def cook_dist(model, X, good):
    # Weight matrix
    W = _weight_matrix(model)

    # Hat matrix
    H = _hat_matrix(X, W)
    hii = np.diag(
        H
    )  # Diagonal values of hat matrix # fit.get_influence().hat_matrix_diag

    # Pearson residuals
    r = model.resid_pearson

    # Cook's distance (formula used by R = (res/(1 - hat))^2 * hat/(dispersion * p))
    # Note: dispersion is 1 since we aren't modeling overdispersion

    resid = good.disp - model.predict(good)
    rss = np.sum(resid ** 2)
    MSE = rss / (good.shape[0] - 2)
    # use the formula from: https://www.mathworks.com/help/stats/cooks-distance.html
    cooks_d = (
        r ** 2 / (2 * MSE) * hii / (1 - hii) ** 2
    )  # (r / (1 - hii)) ** 2 *  / (1 * 2)

    return cooks_d

# ---------------------------------------------------------------------------------------------------
# preprocess utilities
def basic_stats(adata):
    adata.obs['nGenes'], adata.obs['nCounts'] = (adata.X > 0).sum(1), (adata.X).sum(1)
    mito_genes = adata.var_names.str.upper().str.startswith('MT-')
    try:
        adata.obs['pMito'] = (adata[:, mito_genes].X).sum(1).A1 / adata.obs['nCounts'] if issparse(adata.X) else  \
            (adata[:, mito_genes].X).sum(1) / adata.obs['nCounts']
    except:
        raise Exception(f"looks like your var_names may be corrupted (i.e. include nan values)")

def unique_var_obs_adata(adata):
    """Function to make the obs and var attribute's index unique"""
    adata.obs_names_make_unique()
    adata.var_names_make_unique()

    return adata

def layers2csr(adata):
    """Function to make the obs and var attribute's index unique"""
    for i in adata.layers.keys():
        adata.layers[i] = csr_matrix(adata.layers[i]) if not issparse(adata.layers[i]) else adata.layers[i]

    return adata

def merge_adata_attrs(adata_ori, adata, attr):
    if attr == 'var':
        _columns = set(adata.var.columns).difference(adata_ori.var.columns)
        var_df = adata_ori.var.merge(adata.var[_columns], how='left',
                                            left_index=True, right_index=True)
        adata_ori.var = var_df.loc[adata_ori.var.index, :]
    elif attr == 'obs':
        _columns = set(adata.obs.columns).difference(adata_ori.obs.columns)
        obs_df = adata_ori.obs.merge(adata.obs[_columns], how='left',
                                            left_index=True, right_index=True)
        adata_ori.obs = obs_df.loc[adata_ori.obs.index, :]

    return adata_ori

def allowed_layer_raw_names():
    only_splicing = ['spliced', 'unspliced']
    only_labeling = ['new', 'total']
    splicing_and_labeling = ['uu', 'ul', 'su', 'sl']

    return only_splicing, only_labeling, splicing_and_labeling

def allowed_X_layer_names():
    only_splicing = ['X_spliced', 'X_unspliced']
    only_labeling = ['X_new', 'X_total']
    splicing_and_labeling = ['X_uu', 'X_ul', 'X_su', 'X_sl']

    return only_splicing, only_labeling, splicing_and_labeling

def get_layer_keys(adata, layers="all", remove_normalized=True, include_protein=True):
    """Get the list of available layers' keys.
    """
    layer_keys = list(adata.layers.keys())
    if remove_normalized:
        layer_keys = [i for i in layer_keys if not i.startswith("X_")]

    if "protein" in adata.obsm.keys() and include_protein:
        layer_keys.extend(["X", "protein"])
    else:
        layer_keys.extend(["X"])
    layers = (
        layer_keys
        if layers == "all"
        else list(set(layer_keys).intersection(list(layers)))
    )

    layers = list(set(layers).difference(["matrix", "ambiguous", "spanning"]))
    return layers

def get_shared_counts(adata, layers, min_shared_count, type="gene"):
    layers = list(set(layers).difference(["X", "matrix", "ambiguous", "spanning"]))
    layers = np.array(layers)[~pd.DataFrame(layers)[0].str.startswith("X_").values]

    _nonzeros, _sum = None, None
    for layer in layers:
        if issparse(adata.layers[layers[0]]):
            _nonzeros = (
                adata.layers[layer] > 0
                if _nonzeros is None
                else _nonzeros.multiply(adata.layers[layer] > 0)
            )
        else:
            _nonzeros = (
                adata.layers[layer] > 0
                if _nonzeros is None
                else _nonzeros * (adata.layers[layer] > 0)
            )

    for layer in layers:
        if issparse(adata.layers[layers[0]]):
            _sum = (
                _nonzeros.multiply(adata.layers[layer])
                if _sum is None
                else _sum + _nonzeros.multiply(adata.layers[layer])
            )
        else:
            _sum = (
                np.multiply(_nonzeros, adata.layers[layer])
                if _sum is None
                else _sum + np.multiply(_nonzeros, adata.layers[layer])
            )

    if type == "gene":
        return (
            np.array(_sum.sum(0).A1 >= min_shared_count)
            if issparse(adata.layers[layers[0]])
            else np.array(_sum.sum(0) >= min_shared_count)
        )
    if type == "cells":
        return (
            np.array(_sum.sum(1).A1 >= min_shared_count)
            if issparse(adata.layers[layers[0]])
            else np.array(_sum.sum(1) >= min_shared_count)
        )

def clusters_stats(U, S, clusters_uid, cluster_ix, size_limit=40):
    """Calculate the averages per cluster

    If the cluster is too small (size<size_limit) the average of the total is reported instead
    This function is modified from velocyto in order to reproduce velocyto's DentateGyrus notebook.
    """
    U_avgs = np.zeros((S.shape[1], len(clusters_uid)))
    S_avgs = np.zeros((S.shape[1], len(clusters_uid)))
    avgU_div_avgS = np.zeros((S.shape[1], len(clusters_uid)))
    slopes_by_clust = np.zeros((S.shape[1], len(clusters_uid)))

    for i, uid in enumerate(clusters_uid):
        cluster_filter = cluster_ix == i
        n_cells = np.sum(cluster_filter)
        if n_cells > size_limit:
            U_avgs[:, i], S_avgs[:, i] = (
                U[cluster_filter, :].mean(0),
                S[cluster_filter, :].mean(0),
            )
        else:
            U_avgs[:, i], S_avgs[:, i] = U.mean(0), S.mean(0)

    return U_avgs, S_avgs

def get_svr_filter(adata, layer="spliced", n_top_genes=3000, return_adata=False):
    score_name = "score" if layer in ["X", "all"] else layer + "_score"
    valid_idx = np.where(np.isfinite(adata.var.loc[:, score_name]))[0]

    valid_table = adata.var.iloc[valid_idx, :]
    nth_score = np.sort(valid_table.loc[:, score_name])[::-1][
        np.min((n_top_genes - 1, valid_table.shape[0] - 1))
    ]

    feature_gene_idx = np.where(valid_table.loc[:, score_name] >= nth_score)[0][
        :n_top_genes
    ]
    feature_gene_idx = valid_idx[feature_gene_idx]

    if return_adata:
        adata.var.loc[:, "use_for_pca"] = False
        adata.var.loc[adata.var.index[feature_gene_idx], "use_for_pca"] = True
        res = adata
    else:
        filter_bool = np.zeros(adata.n_vars, dtype=bool)
        filter_bool[feature_gene_idx] = True
        res = filter_bool

    return res

def sz_util(adata, layer, round_exprs, method, locfunc, total_layers=None, CM=None, scale_to=None):
    adata = adata.copy()

    if layer == '_total_' and '_total_' not in adata.layers.keys():
        if total_layers is not None:
            if not isinstance(total_layers, list): total_layers = [total_layers]
            if len(set(total_layers).difference(adata.layers.keys())) == 0:
                total = None
                for t_key in total_layers:
                    total = (
                        adata.layers[t_key] if total is None else total + adata.layers[t_key]
                    )
                adata.layers["_total_"] = total

    if layer == "raw":
        CM = adata.raw.X if CM is None else CM
    elif layer == "X":
        CM = adata.X if CM is None else CM
    elif layer == "protein":
        if "protein" in adata.obsm_keys():
            CM = adata.obsm["protein"] if CM is None else CM
        else:
            return None, None
    else:
        CM = adata.layers[layer] if CM is None else CM

    if round_exprs:
        if issparse(CM):
            CM.data = np.round(CM.data, 0)
        else:
            CM = CM.round().astype("int")

    cell_total = CM.sum(axis=1).A1 if issparse(CM) else CM.sum(axis=1)
    cell_total += cell_total == 0  # avoid infinity value after log (0)

    if method in ["mean-geometric-mean-total", 'geometric']:
        sfs = cell_total / (np.exp(locfunc(np.log(cell_total))) if scale_to is None else scale_to)
    elif method == "median":
        sfs = cell_total / (np.nanmedian(cell_total) if scale_to is None else scale_to)
    elif method == "mean":
        sfs = cell_total / (np.nanmean(cell_total) if scale_to is None else scale_to)
    else:
        raise NotImplementedError(f"This method {method} is not supported!")

    return sfs, cell_total

def get_sz_exprs(adata, layer, total_szfactor=None):
    if layer == "raw":
        CM = adata.raw.X
        szfactors = adata.obs[layer + "Size_Factor"].values[:, None]
    elif layer == "X":
        CM = adata.X
        szfactors = adata.obs["Size_Factor"].values[:, None]
    elif layer == "protein":
        if "protein" in adata.obsm_keys():
            CM = adata.obsm[layer]
            szfactors = adata.obs["protein_Size_Factor"].values[:, None]
        else:
            CM, szfactors = None, None
    else:
        CM = adata.layers[layer]
        szfactors = adata.obs[layer + "_Size_Factor"].values[:, None]

    if total_szfactor is not None and total_szfactor in adata.obs.keys():
        szfactors = adata.obs[total_szfactor][:, None]
    else:
        warnings.warn("{total_szfactor} is not None but it is not in adata object.")

    return szfactors, CM

def normalize_util(CM, szfactors, relative_expr, pseudo_expr, norm_method=np.log1p):
    if norm_method == np.log1p: pseudo_expr = 0
    if relative_expr:
        CM = (
            CM.multiply(csr_matrix(1 / szfactors))
            if issparse(CM)
            else CM / szfactors
        )

    if pseudo_expr is None:
        pseudo_expr = 1
    if issparse(CM):
        CM.data = (
            norm_method(CM.data + pseudo_expr)
            if norm_method is not None
            else CM.data
        )
        if norm_method is not None and norm_method.__name__ == 'Freeman_Tukey': CM.data -= 1
    else:
        CM = (
            norm_method(CM + pseudo_expr)
            if norm_method is not None
            else CM
        )

    return CM

def Freeman_Tukey(X, inverse=False):
    if inverse:
        res = np.sqrt(X) + np.sqrt((X + 1))
    else:
        res = (X**2 - 1)**2 / (4 * X**2)

    return res

# ---------------------------------------------------------------------------------------------------
# pca

def pca(adata, CM, n_pca_components=30, pca_key='X', pcs_key='PCs'):

    if adata.n_obs < 100000:
        pca = PCA(n_components=min(n_pca_components, CM.shape[1] - 1), svd_solver="arpack", random_state=0)
        fit = pca.fit(CM.toarray()) if issparse(CM) else pca.fit(CM)
        X_pca = fit.transform(CM.toarray()) if issparse(CM) else fit.transform(CM)
        adata.obsm[pca_key] = X_pca
        adata.uns[pcs_key] = fit.components_.T

        adata.uns["explained_variance_ratio_"] = fit.explained_variance_ratio_
    else:
        # unscaled PCA
        fit = TruncatedSVD(
            n_components=min(n_pca_components + 1, CM.shape[1] - 1), random_state=0
        )
        # first columns is related to the total UMI (or library size)
        X_pca = fit.fit_transform(CM)[:, 1:]
        adata.obsm[pca_key] = X_pca
        adata.uns[pcs_key] = fit.components_.T

        adata.uns["explained_variance_ratio_"] = fit.explained_variance_ratio_[1:]

    return adata, fit, X_pca

def pca_genes(PCs, n_top_genes=100):
    ret = np.zeros(PCs.shape[0], dtype=bool)
    for q in PCs.T:
        qq = np.sort(np.abs(q))[::-1]
        idx = np.abs(q) > qq[n_top_genes]
        ret = np.logical_or(idx, ret)
    return ret

def add_noise_to_duplicates(adata, basis='pca'):
    X_data = adata.obsm['X_' + basis]
    min_val = abs(X_data).min()

    n_obs, n_var = X_data.shape
    while(True):
        _, index = np.unique(X_data, axis=0, return_index=True)
        duplicated_idx = np.setdiff1d(np.arange(n_obs), index)

        if len(duplicated_idx) == 0:
            adata.obsm['X_' + basis] = X_data
            break
        else:
            X_data[duplicated_idx, :] += np.random.normal(0, min_val / 1000, (len(duplicated_idx), n_var))

# ---------------------------------------------------------------------------------------------------
# labeling related

def collapse_adata(adata):
    """Function to collapse the four species data, will be generalized to handle dual-datasets"""
    only_splicing, only_labeling, splicing_and_labeling = allowed_layer_raw_names()

    if np.all([i in adata.layers.keys() for i in splicing_and_labeling]):
        if only_splicing[0] not in adata.layers.keys(): adata.layers[only_splicing[0]] = adata.layers['su'] + adata.layers['sl']
        if only_splicing[1] not in adata.layers.keys(): adata.layers[only_splicing[1]] = adata.layers['uu'] + adata.layers['ul']
        if only_labeling[0] not in adata.layers.keys(): adata.layers[only_labeling[0]] = adata.layers['ul'] + adata.layers['sl']
        if only_labeling[1] not in adata.layers.keys(): adata.layers[only_labeling[1]] = adata.layers[only_labeling[0]] + adata.layers['uu'] + adata.layers['su']

    return adata

def detect_datatype(adata):
    has_splicing, has_labeling, has_protein = False, False, False

    layers = adata.layers.keys()
    if len({'ul', 'sl', 'uu', 'su'}.difference(layers)) == 0:
        has_splicing, has_labeling = True, True
    elif len({'unspliced', 'spliced'}.difference(layers)) == 0:
        has_splicing = True
    elif len({'new', 'total'}.difference(layers)) == 0:
        has_labeling = True

    if "protein" in adata.obsm.keys():
        has_protein = True

    return has_splicing, has_labeling, has_protein

def default_layer(adata):
    has_splicing, has_labeling, _ = detect_datatype(adata)

    if has_splicing:
        if has_labeling:
            if len(set(adata.layers.keys()).intersection(['new', 'total', 'spliced', 'unspliced'])) == 4:
                adata = collapse_adata(adata)
            default_layer = "M_t" if "M_t" in adata.layers.keys() else "X_total" if \
                "X_total" in adata.layers.keys() else "total"
        else:
            default_layer = "M_s" if "M_s" in adata.layers.keys() else "X_spliced" if \
                "X_spliced" in adata.layers.keys() else "spliced"
    else:
        default_layer = "M_t" if "M_t" in adata.layers.keys() else "X_total" if \
            "X_total" in adata.layers.keys() else "total"

    return default_layer

def NTR(adata):
    """calculate the new to total ratio across cells. Note that
    NTR for the first time point in degradation approximates gamma/beta."""

    if len({'new', 'total'}.intersection(adata.layers.keys())) == 2:
        ntr = adata.layers['new'].sum(1) / adata.layers['total'].sum(1)
        ntr = ntr.A1 if issparse(adata.layers['new']) else ntr

        var_ntr = adata.layers['new'].sum(0) / adata.layers['total'].sum(0)
        var_ntr = var_ntr.A1 if issparse(adata.layers['new']) else var_ntr
    elif len({'uu', 'ul', 'su', 'sl'}.intersection(adata.layers.keys())) == 4:
        new = adata.layers['ul'].sum(1) + \
              adata.layers['sl'].sum(1)
        total = new + adata.layers['uu'].sum(1) + \
                adata.layers['su'].sum(1)
        ntr = new / total

        ntr = ntr.A1 if issparse(adata.layers['uu']) else ntr

        new = adata.layers['ul'].sum(0) + \
              adata.layers['sl'].sum(0)
        total = new + adata.layers['uu'].sum(0) + \
                adata.layers['su'].sum(0)
        var_ntr = new / total

        var_ntr = var_ntr.A1 if issparse(adata.layers['uu']) else var_ntr
    elif len({'unspliced', 'spliced'}.intersection(adata.layers.keys())) == 2:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")

            ntr = adata.layers['unspliced'].sum(1) / (adata.layers['unspliced'] + adata.layers['spliced']).sum(1)
            var_ntr = adata.layers['unspliced'].sum(0) / (adata.layers['unspliced'] + adata.layers['spliced']).sum(0)

        ntr = ntr.A1 if issparse(adata.layers['unspliced']) else ntr
        var_ntr = var_ntr.A1 if issparse(adata.layers['unspliced']) else var_ntr
    else:
        ntr, var_ntr = None, None

    return ntr, var_ntr

def scale(adata, layers=None, scale_to_layer=None, scale_to=1e6):
    """scale layers to a particular total expression value, similar to `normalize_expr_data` function."""
    layers = get_layer_keys(adata, layers)
    has_splicing, has_labeling, _ = detect_datatype(adata)

    if scale_to_layer is None:
        scale_to_layer = 'total' if has_labeling else None
        scale = scale_to / adata.layers[scale_to_layer].sum(1)
    else:
        scale = None

    for layer in layers:
        if scale is None:
            scale = scale_to / adata.layers[layer].sum(1)

        adata.layers[layer] = csr_matrix(adata.layers[layer].multiply(scale))

    return adata

# ---------------------------------------------------------------------------------------------------
# ERCC related

def relative2abs(adata,
                 dilution,
                 volume,
                 from_layer=None,
                 to_layers=None,
                 mixture_type=1,
                 ERCC_controls=None,
                 ERCC_annotation=None):
    """Converts FPKM/TPM data to transcript counts using ERCC spike-in. This is based on the relative2abs function from
    monocle 2 (Qiu, et. al, Nature Methods, 2017).

    Parameters
    ----------
        adata: :class:`~anndata.AnnData`
            an Annodata object
        dilution: `float`
            the dilution of the spikein transcript in the lysis reaction mix. Default is 40, 000. The number of spike-in
            transcripts per single-cell lysis reaction was calculated from.
        volume: `float`
            the approximate volume of the lysis chamber (nanoliters). Default is 10
        from_layer: `str` or `None`
            The layer in which the ERCC TPM values will be used as the covariate for the ERCC based linear regression.
        to_layers: `str`, `None` or `list-like`
            The layers that our ERCC based transformation will be applied to.
        mixture_type:
            the type of spikein transcripts from the spikein mixture added in the experiments. By default, it is mixture 1.
            Note that m/c we inferred are also based on mixture 1.
        ERCC_controls:
            the FPKM/TPM matrix for each ERCC spike-in transcript in the cells if user wants to perform the transformation based
            on their spike-in data. Note that the row and column names should match up with the ERCC_annotation and relative_
            exprs_matrix respectively.
        ERCC_annotation:
            the ERCC_annotation matrix from illumina USE GUIDE which will be ued for calculating the ERCC transcript copy
            number for performing the transformation.

    Returns
    -------
        An adata object with the data specified in the to_layers transformed into absolute counts.
    """

    if ERCC_annotation is None:
        ERCC_annotation = pd.read_csv('https://www.dropbox.com/s/cmiuthdw5tt76o5/ERCC_specification.txt?dl=1', sep='\t')

    ERCC_id = ERCC_annotation['ERCC ID']

    ERCC_id = adata.var_names.intersection(ERCC_id)
    if len(ERCC_id) < 10 and ERCC_controls is None:
        raise Exception(f'The adata object you provided has less than 10 ERCC genes.')

    if to_layers is not None:
        to_layers = [to_layers] if to_layers is str else to_layers
        to_layers = list(set(adata.layers.keys()).intersection(to_layers))
        if len(to_layers) == 0:
            raise Exception(f"The layers {to_layers} that will be converted to absolute counts doesn't match any layers"
                            f"from the adata object.")

    mixture_name = "concentration in Mix 1 (attomoles/ul)" if mixture_type == 1 else "concentration in Mix 2 (attomoles/ul)"
    ERCC_annotation['numMolecules'] = ERCC_annotation.loc[:, mixture_name] * (
                volume * 10 ** (-3) * 1 / dilution * 10 ** (-18) * 6.02214129 * 10 ** (23))

    ERCC_annotation['rounded_numMolecules'] = ERCC_annotation['numMolecules'].astype(int)

    if from_layer in [None, 'X']:
        X, X_ercc = (adata.X, adata[:, ERCC_id].X if ERCC_controls is None else ERCC_controls)
    else:
        X, X_ercc = (adata.layers[from_layer], adata[:, ERCC_id] \
            if ERCC_controls is None else ERCC_controls)

    logged = False if X.max() > 100 else True

    if not logged:
        X, X_ercc = (np.log1p(X.A) if issparse(X_ercc) else np.log1p(X), \
                     np.log1p(X_ercc.A) if issparse(X_ercc) else np.log1p(X_ercc))
    else:
        X, X_ercc = (X.A if issparse(X_ercc) else X, X_ercc.A if issparse(X_ercc) else X_ercc)

    y = np.log1p(ERCC_annotation['numMolecules'])

    for i in range(adata.n_obs):
        X_i, X_ercc_i = X[i, :], X_ercc[i, :]

        X_i, X_ercc_i = sm.add_constant(X_i),  sm.add_constant(X_ercc_i)
        res = sm.RLM(y, X_ercc_i).fit()
        k, b = res.params[::-1]

        if to_layers is None:
            X = adata.X
            logged = False if X.max() > 100 else True

            if not logged:
                X_i = np.log1p(X[i, :].A) if issparse(X) else np.log1p(X[i, :])
            else:
                X_i = X[i, :].A if issparse(X) else X[i, :]

            res = k * X_i + b
            res = res if logged else np.expm1(res)
            adata.X[i, :] = csr_matrix(res) if issparse(X) else res
        else:
            for cur_layer in to_layers:
                X = adata.layers[cur_layer]

                logged = False if X.max() > 100 else True
                if not logged:
                    X_i = np.log1p(X[i, :].A) if issparse(X) else np.log1p(X[i, :])
                else:
                    X_i = X[i, :].A if issparse(X) else X[i, :]

                res = k * X_i + b if logged else np.expm1(k * X_i + b)
                adata.layers[cur_layer][i, :] = csr_matrix(res) if issparse(X) else res


#! /usr/bin/env python3

import numpy as np
from scipy.stats import loguniform, truncnorm
import argparse

parser = argparse.ArgumentParser(description="Generate initial (random) sample grid")
parser.add_argument("--fixed-parameter", nargs=2, action="append", help="Fix a parameter's value, e.g. `--fixed-parameter dist 40.0`")
parser.add_argument("--set-limit", nargs=3, action="append", help="Set a parameter's limits to something other than the default, e.g. `--set-limit mej_dyn 0.01, 0.05`")
parser.add_argument("--npts", type=int, default=25000, help="Number of points to use in the grid")
parser.add_argument("--output-file", help="Filename to save grid to")
parser.add_argument("--gaussian-prior", action="append", nargs=3, help="Give a parameter a Gaussian prior, specifying the mean and standard deviation; for example, `--gaussian-prior theta 20.0 5.0`")
args = parser.parse_args()

#
# Prior functions
#

def uniform(llim, rlim, x):
    return 1. / (rlim - llim)

def log_uniform(llim, rlim, x):
    return loguniform.pdf(x, llim, rlim)

def gaussian(llim, rlim, mu, sigma, x):
    return truncnorm.pdf(x, llim, rlim, loc=mu, scale=sigma)

#
# Functions to draw samples from the priors
#
def sample_uniform(llim, rlim, n):
    return np.random.uniform(llim, rlim, size=n)

def sample_log_uniform(llim, rlim, n):
    return loguniform.rvs(llim, rlim, size=n)

def sample_gaussian(llim, rlim, mu, sigma, n):
    return truncnorm.rvs(llim, rlim, loc=mu, scale=sigma, size=n)

# the parameters FIXME: should this be more flexible?
ordered_parameters = ["mej_dyn", "vej_dyn", "mej_wind", "vej_wind", "theta"]

# Parameter limits
limits = {
        "mej_dyn":[0.001, 0.1],
        "vej_dyn":[0.05, 0.3],
        "mej_wind":[0.001, 0.1],
        "vej_wind":[0.05, 0.3],
        "theta":[0., 90.]
}

# If the user specified different limits, change them accordingly
if args.set_limit is not None:
    for [_parameter, _llim, _rlim] in args.set_limit:
        limits[_parameter] = [float(_llim), float(_rlim)]

# Specify each parameter's prior
prior_functions = {
        "mej_dyn":lambda x: log_uniform(*limits["mej_dyn"], x),
        "vej_dyn":lambda x: uniform(*limits["vej_dyn"], x),
        "mej_wind":lambda x: log_uniform(*limits["mej_wind"], x),
        "vej_wind":lambda x: uniform(*limits["vej_wind"], x),
        "theta":lambda x: uniform(*limits["theta"], x)
}

# Specify each parameter's prior sampling function
prior_sampling_functions = {
        "mej_dyn":lambda n: sample_log_uniform(*limits["mej_dyn"], n),
        "vej_dyn":lambda n: sample_uniform(*limits["vej_dyn"], n),
        "mej_wind":lambda n: sample_log_uniform(*limits["mej_wind"], n),
        "vej_wind":lambda n: sample_uniform(*limits["vej_wind"], n),
        "theta":lambda n: sample_uniform(*limits["theta"], n)
}

# Deal with possible fixed parameters
fixed_parameters = {}
if args.fixed_parameter is not None:
    for [_parameter, _val] in args.fixed_parameter:
        fixed_parameters[_parameter] = float(_val)

# Gaussian priors
if args.gaussian_prior is not None:
    for [_parameter, _mu, _sigma] in args.gaussian_prior:
        _mu, _sigma = float(_mu), float(_sigma)
        a, b = (limits[_parameter][0] - _mu) / _sigma, (limits[_parameter][1] - _mu) / _sigma
        prior_functions[_parameter] = lambda x: gaussian(a, b, _mu, _sigma, x)
        prior_sampling_functions[_parameter] = lambda n: sample_gaussian(a, b, _mu, _sigma, n)

#
# Generate the grid
#

grid = np.empty((args.npts, len(ordered_parameters) + 3))

# The first column, for lnL, gets filled in later (by generate_posterior_samples.py), so for now make it 0
grid[:,0] = 0.

# The second and third columns are the prior and sampling prior, respectively, which are the same for the initial grid.
# The joint prior is the product of all the separate priors, so we'll set them to 1 now and multiply them by each parameter's prior in the loop.
grid[:,1] = 0.

# Do the sampling and compute priors
for i, _parameter in enumerate(ordered_parameters):
    grid[:,i + 3] = prior_sampling_functions[_parameter](args.npts) if _parameter not in fixed_parameters.keys() else fixed_parameters[_parameter]
    grid[:,1] += np.log(prior_functions[_parameter](grid[:,i + 3]))
grid[:,2] = grid[:,1]

# Save the grid
np.savetxt(args.output_file, grid, header=("ln(L) ln(p) ln(ps) " + " ".join(ordered_parameters)))