#!/usr/bin/env python3

import logging
import math
import random
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import (
    Iterable,
    Mapping,
    MutableMapping,
    MutableSequence,
    Optional,
    Sequence,
    Set,
    Tuple,
    Union,
)

import numpy as np
import torch
from reagent.ope.estimators.estimator import (
    Estimator,
    EstimatorResult,
    EstimatorSampleResult,
)
from reagent.ope.estimators.types import (
    Action,
    Items,
    Objects,
    Probability,
    Reward,
    Trainer,
    TrainingData,
    TypeWrapper,
    Values,
    ValueType,
)
from reagent.ope.utils import Clamper, RunningAverage
from torch import Tensor

# Main algorithms are from two papers:
#   1. Offline Evaluation of Ranking Policies with Click Models
#      https://arxiv.org/abs/1804.10488
#   2. Off-policy evaluation for slate recommendation
#      https://arxiv.org/abs/1605.04812

# Types for slates
SlateSlotType = Union[int, Tuple[int], float, Tuple[float], np.ndarray, Tensor]
SlateSlot = TypeWrapper[SlateSlotType]

class SlateSlotValues(Values[SlateSlot]):
    """
    Map from a slot to a value
    """

    def _to_key(self, k: int) -> SlateSlot:
        return SlateSlot(k)

class SlateSlots(Items[SlateSlot]):
    """
    List of slot
    """

    def _new_item(self, i: int) -> SlateSlot:
        return SlateSlot(i)

    def fill(
        self,
        values: Union[Mapping[SlateSlot, float], Sequence[float], np.ndarray, Tensor],
    ) -> SlateSlotValues:
        """
        Map slots to given values
        Args:
            values: given values

        Returns:
            Map from slots to given values
        """
        return SlateSlotValues(super().fill(values))

class SlateSlotObjects(Objects[SlateSlot, ValueType]):
    def __init__(
        self,
        values: Union[MutableMapping[SlateSlot, ValueType], MutableSequence[ValueType]],
    ):
        assert (len(values)) > 0
        super().__init__(values)

    def _to_key(self, k: int) -> SlateSlot:
        return SlateSlot(k)

    @property
    def slots(self) -> SlateSlots:
        if self.is_sequence:
            return SlateSlots(len(self._values))
        else:
            return SlateSlots(list(self._key_to_index.keys()))

    @property
    def objects(self) -> Sequence[ValueType]:
        return super().values

    def fill(
        self, values: Sequence[ValueType]
    ) -> Union[Mapping[SlateSlot, ValueType], Sequence[ValueType]]:
        assert len(values) >= len(self._values)
        if self._key_to_index is None:
            return values[: len(self._values)]
        else:
            return {s: v for s, v in zip(self.slots, values[: len(self._values)])}

# type of valid slate candidates, e.g., doc id
SlateItem = Action

class SlateItems(Items[SlateItem]):
    def _new_item(self, i: int) -> SlateItem:
        return SlateItem(i)

class SlateItemValues(Values[SlateItem]):
    def _to_key(self, k: int) -> SlateItem:
        return SlateItem(k)

    @property
    def items(self) -> SlateItems:
        if self.is_sequence:
            return SlateItems(len(self))
        else:
            return SlateItems(super().keys)

class SlateItemFeatures(Objects[SlateItem, Tensor]):
    def __init__(
        self,
        values: Union[Mapping[SlateItem, Tensor], Sequence[Tensor], Tensor, np.ndarray],
    ):
        super().__init__(values)

    def _init_values(
        self,
        values: Union[Mapping[SlateItem, Tensor], Sequence[Tensor], Tensor, np.ndarray],
    ):
        if isinstance(values, Tensor):
            self._values = values.to(dtype=torch.double)
        elif isinstance(values, np.ndarray):
            self._values = torch.as_tensor(values, dtype=torch.double)
        elif isinstance(values, Sequence):
            self._values = torch.stack(values).to(dtype=torch.double)
        elif isinstance(values, Mapping):
            self._key_to_index = dict(zip(values.keys(), range(len(values))))
            self._index_to_key = list(values.keys())
            self._values = torch.stack(list(values.values())).to(dtype=torch.double)
        else:
            raise TypeError(f"Unsupported values type {type(values)}")

    def _to_key(self, k: int) -> SlateItem:
        return SlateItem(k)

    @property
    def items(self) -> SlateItems:
        if self.is_sequence:
            return SlateItems(len(self))
        else:
            return SlateItems(super().keys)

# SlateSlotFeatures = SlateSlotObjects[Tensor]
class SlateSlotFeatures(SlateSlotObjects[Tensor]):
    @property
    def features(self) -> Tensor:
        return torch.stack(self._values)

class Slate(SlateSlotObjects[SlateItem]):
    """
    Class represents a slate: map from slots to items/docs
    """

    def one_hots(self, items: SlateItems, device=None) -> Tensor:
        oh = torch.zeros((len(self), len(items)), dtype=torch.double, device=device)
        for t, i in zip(oh, self._values):
            t[items.index_of(i)] = 1.0
        return oh

    @property
    def items(self) -> Sequence[SlateItem]:
        return super().values

    def slot_values(self, item_values: SlateItemValues) -> SlateSlotValues:
        """
        Map items in the slate to given values
        Args:
            item_values: Map from all items to some values

        Returns:
            List of values in the slate
        """
        if self._key_to_index is None:
            return SlateSlotValues([item_values[i] for i in self._values])
        else:
            return SlateSlotValues({k: item_values[i] for k, i in self._key_to_index})

    def slot_features(self, item_features: SlateItemFeatures) -> SlateSlotFeatures:
        """
        Map items in the slate to given values
        Args:
            item_values: Map from all items to some values

        Returns:
            List of values in the slate
        """
        if self._key_to_index is None:
            return SlateSlotFeatures(
                [item_features[i].detach().clone() for i in self._values]
            )
        else:
            return SlateSlotFeatures(
                {k: item_features[i].detach().clone() for k, i in self._key_to_index}
            )

    def __repr__(self):
        return f"{self.__class__.__name__}{{value[{self._values}]}}"

def make_slate(slots: SlateSlots, items: Sequence[SlateItem]) -> Slate:
    """
    Assign items to slots to make a slate
    """
    assert len(items) >= len(slots)
    if slots.is_sequence:
        return Slate(list(items[: len(slots)]))
    else:
        return Slate(dict(zip(slots, items[: len(slots)])))

class SlateSlotItemValues(SlateSlotObjects[SlateItemValues]):
    def __init__(
        self,
        values: Union[
            MutableMapping[SlateSlot, SlateItemValues], MutableSequence[SlateItemValues]
        ],
    ):
        super().__init__(values)
        self._item_size = len(self._values[0])
        for v in self._values[1:]:
            assert self._item_size == len(v)

    def values_tensor(self, device=None) -> Tensor:
        dist = [v.values for v in self._values]
        return torch.stack(dist).to(device=device)

class SlateSlotItemExpectations(SlateSlotItemValues):
    def expected_rewards(
        self, item_rewards: SlateItemValues, device=None
    ) -> SlateSlotValues:
        """
        Calculate expected relevances of each slot, given each item's
        relevances, under this distribution
        Args:
            item_rewards:
            device:

        Returns:
            Map of slots to their expected relevance
        """
        dist = self.values_tensor(device)
        rewards = item_rewards.values.to(device=device)
        rewards = torch.mm(dist, rewards.unsqueeze(0).t()).squeeze()
        if self.is_sequence:
            return SlateSlotValues(rewards)
        else:
            return SlateSlotValues(dict(zip(self.slots, rewards.tolist())))

    @property
    def expectations(self) -> Sequence[SlateItemValues]:
        return super().values

def make_slot_item_distributions(
    slots: SlateSlots, dists: Sequence[SlateItemValues]
) -> SlateSlotItemExpectations:
    assert len(dists) >= len(slots)
    if slots.is_sequence:
        return SlateSlotItemExpectations(list(dists[: len(slots)]))
    else:
        return SlateSlotItemExpectations(dict(zip(slots, dists[: len(slots)])))

def is_to_calculate_expectation(slate_size: int, item_size: int) -> bool:
    """
    Switch between calculating and sampling expectations, balanced by execution
    time and accuracy
    Return:
        True to calculate
        False to sample
    """
    return (
        slate_size < 4
        or (slate_size == 4 and item_size < 182)
        or (slate_size == 5 and item_size < 47)
        or (slate_size == 6 and item_size < 22)
        or (slate_size == 7 and item_size < 15)
    )

def _calculate_slot_expectation(
    d_out: Tensor,
    probs: Sequence[float],
    buffer: Iterable[Tuple[Set[int], float, float, float]],
) -> Iterable[Tuple[Set[int], float, float, float]]:
    """
    A helper function to calculate items' expectations for a slot
    """
    assert d_out.shape[0] == len(probs)
    next_buffer = []
    for b0, b1, b2, _ in buffer:
        # memory buffer for all ordered combinations so far, list of tuples of
        #   b0: all the items in this ordered combination
        #   b1: cumulative probability of b0
        #   b2: sum of the probabilities of b0
        #   b3: = b1 / (1.0 - b2) cached value for faster computation
        for i, i_prob in enumerate(probs):
            # only add i if it's not already in
            if i in b0:
                continue
            # nb* are next buffer values
            nb2 = b2 + i_prob
            # due to precision errors, sometimes nb2 becomes 1, in this
            # case, discard the combination
            if nb2 < 1.0:
                nb1 = b1 * i_prob / (1.0 - b2)
                next_buffer.append(({*b0, i}, nb1, nb2, nb1 / (1.0 - nb2)))
    for i, i_prob in enumerate(probs):
        p = 0.0
        for b0, _, _, b3 in next_buffer:
            if i in b0:
                continue
            p += b3
        d_out[i] = p * i_prob
    return next_buffer

