from typing import List, Tuple, Union, Mapping, TypeVar, Callable, Iterable, Optional
from typing_extensions import Literal

from copy import copy
from enum import auto
from types import FunctionType, MappingProxyType
from inspect import signature
from pathlib import Path
from functools import wraps
from itertools import combinations

from anndata import AnnData
from cellrank import logging as logg
from cellrank._key import Key
from cellrank.tl._enum import ModeEnum
from cellrank.ul._docs import d, inject_docs
from cellrank.tl._utils import save_fig, _convert_lineage_name, _unique_order_preserving
from cellrank.tl._colors import (
    _get_bg_fg_colors,
    _compute_mean_color,
    _create_categorical_colors,
)

import numpy as np
import pandas as pd
from scipy.stats import entropy
from pandas.api.types import infer_dtype, is_categorical_dtype

import matplotlib.colors as c
import matplotlib.pyplot as plt

ColorLike = TypeVar("ColorLike")
_ERROR_NOT_ITERABLE = "Expected `{}` to be iterable, found type `{}`."
_ERROR_WRONG_SIZE = "Expected `{}` to be of size `{{}}`, found `{{}}`."

_HT_CELLS = 10  # head and tail cells to show
_HTML_REPR_THRESH = 100
_DUMMY_CELL = "...	"
_ORDER = "C"

class PrimingDegree(ModeEnum):  # noqa: D101
    KL_DIVERGENCE = auto()
    ENTROPY = auto()

class Lin(ModeEnum):  # noqa: D101
    REST = auto()
    OTHERS = auto()

class DistanceMeasure(ModeEnum):  # noqa: D101
    COSINE_SIM = auto()
    WASSERSTEIN_DIST = auto()
    KL_DIV = auto()
    JS_DIV = auto()
    MUTUAL_INFO = auto()
    EQUAL = auto()

class NormWeights(ModeEnum):  # noqa: D101
    SCALE = auto()
    SOFTMAX = auto()

class Reduction(ModeEnum):  # noqa: D101
    DIST = auto()
    SCALE = auto()

class LinKind(ModeEnum):  # noqa: D101
    MACROSTATES = auto()
    TERM_STATES = auto()
    ABS_PROBS = auto()

def _at_least_2d(array: np.ndarray, dim: int):
    return np.expand_dims(array, dim) if array.ndim < 2 else array

def wrap(numpy_func: Callable) -> Callable:
    """
    Wrap an numpy function.

    Modifies functionality of some function (e.g. ignoring `.squeeze`, retaining dimensions).

    Parameters
    ----------
    numpy_func
        Function to be wrapped.

    Returns
    -------
    Wrapped function which takes a :class:`cellrank.tl.Lineage` and return :class:`cellrank.tl.Lineage`.
    """

    @wraps(numpy_func)
    def decorator(array, *args, **kwargs):
        if not isinstance(array, Lineage):
            raise TypeError(
                f"Expected array to be of type `Lineage`, found `{type(array).__name__}`."
            )
        if fname == "squeeze":
            return array
        if fname == "array_repr":
            return repr(array)

        if "axis" in kwargs:
            axis = kwargs["axis"]
        elif axis_ix < len(args):
            axis = args[axis_ix]
        else:
            axis = default_axis

        res = np.array(numpy_func(array.X, *args, **kwargs), copy=False)

        # handle expand_dim
        if res.ndim > 2:
            return array

        # handle reductions
        if not res.shape:
            return Lineage(np.array([[res]]), names=[fname], colors=["grey"])
        if res.shape == array.shape:
            return Lineage(res, names=array.names, colors=array.colors)

        res = np.expand_dims(res, axis)

        is_t = int(array._is_transposed)
        if is_t:
            res = res.T

        lin = None
        if res.shape[0] == array.shape[is_t]:
            lin = Lineage(
                res,
                names=[f"{fname} of {', '.join(array.names)}"],
                colors=["grey"],
            )

        if res.shape[1] == array.shape[1 - is_t]:
            lin = Lineage(
                res, names=[f"{fname} of {n}" for n in array.names], colors=array.colors
            )

        if lin is not None:
            return lin.T if is_t else lin

        raise RuntimeError(
            f"Unable to interpret result of function `{fname}` called "
            f"with args `{args}`, kwargs: `{kwargs}`."
        )

    params = signature(numpy_func).parameters
    if "axis" in params:
        axis_ix = list(params.keys()).index("axis") - 1
        default_axis = params["axis"].default
    else:
        axis_ix = 256
        default_axis = None
    assert (
        axis_ix >= 0
    ), f"Expected argument `'axis'` not to be first for function `{numpy_func.__name__}`."

    fname = numpy_func.__name__
    if fname == "amin":
        fname = "min"
    elif fname == "amax":
        fname = "max"

    return decorator

def _register_handled_functions():
    handled_fns = {}

    for attrname in dir(np):
        fn = getattr(np, attrname)
        if isinstance(fn, FunctionType):
            try:
                sig = signature(fn)
                if "axis" in sig.parameters.keys():
                    handled_fns[fn] = wrap(fn)
            except ValueError:
                pass

    handled_fns.pop(np.expand_dims, None)

    handled_fns[np.allclose] = wrap(np.allclose)
    handled_fns[np.array_repr] = wrap(np.array_repr)
    handled_fns[entropy] = wrap(entropy)  # qol change

    return handled_fns

_HANDLED_FUNCTIONS = _register_handled_functions()

class LineageMeta(type):
    """
    Metaclass for Lineaage.

    Registers functions which are handled by us and overloads common attributes, such as `.sum` with these functions.
    """

    __overloaded_functions__ = dict(  # noqa
        sum=np.sum,
        mean=np.mean,
        min=np.min,
        argmin=np.argmin,
        max=np.max,
        argmax=np.argmax,
        std=np.std,
        var=np.var,
        sort=np.sort,
        squeeze=np.squeeze,
        entropy=entropy,
    )

    def __new__(cls, clsname, superclasses, attributedict):  # noqa
        res = type.__new__(cls, clsname, superclasses, attributedict)
        for attrname, fn in LineageMeta.__overloaded_functions__.items():
            wrapped_fn = _HANDLED_FUNCTIONS.get(fn, None)
            if wrapped_fn:
                setattr(res, attrname, wrapped_fn)

        return res

class Lineage(np.ndarray, metaclass=LineageMeta):
    """
    Lightweight :class:`numpy.ndarray` wrapper that adds names and colors.

    Parameters
    ----------
    input_array
        Input array containing lineage probabilities, each lineage being stored in a column.
    names
        Names of the lineages.
    colors
        Colors of the lineages.
    """

    def __new__(
        cls,
        input_array: np.ndarray,
        *,
        names: Iterable[str],
        colors: Optional[Iterable[ColorLike]] = None,
    ) -> "Lineage":
        """Create and return a new object."""

        if not isinstance(input_array, np.ndarray):
            raise TypeError(
                f"Input array must be of type `numpy.ndarray`, found `{type(input_array).__name__!r}`."
            )

        if input_array.ndim == 1:
            input_array = np.expand_dims(input_array, -1)
        elif input_array.ndim > 2:
            raise ValueError(
                f"Input array must be 2-dimensional, found `{input_array.ndim}`."
            )

        if input_array.shape[0] == 0:
            raise ValueError("Expected number cells to be at least 1, found 0.")
        if input_array.shape[1] == 0:
            raise ValueError("Expected number of lineages to be at least 1, found 0.")

        obj = np.array(input_array, copy=True).view(cls)
        obj._n_lineages = obj.shape[1]
        obj._is_transposed = False
        obj.names = names  # these always create a copy, which is a good thing
        obj.colors = colors

        return obj

    def __array_finalize__(self, obj) -> None:
        if obj is None:
            return

        _names = getattr(obj, "_names", None)
        if _names is not None:
            self._names = _names
            self._n_lineages = len(_names)
            self._names_to_ixs = {n: i for i, n in enumerate(self.names)}
        else:
            self._names = None
            self._names_to_ixs = None
            self._n_lineages = getattr(
                obj, "_n_lineages", obj.shape[1] if obj.ndim == 2 else 0
            )

        self._colors = getattr(obj, "colors", None)
        self._is_transposed = getattr(obj, "_is_transposed", False)

    def _mixer(self, rows, mixtures):
        def update_entries(key):
            if key:
                res.append(self[rows, key].X.sum(1))
                # item = (key, rows) if self._is_transposed else (rows, key)
                # res.append(self[item].X.sum(int(not self._is_transposed)))
                names.append(" or ".join(self.names[key]))
                colors.append(_compute_mean_color(self.colors[key]))

        lin_kind = [_ for _ in mixtures if isinstance(_, Lin)]
        if len(lin_kind) > 1:
            raise ValueError(
                f"`Lin` enum is allowed only once in the expression, found `{len(lin_kind)}`."
            )

        keys = [
            tuple(
                self._maybe_convert_names(
                    _convert_lineage_name(mixture), default=mixture
                )
            )
            if isinstance(mixture, str)
            else (mixture,)
            for mixture in mixtures
            if not (isinstance(mixture, Lin))
        ]
        keys = _unique_order_preserving(keys)

        # check the `keys` are unique
        overlap = [set(ks) for ks in keys]
        for c1, c2 in combinations(overlap, 2):
            overlap = c1 & c2
            if overlap:
                raise ValueError(
                    f"Found overlapping keys: `{self.names[list(overlap)]}`."
                )

        seen = set()
        names, colors, res = [], [], []
        for key in map(list, keys):
            seen.update(self.names[key])
            update_entries(key)

        if len(lin_kind) == 1:
            lin_kind = lin_kind[0]
            keys = [i for i, n in enumerate(self.names) if n not in seen]

            if lin_kind == Lin.OTHERS:
                for key in keys:
                    update_entries([key])
            elif lin_kind == Lin.REST:
                update_entries(keys)
                if keys:
                    names[-1] = str(lin_kind)
            else:
                raise NotImplementedError(
                    f"Mixing `{lin_kind}` is not yet implemented."
                )

        res = np.stack(res, axis=-1)
        return Lineage(res, names=names, colors=colors)

    def __getitem__(self, item) -> "Lineage":
        was_transposed = False
        if self._is_transposed:
            was_transposed = True
            self = self.T
            if isinstance(item, tuple):
                item = item[::-1]

        obj = self.__getitem(item)

        return obj.T if was_transposed else obj

    def __getitem(self, item):
        if isinstance(item, tuple):
            if len(item) > 2:
                raise ValueError(
                    f"Expected key to be of length `2`, found `{len(item)}`."
                )

            item = list(item)
            if item[0] is Ellipsis or item[0] is None:
                item[0] = range(self.shape[0])
            if len(item) == 2 and (item[1] is Ellipsis or item[1] is None):
                item[1] = range(self.shape[1])
            item = tuple(item)

        is_tuple_len_2 = (
            isinstance(item, tuple)
            and len(item) == 2
            and isinstance(
                item[0], (int, np.integer, range, slice, tuple, list, np.ndarray)
            )
        )
        if is_tuple_len_2:
            rows, col = item

            if isinstance(col, (int, np.integer, str)):
                col = [col]

            try:
                # slicing an array where row/col are like 2D indices
                if 1 < len(col) == len(rows) and len(rows) == self.shape[0]:
                    col = self._maybe_convert_names(col, make_unique=False)
                    return Lineage(
                        # never remove this expand_dims - it's critical
                        np.expand_dims(self.X[rows, col], axis=-1),
                        names=["mixture"],
                        colors=["#000000"],
                    )
            except TypeError:  # because of range
                pass

            if isinstance(col, (list, tuple, np.ndarray)):
                if any(
                    map(
                        lambda i: isinstance(i, Lin)
                        or (isinstance(i, str) and ("," in i or "or" in i)),
                        col,
                    )
                ):
                    return self._mixer(rows, col)
                col = self._maybe_convert_names(col)
                item = rows, col
        else:
            if isinstance(item, (int, np.integer, str)):
                item = [item]

            col = range(len(self.names))
            if isinstance(item, (tuple, list, np.ndarray)):
                if any(
                    map(
                        lambda i: isinstance(i, Lin)
                        or (isinstance(i, str) and ("," in i or "or" in i)),
                        item,
                    )
                ):
                    return self._mixer(slice(None, None, None), item)
                elif any(map(lambda i: isinstance(i, str), item)):
                    item = (slice(None, None, None), self._maybe_convert_names(item))
                    col = item[1]

        shape, row_order, col_order = None, None, None
        if (
            is_tuple_len_2
            and not isinstance(item[0], slice)
            and not isinstance(item[1], slice)
        ):
            item_0 = (
                np.array(item[0]) if not isinstance(item[0], np.ndarray) else item[0]
            )
            item_1 = (
                np.array(item[1]) if not isinstance(item[1], np.ndarray) else item[1]
            )
            item_0 = _at_least_2d(item_0, -1)
            item_1 = _at_least_2d(item_1, 0)

            if item_1.dtype == bool:
                if item_0.dtype != bool:
                    if not issubclass(item_0.dtype.type, np.integer):
                        raise TypeError(f"Invalid type `{item_0.dtype.type}`.")
                    row_order = (
                        item_0[:, 0]
                        if item_0.shape[0] == self.shape[0]
                        else np.argsort(np.argsort(item_0[:, 0]))
                    )
                    item_0 = _at_least_2d(np.isin(np.arange(self.shape[0]), item_0), -1)
                item = item_0 * item_1
                shape = np.max(np.sum(item, axis=0)), np.max(np.sum(item, axis=1))
                col = np.where(np.all(item_1, axis=0))[0]
            elif item_0.dtype == bool:
                if item_1.dtype != bool:
                    if not issubclass(item_1.dtype.type, np.integer):
                        raise TypeError(f"Invalid type `{item_1.dtype.type}`.")
                    col_order = (
                        item_1[0, :]
                        if item_1.shape[1] == self.shape[1]
                        else np.argsort(np.argsort(item_1[0, :]))
                    )
                    item_1 = _at_least_2d(np.isin(np.arange(self.shape[1]), item_1), 0)

                item = item_0 * item_1
                shape = np.max(np.sum(item, axis=0)), np.max(np.sum(item, axis=1))
                col = np.where(np.all(item_1, axis=0))[0]
            else:
                item = (item_0, item_1)

        obj = super().__getitem__(item)

        if shape is not None:  # keep the resulting shape
            obj = obj.reshape(shape)
        if row_order is not None:
            obj = obj[row_order, :]
        if col_order is not None:
            obj = obj[:, col_order]

        if isinstance(obj, Lineage):
            obj._names = np.atleast_1d(self.names[col])
            obj._colors = np.atleast_1d(self.colors[col])
            if col_order is not None:
                obj._names = obj._names[col_order]
                obj._colors = obj._colors[col_order]
            obj._names_to_ixs = {name: i for i, name in enumerate(obj._names)}

        return obj

    def __array_function__(self, func, types, args, kwargs):
        if func not in _HANDLED_FUNCTIONS:
            return NotImplemented
        # Note: this allows subclasses that don't override
        # __array_function__ to handle MyArray objects
        if not all(issubclass(t, type(self)) for t in types):
            return NotImplemented

        return _HANDLED_FUNCTIONS[func](*args, **kwargs)

    @property
    def names(self) -> np.ndarray:
        """Lineage names. Must be unique."""
        return self._names

    @names.setter
    def names(self, value: Iterable[str]):
        if not isinstance(value, Iterable):
            raise TypeError(_ERROR_NOT_ITERABLE.format("names", type(value).__name__))

        value = [v if isinstance(v, str) else str(v) for v in value]
        value = self._check_shape(value, _ERROR_WRONG_SIZE.format("names"))

        if len(set(value)) != len(value):
            raise ValueError(f"Not all lineage names are unique: `{value}`.")

        self._names = self._prepare_annotation(value)
        self._names_to_ixs = {name: ix for ix, name in enumerate(self.names)}

    @property
    def colors(self) -> np.ndarray:
        """Lineage colors."""
        return self._colors

    @colors.setter
    def colors(self, value: Optional[Iterable[ColorLike]]):
        if value is None:
            value = _create_categorical_colors(self._n_lineages)
        elif not isinstance(value, Iterable):
            raise TypeError(_ERROR_NOT_ITERABLE.format("colors", type(value).__name__))

        value = self._check_shape(value, _ERROR_WRONG_SIZE.format("colors"))
        self._colors = self._prepare_annotation(
            value,
            checker=c.is_color_like,
            transformer=c.to_hex,
            checker_msg="Value `{}` is not a valid color.",
        )

    def _maybe_convert_names(
        self,
        names: Iterable[Union[int, str, bool]],
        is_singleton: bool = False,
        default: Optional[Union[int, str]] = None,
        make_unique: bool = True,
    ) -> Union[int, List[int], List[bool]]:
        if all(map(lambda n: isinstance(n, (bool, np.bool_)), names)):
            return list(names)
        res = []
        for name in names:
            if isinstance(name, str):
                if name in self._names_to_ixs:
                    name = self._names_to_ixs[name]
                elif default is not None:
                    if isinstance(default, str):
                        if default not in self._names_to_ixs:
                            raise KeyError(
                                f"Invalid lineage name: `{name}`. "
                                f"Valid names are: `{list(self.names)}`."
                            )
                        name = self._names_to_ixs[default]
                    else:
                        name = default
                else:
                    raise KeyError(
                        f"Invalid lineage name `{name}`. "
                        f"Valid names are: `{list(self.names)}`."
                    )
            res.append(name)

        if make_unique:
            res = _unique_order_preserving(res)

        return res[0] if is_singleton else res

    def _check_shape(
        self, array: Iterable[Union[str, ColorLike]], msg: str
    ) -> List[Union[str, ColorLike]]:
        array = list(array)
        if len(array) != self._n_lineages:
            raise ValueError(msg.format(self._n_lineages, len(array)))

        return array

    def _prepare_annotation(
        self,
        array: List[str],
        checker: Optional[Callable] = None,
        transformer: Optional[Callable] = None,
        checker_msg: Optional[str] = None,
    ) -> np.ndarray:
        if checker:
            assert checker_msg, "Please provide a message when `checker` is not `None`."
            for v in array:
                if not checker(v):
                    raise ValueError(checker_msg.format(v))

        if transformer is not None:
            array = np.array([transformer(v) for v in array])

        return np.array(array)

    @property
    def X(self) -> np.ndarray:
        """Convert self to numpy array, losing names and colors."""
        return np.array(self, copy=False)

    @property
    def T(self):
        """Transpose of self."""
        obj = self.transpose()
        obj._is_transposed = not self._is_transposed
        return obj

    def view(self, dtype=None, type=None) -> "LineageView":
        """Return a view of self."""
        return LineageView(self)

    def __repr__(self) -> str:
        return f'{super().__repr__()[:-1]},\n  names([{", ".join(self.names)}]))'

    def __str__(self):
        return f'{super().__str__()}\n names=[{", ".join(self.names)}]'

    @property
    def _fmt(self) -> Callable:
        if np.issubdtype(self.dtype, float):
            return "{:.06f}".format
        return "{}".format

    def _repr_html_(self) -> str:
        def format_row(r):
            rng = (
                range(self.shape[1])
                if not self._is_transposed
                or (self._is_transposed and self.shape[1] <= _HTML_REPR_THRESH)
                else list(range(_HT_CELLS))
                + [...]
                + list(range(self.shape[1] - _HT_CELLS - 1, self.shape[1] - 1))
            )

            cells = "".join(
                f"" f"{self._fmt(self.X[r, c])}" f"	"
                if isinstance(c, int)
                else _DUMMY_CELL
                for c in rng
            )
            return f"
{(names[r] if self._is_transposed else '') + cells}
"

        def dummy_row():
            values = "".join(_DUMMY_CELL for _ in range(self.shape[1]))
            return f"
{values}
"

        if self.names is None or self.colors is None:
            raise RuntimeError(
                f"Name or colors are `None`. This can happen when running `array.view({type(self).__name__}`."
            )

        if self.ndim != 2:
            return repr(self)

        styles = [
            f"'background-color: {bg}; color: {fg}; text-align: center; word-wrap: break-word; max-width: 100px'"
            for bg, fg in map(_get_bg_fg_colors, self.colors)
        ]
        names = [f"{n}	" for n, style in zip(self.names, styles)]
        header = f"
{''.join(names)}
"

        if self.shape[0] > _HTML_REPR_THRESH:
            body = "".join(format_row(i) for i in range(_HT_CELLS))
            body += dummy_row()
            body += "".join(
                format_row(i)
                for i in range(self.shape[0] - _HT_CELLS - 1, self.shape[0] - 1)
            )
        else:
            body = "".join(format_row(i) for i in range(self.shape[0]))

        cells = "cells" if self.shape[0] > 1 else "cell"
        lineages = "lineages" if self.shape[1] > 1 else "lineage"
        if self._is_transposed:
            cells, lineages = lineages, cells
        metadata = f"
{self.shape[0]} {cells} x {self.shape[1]} {lineages}
"

        if self._is_transposed:
            header = ""
        return (
            f"
"
            f"
{header}{body}
{metadata}"
            f"
"
        )

    def __format__(self, format_spec):
        if self.shape == (1, 1):
            return format_spec.format(self.X[0, 0])
        if self.shape == (1,):
            return format_spec.format(self.X[0])
        return NotImplemented

    def __setstate__(self, state):
        *state, names, colors, is_t = state
        names = names[-1]
        colors = colors[-1]

        self._names = np.empty(names[1])
        self._colors = np.empty(colors[1])

        super().__setstate__(tuple(state))
        self._names.__setstate__(tuple(names))
        self._colors.__setstate__(tuple(colors))

        self._is_transposed = is_t
        self._n_lineages = len(self.names)
        self._names_to_ixs = {name: ix for ix, name in enumerate(self.names)}

    def __reduce__(self):
        res = list(super().__reduce__())

        names = self.names.__reduce__()
        colors = self.colors.__reduce__()
        res[-1] += (names, colors, self._is_transposed)

        return tuple(res)

    def copy(self, _="C") -> "Lineage":
        """Return a copy of itself."""
        obj = Lineage(
            self.T if self._is_transposed else self,
            names=np.array(self.names, copy=True, order=_ORDER),
            colors=np.array(self.colors, copy=True, order=_ORDER),
        )
        return obj.T if self._is_transposed else obj

    def __copy__(self):
        return self.copy()

    @d.get_full_description(base="lin_pd")
    @d.get_sections(base="lin_pd", sections=["Parameters", "Returns"])
    def priming_degree(
        self,
        method: Literal["kl_divergence", "entropy"] = "kl_divergence",
        early_cells: Optional[np.ndarray] = None,
    ) -> np.ndarray:
        """
        Compute the degree of lineage priming.

        It returns a score in `[0, 1]` where `0` stands for naive and `1` stands for committed.

        Parameters
        ----------
        method
            The method used to compute the degree of lineage priming. Valid options are:

                - `'kl_divergence'` - as in :cite:`velten:17`, computes KL-divergence between the fate probabilities of
                  a cell and the average fate probabilities. Computation of average fate probabilities can be restricted
                  to a set of user-defined ``early_cells``.
                - `'entropy'` - as in :cite:`setty:19`, computes entropy over a cell's fate probabilities.
        early_cells
            Cell IDs or a mask marking early cells. If `None`, use all cells.
            Only used when ``method = 'kl_divergence'``.

        Returns
        -------
        The priming degree.
        """
        early_cells = (
            np.ones((len(self),), dtype=np.bool_)
            if early_cells is None
            else np.asarray(early_cells)
        )
        if not np.issubdtype(early_cells.dtype, np.bool_):
            early_cells = np.unique(early_cells)

        method = PrimingDegree(method)
        probs = self.X
        early_subset = probs[early_cells, :]
        if not len(early_subset):
            raise ValueError("No early cells have been specified.")

        with np.errstate(divide="ignore", invalid="ignore"):
            if method == PrimingDegree.KL_DIVERGENCE:
                probs = np.nan_to_num(
                    np.sum(
                        probs * np.log2(probs / np.mean(early_subset, axis=0)), axis=1
                    ),
                    nan=1.0,
                    copy=False,
                )
            elif method == PrimingDegree.ENTROPY:
                probs = entropy(probs, axis=1)
                probs = np.max(probs) - probs
            else:
                raise NotImplementedError(f"Method `{method}` is not yet implemented")

            minn, maxx = np.min(probs), np.max(probs)
            return (probs - minn) / (maxx - minn)

    @d.dedent
    def plot_pie(
        self,
        reduction: Callable,
        title: Optional[str] = None,
        legend_loc: Optional[str] = "on data",
        legend_kwargs: Mapping = MappingProxyType({}),
        figsize: Optional[Tuple[float, float]] = None,
        dpi: Optional[float] = None,
        save: Optional[Union[Path, str]] = None,
        **kwargs,
    ) -> None:
        """
        Plot a pie chart visualizing aggregated lineage probabilities.

        Parameters
        ----------
        reduction
            Function that will be applied lineage-wise.
        title
            Title of the figure.
        legend_loc
            Location of the legend. If `None`, it is not shown.
        legend_kwargs
            Keyword arguments for :func:`matplotlib.axes.Axes.legend`.
        %(plotting)s

        Returns
        -------
        %(just_plots)s
        """

        if len(self.names) == 1:
            raise ValueError("Cannot plot pie chart for only 1 lineage.")

        fig, ax = plt.subplots(figsize=figsize, dpi=dpi)

        if "autopct" not in kwargs:
            autopct_found = False
            autopct = (
                "{:.1f}%".format
            )  # we don't really care, we don't shot the pct, but the value
        else:
            autopct_found = True
            autopct = kwargs.pop("autopct")

        if title is None:
            title = reduction.__name__ if hasattr(reduction, "__name__") else None

        reduction = reduction(self, axis=int(self._is_transposed)).X.squeeze()
        reduction_norm = reduction / np.sum(reduction)

        wedges, texts, *autotexts = ax.pie(
            reduction_norm.squeeze(),
            labels=self.names if legend_loc == "on data" else None,
            autopct=autopct,
            wedgeprops={"edgecolor": "w"},
            colors=self.colors,
            **kwargs,
        )

        # if autopct is not None
        if len(autotexts):
            autotexts = autotexts[0]
            for name, at in zip(self.names, autotexts):
                ix = self._names_to_ixs[name]
                at.set_color(_get_bg_fg_colors(self.colors[ix])[1])
                if not autopct_found:
                    at.set_text(f"{reduction[ix]:.4f}")

        if legend_loc not in (None, "none", "on data"):
            ax.legend(
                wedges,
                self.names,
                title="lineages",
                loc=legend_loc,
                **legend_kwargs,
            )

        ax.set_title(title)
        ax.set_aspect("equal")

        if save is not None:
            save_fig(fig, save)

    @d.dedent
    @inject_docs(m=Reduction, dm=DistanceMeasure, nw=NormWeights)
    def reduce(
        self,
        *keys: str,
        mode: Literal["dist", "scale"] = Reduction.DIST,
        dist_measure: Literal[
            "cosine_sim", "wasserstein_dist", "kl_div", "js_div", "mutual_info", "equal"
        ] = DistanceMeasure.MUTUAL_INFO,
        normalize_weights: Literal["scale", "softmax"] = NormWeights.SOFTMAX,
        softmax_scale: float = 1,
        return_weights: bool = False,
    ) -> Union["Lineage", Tuple["Lineage", Optional[pd.DataFrame]]]:
        """
        Subset states and normalize them so that they again sum to 1.

        Parameters
        ----------
        keys
            List of keys that define the states, to which this object will be reduced by projecting the values
            of the other states.
        mode
            Reduction mode to use. Valid options are:

                - `{m.DIST!r}` - use a distance measure ``dist_measure`` to compute weights.
                - `{m.SCALE!r}` - just rescale the values.
        dist_measure
            Used to quantify similarity between query and reference states. Valid options are:

                - `{dm.COSINE_SIM!r}` - cosine similarity.
                - `{dm.WASSERSTEIN_DIST!r}` - Wasserstein distance.
                - `{dm.KL_DIV!r}` - Kullback–Leibler divergence.
                - `{dm.JS_DIV!r}` - Jensen–Shannon divergence.
                - `{dm.MUTUAL_INFO!r}` - mutual information.
                - `{dm.EQUAL!r}` - equally redistribute the mass among the rest.

            Only use when ``mode = {m.DIST!r}``.
        normalize_weights
            How to row-normalize the weights. Valid options are:

                - `{nw.SCALE!r}` - divide by the sum.
                - `{nw.SOFTMAX!r}`- use a softmax.

            Only use when ``mode = {m.DIST!r}``.
        softmax_scale
            Scaling factor in the softmax, used for normalizing the weights to sum to `1`.
        return_weights
            If `True`, a :class:`pandas.DataFrame` of the weights used for the projection is also returned.
            If ``mode = {m.SCALE!r}``, the weights will be `None`.

        Returns
        -------
        The lineage object, reduced to the %(initial_or_terminal)s states.
        The weights used for the projection of shape ``(n_query, n_reference)``, if ``return_weights = True``.
        """

        mode = Reduction(mode)
        dist_measure = DistanceMeasure(dist_measure)
        normalize_weights = NormWeights(normalize_weights)

        if self._is_transposed:
            raise RuntimeError("This method works only on non-transposed lineages.")

        if not len(keys):
            raise ValueError("Unable to perform the reduction, no keys specified.")

        # check the lineage object
        if not np.allclose(np.sum(self, axis=1).X, 1.0):
            raise ValueError("Memberships do not sum to one row-wise.")

        if len(keys) == 1:
            tmp = self[:, keys]
            return Lineage(
                np.ones((self.shape[0], 1), dtype=self.dtype),
                names=tmp.names,
                colors=tmp.colors,
            )

        # check input parameters
        if return_weights and mode == Reduction.SCALE:
            logg.warning(
                f"If `mode={mode!r}`, no weights are computed. Returning `None`"
            )

        reference = self[:, keys]
        rest = [
            k for k in self.names if all(map(lambda rk: k not in rk, reference.names))
        ]
        if not len(rest):
            logg.warning(
                "Unable to perform reduction because all keys have been selected. Returning combined object only"
            )
            return (reference.copy(), None) if return_weights else reference.copy()

        query = self[:, rest]

        if mode == Reduction.SCALE:
            reference = _row_normalize(reference)
        elif mode == Reduction.DIST:
            # compute a set of weights of shape (n_query x n_reference)
            if dist_measure == DistanceMeasure.COSINE_SIM:
                weights = _cosine_sim(reference.X, query.X)
            elif dist_measure == DistanceMeasure.WASSERSTEIN_DIST:
                weights = _wasserstein_dist(reference.X, query.X)
            elif dist_measure == DistanceMeasure.KL_DIV:
                weights = _kl_div(reference.X, query.X)
            elif dist_measure == DistanceMeasure.JS_DIV:
                weights = _js_div(reference.X, query.X)
            elif dist_measure == DistanceMeasure.MUTUAL_INFO:
                weights = _mutual_info(reference.X, query.X)
            elif dist_measure == DistanceMeasure.EQUAL:
                weights = _row_normalize(np.ones((query.shape[1], reference.shape[1])))
            else:
                raise NotImplementedError(
                    f"Distance measure `{dist_measure}` is not yet implemented."
                )

            # make some checks on the weights
            if weights.shape != (query.shape[1], reference.shape[1]):
                raise ValueError(
                    f"Expected weight matrix to be of shape `({query.shape[1]}, {reference.shape[1]})`, "
                    f"found `{weights.shape}`."
                )
            if not np.isfinite(weights).all():
                raise ValueError(
                    "Weights matrix contains elements that are not finite."
                )
            if (weights < 0).any():
                raise ValueError("Weights matrix contains negative elements.")

            if (weights == 0).any():
                logg.warning("Weights matrix contains exact zeros.")

            # normalize the weights to row-sum to one
            if normalize_weights == NormWeights.SCALE:
                weights_n = _row_normalize(weights)
            elif normalize_weights == NormWeights.SOFTMAX:
                weights_n = _softmax(_row_normalize(weights), softmax_scale)
            else:
                raise NotImplementedError(
                    f"Normalization method `{normalize_weights}` is yet implemented."
                )

            # check that the weights row-sum to one now
            if not np.allclose(weights_n.sum(1), 1.0):
                raise ValueError("Weights do not sum to 1 row-wise.")

            # use the weights to re-distribute probability mass form query to reference
            for i, w in enumerate(weights_n):
                reference += np.dot(query[:, i].X, w[None, :])
        else:
            raise NotImplementedError(
                f"Reduction mode `{mode}` is not yet implemented."
            )

        # check that the lineages row-sum to one now
        if not np.allclose(reference.sum(1), 1.0):
            raise ValueError("Reduced lineage rows do not sum to 1.")

        # potentially create a weights-df and return everything
        if return_weights:
            if mode == Reduction.DIST:
                return (
                    reference,
                    pd.DataFrame(
                        data=weights_n, columns=reference.names, index=query.names
                    ),
                )
            return reference, None

        return reference

    @classmethod
    @d.dedent
    @inject_docs(lk=LinKind)
    def from_adata(
        cls,
        adata: AnnData,
        backward: bool = False,
        kind: Literal["macrostates", "term_states", "abs_probs"] = LinKind.ABS_PROBS,
    ) -> "Lineage":
        """
        Reconstruct :class:`cellrank.tl.Lineage` from :class:`anndata.AnnData`.

        Parameters
        ----------
        %(adata)s
        %(backward)s
        kind
            Which kind of object to reconstruct. Valid options are:

                - `{lk.MACROSTATES!r}`- macrostates memberships from :class:`cellrank.tl.estimators.GPCCA`.
                - `{lk.TERM_STATES!r}`- terminal states memberships from :class:`cellrank.tl.estimators.GPCCA`.
                - `{lk.ABS_PROBS!r}`- the absorption probabilities.

        Returns
        -------
        The reconstructed lineage object.
        """
        kind = LinKind(kind)
        if kind == LinKind.MACROSTATES:
            nkey = Key.obs.macrostates(backward)
            key = Key.obsm.memberships(nkey)
        elif kind == LinKind.TERM_STATES:
            nkey = Key.obs.term_states(backward)
            key = Key.obsm.memberships(nkey)
        elif kind == LinKind.ABS_PROBS:
            nkey = Key.obs.term_states(backward)
            key = Key.obsm.abs_probs(backward)
        else:
            raise NotImplementedError(f"Lineage kind `{kind}` is not yet implemented.")

        ckey = Key.uns.colors(nkey)

        data: Optional[Union[np.ndarray, Lineage]] = copy(adata.obsm.get(key, None))
        if data is None:
            raise KeyError(f"Unable to find lineage data in `adata.obsm[{key!r}]`.")
        if isinstance(data, Lineage):
            return data
        if data.ndim != 2:
            raise ValueError(f"Expected 2 dimensional data, found `{data.ndim}`.")

        states = adata.obs.get(nkey, None)
        if states is None:
            logg.warning(
                f"Unable to find states in `adata.obs[{nkey!r}]`. Using default names"
            )
        elif not is_categorical_dtype(states):
            logg.warning(
                f"Expected `adata.obs[{key!r}]` to be `categorical`, "
                f"found `{infer_dtype(adata.obs[nkey])}`. Using default names"
            )
        else:
            states = list(states.cat.categories)
            if len(states) != data.shape[1]:
                logg.warning(
                    f"Expected to find `{data.shape[1]}` names, found `{len(states)}`. "
                    f"Using default names"
                )
        if states is None or len(states) != data.shape[1]:
            states = [str(i) for i in range(data.shape[1])]

        colors = adata.uns.get(ckey, None)
        if colors is None:
            logg.warning(
                f"Unable to find colors in `adata.uns[{ckey!r}]`. "
                f"Using default colors"
            )
        elif len(colors) != data.shape[1]:
            logg.warning(
                f"Expected to find `{data.shape[1]}` colors, found `{len(colors)}`. "
                f"Using default colors"
            )
            colors = None

        return Lineage(data, names=states, colors=colors)