class SlateItemProbabilities(SlateItemValues):
    """
    Probabilities of each item being selected into the slate
    """

    def __init__(
        self,
        values: Union[Mapping[SlateItem, float], Sequence[float], np.ndarray, Tensor],
        greedy: bool = False,
    ):
        super().__init__(values)
        self._greedy = greedy
        self._slot_item_expectations = None

    def _to_key(self, k: int) -> SlateItem:
        return SlateItem(k)

    def _reset(self):
        super()._reset()
        self._slot_item_expectations = None

    def slate_probability(self, slate: Slate) -> Probability:
        """
        Calculate probability of a slate under this distribution
        Args:
            slate:

        Returns:
            probability
        """
        if self._greedy:
            items = super().greedy(len(slate))
            for i1, i2 in zip(items, slate.items):
                if i1 != i2:
                    return 0.0
            return 1.0
        else:
            clamped = torch.clamp(self._values, 0.0)
            indices = [self.index_of(item) for _, item in slate]
            probs = clamped[indices]
            sums = clamped[indices]
            clamped[indices] = 0.0
            sums = sums.flip(0).cumsum(0).flip(0) + clamped.sum()
            return Probability((probs / sums).prod().item())

    def slot_item_expectations(self, slots: SlateSlots) -> SlateSlotItemExpectations:
        slate_size = len(slots)
        if (
            self._slot_item_expectations is not None
            and len(self._slot_item_expectations) >= slate_size
        ):
            return self._slot_item_expectations
        item_size = len(self)
        assert item_size >= slate_size
        if self._greedy:
            self._slot_item_expectations = make_slot_item_distributions(
                slots,
                [
                    self.replace(torch.zeros(item_size, dtype=torch.double))
                    for _ in range(len(self))
                ],
            )
            sorted_items, _ = self.sort()
            for item, ds in zip(
                sorted_items, self._slot_item_expectations.expectations
            ):
                ds[item] = 1.0
        else:
            self._normalize()
            if is_to_calculate_expectation(len(slots), len(self)):
                self._calculate_expectations(slots)
            else:
                self._sample_expectations(slots, 20000)
        return self._slot_item_expectations

    def _sample_expectations(self, slots: SlateSlots, num_samples: int):
        slate_size = len(slots)
        item_size = len(self)
        dm = torch.zeros((slate_size, item_size), dtype=torch.double)
        ri = torch.arange(slate_size)
        ws = self._probabilities.repeat((num_samples, 1))
        for _ in range(item_size):
            samples = torch.multinomial(ws, slate_size)
            for sample in samples:
                dm[ri, sample] += 1
        dm /= num_samples * item_size
        self._slot_item_expectations = make_slot_item_distributions(
            slots, [self.replace(vs) for vs in dm]
        )

    def _calculate_expectations(self, slots: SlateSlots):
        """
        A brute-force way to calculate each item's expectations at each slot by
        going through all l-choose-m (l!/(l-m)!) possible slates.
        """
        slate_size = len(slots)
        item_size = len(self)
        dm = torch.zeros((slate_size, item_size), dtype=torch.double)
        dm[0] = self._probabilities
        buffer = [({}, 1.0, 0.0, 1.0)]
        probs = self._probabilities.tolist()
        for d in dm[1:]:
            buffer = _calculate_slot_expectation(d, probs, buffer)
        self._slot_item_expectations = make_slot_item_distributions(
            slots, [self.replace(vs) for vs in dm]
        )

    def sample_slate(self, slots: SlateSlots) -> Slate:
        slate_size = len(slots)
        if self._greedy:
            items = super().greedy(slate_size)
        else:
            items = super().sample(slate_size)
        if slate_size == 1:
            items = [items]
        return make_slate(slots, items)

    @property
    def is_deterministic(self) -> bool:
        return self._greedy

    def slate_space(
        self, slots: SlateSlots, max_size: int = -1
    ) -> Iterable[Tuple[Sequence[SlateItem], float]]:
        """Return all possible slates and their probabilities

        The algorithm is similar to :func:`~_calculate_expectations`, but has
        less value to cache thus save both space and computation
        Args:
            slots: slots to be filled
            max_size: max number of samples to be returned
                      <= 0 return all samples
        """
        slate_size = len(slots)
        item_size = len(self)
        assert item_size >= slate_size
        if self._greedy:
            items = super().greedy(slate_size)
            return [(items, 1.0)]
        else:
            buffer = [([], 1.0, 0.0)]
            probs = self._probabilities.tolist()
            for _ in range(slate_size):
                next_buffer = []
                for b0, b1, b2 in buffer:
                    # memory buffer for all ordered combinations so far, list of tuples of
                    #   b0: all the items in this ordered combination
                    #   b1: cumulative probability of b0
                    #   b2: sum of the probabilities of b0
                    for i, i_prob in enumerate(probs):
                        if i in b0:
                            continue
                        nb2 = b2 + i_prob
                        if nb2 < 1.0:
                            nb1 = b1 * i_prob / (1.0 - b2)
                            next_buffer.append(([*b0, i], nb1, nb2))
                if max_size <= 0 or max_size > len(next_buffer):
                    buffer = next_buffer
                else:
                    buffer = random.sample(next_buffer, max_size)
            return [([SlateItem(i) for i in b[0]], b[1]) for b in buffer]

class SlateSlotItemProbabilities(SlateSlotItemValues):
    def __init__(
        self,
        values: Union[
            MutableMapping[SlateSlot, SlateItemValues], MutableSequence[SlateItemValues]
        ],
        greedy: bool = False,
    ):
        super().__init__(values)
        self._greedy = greedy
        self._slot_item_distributions = None
        self._slot_item_expectations = None

    def slate_probability(self, slate: Slate) -> Probability:
        """
        Calculate probability of a slate under this distribution
        Args:
            slate:

        Returns:
            probability
        """
        assert len(slate) <= len(self)
        if self._greedy:
            for slot, item in slate:
                probs = self[slot]
                its, _ = probs.sort()
                if its[0] != item:
                    return 0.0
            return 1.0
        else:
            p = 1.0
            last_items = []
            for slot, item in slate:
                item_probs = self[slot]
                w = 1.0
                for last_item in last_items:
                    w -= item_probs.probability(last_item)
                if math.fabs(w - 0.0) < 1.0e-10:
                    return 0.0
                p *= item_probs.probability(item) / w
                last_items.append(item)
            return p

    def slot_item_expectations(self, samples: int = 20000) -> SlateSlotItemExpectations:
        slate_size = len(self.slots)
        if (
            self._slot_item_expectations is not None
            and len(self._slot_item_expectations) >= slate_size
        ):
            return self._slot_item_expectations
        item_size = len(self._values[0])
        assert item_size >= slate_size
        ps = self.values_tensor()
        if self._greedy:
            dists = []
            for i, value in zip(range(slate_size), self._values):
                item = ps[i].argmax().item()
                dist = torch.zeros(item_size, dtype=torch.double)
                dist[item] = 1.0
                dists.append(value.replace(dist))
                ps[torch.arange(i + 1, slate_size), item] = 0.0
            self._slot_item_expectations = make_slot_item_distributions(
                self.slots, dists
            )
        else:
            if is_to_calculate_expectation(slate_size, item_size):
                self._calculate_expectations()
            else:
                self._sample_expectations(samples * item_size)
        return self._slot_item_expectations

    def _sample_expectations(self, num_samples: int):
        slate_size = len(self.slots)
        item_size = len(self._values[0])
        dm = torch.zeros((slate_size, item_size), dtype=torch.double)
        ri = torch.arange(slate_size)
        for _ in range(num_samples):
            ps = self.values_tensor()
            sample = []
            for i in range(slate_size):
                item = ps[i].multinomial(1)
                sample.append(item)
                ps[torch.arange(i + 1, slate_size), item] = 0.0
            dm[ri, sample] += 1
        dm /= num_samples
        self._slot_item_expectations = make_slot_item_distributions(
            self.slots, [ivs.replace(vs) for ivs, vs in zip(self._values, dm)]
        )

    def _calculate_expectations(self):
        slate_size = len(self.slots)
        item_size = len(self._values[0])
        dm = torch.zeros((slate_size, item_size), dtype=torch.double)
        prob_list = []
        for v in self._values:
            v._normalize()
            prob_list.append(v._probabilities.detach().clone())
        dm[0] = prob_list[0]
        buffer = [({}, 1.0, 0.0, 1.0)]
        for d, probs in zip(dm[1:], prob_list[1:]):
            buffer = _calculate_slot_expectation(d, probs.tolist(), buffer)
        self._slot_item_expectations = make_slot_item_distributions(
            self.slots, [its.replace(vs) for its, vs in zip(self._values, dm)]
        )

    def sample_slate(self, slots: SlateSlots) -> Slate:
        slate_size = len(slots)
        ps = self.values_tensor()
        items = []
        if self._greedy:
            for i, value in zip(range(slate_size), self._values):
                item = ps[i].argmax().item()
                items.append(value.items[item])
                ps[torch.arange(i + 1, slate_size), item] = 0.0
        else:
            for i, value in zip(range(slate_size), self._values):
                item = ps[i].multinomial(1).item()
                items.append(value.items[item])
                ps[torch.arange(i + 1, slate_size), item] = 0.0
        return make_slate(slots, items)

class RewardDistribution(ABC):
    """
    Return customized probability distribution according to rewards
    """

    def __init__(self, deterministic: bool = False):
        self._deterministic = deterministic

    @abstractmethod
    def distribution(self, rewards: Tensor) -> Tensor:
        pass

    def __call__(self, rewards: SlateItemValues) -> SlateItemProbabilities:
        dist = self.distribution(rewards.values)
        return SlateItemProbabilities(rewards.items.fill(dist), self._deterministic)

    @property
    @abstractmethod
    def name(self) -> str:
        pass

class PassThruDistribution(RewardDistribution):
    """
    No-op distribution, probability determined by reward
    """

    def distribution(self, rewards: Tensor) -> Tensor:
        return rewards.detach().clone()

    @property
    def name(self) -> str:
        return f"{self._deterministic}"

    def __repr__(self):
        return f"PassThruDistribution[deterministic={self._deterministic}]"

class RankingDistribution(RewardDistribution):
    """
    Ranking distribution according to https://arxiv.org/abs/1605.04812
    """

    def __init__(self, alpha: float = -1.0, deterministic: bool = False):
        super().__init__(deterministic)
        self._alpha = alpha

    def distribution(self, rewards: Tensor) -> Tensor:
        dist = rewards.detach().clone()
        if self._alpha >= 0:
            _, ids = torch.sort(rewards, descending=True)
            rank = torch.arange(1, ids.shape[0] + 1, dtype=torch.double)
            dist[ids] = torch.pow(
                2.0, (-1.0 * (self._alpha * torch.log2(rank)).floor_())
            )
        return dist

    @property
    def name(self) -> str:
        return f"ranking_{self._alpha}_{self._deterministic}"

    def __repr__(self):
        return (
            f"RankingDistribution[alpha={self._alpha}"
            f",deterministic={self._deterministic}]"
        )

class FrechetDistribution(RewardDistribution):
    """
    Frechet distribution
    """

    def __init__(self, shape: float, deterministic: bool = False):
        super().__init__(deterministic)
        self._shape = shape

    def distribution(self, rewards: Tensor) -> Tensor:
        return torch.pow(rewards, self._shape)

    @property
    def name(self) -> str:
        return f"frechet_{self._shape}_{self._deterministic}"

    def __repr__(self):
        return (
            f"FrechetDistribution[shape={self._shape}]"
            f",deterministic={self._deterministic}]"
        )

SlateQueryType = Union[int, Tuple[int], float, Tuple[float], np.ndarray, Tensor]
SlateQuery = TypeWrapper[SlateQueryType]

@dataclass(frozen=True)
class SlateContext:
    query: SlateQuery
    slots: SlateSlots
    params: object = None

class SlatePolicy(ABC):
    """
    Policy interface
    """

    def __init__(self, device=None):
        self.device = device

    @abstractmethod
    def _query(self, context: SlateContext) -> SlateItemProbabilities:
        pass

    def __call__(self, context: SlateContext) -> SlateItemProbabilities:
        return self._query(context)

class SlateMetric:
    """
    Metric calculator for a slate: weights (dot) rewards

    Base class is just sum of the all item rewards
    """

    def __init__(self, device=None):
        self._device = device

    def calculate_reward(
        self,
        slots: SlateSlots,
        rewards: SlateSlotValues = None,
        slot_values: SlateSlotValues = None,
        slot_weights: SlateSlotValues = None,
    ) -> float:
        if slot_values is None:
            slot_values = self.slot_values(rewards)
        values = slot_values.values.to(device=self._device)
        if slot_weights is None:
            slot_weights = self.slot_weights(slots)
        weights = slot_weights.values.to(device=self._device)
        return torch.tensordot(values, weights, dims=([0], [0])).item()

    def __call__(self, slots: SlateSlots, rewards: SlateSlotValues) -> float:
        return self.calculate_reward(slots, rewards)

    def slot_weights(self, slots: SlateSlots) -> SlateSlotValues:
        return slots.fill([1.0] * len(slots))

    def slot_values(self, rewards: SlateSlotValues) -> SlateSlotValues:
        return rewards

class DCGSlateMetric(SlateMetric):
    _weights: Tensor = None

    def _get_discount(self, slate_size: int) -> Tensor:
        if (
            DCGSlateMetric._weights is None
            or DCGSlateMetric._weights.shape[0] < slate_size
            or DCGSlateMetric._weights.device != self._device
        ):
            DCGSlateMetric._weights = torch.reciprocal(
                torch.log2(
                    torch.arange(
                        2, slate_size + 2, dtype=torch.double, device=self._device
                    )
                )
            )
        return DCGSlateMetric._weights[:slate_size]

    def slot_weights(self, slots: SlateSlots) -> SlateSlotValues:
        return slots.fill(self._get_discount(len(slots)))

    def slot_values(self, rewards: SlateSlotValues) -> SlateSlotValues:
        return rewards.replace(torch.pow(2.0, rewards.values) - 1.0)

class NDCGSlateMetric(DCGSlateMetric):
    def __init__(self, item_rewards: SlateItemValues, device=None):
        super().__init__(device)
        self._sorted_items, _ = item_rewards.sort()
        self._item_rewards = item_rewards
        self._idcg = {}

    def slot_weights(self, slots: SlateSlots) -> SlateSlotValues:
        slate_size = len(slots)
        assert len(self._sorted_items) >= slate_size
        if slate_size not in self._idcg:
            i_slate = make_slate(slots, self._sorted_items[:slate_size])
            idcg = super().calculate_reward(
                slots,
                i_slate.slot_values(self._item_rewards),
                None,
                super().slot_weights(slots),
            )
            self._idcg[slate_size] = idcg
        else:
            idcg = self._idcg[slate_size]
        return slots.fill(
            torch.zeros(slate_size, dtype=torch.double)
            if idcg == 0
            else self._get_discount(slate_size) / idcg
        )

class ERRSlateMetric(SlateMetric):
    def __init__(self, max_reward: float, device=None):
        super().__init__(device)
        self._max_reward = max_reward

    def slot_weights(self, slots: SlateSlots) -> SlateSlotValues:
        return slots.fill([1.0 / (r + 1) for r in range(len(slots))])

    def slot_values(self, rewards: SlateSlotValues) -> SlateSlotValues:
        d = torch.tensor(self._max_reward, device=self._device).pow(2.0)
        r = (torch.pow(2.0, rewards.values.clamp(0.0, self._max_reward)) - 1.0) / d
        p = 1.0
        err = torch.zeros(len(rewards), dtype=torch.double, device=self._device)
        for i in range(len(rewards)):
            ri = r[i]
            err[i] = p * ri
            p = p * (1.0 - ri.item())
        return rewards.replace(err)

class SlateModel(ABC):
    """
    Model providing item relevance/reward, slot examination (click) distribution
    """

    @abstractmethod
    def item_rewards(self, context: SlateContext) -> SlateItemValues:
        """
        Returns each item's relevance under the context
        Args:
            context:

        Returns:
            Item relevances
        """
        pass

    def slot_probabilities(self, context: SlateContext) -> SlateSlotValues:
        """
        Returns each slot/positions's probability independent of showing item,
        used in PBM estimator
        Args:
            context:

        Returns:

        """
        return context.slots.fill(torch.ones(len(context.slots), dtype=torch.double))

@dataclass(frozen=True)
class LogSample:
    context: SlateContext
    metric: SlateMetric
    log_slate: Slate
    log_reward: Reward
    _log_slate_probability: Probability = float("nan")
    # probability for each item being places at each slot
    _log_slot_item_probabilities: Optional[SlateSlotItemProbabilities] = None
    # item probability distribution from behavior policy
    _log_item_probabilities: Optional[SlateItemProbabilities] = None
    _tgt_slate_probability: Probability = float("nan")
    _tgt_slot_item_probabilities: Optional[SlateSlotItemProbabilities] = None
    # item probability distribution from target policy
    _tgt_item_probabilities: Optional[SlateItemProbabilities] = None
    # gt_item_rewards: Optional[SlateItemValues] = None
    # pre-calculated ground truth for target policy
    ground_truth_reward: Reward = float("nan")
    # context dependent slot weights (e.g. DCG or ERR weights), used by PBM
    slot_weights: Optional[SlateSlotValues] = None
    # item/action independent examination probabilities of each slot, used by PBM
    slot_probabilities: Optional[SlateSlotValues] = None
    # features associated with the slate, to train direct model
    item_features: SlateItemFeatures = None

    def validate(self):
        slate_size = len(self.context.slots)
        item_size = len(self.items)
        assert len(self.log_slate) == slate_size
        assert (
            math.isnan(self._log_slate_probability)
            or self._log_slate_probability <= 1.0
        )
        assert (
            math.isnan(self._tgt_slate_probability)
            or self._tgt_slate_probability <= 1.0
        )
        assert (
            self._log_slot_item_probabilities is None
            or len(self._log_slot_item_probabilities) == slate_size
        )
        assert (
            self._log_item_probabilities is None
            or len(self._log_item_probabilities) == item_size
        )
        assert (
            self._tgt_slot_item_probabilities is None
            or len(self._tgt_slot_item_probabilities) == slate_size
        )
        assert (
            self._tgt_item_probabilities is None
            or len(self._tgt_item_probabilities) == item_size
        )
        assert self.slot_weights is None or len(self.slot_weights) == slate_size
        assert (
            self.slot_probabilities is None
            or len(self.slot_probabilities) == slate_size
        )

    def log_slot_item_expectations(
        self, slots: SlateSlots
    ) -> Optional[SlateSlotItemExpectations]:
        if self._log_slot_item_probabilities is not None:
            return self._log_slot_item_probabilities.slot_item_expectations()
        if self._log_item_probabilities is not None:
            return self._log_item_probabilities.slot_item_expectations(slots)
        return None

    def log_slate_probability(self, slate: Optional[Slate] = None) -> float:
        if not math.isnan(self._log_slate_probability):
            return self._log_slate_probability
        if slate is None:
            slate = self.log_slate
        if self._log_slot_item_probabilities is not None:
            return self._log_slot_item_probabilities.slate_probability(slate)
        if self._log_item_probabilities is not None:
            return self._log_item_probabilities.slate_probability(slate)
        return 0.0

    def tgt_slot_expectations(
        self, slots: SlateSlots
    ) -> Optional[SlateSlotItemExpectations]:
        if self._tgt_slot_item_probabilities is not None:
            return self._tgt_slot_item_probabilities.slot_item_expectations()
        if self._tgt_item_probabilities is not None:
            return self._tgt_item_probabilities.slot_item_expectations(slots)
        return None

    def tgt_slate_probability(self) -> float:
        if not math.isnan(self._tgt_slate_probability):
            return self._tgt_slate_probability
        if self._tgt_slot_item_probabilities is not None:
            return self._tgt_slot_item_probabilities.slate_probability(self.log_slate)
        if self._tgt_item_probabilities is not None:
            return self._tgt_item_probabilities.slate_probability(self.log_slate)
        return 0.0

    def tgt_slate_space(
        self, slots: SlateSlots
    ) -> Iterable[Tuple[Sequence[SlateItem], float]]:
        if self._tgt_item_probabilities is not None:
            return self._tgt_item_probabilities.slate_space(slots)
        return []

    @property
    def items(self) -> SlateItems:
        if self._log_slot_item_probabilities is not None:
            return self._log_slot_item_probabilities._values[0].items
        if self._log_item_probabilities is not None:
            return self._log_item_probabilities.items
        return SlateItems(0)

@dataclass(frozen=True)
class SlateEstimatorInput:
    samples: Sequence[LogSample]

    def validate(self):
        for s in self.samples:
            s.validate()

class SlateEstimator(Estimator):
    @abstractmethod
    def _evaluate_sample(
        self, sample: LogSample, logger: logging.Logger
    ) -> Optional[EstimatorSampleResult]:
        pass

class DMEstimator(SlateEstimator):
    """
    Direct Method estimator
    """

    def __init__(self, trainer: Trainer, training_sample_ratio: float, device=None):
        super().__init__(device)
        self._trainer = trainer
        self._training_sample_ratio = training_sample_ratio

    def _train_model(
        self, samples: Sequence[LogSample], logger: logging.Logger
    ) -> Optional[Iterable[LogSample]]:
        if self._trainer is None:
            logger.error("Target model trainer is none, DM is not available")
            return None
        self._trainer.reset()
        logger.info("  training direct model...")
        st = time.perf_counter()
        sample_size = len(samples)
        if self._training_sample_ratio > 0.0 and self._training_sample_ratio < 1.0:
            training_samples = range(int(sample_size * self._training_sample_ratio))
        else:
            training_samples = range(sample_size)
        train_x = []
        train_y = []
        vali_mask = [True] * len(samples)
        for i in training_samples:
            sample = samples[i]
            if sample.item_features is None:
                continue
            slate_features = sample.log_slate.slot_features(sample.item_features)
            train_x.append(slate_features.features.flatten())
            train_y.append(sample.log_reward)
            vali_mask[i] = False
        if len(train_x) == 0:
            logger.error("Slate features not provided, DM is not available")
            return None
        train_x = torch.stack(train_x)
        train_y = torch.tensor(train_y, dtype=torch.double, device=train_x.device)
        vali_x = []
        vali_y = []
        evaluate_samples = []
        for mask, sample in zip(vali_mask, samples):
            if not mask or sample.item_features is None:
                continue
            slate_features = sample.log_slate.slot_features(sample.item_features)
            vali_x.append(slate_features.features.flatten())
            vali_y.append(sample.log_reward)
            evaluate_samples.append(sample)
        if len(vali_x) == 0:
            vali_x = train_x.detach().clone()
            vali_y = train_y.detach().clone()
            evaluate_samples = samples
        else:
            vali_x = torch.stack(vali_x)
            vali_y = torch.tensor(vali_y, dtype=torch.double, device=vali_x.device)
        training_data = TrainingData(train_x, train_y, None, vali_x, vali_y, None)
        self._trainer.train(training_data)
        logger.info(f"  training direct model done: {time.perf_counter() - st}s")

        return evaluate_samples

    def _evaluate_sample(
        self, sample: LogSample, logger: logging.Logger
    ) -> Optional[EstimatorSampleResult]:
        slots = sample.context.slots
        tgt_slate_space = sample.tgt_slate_space(slots)
        features = []
        probs = []
        for items, prob in tgt_slate_space:
            slate = make_slate(slots, items)
            slate_features = slate.slot_features(sample.item_features)
            features.append(slate_features.features.flatten())
            probs.append(prob)
        preds = self._trainer.predict(torch.stack(features), device=self._device)
        tgt_reward = torch.dot(
            preds.scores, torch.tensor(probs, dtype=torch.double, device=self._device)
        )
        return EstimatorSampleResult(
            sample.log_reward,
            tgt_reward.item(),
            sample.ground_truth_reward,
            float("nan"),
        )

    def evaluate(
        self, input: SlateEstimatorInput, *kwargs
    ) -> Optional[EstimatorResult]:
        input.validate()
        logger = Estimator.logger()
        samples = self._train_model(input.samples, logger)
        if samples is None:
            return None

        log_avg = RunningAverage()
        tgt_avg = RunningAverage()
        gt_avg = RunningAverage()
        for sample in samples:
            result = self._evaluate_sample(sample, logger)
            if result is None:
                continue
            log_avg.add(result.log_reward)
            tgt_avg.add(result.target_reward)
            gt_avg.add(result.ground_truth_reward)
        return EstimatorResult(
            log_avg.average, tgt_avg.average, gt_avg.average, tgt_avg.count
        )

    def __repr__(self):
        return (
            f"DMEstimator(trainer({self._trainer.name})"
            f",ratio({self._training_sample_ratio}),device({self._device}))"
        )