class LineageView(Lineage):
    """View of :class:`cellrank.tl.Lineage`."""

    def __new__(cls, lineage: Lineage) -> "LineageView":
        """Create a LineageView."""
        if not isinstance(lineage, Lineage):
            raise TypeError(
                f"Cannot create a `{cls.__name__}` of `{type(lineage).__name__}`."
            )

        view = np.array(lineage, copy=False).view(cls)
        view._owner = lineage
        view._names = lineage.names
        view._n_lineages = len(view.names)
        view._names_to_ixs = lineage._names_to_ixs
        view._colors = lineage.colors
        view._is_transposed = lineage._is_transposed

        return view

    @property
    def names(self) -> np.ndarray:
        """Lineage names."""
        return super().names

    @property
    def owner(self) -> Lineage:
        """Return the lineage associated with this view."""
        return self._owner

    @names.setter
    def names(self, _):
        raise RuntimeError(f"Unable to set names of `{type(self).__name__}`.")

    @property
    def colors(self) -> np.ndarray:
        """Lineage colors."""
        return super().colors

    @colors.setter
    def colors(self, _):
        raise RuntimeError(f"Unable to set colors of `{type(self).__name__}`.")

    def view(self, dtype=None, type=None) -> "LineageView":
        """Return self."""
        return self

    def copy(self, _="C") -> Lineage:
        """Return a copy of itself."""
        was_trasposed = False
        if self._is_transposed:
            self = self.T
            was_trasposed = True

        obj = Lineage(
            self,
            names=np.array(self.names, copy=True, order=_ORDER),
            colors=np.array(self.colors, copy=True, order=_ORDER),
        )
        return obj.T if was_trasposed else obj