class IPSEstimator(SlateEstimator):
    def __init__(
        self, weight_clamper: Clamper = None, weighted: bool = True, device=None
    ):
        super().__init__(device)
        self._weight_clamper = (
            weight_clamper if weight_clamper is not None else Clamper()
        )
        self._weighted = weighted

    def _evaluate_sample(
        self, sample: LogSample, logger: logging.Logger
    ) -> Optional[EstimatorSampleResult]:
        tgt_prob = sample.tgt_slate_probability()
        log_prob = sample.log_slate_probability(sample.log_slate)
        if tgt_prob == log_prob:
            weight = 1.0
        elif tgt_prob <= 0.0:
            weight = 0.0
        elif log_prob <= 0.0:
            return None
        else:
            weight = self._weight_clamper(tgt_prob / log_prob)
        return EstimatorSampleResult(
            sample.log_reward,
            sample.log_reward * weight,
            sample.ground_truth_reward,
            weight,
        )

    def evaluate(
        self, input: SlateEstimatorInput, *kwargs
    ) -> Optional[EstimatorResult]:
        input.validate()
        logger = Estimator.logger()
        log_avg = RunningAverage()
        tgt_avg = RunningAverage()
        acc_weight = RunningAverage()
        gt_avg = RunningAverage()
        zw = 0
        for sample in input.samples:
            result = self._evaluate_sample(sample, logger)
            if result is None:
                zw += 1
                continue
            log_avg.add(result.log_reward)
            tgt_avg.add(result.target_reward)
            gt_avg.add(result.ground_truth_reward)
            acc_weight.add(result.weight)
            if result.weight == 0.0:
                zw += 1
        logging.info(
            f"IPSEstimator invalid sample pct: {zw * 100 / len(input.samples)}%"
        )
        if tgt_avg.count == 0:
            return None
        if self._weighted:
            estimated = tgt_avg.total / acc_weight.total
            return EstimatorResult(
                log_avg.average, estimated, gt_avg.average, acc_weight.average
            )
        else:
            return EstimatorResult(
                log_avg.average, tgt_avg.average, gt_avg.average, tgt_avg.count
            )

    def __repr__(self):
        return (
            f"IPSEstimator(weight_clamper({self._weight_clamper})"
            f",weighted({self._weighted}),device({self._device}))"
        )

class DoublyRobustEstimator(DMEstimator):
    def __init__(
        self,
        trainer: Trainer,
        training_sample_ratio: float,
        weight_clamper: Clamper = None,
        weighted: bool = False,
        device=None,
    ):
        super().__init__(trainer, training_sample_ratio, device)
        self._weight_clamper = (
            weight_clamper if weight_clamper is not None else Clamper()
        )
        self._weighted = weighted

    def _evaluate_sample(
        self, sample: LogSample, logger: logging.Logger
    ) -> Optional[EstimatorSampleResult]:
        slots = sample.context.slots
        if self._trainer.is_trained:
            tgt_slate_space = sample.tgt_slate_space(slots)
            features = []
            probs = []
            for items, prob in tgt_slate_space:
                slate = make_slate(slots, items)
                slate_features = slate.slot_features(sample.item_features)
                features.append(slate_features.features.flatten())
                probs.append(prob)
            preds = self._trainer.predict(torch.stack(features), device=self._device)
            dm_reward = torch.dot(
                preds.scores,
                torch.tensor(probs, dtype=torch.double, device=self._device),
            ).item()
            log_slate_feature = sample.log_slate.slot_features(sample.item_features)
            pred = self._trainer.predict(
                torch.unsqueeze(log_slate_feature.features.flatten(), dim=0),
                device=self._device,
            )
            log_dm_reward = pred.scores[0].item()
        else:
            dm_reward = 0.0
            log_dm_reward = 0.0
        tgt_prob = sample.tgt_slate_probability()
        log_prob = sample.log_slate_probability(sample.log_slate)
        if tgt_prob == log_prob:
            weight = 1.0
        elif tgt_prob <= 0.0:
            weight = 0.0
        elif log_prob <= 0.0:
            return None
        else:
            weight = self._weight_clamper(tgt_prob / log_prob)
        target_reward = (sample.log_reward - log_dm_reward) * weight + dm_reward
        return EstimatorSampleResult(
            sample.log_reward, target_reward, sample.ground_truth_reward, weight
        )

    def evaluate(
        self, input: SlateEstimatorInput, *kwargs
    ) -> Optional[EstimatorResult]:
        input.validate()
        logger = Estimator.logger()
        samples = self._train_model(input.samples, logger)
        if samples is None:
            samples = input.samples

        log_avg = RunningAverage()
        tgt_avg = RunningAverage()
        acc_weight = RunningAverage()
        gt_avg = RunningAverage()
        for sample in samples:
            result = self._evaluate_sample(sample, logger)
            if result is None:
                continue
            log_avg.add(result.log_reward)
            tgt_avg.add(result.target_reward)
            acc_weight.add(result.weight)
            gt_avg.add(result.ground_truth_reward)
        if self._weighted:
            estimated = tgt_avg.total / acc_weight.total
            return EstimatorResult(
                log_avg.average, estimated, gt_avg.average, acc_weight.average
            )
        else:
            return EstimatorResult(
                log_avg.average, tgt_avg.average, gt_avg.average, tgt_avg.count
            )

    def __repr__(self):
        return (
            f"DoublyRobustEstimator(trainer({self._trainer.name})"
            f",ratio({self._training_sample_ratio})"
            f",weight_clamper({self._weight_clamper})"
            f",weighted({self._weighted}),device({self._device}))"
        )

class PseudoInverseEstimator(SlateEstimator):
    """
    Estimator from reference 2
    """

    def __init__(
        self, weight_clamper: Clamper = None, weighted: bool = True, device=None
    ):
        super().__init__(device)
        self._weight_clamper = (
            weight_clamper if weight_clamper is not None else Clamper()
        )
        self._weighted = weighted

    def _evaluate_sample(
        self, sample: LogSample, logger: logging.Logger
    ) -> Optional[EstimatorSampleResult]:
        log_slot_expects = sample.log_slot_item_expectations(sample.context.slots)
        if log_slot_expects is None:
            logger.warning(f"Log slot distribution not available")
            return None
        tgt_slot_expects = sample.tgt_slot_expectations(sample.context.slots)
        if tgt_slot_expects is None:
            logger.warning(f"Target slot distribution not available")
            return None
        log_indicator = log_slot_expects.values_tensor(self._device)
        tgt_indicator = tgt_slot_expects.values_tensor(self._device)
        lm = len(sample.context.slots) * len(sample.items)
        gamma = torch.as_tensor(
            np.linalg.pinv(
                torch.mm(
                    log_indicator.view((lm, 1)), log_indicator.view((1, lm))
                ).numpy()
            )
        )
        # torch.pinverse is not very stable
        # gamma = torch.pinverse(
        #     torch.mm(log_indicator.view((lm, 1)), log_indicator.view((1, lm)))
        # )
        ones = sample.log_slate.one_hots(sample.items, self._device)
        weight = self._weight_clamper(
            torch.mm(tgt_indicator.view((1, lm)), torch.mm(gamma, ones.view((lm, 1))))
        ).item()
        return EstimatorSampleResult(
            sample.log_reward,
            sample.log_reward * weight,
            sample.ground_truth_reward,
            weight,
        )

    def evaluate(
        self, input: SlateEstimatorInput, *kwargs
    ) -> Optional[EstimatorResult]:
        input.validate()
        logger = Estimator.logger()
        log_avg = RunningAverage()
        tgt_avg = RunningAverage()
        acc_weight = RunningAverage()
        gt_avg = RunningAverage()
        zw = 0
        for sample in input.samples:
            result = self._evaluate_sample(sample, logger)
            if result is None:
                zw += 1
                continue
            log_avg.add(result.log_reward)
            tgt_avg.add(result.target_reward)
            gt_avg.add(result.ground_truth_reward)
            acc_weight.add(result.weight)
            if result.weight == 0.0:
                zw += 1
            if tgt_avg.count % 1000 == 0:
                logger.info(f"  PseudoInverseEstimator: processed {tgt_avg.count}")
        logging.info(
            f"PseudoInverseEstimator invalid sample pct: {zw * 100 / len(input.samples)}%"
        )
        if tgt_avg.count == 0:
            return None
        if self._weighted:
            estimated = tgt_avg.total / acc_weight.total
            return EstimatorResult(
                log_avg.average, estimated, gt_avg.average, acc_weight.average
            )
        else:
            return EstimatorResult(
                log_avg.average, tgt_avg.average, gt_avg.average, tgt_avg.count
            )

    def __repr__(self):
        return (
            f"PseudoInverseEstimator(weight_clamper({self._weight_clamper})"
            f",weighted({self._weighted}),device({self._device}))"
        )

class PBMEstimator(SlateEstimator):
    """
    Estimator from reference 1: Position-Based Click Model
    """

    def __init__(
        self, weight_clamper: Clamper = None, weighted: bool = True, device=None
    ):
        super().__init__(device)
        self._weight_clamper = (
            weight_clamper if weight_clamper is not None else Clamper()
        )
        self._weighted = weighted

    def _evaluate_sample(
        self, sample: LogSample, logger: logging.Logger
    ) -> Optional[EstimatorSampleResult]:
        log_slot_expects = sample.log_slot_item_expectations(sample.context.slots)
        if log_slot_expects is None:
            logger.warning(f"  Log slot distribution not available")
            return None
        tgt_slot_expects = sample.tgt_slot_expectations(sample.context.slots)
        if tgt_slot_expects is None:
            logger.warning(f"  Target slot distribution not available")
            return None
        slate_size = len(sample.context.slots)
        slot_weights = sample.slot_weights
        if slot_weights is None:
            slot_weights = SlateSlotValues(torch.ones(slate_size, dtype=torch.double))
        weights = slot_weights.values.to(device=self._device)
        if sample.slot_probabilities is not None:
            weights *= sample.slot_probabilities.values
        h = torch.zeros(slate_size, dtype=torch.double, device=self._device)
        p = torch.zeros(slate_size, dtype=torch.double, device=self._device)
        i = 0
        for slot, item in sample.log_slate:
            h[i] = tgt_slot_expects[slot][item]
            p[i] = log_slot_expects[slot][item]
            i += 1
        nu = torch.tensordot(h, weights, dims=([0], [0]))
        de = torch.tensordot(p, weights, dims=([0], [0]))
        if nu == de:
            weight = 1.0
        elif nu == 0:
            weight = 0.0
        elif de == 0:
            return None
        else:
            weight = self._weight_clamper(nu / de)
        return EstimatorSampleResult(
            sample.log_reward,
            sample.log_reward * weight,
            sample.ground_truth_reward,
            weight,
        )

    def evaluate(
        self, input: SlateEstimatorInput, *kwargs
    ) -> Optional[EstimatorResult]:
        input.validate()
        logger = Estimator.logger()
        log_avg = RunningAverage()
        tgt_avg = RunningAverage()
        acc_weight = RunningAverage()
        gt_avg = RunningAverage()
        zw = 0
        for sample in input.samples:
            result = self._evaluate_sample(sample, logger)
            if result is None:
                zw += 1
                continue
            log_avg.add(result.log_reward)
            tgt_avg.add(result.target_reward)
            gt_avg.add(result.ground_truth_reward)
            acc_weight.add(result.weight)
            if result.weight == 0.0:
                zw += 1
            if tgt_avg.count % 1000 == 0:
                logger.info(f"  PBMEstimator: processed {tgt_avg.count}")
        logging.info(
            f"PBMEstimator invalid sample pct: {zw * 100 / len(input.samples)}%"
        )
        if tgt_avg.count == 0:
            return None
        if self._weighted:
            estimated = tgt_avg.total / acc_weight.total
            return EstimatorResult(
                log_avg.average, estimated, gt_avg.average, acc_weight.average
            )
        else:
            return EstimatorResult(
                log_avg.average, tgt_avg.average, gt_avg.average, tgt_avg.count
            )

    def __repr__(self):
        return (
            f"PBMEstimator(weight_clamper({self._weight_clamper})"
            f",weighted({self._weighted}),device({self._device}))"
        )


from __future__ import absolute_import, unicode_literals, print_function

import os, re, math
from ctypes import *
import numpy as N
from numpy.ctypeslib import as_array

# don't bother with parsing error
try:
    lib = cdll.LoadLibrary('libnest3.so')
except:
    lib = cdll.LoadLibrary(os.path.dirname(__file__) + '/libnest3.so')

# if we want to do OS X version detection:
# import platform
# if platform.system() == 'Darwin'
# '.'.join(platform.mac_ver().split('.')[:2]) --> 10.X

# libstempo.multinest.run borrows heavily from Johannes Buchner's pymultinest;
# it requires MultiNest v3.2 patched with cwrapper.f90

def run(LogLikelihood,
    Prior,
    n_dims,
    n_params = None,
    n_clustering_params = None, wrapped_params = None,
    importance_nested_sampling = True,
    multimodal = True, const_efficiency_mode = False, n_live_points = 400,
    evidence_tolerance = 0.5, sampling_efficiency = 0.8,
    n_iter_before_update = 100, null_log_evidence = -1e90,
    max_modes = 100, mode_tolerance = -1e90,
    outputfiles_basename = "./multinest-", seed = -1, verbose = False,
    resume = True, context = None, write_output = True, log_zero = -1e100,
    max_iter = 0, init_MPI = True, dump_callback = None):
    """
    Runs MultiNest

    The most important parameters are the two log-probability functions Prior
    and LogLikelihood. They are called by MultiNest.

    Prior should transform the unit cube into the parameter cube. Here
    is an example for a uniform prior::

        def Prior(cube, ndim, nparams):
            for i in range(ndim):
                cube[i] = cube[i] * 10 * math.pi

    The LogLikelihood function gets this parameter cube and should
    return the logarithm of the likelihood.
    Here is the example for the eggbox problem::

        def Loglike(cube, ndim, nparams):
            chi = 1.

            for i in range(ndim):
                chi *= math.cos(cube[i] / 2.)
            return math.pow(2. + chi, 5)

    Some of the parameters are explained below. Otherwise consult the
    MultiNest documentation.

    @param importance_nested_sampling:
        If True, Multinest will use Importance Nested Sampling (INS). Read http://arxiv.org/abs/1306.2144
        for more details on INS. Please read the MultiNest README file before using the INS in MultiNest v3.0.

    @param n_params:
        Total no. of parameters, should be equal to ndims in most cases
        but if you need to store some additional
        parameters with the actual parameters then you need to pass
        them through the likelihood routine.

    @param sampling_efficiency:
        defines the sampling efficiency. 0.8 and 0.3 are recommended
        for parameter estimation & evidence evalutation
        respectively.
        use 'parameter' or 'model' to select the respective default
        values

    @param mode_tolerance:
        MultiNest can find multiple modes & also specify which samples belong to which mode. It might be
        desirable to have separate samples & mode statistics for modes with local log-evidence value greater than a
        particular value in which case Ztol should be set to that value. If there isn't any particularly interesting
        Ztol value, then Ztol should be set to a very large negative number (e.g. -1e90).

    @param evidence_tolerance:
        A value of 0.5 should give good enough accuracy.

    @param n_clustering_params:
        If mmodal is T, MultiNest will attempt to separate out the
        modes. Mode separation is done through a clustering
        algorithm. Mode separation can be done on all the parameters
        (in which case nCdims should be set to ndims) & it
        can also be done on a subset of parameters (in which case
        nCdims < ndims) which might be advantageous as
        clustering is less accurate as the dimensionality increases.
        If nCdims < ndims then mode separation is done on
        the first nCdims parameters.

    @param null_log_evidence:
        If mmodal is T, MultiNest can find multiple modes & also specify
        which samples belong to which mode. It might be
        desirable to have separate samples & mode statistics for modes
        with local log-evidence value greater than a
        particular value in which case nullZ should be set to that
        value. If there isn't any particulrly interesting
        nullZ value, then nullZ should be set to a very large negative
        number (e.g. -1.d90).

    @param init_MPI:
        initialize MPI routines?, relevant only if compiling with MPI

    @param log_zero:
        points with loglike < logZero will be ignored by MultiNest

    @param max_iter:
        maximum number of iterations. 0 is unlimited.

    @param write_output:
        write output files? This is required for analysis.

    @param dump_callback:
        a callback function for dumping the current status

    """

    if n_params == None:
        n_params = n_dims
    if n_clustering_params == None:
        n_clustering_params = n_dims
    if wrapped_params == None:
        wrapped_params = [0] * n_dims

    WrappedType = c_int * len(wrapped_params)
    wraps = WrappedType(*wrapped_params)

    if sampling_efficiency == 'parameter':
        sampling_efficiency = 0.8
    if sampling_efficiency == 'model':
        sampling_efficiency = 0.3

    # MV 20130923

    loglike_type = CFUNCTYPE(c_double,
                             POINTER(c_double),c_int,c_int,c_void_p)

    dumper_type  = CFUNCTYPE(c_void_p,
                             c_int,c_int,c_int,
                             POINTER(c_double),POINTER(c_double),POINTER(c_double),
                             c_double,c_double,c_double,c_void_p)

    if hasattr(LogLikelihood,'loglike') and hasattr(Prior,'remap') and hasattr(Prior,'prior'):
        def loglike(cube,ndim,nparams,nullcontext):
            # we're not using context with libstempo.like objects

            pprior = Prior.premap(cube)

            # mappers are supposed to throw a ValueError if they get out of range
            try:
                pars = Prior.remap(cube)
            except ValueError:
                return -N.inf

            prior = pprior * Prior.prior(pars)

            return -N.inf if not prior else math.log(prior) + LogLikelihood.loglike(pars)
    else:
        def loglike(cube,ndim,nparams,nullcontext):
            # it's actually easier to use the context, if any, at the Python level
            # and pass a null pointer to MultiNest...

            args = [cube,ndim,nparams] + ([] if context is None else context)

            if Prior:
                Prior(*args)

            return LogLikelihood(*args)

    def dumper(nSamples,nlive,nPar,
               physLive,posterior,paramConstr,
               maxLogLike,logZ,logZerr,nullcontext):

        if dump_callback:
            # It's not clear to me what the desired PyMultiNest dumper callback
            # syntax is... but this should pass back the right numpy arrays,
            # without copies. Untested!
            pc =  as_array(paramConstr,shape=(nPar,4))

            dump_callback(nSamples,nlive,nPar,
                          as_array(physLive,shape=(nPar+1,nlive)).T,
                          as_array(posterior,shape=(nPar+2,nSamples)).T,
                          (pc[0,:],pc[1,:],pc[2,:],pc[3,:]),    # (mean,std,bestfit,map)
                          maxLogLike,logZ,logZerr)

    # MV 20130923: currently we support only multinest 3.2 (24 parameters),
    # but it would not be a problem to build up the parameter list dynamically

    lib.run(c_bool(importance_nested_sampling),c_bool(multimodal),c_bool(const_efficiency_mode),
            c_int(n_live_points),c_double(evidence_tolerance),
            c_double(sampling_efficiency),c_int(n_dims),c_int(n_params),
            c_int(n_clustering_params),c_int(max_modes),
            c_int(n_iter_before_update),c_double(mode_tolerance),
            create_string_buffer(outputfiles_basename.encode()),    # MV 20130923: need a regular C string
            c_int(seed),wraps,
            c_bool(verbose),c_bool(resume),
            c_bool(write_output),c_bool(init_MPI),
            c_double(log_zero),c_int(max_iter),
            loglike_type(loglike),dumper_type(dumper),
            c_void_p(0))

class multinestdata(dict):
    pass

class multinestpar(object):
    pass

# where are the multinest files?
def _findfiles(multinestrun,dirname,suffix='-post_equal_weights.dat'):
    # try chains/multinestrun-...
    #     chains/multinestrun/multinestrun-...
    root = [dirname + '/',dirname + '/' + multinestrun]

    # and if multinestrun is something like pulsar-model,
    # try chains/pulsar/model/pulsar-model-...
    if '-' in multinestrun:
        tokens = multinestrun.split('-')[:-1]
        pulsar, model = '-'.join(tokens[:-1]), tokens[-1]
        root.append(dirname + '/' + pulsar + '/' + model)

    return filter(lambda r: os.path.isfile(r + '/' + multinestrun + suffix),root)

def _getcomment(ret,filename):
    try:
        ret.comment = open(filename,'r').read()
    except IOError:
        pass