def _remove_zero_rows(a: np.ndarray, b: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    if a.shape[0] != b.shape[0]:
        raise ValueError("Lineage objects have unequal cell numbers")

    bool_a = (a == 0).any(axis=1)
    bool_b = (b == 0).any(axis=1)
    mask = ~np.logical_or(bool_a, bool_b)

    logg.warning(
        f"Removed {a.shape[0] - np.sum(mask)} rows because they contained zeros"
    )

    return a[mask, :], b[mask, :]

def _softmax(X, beta: float = 1):
    return np.exp(X * beta) / np.expand_dims(np.sum(np.exp(X * beta), axis=1), -1)

def _row_normalize(X):
    if isinstance(X, Lineage):
        return X / X.sum(1)  # Lineage is shape-preserving
    return X / np.expand_dims(X.sum(1), -1)

def _col_normalize(X, norm_ord=2):
    from numpy.linalg import norm

    return X / norm(X, ord=norm_ord, axis=0)

def _cosine_sim(reference, query):
    # the cosine similarity is symmetric

    # normalize these to have 2-norm 1
    reference_n, query_n = _col_normalize(reference, 2), _col_normalize(query, 2)

    return (reference_n.T @ query_n).T

def _point_wise_distance(reference, query, distance):
    # utility function for all point-wise distances/divergences
    # take care of rows that contain zeros
    reference_no_zero, query_no_zero = _remove_zero_rows(reference, query)

    # normalize these to be valid probability distributions (column-wise)
    reference_n, query_n = (
        _col_normalize(reference_no_zero, 1),
        _col_normalize(query_no_zero, 1),
    )

    # loop over query and reference columns and compute pairwise wasserstein distances
    weights = np.zeros((query.shape[1], reference.shape[1]))
    for i, q_d in enumerate(query_n.T):
        for j, r_d in enumerate(reference_n.T):
            weights[i, j] = 1.0 / distance(q_d, r_d)

    return weights

def _wasserstein_dist(reference, query):
    # the wasserstein distance is symmetric
    from scipy.stats import wasserstein_distance

    return _point_wise_distance(reference, query, wasserstein_distance)

def _kl_div(reference, query):
    # the KL divergence is not symmetric
    from scipy.stats import entropy

    return _point_wise_distance(reference, query, entropy)

def _js_div(reference, query):
    # the js divergence is symmetric
    from scipy.spatial.distance import jensenshannon

    return _point_wise_distance(reference, query, jensenshannon)

def _mutual_info(reference, query):
    # mutual information is not symmetric. We don't need to normalise the vectors, it's invariant under scaling.
    from sklearn.feature_selection import mutual_info_regression

    weights = np.zeros((query.shape[1], reference.shape[1]))
    for i, target in enumerate(query.T):
        weights[i, :] = mutual_info_regression(reference, target)

    return weights


# -*- coding: utf-8 -*-
"""
SVGPATH2MPL
~~~~~~~~~~~
Parse SVG path definition strings into matplotlib Path objects.
A path in SVG is defined by a 'path' element which contains a
``d="(path data)"`` attribute that contains moveto, line, curve (both
cubic and quadratic Béziers), arc and closepath instructions. See the SVG
Path specification at <https://www.w3.org/TR/SVG/paths.html>.

:copyright: (c) 2016, Nezar Abdennur.
:license: BSD.

"""
from __future__ import division, print_function
from math import sin, cos, sqrt, degrees, radians, acos
import re

from matplotlib.path import Path
import matplotlib.transforms as transforms
import numpy as np

__version__ = '1.0.0'
__all__ = ['parse_path']

COMMANDS = set('MmZzLlHhVvCcSsQqTtAa')
UPPERCASE = set('MZLHVCSQTA')

COMMAND_RE = re.compile(r"([MmZzLlHhVvCcSsQqTtAa])")
FLOAT_RE = re.compile(r"[-+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?")

COMMAND_CODES = {
    'M': (Path.MOVETO,),    # moveto
    'L': (Path.LINETO,),    # line
    'H': (Path.LINETO,),    # shorthand for horizontal line
    'V': (Path.LINETO,),    # shorthand for vertical line
    'Q': (Path.CURVE3,)*2,  # quadratic bezier
    'T': (Path.CURVE3,)*2,  # shorthand for smooth quadratic bezier
    'C': (Path.CURVE4,)*3,  # cubic bezier
    'S': (Path.CURVE4,)*3,  # shorthand for smooth cubic bezier
    'Z': (Path.CLOSEPOLY,), # closepath
    'A': None               # arc
}

PARAMS = {
    'M': 2, # moveto
    'L': 2, # line
    'H': 1, # shorthand for horizontal line
    'V': 1, # shorthand for vertical line
    'Q': 4, # quadratic bezier
    'T': 4, # shorthand for smooth quadratic bezier
    'C': 6, # cubic bezier
    'S': 6, # shorthand for smooth cubic bezier
    'Z': 0, # closepath
    'A': 7  # arc
}

def endpoint_to_center(start, radius, rotation, large, sweep, end):
    """
    Translates the "endpoint" parameterization of an elliptical arc used by
    the SVG spec to the "center" parameterization.

    Parameters
    ----------
    start : complex
        Starting point (x1, y1).
    radius : complex
        Two elliptical radii (rx, ry).
    rotation : float
        Angle from the x-axis of the current coordinate system to the x-axis of
        the ellipse.
    large : bool
        False if an arc spanning < 180 degrees is to be drawn, True if an arc
        spanning >= 180 degrees is to be drawn.
    sweep : bool
        If sweep-flag is True, then the arc will be drawn in a "positive-angle"
        direction from start to end.
    end : complex
        End point (x2, y2).

    Returns
    -------
    radius : complex
        Radii of the ellipse (rx, ry), potentially corrected if out-of-range.
    center : complex
        Center of the ellipse (xc, yc).
    theta1, theta2 : float
        Start and end angles of an arc on the unit circle prior to being
        stretched and rotated into an elliptical arc.

    Notes
    -----
    See [1]_, `svg.path <https://github.com/regebro/svg.path>`_, and this
    `SO question <https://stackoverflow.com/q/197649>`_.

    References
    ----------

    .. [1] http://www.w3.org/TR/SVG/implnote.html#ArcImplementationNotes

    """
    # Step 1: compute x1prim, y1prim
    cosr = cos(radians(rotation))
    sinr = sin(radians(rotation))
    dx = (start.real - end.real) / 2
    dy = (start.imag - end.imag) / 2
    x1prim = cosr * dx + sinr * dy
    y1prim = -sinr * dx + cosr * dy
    x1prim_sq = x1prim * x1prim
    y1prim_sq = y1prim * y1prim

    rx = abs(radius.real)
    ry = abs(radius.imag)
    rx_sq = rx * rx
    ry_sq = ry * ry

    # Correct out of range radii
    radius_scale = (x1prim_sq / rx_sq) + (y1prim_sq / ry_sq)
    if radius_scale > 1:
        radius_scale = sqrt(radius_scale)
        rx *= radius_scale
        ry *= radius_scale
        rx_sq = rx * rx
        ry_sq = ry * ry
    else:
        # SVG spec only scales UP
        radius_scale = 1
    radius = rx + ry * 1j

    # Step 2: compute cxprim, cyprim
    t1 = rx_sq * y1prim_sq
    t2 = ry_sq * x1prim_sq
    c = sqrt(abs((rx_sq * ry_sq - t1 - t2) / (t1 + t2)))
    if large == sweep:
        c = -c
    cxprim = c * rx * y1prim / ry
    cyprim = -c * ry * x1prim / rx

    # Step 3: compute the center from cxprim, cyprim
    center = complex(
        (cosr * cxprim - sinr * cyprim) + ((start.real + end.real) / 2),
        (sinr * cxprim + cosr * cyprim) + ((start.imag + end.imag) / 2),
    )

    # Step 4: compute theta and delta_theta
    ux = (x1prim - cxprim) / rx
    uy = (y1prim - cyprim) / ry
    vx = (-x1prim - cxprim) / rx
    vy = (-y1prim - cyprim) / ry
    n = sqrt(ux * ux + uy * uy)
    p = ux
    theta = degrees(acos(p / n))
    if uy < 0:
        theta = -theta
    theta = theta % 360

    n = sqrt((ux * ux + uy * uy) * (vx * vx + vy * vy))
    p = ux * vx + uy * vy
    d = p / n
    # In certain cases the above calculation can through inaccuracies
    # become just slightly out of range, f ex -1.0000000000000002.
    d = np.clip(d, -1.0, 1.0)

    delta = degrees(acos(d))
    if (ux * vy - uy * vx) < 0:
        delta = -delta
    delta = delta % 360
    if not sweep:
        if delta > 0:
            delta -= 360

    return radius, center, theta, theta + delta

def arc_path(start, radius, rotation, large, sweep, end):
    """
    Generate an elliptical arc path given an endpoint parameterization.

    Uses matplotlib to draw the arc using quadratic Bezier curves.

    Parameters
    ----------
    start : complex
        Starting point (x1, y1).
    radius : complex
        Two elliptical radii (rx, ry).
    rotation : float
        Angle from the x-axis of the current coordinate system to the x-axis of
        the ellipse.
    large : bool
        False if an arc spanning < 180 degrees is to be drawn, True if an arc
        spanning >= 180 degrees is to be drawn.
    sweep : bool
        If sweep-flag is True, then the arc will be drawn in a "positive-angle"
        direction from start to end.
    end : complex
        End point (x2, y2).

    Returns
    -------
    codes : array (n,)
        Command codes
    verts : array (n,2)
        Vertices

    Notes
    -----
    We first perform a conversion to center parameterization and generate
    a circular arc at the origin. Then we apply scaling, translation and
    rotation.

    One can think of an ellipse as a circle that has been stretched and then
    rotated. Start by making an arc along the unit circle from `theta1` to
    `theta2`, centered at `center`. Then scale the circle along the x and y
    axes according to the given radii. Finally, rotate the arc around the
    center through the given angle `rotation`.

    """
    radius, center, theta1, theta2 = endpoint_to_center(
        start,
        radius,
        rotation,
        large,
        sweep,
        end
    )

    # Create an arc on the unit circle
    # Matplotlib does this using CURVE4 operations
    # https://matplotlib.org/stable/_modules/matplotlib/path.html#Path.arc
    if theta2 > theta1:
        arc = Path.arc(theta1=theta1, theta2=theta2)
        reverse_path = False
    else:
        arc = Path.arc(theta1=theta2, theta2=theta1)
        reverse_path = True

    # Transform it into an elliptical arc:
    # * scale the minor and major axes
    # * translate it to the center
    # * rotate the x-axis of the ellipse from the x-axis of the current
    #   coordinate system
    trans = (
        transforms.Affine2D()
        .scale(radius.real, radius.imag)
        .translate(center.real, center.imag)
        .rotate_deg_around(center.real, center.imag, rotation)
    )
    arc = trans.transform_path(arc)
    codes = np.array(arc.codes)
    verts = np.array(arc.vertices)

    # Make sure we are drawing from start to end
    if reverse_path:
        verts = verts[::-1, :]

    # Change the initial MOVETO operation into a LINETO to connect to the
    # previous path
    codes[0] = Path.LINETO

    return codes, verts

def _tokenize_path(pathdef):
    for x in COMMAND_RE.split(pathdef):
        if x in COMMANDS:
            yield x
        for token in FLOAT_RE.findall(x):
            yield token

def _next_pos(elements):
    return float(elements.pop()) + float(elements.pop()) * 1j

def _parse_path(pathdef, current_pos):
    # In the SVG specs, initial movetos are absolute, even if specified as 'm'.
    # This is the default behavior here as well. But if you pass in a
    # current_pos variable, the initial moveto will be relative to that
    # current_pos. This is useful.
    elements = list(_tokenize_path(pathdef))
    # Reverse for easy use of .pop()
    elements.reverse()

    start_pos = None
    command = None

    while elements:

        # 1. Determine the current command

        if elements[-1] in COMMANDS:
            # New command.
            last_command = command  # Used by S and T
            command = elements.pop()
            absolute = command in UPPERCASE
            command = command.upper()
        else:
            # Implicit command.
            # If this element starts with numbers, it is an implicit command
            # and we don't change the command. Check that it's allowed:
            if command is None:
                raise ValueError(
                    "Unallowed implicit command in {}, position {}".format(
                        pathdef, len(pathdef.split()) - len(elements))
                    )
            last_command = command  # Used by S and T

        # 2. Parse the current command

        # MOVETO
        if command == 'M':
            pos = _next_pos(elements)
            if absolute:
                current_pos = pos
            else:
                current_pos += pos

            # when M is called, reset start_pos
            # This behavior of Z is defined in svg spec:
            # http://www.w3.org/TR/SVG/paths.html#PathDataClosePathCommand
            start_pos = current_pos

            yield COMMAND_CODES['M'], [(current_pos.real, current_pos.imag)]

            # Implicit moveto commands are treated as lineto commands.
            # So we set command to lineto here, in case there are
            # further implicit commands after this moveto.
            command = 'L'

        # CLOSEPATH
        elif command == 'Z':
            # path closure
            if current_pos != start_pos:
                verts = [(start_pos.real, start_pos.imag)]
                yield COMMAND_CODES['L'], verts

            # mpl.Path: a point is required but ignored
            verts = [(start_pos.real, start_pos.imag)]
            yield COMMAND_CODES['Z'], verts

            current_pos = start_pos
            start_pos = None
            command = None  # You can't have implicit commands after closing.

        # LINETO
        elif command == 'L':
            pos = _next_pos(elements)
            if not absolute:
                pos += current_pos
            verts = [(pos.real, pos.imag)]
            yield COMMAND_CODES['L'], verts
            current_pos = pos

        # HORIZONTAL_PATHTO
        elif command == 'H':
            x = elements.pop()
            pos = float(x) + current_pos.imag * 1j
            if not absolute:
                pos += current_pos.real
            verts = [(pos.real, pos.imag)]
            yield COMMAND_CODES['H'], verts
            current_pos = pos

        # VERTICAL_PATHTO
        elif command == 'V':
            y = elements.pop()
            pos = current_pos.real + float(y) * 1j
            if not absolute:
                pos += current_pos.imag * 1j
            verts = [(pos.real, pos.imag)]
            yield COMMAND_CODES['V'], verts
            current_pos = pos

        # CUBIC_BEZIER
        elif command == 'C':
            control1 = _next_pos(elements)
            control2 = _next_pos(elements)
            end = _next_pos(elements)
            if not absolute:
                control1 += current_pos
                control2 += current_pos
                end += current_pos
            verts = [
                 (control1.real, control1.imag),
                 (control2.real, control2.imag),
                 (end.real, end.imag)
            ]
            yield COMMAND_CODES['C'], verts
            current_pos = end

        # SMOOTH_CUBIC_BEZIER
        elif command == 'S':
            # Smooth curve. First control point is the "reflection" of
            # the second control point in the previous path.
            if last_command not in 'CS':
                # If there is no previous command or if the previous command
                # was not an C, c, S or s, assume the first control point is
                # coincident with the current point.
                control1 = current_pos
            else:
                # The first control point is assumed to be the reflection of
                # the second control point on the previous command relative
                # to the current point.
                last_control = control2
                control1 = current_pos + current_pos - last_control
            control2 = _next_pos(elements)
            end = _next_pos(elements)
            if not absolute:
                control2 += current_pos
                end += current_pos
            verts = [
                (control1.real, control1.imag),
                (control2.real, control2.imag),
                (end.real, end.imag)
            ]
            yield COMMAND_CODES['S'], verts
            current_pos = end

        # QUADRATIC_BEZIER
        elif command == 'Q':
            control = _next_pos(elements)
            end = _next_pos(elements)
            if not absolute:
                control += current_pos
                end += current_pos
            verts = [
                (control.real, control.imag),
                (end.real, end.imag)
            ]
            yield COMMAND_CODES['Q'], verts
            current_pos = end

        # SMOOTH_QUADRATIC_BEZIER
        elif command == 'T':
            # Smooth curve. Control point is the "reflection" of
            # the second control point in the previous path.
            if last_command not in 'QT':
                # If there is no previous command or if the previous command
                # was not an Q, q, T or t, assume the first control point is
                # coincident with the current point.
                control = current_pos
            else:
                # The control point is assumed to be the reflection of
                # the control point on the previous command relative
                # to the current point.
                last_control = control
                control = current_pos + current_pos - last_control
            end = _next_pos(elements)
            if not absolute:
                end += current_pos
            verts = [
                (control.real, control.imag),
                (end.real, end.imag)
            ]
            yield COMMAND_CODES['T'], verts
            current_pos = end

        # ELLIPTICAL_ARC
        elif command == 'A':
            radius = _next_pos(elements)
            rotation = float(elements.pop())
            large = float(elements.pop())
            sweep = float(elements.pop())
            end = _next_pos(elements)

            if not absolute:
                end += current_pos

            if current_pos == end:
                # This is equivalent of omitting the segment, so do nothing
                continue
            elif radius.real == 0 or radius.imag == 0:
                # This should be treated as a straight line
                verts = [(end.real, end.imag)]
                yield COMMAND_CODES['L'], verts
            else:
                codes, verts = arc_path(
                    current_pos, radius, rotation, large, sweep, end
                )
                yield codes, verts

            current_pos = end

def parse_path(pathdef, current_pos=0 + 0j):
    """
    Parse an SVG path definition string into a matplotlib Path object.

    Parameters
    ----------
    pathdef : str
        SVG path 'd' attribute, e.g. 'M 100 100 L 300 100 L 200 300 z'.
    current_pos : complex, optional
        Coordinates of the starting position of the path, given as a complex
        number. When provided, an initial moveto operation will be intepreted
        as relative to this position even if given as M.

    Returns
    -------
    :class:`matplotlib.path.Path` instance

    See also
    --------
    matplotlib.path.Path
    matplotlib.patches.PathPatch
    matplotlib.collections.PathCollection
    matplotlib.transforms

    """
    codes = []
    verts = []
    for c, v in _parse_path(pathdef, current_pos):
        codes.extend(c)
        verts.extend(v)
    return Path(verts, codes)


# The original code requires pylcaio release v1.0, from Agez, found at:
# https://github.com/MaximeAgez/pylcaio

import pandas as pd
import numpy as np
import os, sys
import pickle
import logging
import pylcaio

fp_results = os.path.join(os.path.curdir, 'results')
fp_output = os.path.join(os.path.curdir, 'output')

def hybrid_emission_factors(trade_only, year):
    Analysis = pylcaio.Analysis('ecoinvent3.5','exiobase3', method_double_counting='STAM')

    # For making human-readable indices; consists of all ecoinvent process metadata
    labels = Analysis.PRO_f

    # Make DataFrame for whole ecoinvent requirements matrix with human-readable indices
    Aff_labels = (labels.join(Analysis.A_ff,how='inner'))
    Aff_labels = (Aff_labels).set_index(keys=['activityName', 'geography', 'productName'], append=True)
    Aff_labels = Aff_labels.iloc[:,19:] # remove extraeneous metadata
    Aff_labels.columns = Aff_labels.index

    # Calculate life cycle impacts
    I = pd.DataFrame(np.eye(len(Analysis.A_ff)), Analysis.A_ff.index, Analysis.A_ff.columns)
    X = pd.DataFrame(np.linalg.solve(I-Analysis.A_ff, I), Analysis.A_ff.index, I.columns) # pro x pro
    D = Analysis.C_f.fillna(0).dot(Analysis.F_f).dot(X) # pro x imp

    D = D.T
    D_labels = labels.join(D, how='inner')
    D_labels.set_index(keys=['activityName', 'geography', 'productName'], append=True, inplace=True)
    D_labels = D_labels.iloc[:, 19:]
    D_labels = D_labels.T

    # Other possible impact methods
    imp_list = ['EDIP; environmental impact; global warming, GWP 100a; kg CO2-Eq', 'ReCiPe Midpoint (H); climate change; GWP100; kg CO2-Eq', 'ReCiPe Midpoint (H) V1.13; climate change; GWP100; kg CO2-Eq', 'IPCC 2013; climate change; GTP 100a; kg CO2-Eq']
    D_labels_comp = D_labels.loc[imp_list]
    D_labels = D_labels.loc['EDIP; environmental impact; global warming, GWP 100a; kg CO2-Eq']

    # the footprint of all hybridized processes from ecoinvent, with their original footprint
    fp_hybrid = os.path.join(os.path.curdir, 'data', 'Results_full_database_STAM_2011.xlsx')
    hybrid_results = pd.read_excel(fp_hybrid, 'GWP_only_hyb', usecols='A:E,G', index_col=[0,1,2,3,4])

    # ## List of ENTSO-E countries
    country_list = ['AL','AT','BA','BE','BG','CH','CZ','DE','DK','EE',
                    'ES','FI','FR','GR','HR','HU','IE','IT','LT','LU',
                    'LV','ME','MK','NL','NO','PL','PT','RO','RS','SE',
                    'SI','SK','GB','TR','GE','CY','HR','NI','UK']

    # Find all electricity processes
    set(labels.loc[labels['activityName'].str.contains('electricity')]['activityName'])

    # Isolate high voltage (production) mixes and low voltage mixes (which includes solar PV)
    # note that pumped storage hydropower in CH gets excluded as it has 0 contribution to the production mix (for whatever reason)
    market_mixes = labels.loc[labels['activityName'].str.contains('electricity, high voltage, production mix')] # production mixes, do not include solar or waste (or Swiss pumped hydro)
    waste_mixes = labels.loc[labels['activityName'].str.contains('electricity, from municipal waste incineration to generic market for electricity, medium voltage')]  # production mixes, do not include solar
    solar_mixes = labels.loc[labels['activityName'].str.contains('electricity, low voltage')]  # use low voltage mixes to get shares of solar photovoltaic

    # ### Filter for high-voltage production mixes
    # First, retrieve electricity mixes, and reduce un-needed processes
    df = Aff_labels[market_mixes.index]
    df_waste = Aff_labels.loc[waste_mixes.index]
    df_solar = Aff_labels[solar_mixes.index]

    df = df.loc[:, df.columns.isin(country_list, level=2)]

    # manual workaround for keeping CH pumped hydro after removing unneeded rows
    # (since this ecoinvent process apparently does not contribute to any electricity mixes in ecoinvent)
    ind =('69ccf019-081e-42e3-b7f5-879dc7dde86a_66c93e71-f32b-4591-901c-55395db5c132', 'electricity production, hydro, pumped storage', 'CH', 'electricity, high voltage')  # Activityid for Swiss pumped hydro
    col = df.columns.get_loc(('54ed269b-e329-469c-902e-a8991ee7d24a_66c93e71-f32b-4591-901c-55395db5c132', 'electricity, high voltage, production mix', 'CH', 'electricity, high voltage'))
    index = df.index.get_loc(ind)
    df.iloc[index, col] = 1e-7

    df = df.replace(0, np.nan).dropna(how='all', axis=0)

    # ### Filter for medium-voltage mixes ( to find waste incineration shares)
    df_waste = df_waste.loc[:, df_waste.columns.isin(country_list, level=2)]
    df_waste = df_waste.replace(to_replace=0, value=np.nan).dropna(how='all', axis=0)

    # check unique waste incineration activities
    set(df_waste.index.get_level_values('activityName'))

    # ### Filter for low-voltage mixes (to find solar photovoltaic shares)
    df_solar = df_solar.loc[:, df_solar.columns.isin(country_list, level=2)]
    df_solar = df_solar.replace(to_replace=0, value=np.nan).dropna(how='all', axis=0)

    # ### Develop correspondence between ecoinvent and ENTSO-E technology categories

    # Correspondence from ENTSO-E technology categories to keywords for searching in ecoinvent processes
    # Note that the 'fossil oil shale', 'other' and 'other renewable' do not have ecoinvent equivalents;
    # these are dealt with later
    tec_dict = {'Biomass': ['biogas', 'wood chips'],
                'Fossil Brown coal/Lignite': 'lignite',
                'Fossil Coal-derived gas': 'coal gas',
                'Fossil Gas': 'natural gas',
                'Fossil Hard coal': 'hard coal',
                'Fossil Oil': ' oil',
                'Fossil Oil shale': '--',
                'Fossil Peat': 'peat',
                'Geothermal': 'geothermal',
                'Hydro Pumped Storage': 'pumped storage',
                'Hydro Run-of-river and poundage': 'run-of-river',
                'Hydro Water Reservoir': 'reservoir',
                'Nuclear': 'nuclear',
                'Other': '---',
                'Other renewable': '----',
                'Solar': ['solar thermal', 'solar tower', 'photovoltaic'],
                'Waste': 'waste incineration',
                'Wind Offshore': 'offshore',
                'Wind Onshore': 'onshore'
    }

    # reverse correspondence from ecoinvent search keywords to ENTSO-E categories
    temp_tecdict = { 'biogas':'Biomass',
                    'wood chips':'Biomass',
                'lignite': 'Fossil Brown coal/Lignite' ,
                'coal gas': 'Fossil Coal-derived gas',
                'natural gas': 'Fossil Gas' ,
                'hard coal': 'Fossil Hard coal' ,
                ' oil': 'Fossil Oil',
                '--' : 'Fossil Oil shale',
                'peat': 'Fossil Peat' ,
                'geothermal':'Geothermal',
                'pumped storage': 'Hydro Pumped Storage',
                'run-of-river': 'Hydro Run-of-river and poundage',
                'reservoir': 'Hydro Water Reservoir',
                'nuclear' : 'Nuclear',
                '---':'Other' ,
                '----': 'Other renewable',
                'solar thermal': 'Solar',
                'solar tower': 'Solar',
                'photovoltaic': 'Solar',
                'waste incineration': 'Waste', #  'electricity, from municipal waste incineration to generic market for electricity, medium voltage',
                'offshore': 'Wind Offshore',
                'onshore':'Wind Onshore'
    }

    # flatten tec list
    tec_list = list(tec_dict.values())
    for item in tec_list:
        if type(item)==list:
            for subitem in item:
                tec_list.append(subitem)
            tec_list.remove(item)

    # Get list of solar PV processes in ecoinvent
    d = list(set(df_solar.index.get_level_values(level=1)))
    pv = [entry for entry in d if 'electricity production' in entry]  # all low-voltage production technologies

    # Fetch waste incineration process names
    waste = ['electricity, from municipal waste incineration to generic market for electricity, medium voltage']

    # Make list of electricity generation technologies (both high voltage and low-voltage, i.e., solar PV)
    ecoinvent_tecs = list(set(df.index.get_level_values(level=1))) + pv + waste

    # Quality assurance; make sure all ecoinvent technologies are covered by the keyword search, and keyword matches are correct
    # also, check number of matches for each keyword
    matches = []
    match_keys = []
    num_matches = dict((keyword,0) for keyword in tec_list)

    for tec in ecoinvent_tecs:
        for keyword in tec_list:
            if keyword in tec and tec not in matches:
#                print(f'{keyword} is in {tec}')
                matches.append(tec)
                match_keys.append(temp_tecdict[keyword])
                num_matches[keyword] += 1

    print(num_matches)

    # Find ecoinvent tecs that were not matched
    if len(ecoinvent_tecs)-len(matches) > 0:
        print('ecoinvent process(es) not matched:')
        print(set(ecoinvent_tecs).difference(set(matches)))

    # Correspondence of ENTSO-E categories and specific ecoinvent processes
    from collections import defaultdict

    ei_tec_dict = defaultdict(list) # keys as entso-e categories, values are corresponding ei
    rev_ei_tec_dict = defaultdict(list) # vice-versa; keys as ei, values are entso-e categories
    for i, j in zip(match_keys, matches):
        ei_tec_dict[i].append(j)
        rev_ei_tec_dict[j].append(i)

    # Build correspondence table for Supplementary Information
    tec_index = pd.Index(list(ei_tec_dict.keys()))
    tec_values = list(ei_tec_dict.values())
    ei_tec_table = pd.DataFrame(tec_values, index=tec_index).stack()
    ei_tec_table.index = ei_tec_table.index.droplevel(level=1)
    ei_tec_table.sort_index(axis=0, inplace=True)

    # ### Calculate shares of ecoinvent processes
    # For multi-process technology categories, e.g., solar, biomass...

    # combine high-, medium- and low-voltage electricity mixes for full tech resolution
    f = pd.concat([df, df_waste, df_solar], axis=1)

    rev_ei_tec_dict = {k: v for k, v in rev_ei_tec_dict.items() if not not v}

    temp_list = f.index.get_level_values('activityName')
    entso_list = []
    f_check = []

    # For every activity in ecoinvent,if they are related to electricity production, add to entso_list
    # If not, drop from the matrix; they are irrelevant
    for entry in temp_list:
        try:
            entso_list.append(rev_ei_tec_dict[entry])
        except:
            f_check.append(entry)

    f.drop(index=set(f_check), level=1, axis=0, inplace=True)

    # inputs to electricity mixes not included in the calculations of emission factors
    set(f_check)

    entso_list = [item for sublist in entso_list for item in sublist] #flatten list

    entso_index = pd.Index(tuple(entso_list), name='entso_e')
    f.set_index(entso_index, append=True, inplace=True)

    # ### Get shares of each ecoinvent process for each ENTSO-E technology category

    f_gpby = f.groupby(level='entso_e')

    # Get shares of each ecoinvent process in each ENTSO-E technology category;
    # for calculating weighted average for emission factor
    ei_process_shares = f.divide(f_gpby.sum(), level=4, axis=1)

    hybrid_temp = hybrid_results.reset_index(level=[1,2,3,4], drop=True)
    g_temp = ei_process_shares.reset_index(level=[1,2,3,4], drop=True)

    # ### Calculate weighted average, hybridized emissions factor for each technology

    ef = ei_process_shares.mul(hybrid_temp.iloc[:, 0], axis=0, level=0)

    ef_aggregated = ef.groupby(level='entso_e').sum()

    ef_countries = ef_aggregated.sum(axis=1, level='geography')
    ef_countries.replace(0, np.nan, inplace=True)
    ef_countries.dropna(axis=1, how='all', inplace=True)
    ef_countries.sort_index(inplace=True)

    ef_countries = ef_countries.T
    ef_countries.sort_index(axis=0, inplace=True)

    entso_fp = os.path.join(fp_output, 'entsoe', 'ENTSO_production_volumes_' + str(year) + '.csv')

    entso_e_production = pd.read_csv(entso_fp, header=0, index_col=[0])
    entso_e_production.replace(0, np.nan,inplace=True)
    entso_mask = entso_e_production.isna().sort_index()
    ef_mask = ef_countries.sort_index().isna()

    entso_mask.drop(columns=['Fossil Oil shale', 'Other', 'Other renewable', 'Marine'], inplace=True)

    entso_mask.index.rename('geography', inplace=True)
    entso_mask.columns.rename('entso_e', inplace=True)

    # Countries for which we have production data, but no hybridized emission factors
    no_ef_countries = set(entso_e_production.index) - set(ef_countries.index)
    print('Countries with production data, but no hybridized emission factors:')
    print(no_ef_countries)
    logging.warning(f'{no_ef_countries} have production data from ENTSO_E but no hybridized emission factors')

    list(set(ef_countries.columns) - set(entso_e_production.index))

    check_df = entso_mask.eq(ef_mask, axis=0)

    # Find countries missing from ecoinvent
    missing_countries = []
    for country, row in check_df.iterrows():
        if not row.any():
            missing_countries.append(country)
    # Optionally, drop them from mask
    check_df.drop(index = missing_countries, inplace=True)

    # Find the technology-country pairs that are missing hybridized emission factors
    countries_missing_ef = {}
    for tec, col in check_df.iteritems():
        temp = check_df.index[check_df[tec] == False].tolist()
        print(f'For {tec}, regionalized, hybridized emission factors are missing for {temp}. Using technological average from available regions instead.')
        temp2 = []
        for country in temp:
            if ef_mask.loc[country, tec]:
                temp2.append(country)
        countries_missing_ef[tec] = temp2#.append(country)
    countries_missing_ef

    no_ef = pd.DataFrame({k:pd.Series(v[:13]) for k, v in countries_missing_ef.items()})

    # ## Calculate LC emissions for production mix for countries with missing production data

#    trade_only = ['AL','BY','HR','LU','MD','MT','RU','TR','UA']
    low_volt = labels.loc[labels['activityName'] == 'market for electricity, low voltage']
    trade_only_mixes_lv = low_volt.loc[low_volt['geography'].isin(trade_only)].index
    lv_mix_ef = D_labels.loc[trade_only_mixes_lv]  # low voltage mixes, not used

    high_volt = labels.loc[labels['activityName'] == 'market for electricity, high voltage']
    trade_only_mixes_hv = high_volt.loc[high_volt['geography'].isin(trade_only)].index
    hv_mix_ef = D_labels.loc[trade_only_mixes_hv]

    lv_mix_ef_comp = D_labels_comp[trade_only_mixes_lv]  # low voltage mixes, not used
    hv_mix_ef_comp = D_labels_comp[trade_only_mixes_hv]

    trade_mixes_comp = lv_mix_ef_comp.join(hv_mix_ef_comp)

    labels.loc[labels['activityName'] == 'electricity, from municipal waste incineration to generic market for electricity, medium voltage']

    # ## Export results
    # Convert emission factors to g CO2-eq/kWh
    ef_countries = ef_countries * 1000
    lv_mix_ef = lv_mix_ef * 1000  # low voltage mixes, not used
    hv_mix_ef = hv_mix_ef * 1000
    trade_mixes_comp = trade_mixes_comp * 1000

    # Calculate mean technology footprints as quick check
    ef_aggregated.replace(0, np.nan).mean(axis=1)

    fp_results = os.path.join(os.path.curdir,'results')

    with pd.ExcelWriter(os.path.join(fp_results, 'hybrid_emission_factors_final_' + str(year) + '.xlsx')) as writer:
        ef_countries.to_excel(writer, sheet_name='country_emission_factors')
        no_ef.to_excel(writer, sheet_name='missing_emission_factors')
        lv_mix_ef.to_excel(writer, sheet_name='ecoinvent ef')
        hv_mix_ef.to_excel(writer, sheet_name='ecoinvent ef_hv')
        trade_mixes_comp.to_excel(writer, sheet_name='trade_mixes_from_ei')
        ei_tec_table.to_excel(writer, sheet_name='correspondence table')
        ef_aggregated.to_excel(writer, sheet_name='mean tech fps')

    # Export to .csv for use in BEV_footprints_calculations
    # append extra countries to file
    try:
        existing_tradeonly = pd.read_csv(os.path.join(fp_output, 'ecoinvent_ef_hv.csv'))['geography'].values
        add_ef_countries = set(hv_mix_ef.index.get_level_values('geography')) - set(existing_tradeonly)  # countries not already in csv file
        if add_ef_countries:
            hv_mix_ef_exp = hv_mix_ef.loc[:,:,list(add_ef_countries),:]
            hv_mix_ef_exp.to_csv(os.path.join(fp_output, 'ecoinvent_ef_hv.csv'), mode='a', header=False)  # used in BEV_footprints_calculations
    except FileNotFoundError:
        print('Previous ecoinvent file does not exist. Making file.')
        hv_mix_ef.to_csv(os.path.join(fp_output, 'ecoinvent_ef_hv.csv'))

    ef_countries.to_csv(os.path.join(fp_output, 'country_emission_factors.csv'))  # used in clean_impact_factors
    no_ef.to_csv(os.path.join(fp_output, 'missing_emission_factors.csv'))  # used in clean_impact_factors
#    lv_mix_ef.to_csv(os.path.join(fp_results, 'ecoinvent_ef_lv.csv'))

    return ef_countries, no_ef, countries_missing_ef

def clean_impact_factors(year, trade_only, calc_hybrid):
    entsoe_fp = os.path.join(fp_output, 'entsoe')
    with open(os.path.join(entsoe_fp, 'gen_final_' + str(year) + '.pkl'), 'rb') as handle:
        gen_df = pickle.load(handle)

    with open(os.path.join(entsoe_fp,'trade_final_' + str(year) + '.pkl'), 'rb') as handle:
        trade_df = pickle.load(handle)

    if calc_hybrid:
        try:
            print('\n Calculating hybridized emission factors from pyLCAIO')
            # gen_df not modified in hybrid_emission_factors, just used to determine relevant labels
            ef, missing_factors, countries_missing_ef = hybrid_emission_factors(trade_only, year)
        except Exception as e:
            print('\n Calculating from pyLCAIO failed. Importing previously calculated emission factors instead')
            print(e)
            # import ready-calculated emission factors if no access to pylcaio object
            ef = pd.read_csv(os.path.join(fp_output, 'country_emission_factors.csv'), index_col=[0])
            missing_factors = pd.read_csv(os.path.join(fp_output, 'missing_emission_factors.csv'), index_col=[0])
    else:
        print('\n Loading previously calculated emission factors')
        ef = pd.read_csv(os.path.join(fp_output, 'country_emission_factors.csv'), index_col=[0])
        missing_factors = pd.read_csv(os.path.join(fp_output, 'missing_emission_factors.csv'), index_col=[0])

    countries_missing_ef = dict((k, list(v.dropna())) for k, v in missing_factors.iteritems())  # make dict of tec: [countries]
    countries_missing_ef = {key:value for key, value in countries_missing_ef.items() if len(value)}  # remove entries with empty values

    """ Fix missing emission factors """
    # Note that these disregard whether or not there is actual production
    ef['Fossil Oil shale'] = ef['Fossil Oil'] # Proxy shale oil with conventional oil
    ef['Waste'] = 0  # Allocate incineration emissions to waste treatment as per ecoinvent
    ef['Marine'] = 0

    # Put in proxy for missing regions and technologies
    missing_factors.dropna(how='all', axis=1, inplace=True)
    missing_factors_dict = {}
    # build dict of missing technology-country factors
    for tec, col in missing_factors.iteritems():
        missing_factors_dict[tec] = list(col.dropna(how='any', axis=0).values)

    # countries with no geo-specific factors
    missing_countries = list(set(gen_df.index) - set(ef.index))
    for country in missing_countries:
        ef = ef.append(pd.Series(name=country))

    temp_gen_df = gen_df.loc[missing_countries]
    temp_gen_df = temp_gen_df.replace(0, np.nan).dropna(how='all', axis=1)
    temp_dict = {}
    for tec, col in temp_gen_df.items():
        temp = temp_gen_df[tec] > 0
        temp_ind = temp.index
        temp_dict[tec] = temp_ind[temp.values].to_list()

    """ Merge dictionaries and keep values of common keys in list"""
    def mergeDict(dict1, dict2):
        dict3 = {**dict1, **dict2}
        for key, value in dict3.items():
                if key in dict1 and key in dict2:
                    dict3[key] = (dict1[key] + dict2[key])
        return dict3

    # Merge dictionaries and add values of common keys in a list
    missing_factors_dict2 = mergeDict(missing_factors_dict, temp_dict)
    missing_factors_dict2

    # Use continental arithmetic average of same technology
    for key, countries in missing_factors_dict2.items():
        if not key=='Other' and not key=='Other renewable': # these are dealt with below
            ef[key].loc[countries] = ef[key].mean()

    # Make 'other' and 'other renewable' approximations
    renewable_dict = {'renewable': ['Biomass', 'Geothermal', 'Hydro Pumped Storage', 'Hydro Run-of-river and poundage',
                                    'Hydro Water Reservoir','Solar','Waste','Wind Onshore', 'Wind Offshore','Marine'],
                    'non-renewable': ['Fossil Gas','Fossil Hard coal','Fossil Oil','Fossil Brown coal/Lignite','Nuclear','Fossil Coal-derived gas','Fossil Oil shale','Fossil Peat']}

    other_tec = gen_df.loc[:,renewable_dict['non-renewable']]
    other_wf = other_tec.div(other_tec.sum(axis=1), axis=0) # determine generation shares for each country

    other_renew = gen_df.loc[:,renewable_dict['renewable']]
    other_renew_wf = other_renew.div(other_renew.sum(axis=1), axis=0)

    ef['Other'] = (ef*other_wf).sum(axis=1)
    ef['Other renewable'] = (ef*other_renew_wf).sum(axis=1)

    # Sort the indices of production, trade and emission factor matrices before export
    def sort_indices(df):
        df.sort_index(axis=0, inplace=True)
        df.sort_index(axis=1, inplace=True)

    sort_indices(trade_df)
    sort_indices(gen_df)
    sort_indices(ef)

    # Export all data
    # .csv for use in BEV_footprints_calculations.py

    ef.to_csv(os.path.join(fp_output,'final_emission_factors_' + str(year) + '.csv'))

    with pd.ExcelWriter(os.path.join(fp_results, 'entsoe_export_final_' + str(year) + '.xlsx')) as writer:
        trade_df.to_excel(writer, 'trade')
        gen_df.to_excel(writer, 'generation')
        ef.to_excel(writer,'new ef')

    return countries_missing_ef