def _getmeta(ret,filename):
    try:
        meta = N.load(filename)
    except IOError:
        return

    ret.parnames  = list(meta['name'])
    ret.tempopars = list(meta['val'])   # somewhat legacy?
    ret.tempo = {}

    ml = N.argmax(ret.data[:,-1])

    for i,par in enumerate(ret.parnames):
        ret[par] = multinestpar()

        try:
            ret[par].val, ret[par].err = N.mean(ret.data[:,i]) + meta['offset'][i], math.sqrt(N.var(ret.data[:,i]))
            ret[par].offset = meta['offset'][i]
        except ValueError:
            ret[par].val, ret[par].err = N.mean(ret.data[:,i]), math.sqrt(N.var(ret.data[:,i]))

        if 'ml' in meta.dtype.names:
            ret[par].ml = meta['ml'][i]
        else:
            ret[par].ml = ret.data[ml,i] + (meta['offset'][i] if 'offset' in meta.dtype.names else 0)

        ret.tempo[par] = multinestpar()
        ret.tempo[par].val, ret.tempo[par].err = meta['val'][i], meta['err'][i]

def load_mcmc(mcrun,dirname='.'):
    root = _findfiles(mcrun,dirname,'-chain.npy')

    ret = multinestdata()
    ret.dirname = root[0]

    alldata = N.load('{0}/{1}-chain.npy'.format(root[0],mcrun))

    # keep all the steps
    ret.data = alldata[:,:]

    _getmeta(ret,'{0}/{1}-meta.npy'.format(root[0],mcrun))
    _getcomment(ret,'{0}/{1}-comment.txt'.format(root[0],mcrun))

    return ret

def load_emcee(emceerun,dirname='.',chains=False):
    root = _findfiles(emceerun,dirname,'-chain.npy')

    ret = multinestdata()
    ret.dirname = root[0]

    alldata = N.load('{0}/{1}-chain.npy'.format(root[0],emceerun))

    # keep the last iteration of the walker cloud
    ret.data = alldata[:,-1,:]

    if chains:
        ret.chains = alldata

    _getmeta(ret,'{0}/{1}-meta.npy'.format(root[0],emceerun))
    _getcomment(ret,'{0}/{1}-comment.txt'.format(root[0],emceerun))

    return ret

def load(multinestrun,dirname='.'):
    root = _findfiles(multinestrun,dirname,'-post_equal_weights.dat')

    if not root:
        # try to find a tar.gz archive
        import tempfile, tarfile
        root = _findfiles(multinestrun,dirname,'.tar.gz')
        tar = tarfile.open('{0}/{1}.tar.gz'.format(root[0],multinestrun),mode='r|gz')
        root = [tempfile.mkdtemp(prefix='/tmp/')]
        tar.extractall(path=root[0])

    ret = multinestdata()
    ret.dirname = root[0]

    # get data
    ret.data = N.loadtxt('{0}/{1}-post_equal_weights.dat'.format(root[0],multinestrun))[:,:-1]

    # get evidence
    try:
        lines = open('{0}/{1}-stats.dat'.format(root[0],multinestrun),'r').readlines()
        try:
            ret.ev = float(re.search(r'Global Evidence:\s*(\S*)\s*\+/-\s*(\S*)',lines[0]).group(1))
        except:
            ret.ev = float(re.search(r'Global Log-Evidence           :\s*(\S*)\s*\+/-\s*(\S*)',lines[0]).group(1))
    except IOError:
        pass

    # get metadata
    _getmeta(ret,'{0}/{1}-meta.npy'.format(root[0],multinestrun))
    _getcomment(ret,'{0}/{1}-comment.txt'.format(root[0],multinestrun))

    if root[0][:4] == '/tmp':
        import shutil
        shutil.rmtree(root[0])

    return ret

def compress(rootname):
    import sys, os, glob

    dirname, filename = os.path.dirname(rootname), os.path.basename(rootname)

    if filename[-1] == '-':
        filename = filename[:-1]

    files = [filename + '-' + ending for ending in ('.txt','phys_live.points','stats.dat','ev.dat',
                                                    'post_equal_weights.dat','summary.txt','live.points',
                                                    'post_separate.dat','meta.npy','resume.dat','comment.txt')]

    cd = os.getcwd()
    os.chdir(dirname)

    os.system('tar zcf {0}.tar.gz {1}'.format(filename,' '.join(files)))

    files_exclude = [filename + '-' + ending for ending in ('IS.iterinfo','IS.points','IS.ptprob')]

    for f in files + files_exclude:
        if os.path.isfile(f):
            os.unlink(f)

    os.chdir(cd)


# ######################################################################
# Original code:                                                       #
# @author: Robert B. Von Dreele and Brian Toby                         #
# General Structure Analysis System - II (GSAS-II)                     #
# https://subversion.xor.aps.anl.gov/trac/pyGSAS                       #
# Copyright 2010, UChicago Argonne, LLC, Operator of                   #
# Argonne National Laboratory All rights reserved.                     #
#                                                                      #
# Copyright (c) 2014, Brookhaven Science Associates, Brookhaven        #
# National Laboratory. All rights reserved.                            #
#                                                                      #
# Redistribution and use in source and binary forms, with or without   #
# modification, are permitted provided that the following conditions   #
# are met:                                                             #
#                                                                      #
# * Redistributions of source code must retain the above copyright     #
#   notice, this list of conditions and the following disclaimer.      #
#                                                                      #
# * Redistributions in binary form must reproduce the above copyright  #
#   notice this list of conditions and the following disclaimer in     #
#   the documentation and/or other materials provided with the         #
#   distribution.                                                      #
#                                                                      #
# * Neither the name of the Brookhaven Science Associates, Brookhaven  #
#   National Laboratory nor the names of its contributors may be used  #
#   to endorse or promote products derived from this software without  #
#   specific prior written permission.                                 #
#                                                                      #
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS  #
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT    #
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS    #
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE       #
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,           #
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES   #
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR   #
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)   #
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,  #
# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OTHERWISE) ARISING   #
# IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE   #
# POSSIBILITY OF SUCH DAMAGE.                                          #
########################################################################

"""
This is the module for reading files created in GSAS file formats
https://subversion.xor.aps.anl.gov/trac/pyGSAS
"""
from __future__ import absolute_import, division, print_function
import os
import numpy as np

def gsas_reader(file):
    """
    Parameters
    ----------
    file: str
        GSAS powder data file

    Returns
    --------
     tth : ndarray
        twotheta values (degrees) shape (N, ) array

    intensity : ndarray
        intensity values shape (N, ) array

    err : ndarray
        error value of intensity shape(N, ) array
    """

    if os.path.splitext(file)[1] != ".gsas":
        raise IOError("Provide a file with diffraction data saved in GSAS,"
                      " file extension has to be .gsas ")

    # find the file mode, could be 'std', 'esd', 'fxye'
    with open(file, 'r') as fi:
        S = fi.readlines()[1]
        mode = S.split()[9]

    try:
        tth, intensity, err = _func_look_up[mode](file)
    except KeyError:
        raise ValueError("Provide a correct mode of the GSAS file, "
                         "file modes could be in 'STD', 'ESD', 'FXYE' ")

    return tth, intensity, err

def _get_fxye_data(file):
    """
    Parameters
    ----------
    file: str
        GSAS powder data file

    Return
    ------
    tth : ndarray
        twotheta values (degrees) shape (N, ) array

    intensity : ndarray
        intensity values shape (N, ) array

    err : ndarray
        error value of intensity shape(N, ) array

    """
    tth = []
    intensity = []
    err = []

    with open(file, 'r') as fi:
        S = fi.readlines()[2:]
        for line in S:
            vals = line.split()

            tth.append(float(vals[0]))
            f = float(vals[1])
            s = float(vals[2])

            if f <= 0.0:
                intensity.append(0.0)
            else:
                intensity.append(float(vals[1]))

            if s > 0.0:
                err.append(1.0/float(vals[2])**2)
            else:
                err.append(0.0)

    return [np.array(tth), np.array(intensity), np.array(err)]

def _get_esd_data(file):
    """
    Parameters
    ----------
    file: str
        GSAS powder data file

    Return
    ------
    tth : ndarray
        twotheta values (degrees) shape (N, ) array

    intensity : ndarray
        intensity values shape (N, ) array

    err : ndarray
        error value of intensity shape(N, ) array

    """
    tth = []
    intensity = []
    err = []

    with open(file, 'r') as fi:
        S = fi.readlines()[1:]

        # convert from centidegrees to degrees
        start = float(S[0].split()[5])/100.0
        step = float(S[0].split()[6])/100.0

        j = 0
        for line in S[1:]:
            for i in range(0, 80, 16):
                xi = start + step*j
                yi = _sfloat(line[i: i + 8])
                ei = _sfloat(line[i + 8: i + 16])
                tth.append(xi)

                if yi > 0.0:
                    intensity.append(yi)
                else:
                    intensity.append(0.0)

                if ei > 0.0:
                    err.append(1.0/ei**2)
                else:
                    err.append(0.0)
                j += 1
    return [np.array(tth), np.array(intensity), np.array(err)]

def _get_std_data(file):
    """
    Parameters
    ----------
    file: str
        GSAS powder data file

    Return
    ------
    tth : ndarray
        twotheta values (degrees) shape (N, ) array

    intensity : ndarray
        intensity values shape (N, ) array

    err : ndarray
        error value of intensity shape(N, ) array

    """
    tth = []
    intensity = []
    err = []

    with open(file, 'r') as fi:
        S = fi.readlines()[1:]

        # convert from centidegrees to degrees
        start = float(S[0].split()[5])/100.0
        step = float(S[0].split()[6])/100.0

        # number of data values(two theta or intensity)
        nch = float(S[0].split()[2])

        j = 0
        for line in S[1:]:
            for i in range(0, 80, 8):
                xi = start + step*j
                ni = max(_sint(line[i: i + 2]), 1)
                yi = max(_sfloat(line[i + 2: i + 8]), 0.0)
                if yi:
                    vi = yi/ni
                else:
                    yi = 0.0
                    vi = 0.0
                if j < nch:
                    tth.append(xi)
                    if vi <= 0.:
                        intensity.append(0.)
                        err.append(0.)
                    else:
                        intensity.append(yi)
                        err.append(1.0/vi)
                j += 1
    return [np.array(tth), np.array(intensity), np.array(err)]

# find the which function to use according to mode of the GSAS file
# mode could be "STD", "ESD" or "FXYE"
_func_look_up = {'STD': _get_std_data, 'ESD': _get_esd_data,
                 'FXYE': _get_fxye_data}

def _sfloat(S):
    """
    convert a string to a float, treating an all-blank string as zero
    Parameter
    ---------
    S : str
        string that need to be converted as float treating an
        all-blank string as zero

    Returns
    -------
    float or zero
    """
    if S.strip():
        return float(S)
    else:
        return 0.0

def _sint(S):
    """
    convert a string to an integer, treating an all-blank string as zero
    Parameter
    ---------
    S : str
        string that need to be converted as integer treating an all-blank
        strings as zero

    Returns
    -------
    integer or zero
    """
    if S.strip():
        return int(S)
    else:
        return 0


"""

Grid.py

Author: Jordan Mirocha
Affiliation: University of Colorado at Boulder
Created on: Thu Sep 20 14:18:27 2012

Description:

"""

import copy, types
import numpy as np
from collections import Iterable
from ..util import parse_kwargs, rebin
from ..physics.Hydrogen import Hydrogen
from ..physics.Cosmology import Cosmology
from ..physics.Constants import k_B, cm_per_kpc, s_per_myr, m_H
from ..physics.CrossSections import PhotoIonizationCrossSection

try:
    import chianti.core as cc
    import chianti.util as util
    have_chianti = True
except ImportError:
    from ..util import fake_chianti
    util = fake_chianti()
    have_chianti = False

try:
    from periodic.table import element as ELEMENT
except ImportError:
    from ..util import ELEMENT

tiny_number = 1e-8  # A relatively small species fraction

class Grid(object):
    def __init__(self, dims=64, length_units=cm_per_kpc, start_radius=0.01,
        approx_Salpha=1, approx_lya=0, logarithmic_grid=False):
        """
        Initialize grid object.

        Parameters
        ----------
        dims : int
            Number of resolution elements in grid.
        length_units : float
            Size of domain in centimeters.
        start_radius : float
            Radius (in code units) within which to ignore.

        """

        self.dims = int(dims)
        self.length_units = length_units
        self.start_radius = start_radius
        self.approx_Salpha = approx_Salpha
        self.approx_lya = approx_lya
        self.log_grid = logarithmic_grid

        # Compute cell centers and edges
        if logarithmic_grid:
            self.r_edg = self.r = \
                np.logspace(np.log10(self.R0), np.log10(length_units),
                self.dims + 1)
        else:
            self.r_edg = self.r = \
                np.linspace(self.R0, length_units, self.dims + 1)

        # Compute interior cell walls, spacing, and mid-points
        self.r_int = self.r_edg[0:-1]
        self.dr = np.diff(self.r_edg)
        self.r_mid = rebin(self.r_edg)

        self.zi = 0

    @property
    def zeros_absorbers(self):
        return np.zeros(self.N_absorbers)

    @property
    def zeros_absorbers2(self):
        return np.zeros([self.N_absorbers] * 2)

    @property
    def zeros_grid_x_absorbers(self):
        return np.zeros([self.dims, self.N_absorbers])

    @property
    def zeros_grid_x_absorbers2(self):
        return np.zeros([self.dims, self.N_absorbers, self.N_absorbers])

    @property
    def R0(self):
        """ Start radius in length_units. """
        return self.start_radius * self.length_units

    @property
    def Vsh(self):
        """ Shell volume in length_units**3. """
        if not hasattr(self, '_Vsh_all'):
            self._Vsh_all = self.ShellVolume(self.r_edg[0:-1], self.dr)

        return self._Vsh_all

    @property
    def neutrals(self):
        """ Return list of all neutral species. """
        if not hasattr(self, '_neutral_species'):
            self._neutral_species = []
            for element in self.elements:
                self._neutral_species.append('%s_1' % element)

        return self._neutral_species

    @property
    def ions(self):
        """ Return list of all ionized species. """
        if not hasattr(self, '_ionized_species'):
            neutrals = self.neutrals
            self._ionized_species = []
            for ion in self.all_ions:
                if ion in neutrals:
                    continue

                self._ionized_species.append(ion)

        return self._ionized_species

    @property
    def absorbers(self):
        """ Return list of absorbers (don't include electrons). """
        if not hasattr(self, '_absorbing_species'):
            self._absorbing_species = copy.copy(self.neutrals)
            for parent in self.ions_by_parent:
                self._absorbing_species.extend(self.ions_by_parent[parent][1:-1])

        return self._absorbing_species

    @property
    def N_absorbers(self):
        """ Return number of absorbing species. """
        if not hasattr(self, 'self._num_of_absorbers'):
            absorbers = self.absorbers
            self._num_of_absorbers = int(len(absorbers))

        return self._num_of_absorbers

    @property
    def metals(self):
        """ Return list of anything that is not hydrogen or helium. """
        if not hasattr(self, '_metals'):
            self._metals = []
            self._metal_ions = []
            for element in self.ions_by_parent:
                if element in ['h', 'he']:
                    continue

                self._metals.append(element)
                for ion in self.ions_by_parent[element]:
                    self._metal_ions.append(ion)

        return self._metals

    @property
    def metal_ions(self):
        """ Return list of all metal (not H or He) ions."""
        if not hasattr(self, '_metal_ions'):
            all_metals = self.metals

        return self._metal_ions

    @property
    def species_abundances(self):
        """
        Return dictionary containing abundances of parent
        elements of all ions.
        """
        if not hasattr(self, '_species_abundances'):
            self._species_abundances = {}
            for ion in self.ions_by_parent:
                for state in self.ions_by_parent[ion]:
                    self._species_abundances[state] = \
                        self.element_abundances[self.elements.index(ion)]

        return self._species_abundances

    @property
    def species(self):
        if not hasattr(self, '_species'):
            self._species = []
            for parent in self.ions_by_parent:
                for ion in self.ions_by_parent[parent]:
                    self._species.append(ion)

        return self._species

    @property
    def types(self):
        """
        Return list (matching evolving_fields) with integers describing
        species type:
            0 = neutral
           +1 = ion
           -1 = other
        """

        if not hasattr(self, '_species_types'):
            self._species_types = []
            for species in self.evolving_fields:
                if species in self.neutrals:
                    self._species_types.append(0)
                elif species in self.ions:
                    self._species_types.append(1)
                else:
                    self._species_types.append(-1)

        return self._species_types

    @property # MUST GENERALIZE THIS
    def ioniz_thresholds(self):
        """
        Return dictionary containing ionization threshold energies (in eV)
        for all absorbers.
        """

        if not hasattr(self, '_ioniz_thresholds'):
            self._ioniz_thresholds = {}
            #for absorber in self.absorbers:
            #if absorber == 'h_1':
            self._ioniz_thresholds['h_1'] = 13.6
            #elif absorber == 'he_1':
            self._ioniz_thresholds['he_1'] = 24.4
            #elif absorber == 'he_2':
            self._ioniz_thresholds['he_2'] = 54.4

        return self._ioniz_thresholds

    @property # MUST GENERALIZE THIS
    def bf_cross_sections(self):
        """
        Return dictionary containing functions that compute the bound-free
        absorption cross-sections for all absorbers.
        """

        if not hasattr(self, 'all_xsections'):
            self._bf_xsections = {}
            #for absorber in self.absorbers:
                #ion = cc.continuum(absorber)
                #ion.vernerCross(energy = np.logspace(1, 5, 1000))
                #if absorber == 'h_1':
            self._bf_xsections['h_1'] = lambda E: \
                PhotoIonizationCrossSection(E, species=0)
                #elif absorber == 'he_1':
            self._bf_xsections['he_1'] = lambda E: \
                PhotoIonizationCrossSection(E, species=1)
                #elif absorber == 'he_2':
            self._bf_xsections['he_2'] = lambda E: \
                PhotoIonizationCrossSection(E, species=2)

        return self._bf_xsections

    @property
    def x_to_n(self):
        """
        Return dictionary containing conversion factor between species
        fraction and number density for all species.
        """
        if not hasattr(self, '_x_to_n_converter'):
            self._x_to_n_converter = {}
            for ion in self.all_ions:
                self._x_to_n_converter[ion] = self.n_ref \
                    * self.species_abundances[ion]

        return self._x_to_n_converter

    @property
    def expansion(self):
        if not hasattr(self, '_expansion'):
            self.set_physics()
        return self._expansion

    @property
    def isothermal(self):
        if not hasattr(self, '_isothermal'):
            self.set_physics()
        return self._isothermal

    @property
    def secondary_ionization(self):
        if not hasattr(self, '_secondary_ionization'):
            self.set_physics()
        return self._secondary_ionization

    @property
    def compton_scattering(self):
        if not hasattr(self, '_compton_scattering'):
            self.set_physics()
        return self._compton_scattering

    @property
    def recombination(self):
        if not hasattr(self, '_recombination'):
            self.set_physics()
        return self._recombination

    @property
    def clumping_factor(self):
        if not hasattr(self, '_clumping_factor'):
            self.set_physics()
        return self._clumping_factor

    @property
    def hydr(self):
        if not hasattr(self, '_hydr'):
            self._hydr = Hydrogen(self.cosm, approx_Salpha=self.approx_Salpha,
                approx_lya=self.approx_lya)
        return self._hydr

    @property
    def cosm(self):
        if not hasattr(self, '_cosm'):
            self._cosm = Cosmology()
        return self._cosm

    def initialize(self, data):
        self.set_chemistry()

    def set_physics(self, isothermal=False, compton_scattering=False,
        secondary_ionization=0, expansion=False, recombination='B',
        clumping_factor=1.0):
        self._isothermal = isothermal
        self._compton_scattering = compton_scattering
        self._secondary_ionization = secondary_ionization
        self._expansion = expansion
        self._recombination = recombination

        if type(clumping_factor) is not types.FunctionType:
            self._clumping_factor = lambda z: clumping_factor
        else:
            self._clumping_factor = clumping_factor

        if self._expansion:
            self.set_cosmology()

    @property
    def in_bubbles(self):
        if not hasattr(self, '_in_bubbles'):
            self.set_recombination_rate()

        return self._in_bubbles

    def set_recombination_rate(self, in_bubbles=False):
        self._in_bubbles = in_bubbles

    def set_cosmology(self, initial_redshift=1e3, OmegaMatterNow=0.272,
        OmegaLambdaNow=0.728, OmegaBaryonNow=0.044, HubbleParameterNow=0.702,
        HeliumAbundanceByNumber=0.08, CMBTemperatureNow=2.725,
        approx_highz=False):

        self.zi = initial_redshift
        self._cosm = Cosmology(OmegaMatterNow=OmegaMatterNow,
            OmegaLambdaNow=OmegaLambdaNow, OmegaBaryonNow=OmegaBaryonNow,
            HubbleParameterNow=HubbleParameterNow,
            HeliumAbundanceByNumber=HeliumAbundanceByNumber,
            CMBTemperatureNow=CMBTemperatureNow,
            approx_highz=approx_highz)

    def set_chemistry(self, Z=1, abundances=1.0, energy=False):
        """
        Initialize chemistry.

        This routine sets the chemical composition of the medium being
        simulated.

        Parameters
        ----------
        Z : int, list
            Atomic number(s) of elements to include in calculation.
        abundances : float, list, str
            Abundance(s) (relative to hydrogen) of elements.
            If chiantiPy is installed, can be a string. Some acceptable
            abundance strings are:
                'cosmic', 'sun_photospheric', 'sun_coronal'
        energy : bool
            Solve for internal energy or temperature? (not sure if this is
            used anywhere)

        Example
        -------
        grid = Grid(dims=32)
        grid.set_chemistry(Z=[1,2], abundances='cosmic')

        """

        if type(Z) is not list:
            Z = [Z]

        if type(abundances) not in [str, list]:
            abundances = [abundances]

        self.abundances = abundances

        self.Z = np.array(Z)
        self.ions_by_parent = {} # Ions sorted by parent element in dictionary
        self.parents_by_ion = {} # From ion name, determine parent element
        self.elements = []       # Just a list of element names
        self.all_ions = []       # All ion species
        self.tracked_fields = [] # Anything we're keeping track of
        self.evolving_fields = []# Anything with an ODE we'll later solve

        for i, element in enumerate(self.Z):
            element_name = util.z2element(element)

            self.ions_by_parent[element_name] = []
            self.elements.append(element_name)
            for ion in xrange(element + 1):
                name = util.zion2name(element, ion + 1)
                self.all_ions.append(name)
                self.ions_by_parent[element_name].append(name)
                self.parents_by_ion[name] = element_name
                self.tracked_fields.append(name)
                self.evolving_fields.append(name)

        self.solve_ge = False
        self.evolving_fields.append('e')
        if not self.isothermal:
            if energy:
                self.solve_ge = True
                self.evolving_fields.append('ge')
            else:
                self.evolving_fields.append('Tk')

        # Create blank data fields
        if not hasattr(self, 'data'):
            self.data = {}
            for field in self.tracked_fields:
                self.data[field] = np.zeros(self.dims)

        # Read abundances from chianti
        if type(abundances) is str and have_chianti:
            self.abundances_by_number = util.abundanceRead(abundance)['abundance']
            self.element_abundances = []
            for i, Z in enumerate(self.Z):
                self.element_abundances.append(self.abundances_by_number[Z-1])
        elif type(abundances) is str:
            raise ValueError('If chianti is not installed, must supply abundances by number.')
        else:
            self.abundances_by_number = self.abundances
            self.element_abundances = []
            for i, Z in enumerate(self.Z):
                self.element_abundances.append(self.abundances_by_number[i])

        # Initialize mapping between q-vector and physical quantities (dengo)
        self._set_qmap()

    def set_density(self, rho0=None):
        """
        Initialize gas density, and from that, the hydrogen number density.

        Setting the gas density is necessary for computing the hydrogen
        number density, which normalizes fractional abundances of elements
        to proper number densities of all species.

        Parameters
        ----------
        rho0 : float, array
            Density of medium in g / cm**3. Can be a float (uniform medium),
            or an array of values the same size as the grid itself.

        """

        if isinstance(rho0, Iterable):
            self.data['rho'] = rho0
        else:
            self.data['rho'] = rho0 * np.ones(self.dims)

        if len(self.Z) == 1:
            self.n_H = self.data['rho'] / m_H
            self.n_He = 0.0
        elif len(self.Z) == 2:
            if 2 not in self.Z:
                raise ValueError('Only know how to do H+He gas.')

            self.n_H = (1. - self.cosm.Y) * self.data['rho'] / m_H
            self.n_He = self.cosm.Y * self.data['rho'] / 4. / m_H

        self.n_ref = self.n_H

        #if len(self.Z) == 1:
        #    if self.Z == np.ones(1):
        #        self.abundances_by_number = self.element_abundances = np.ones(1)
        #        if self.expansion:
        #            self.n_H = \
        #                (1. - self.cosm.Y) * self.data['rho'] / m_H
        #        else:
        #            self.n_H = self.data['rho'] / m_H
        #    else:
        #        self.n_H = self.data['rho'] / m_H \
        #            / ELEMENT(self.elements[0]).mass
        #
        #    self.n_ref = copy.deepcopy(self.n_H)
        #    return
        #
        #
        #
        ## Set hydrogen number density (which normalizes all other species)
        #X = 0.0
        #for i in xrange(len(self.abundances_by_number) - 1):
        #    name = util.z2element(i + 1)
        #    if not name.strip():
        #        continue
        #
        #    X += self.abundances_by_number[i] * ELEMENT(name).mass
        #
        ## Set reference number density
        #if 'h' in self.elements:
        #    self.n_H = self.n_ref = self.data['rho'] / m_H / X
        #else:
        #    self.n_H = self.n_ref = self.data['rho'] / m_H \
        #        / ELEMENT(self.elements[0]).mass

    def set_temperature(self, T0):
        """
        Set initial temperature in grid.

        Parameters
        ----------
        T0 : float, array
            Initial temperature in grid. Can be constant value (corresponding
            to uniform medium), or an array of values like the grid.
        """

        if isinstance(T0, Iterable):
            self.data['Tk'] = T0
        else:
            self.data['Tk'] = T0 * np.ones(self.dims)

    def set_ionization(self, Z=None, x=None, state=None, perturb=0):
        """
        Set initial ionization state.  If Z is None, assume constant ion fraction
        of 1 / (1 + Z) for all elements.  Can be overridden by 'state', which can be
        'equilibrium', and maybe eventually other options (e.g. perturbed out of
        equilibrium slightly, perhaps).
        """

        if x is not None:
            self.data[util.zion2name(Z, 1)].fill(1. - x)
            self.data[util.zion2name(Z, 2)].fill(x)

        elif state == 'equilibrium':
            np.seterr(all = 'ignore')   # This tends to produce divide by zero errors
            for Z in self.Z:
                eq = cc.ioneq(Z, self.data['Tk'])

                for i in xrange(1 + Z):
                    mask = np.isnan(eq.Ioneq[i])
                    name = util.zion2name(Z, i + 1)
                    self.data[name][:] = eq.Ioneq[i]
                    self.data[name][mask] = np.ones_like(mask[mask == True])
                    # For some reason chianti sometimes gives nans where
                    # the neutral fraction (in oxygen at least) should be 1.
                    # It only happens when cc.ioneq is given an array of temps,
                    # i.e. everything is fine if you loop over T but that's
                    # way slower.
                    if perturb > 0:
                        tmp = self.data[name] * np.random.normal(loc=1.0,
                            scale=perturb, size=self.dims)
                        tmp[tmp < tiny_number] = tiny_number
                        self.data[name] = copy.copy(tmp)

                # Renormalize
                if perturb > 0:
                    C = 0
                    for i in xrange(1 + Z):
                        name = util.zion2name(Z, i + 1)
                        C += self.data[name]

                    for i in xrange(1 + Z):
                        name = util.zion2name(Z, i + 1)
                        self.data[name] /= C

            np.seterr(all=None)

        elif state == 'neutral':
            for Z in self.Z:

                N = len(self.ions_by_parent[util.z2element(Z)]) - 1

                for i in xrange(1 + Z):
                    name = util.zion2name(Z, i + 1)

                    if i == 0:
                        self.data[name] = np.ones(self.dims) - N * tiny_number
                    else:
                        self.data[name] = tiny_number * np.ones(self.dims)

        else:
            for species in self.all_ions:
                self.data[species].fill(1. \
                    / (1. + util.convertName(species)['Z']))

        # Set electron density
        self._set_electron_fraction()

        if self.solve_ge:
            self.set_gas_energy()

    def set_gas_energy(self):
        """
        Store ge.
        """

        self.data['n'] = self.particle_density(self.data)
        self.data['ge'] = 1.5 * self.data['n'] * k_B * self.data['Tk']

    def set_ics(self, data):
        """
        Simple way of setting all initial conditions at once with a data
        dictionary.
        """

        self.data = {}
        for key in data.keys():
            if type(data[key]) is float:
                self.data[key] = data[key]
                continue

            self.data[key] = data[key].copy()

    def make_clump(self, position = None, radius = None, overdensity = None,
        temperature = None, ionization = None, profile = None):
        """ Create a clump! """

        # Figure out where the clump is
        gridarr = np.linspace(0, 1, self.dims)
        isclump = (gridarr >= (position - radius)) \
                & (gridarr <= (position + radius))

        # First, modify density and temperature
        if profile == 0:
            self.data['rho'][isclump] *= overdensity
            self.data['n'][isclump] *= overdensity
            self.n_H[isclump] *= overdensity
            self.n_ref[isclump] *= overdensity
            self.data['Tk'][isclump] = temperature
        #if profile == 1:
        #    self.data['rho'] += self.data['rho'] * overdensity \
        #        * np.exp(-(gridarr - position)**2 / 2. / radius**2)
        #    self.n_H += self.n_H * overdensity \
        #        * np.exp(-(gridarr - position)**2 / 2. / radius**2)
        #    self.data['T'] -= self.data['T'] * overdensity \
        #        * np.exp(-(gridarr - position)**2 / 2. / radius**2)

        # Need to think more about Gaussian clump T, x.

        # Ionization state - could generalize this more
        for neutral in self.neutrals:
            self.data[neutral][isclump] = 1. - ionization
        for ion in self.ions:
            self.data[ion][isclump] = ionization

        # Reset electron density, particle density, and gas energy
        self._set_electron_fraction()

        if hasattr(self, '_x_to_n_converter'):
            del self._x_to_n_converter

    def _set_electron_fraction(self):
        """
        Set electron density - must have run set_density beforehand.
        """

        self.data['e'] = np.zeros(self.dims)
        for i, Z in enumerate(self.Z):
            for j in np.arange(1, 1 + Z):   # j = number of electrons donated by ion j + 1
                x_i_jp1 = self.data[util.zion2name(Z, j + 1)]
                self.data['e'] += j * x_i_jp1 * self.n_ref \
                    * self.element_abundances[i]

        self.data['e'] /= self.n_H

    def particle_density(self, data, z=0):
        """
        Compute total particle number density.
        """

        n = data['e'].copy()
        #for ion in self.all_ions:
        #    n += data[ion] * self.x_to_n[ion] * (1. + z)**3 \
        #        / (1. + self.zi)**3

        if self.expansion:
            n *= self.cosm.nH(z)

            n += self.cosm.nH(z)

            if 2 in self.Z:
                n += self.cosm.nHe(z)

        else:
            n *= self.n_H

            n += self.n_H

            if 2 in self.Z:
                n += self.n_H * self.abundances[1]

        return n

    def electron_fraction(self, data, z):
        de = np.zeros(self.dims)
        for i, Z in enumerate(self.Z):
            for j in np.arange(1, 1 + Z):   # j = number of electrons donated by ion j + 1
                x_i_jp1 = data[util.zion2name(Z, j + 1)]
                de += j * x_i_jp1 * self.n_ref * (1. + z)**3 / (1. + self.zi)**3 \
                    * self.element_abundances[i]

        return de / self.n_H

    def ColumnDensity(self, data):
        """ Compute column densities for all absorbing species. """

        N = {}
        Nc = {}
        logN = {}
        for absorber in self.absorbers:
            Nc[absorber] = self.dr * data[absorber] * self.x_to_n[absorber]
            N[absorber] = np.cumsum(Nc[absorber])
            logN[absorber] = np.log10(N[absorber])

        return N, logN, Nc

    def _set_qmap(self):
        """
        The vector 'q' is an array containing the values of all ion fractions and the
        gas energy.  This routine sets up the mapping between elements in q and the
        corrresponding physical quantities.

        Will be in order of increasing Z, then de, then ge.
        """

        self.qmap = []
        for species in self.evolving_fields:
            self.qmap.append(species)

    def ShellVolume(self, r, dr):
        """
        Return volume of shell at distance r, thickness dr.
        """

        return 4. * np.pi * ((r + dr)**3 - r**3) / 3.