#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
# graph_tool -- a general graph manipulation python module
#
# Copyright (C) 2006-2018 Tiago de Paula Peixoto <tiago@skewed.de>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import division, absolute_import, print_function
import sys
if sys.version_info < (3,):
    range = xrange

from .. import _degree, _prop, Graph, GraphView, libcore, _get_rng, PropertyMap, \
    conv_pickle_state

import random
from numpy import *
import numpy
from collections import defaultdict

from .. import group_vector_property, ungroup_vector_property

from .. dl_import import dl_import
dl_import("from . import libgraph_tool_inference as libinference")

from . blockmodel import *
from . blockmodel import _bm_test

class OverlapBlockState(BlockState):
    r"""The overlapping stochastic block model state of a given graph.

    Parameters
    ----------
    g : :class:`~graph_tool.Graph`
        Graph to be modelled.
    b : :class:`~graph_tool.PropertyMap` or :class:`numpy.ndarray` (optional, default: ``None``)
        Initial block labels on the vertices or half-edges. If not supplied, it
        will be randomly sampled.
        If the value passed is a vertex property map, it will be assumed to be a
        non-overlapping partition of the vertices. If it is an edge property
        map, it should contain a vector for each edge, with the block labels at
        each end point (sorted according to their vertex index, in the case of
        undirected graphs, otherwise from source to target). If the value is an
        :class:`numpy.ndarray`, it will be assumed to correspond directly to a
        partition of the list of half-edges.
    B : ``int`` (optional, default: ``None``)
        Number of blocks (or vertex groups). If not supplied it will be obtained
        from the parameter ``b``.
    recs : list of :class:`~graph_tool.PropertyMap` instances (optional, default: ``[]``)
        List of real or discrete-valued edge covariates.
    rec_types : list of edge covariate types (optional, default: ``[]``)
        List of types of edge covariates. The possible types are:
        ``"real-exponential"``, ``"real-normal"``, ``"discrete-geometric"``,
        ``"discrete-poisson"`` or ``"discrete-binomial"``.
    rec_params : list of ``dict`` (optional, default: ``[]``)
        Model hyperparameters for edge covariates. This should a list of
        ``dict`` instances. See :class:`~graph_tool.inference.blockmodel.BlockState` for
        more details.
    clabel : :class:`~graph_tool.PropertyMap` (optional, default: ``None``)
        Constraint labels on the vertices. If supplied, vertices with different
        label values will not be clustered in the same group.
    deg_corr : ``bool`` (optional, default: ``True``)
        If ``True``, the degree-corrected version of the blockmodel ensemble will
        be assumed, otherwise the traditional variant will be used.
    allow_empty : ``bool`` (optional, default: ``False``)
        If ``True``, partition description length computed will allow for empty
        groups.
    max_BE : ``int`` (optional, default: ``1000``)
        If the number of blocks exceeds this number, a sparse representation of
        the block graph is used, which is slightly less efficient, but uses less
        memory,
    """

    def __init__(self, g, b=None, B=None, recs=[], rec_types=[], rec_params=[],
                 clabel=None, pclabel=None, deg_corr=True, allow_empty=False,
                 max_BE=1000, **kwargs):

        kwargs = kwargs.copy()

        # determine if there is a base graph, and overlapping structure
        self.base_g = kwargs.pop("base_g", None)

        # overlapping information
        node_index = kwargs.pop("node_index", None)
        node_in_degs = kwargs.pop("node_in_degs", None)
        node_out_degs = kwargs.pop("node_out_degs", None)
        half_edges = kwargs.pop("half_edges", None)
        eindex = kwargs.pop("eindex", None)

        if node_index is not None and self.base_g is None:
            raise ValueError("Must specify base graph if node_index is specified...")

        # create overlapping structure
        if node_index is None:
            # keep base graph
            self.base_g = g

            if len(recs) == 0:
                rec = self.base_g.new_ep("vector<double>")
            else:
                recs = [x.copy("double") for x in recs]
                rec = group_vector_property(recs)

            # substitute provided graph by its half-edge graph
            g, b, node_index, half_edges, eindex, rec = \
                                                half_edge_graph(g, b, B, rec)

            if len(recs) > 0:
                recs = ungroup_vector_property(rec, range(len(recs)))

        # create half edges set if absent
        if half_edges is None:
            half_edges = self.base_g.new_vertex_property("vector<int64_t>")
            libinference.get_nodeset_overlap(g._Graph__graph,
                                             _prop("v", g, node_index),
                                             _prop("v", self.base_g, half_edges))

        self.overlap = True
        self.node_index = node_index
        self.half_edges = half_edges
        self.eindex = eindex

        # configure the main graph and block model parameters
        self.g = g

        self.deg_corr = deg_corr

        self.is_edge_weighted = False
        self.is_vertex_weighted = False
        self.is_weighted = False

        if b is None:
            # create a random partition into B blocks.
            B = min(B, self.g.num_vertices())
            ba = random.randint(0, B, self.g.num_vertices())
            ba[:B] = arange(B)        # avoid empty blocks
            if B < self.g.num_vertices():
                random.shuffle(ba)
            b = g.new_vertex_property("int")
            b.fa = ba
            self.b = b
        else:
            # if a partition is available, we will incorporate it.
            # in the overlapping case
            # at this point, *b* must correspond to the partition of
            # *half-edges*
            if isinstance(b, numpy.ndarray):
                self.b = g.new_vertex_property("int")
                self.b.fa = b
            else:
                b = b.copy(value_type="int")
                b = g.own_property(b)
                self.b = b
            if B is None:
                B = int(self.b.fa.max()) + 1

        if self.b.fa.max() >= B:
            raise ValueError("Maximum value of b is larger or equal to B!")

        self.rec = [self.g.own_property(p) for p in recs]
        for i in range(len(self.rec)):
            if self.rec[i].value_type() != "double":
                self.rec[i] = self.rec[i].copy("double")
        self.drec = kwargs.pop("drec", None)
        if self.drec is None:
            self.drec = []
            for rec in self.rec:
                self.drec.append(self.g.new_ep("double", rec.fa ** 2))
        else:
            self.drec = [self.g.own_property(p) for p in self.drec]

        rec_types = list(rec_types)
        rec_params = list(rec_params)

        # if len(rec_params) < len(rec_types):
        #     rec_params += [{} for i in range((len(rec_types) -
        #                                       len(rec_params)))]

        if len(self.rec) > 0 and rec_types[0] != libinference.rec_type.count:
            rec_types.insert(0, libinference.rec_type.count)
            rec_params.insert(0, {})
            self.rec.insert(0, self.g.new_ep("double", 1))
            self.drec.insert(0, self.g.new_ep("double"))

        # Construct block-graph
        self.bg = get_block_graph(g, B, self.b, rec=self.rec, drec=self.drec)
        self.bg.set_fast_edge_removal()

        self.mrs = self.bg.ep["count"]
        self.wr = self.bg.vp["count"]

        self.mrp = self.bg.degree_property_map("out", weight=self.mrs)

        if g.is_directed():
            self.mrm = self.bg.degree_property_map("in", weight=self.mrs)
        else:
            self.mrm = self.mrp

        self.B = B

        self.candidate_blocks = Vector_size_t()
        self.candidate_blocks.extend(arange(self.B, dtype="int"))

        if pclabel is not None:
            if isinstance(pclabel, PropertyMap):
                self.pclabel = self.g.own_property(pclabel).copy("int")
            else:
                self.pclabel = self.g.new_vp("int")
                self.pclabel.fa = pclabel
        else:
            self.pclabel = self.g.new_vp("int")

        if clabel is not None:
            if isinstance(clabel, PropertyMap):
                self.clabel = self.g.own_property(clabel).copy("int")
            else:
                self.clabel = self.g.new_vp("int")
                self.clabel.fa = clabel
        elif self.pclabel.fa.max() > 0:
            self.clabel = self.pclabel
        else:
            self.clabel = self.g.new_vp("int")

        self.bclabel = self.get_bclabel()
        self.hclabel = self.bg.new_vp("int")

        BlockState._init_recs(self, self.rec, rec_types, rec_params)
        self.recdx = libcore.Vector_double(len(self.rec))
        self.Lrecdx = kwargs.pop("Lrecdx", None)
        if self.Lrecdx is None:
            self.Lrecdx = libcore.Vector_double(len(self.rec)+1)
            self.Lrecdx[0] = -1
        self.Lrecdx.resize(len(self.rec)+1)
        self.epsilon = kwargs.pop("epsilon", None)
        if self.epsilon is None:
            self.epsilon = libcore.Vector_double(len(self.rec))
            for i in range(len(self.rec)):
                idx = self.rec[i].a != 0
                if numpy.any(idx):
                    self.epsilon[i] = abs(self.rec[i].a[idx]).min() / 10

        self.max_BE = max_BE

        self.use_hash = self.B > self.max_BE

        self.allow_empty = True

        self._abg = self.bg._get_any()
        self._state = libinference.make_overlap_block_state(self, _get_rng())

        if deg_corr:
            init_q_cache(max(self.get_E(), self.get_N()) + 1)

        self._entropy_args = dict(adjacency=True, deg_entropy=True, dl=True,
                                  partition_dl=True, degree_dl=True,
                                  degree_dl_kind="distributed", edges_dl=True,
                                  dense=False, multigraph=True, exact=True,
                                  recs=True, recs_dl=True, beta_dl=1.)

        self._coupled_state = None

        vweight = kwargs.pop("vweight", "unity")
        eweight = kwargs.pop("eweight", "unity")

        if vweight != "unity":
            kwargs["vweight"] = vweight
        if eweight != "unity":
            kwargs["eweight"] = eweight

        self.ignore_degrees = kwargs.pop("ignore_degrees", self.g.new_vp("bool"))

        if len(kwargs) > 0:
            warnings.warn("unrecognized keyword arguments: " +
                          str(list(kwargs.keys())))

    def __repr__(self):
        return "<OverlapBlockState object with %d blocks,%s%s for graph %s, at 0x%x>" % \
            (self.B, " degree corrected," if self.deg_corr else "",
             ((" with %d edge covariate%s," % (len(self.rec_types) - 1,
                                               "s" if len(self.rec_types) > 2 else ""))
              if len(self.rec_types) > 0 else ""),
             str(self.base_g), id(self))

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

    def __deepcopy__(self, memo):
        g = copy.deepcopy(self.g, memo)
        node_index = g.own_property(copy.deepcopy(self.node_index, memo))
        eindex = g.own_property(copy.deepcopy(self.eindex, memo))
        half_edges = base_g.own_property(copy.deepcopy(self.half_edges, memo))
        base_g = copy.deepcopy(self.base_g, memo)
        return self.copy(g=g, node_index=node_index, eindex=eindex,
                         half_edges=half_edges, base_g=base_g)

    def copy(self, g=None, b=None, B=None, deg_corr=None, clabel=None,
             pclabel=None, **kwargs):
        r"""Copies the block state. The parameters override the state properties, and
         have the same meaning as in the constructor. If ``overlap=False`` an
         instance of :class:`~graph_tool.inference.blockmodel.BlockState` is returned. This
         is by default a shallow copy."""

        state = OverlapBlockState(self.g if g is None else g,
                                  b=self.b if b is None else b,
                                  B=(self.B if b is None else None) if B is None else B,
                                  clabel=self.clabel.fa if clabel is None else clabel,
                                  pclabel=self.pclabel if pclabel is None else pclabel,
                                  deg_corr=self.deg_corr if deg_corr is None else deg_corr,
                                  recs=kwargs.pop("recs", self.rec),
                                  drec=kwargs.pop("drec", self.drec),
                                  rec_types=kwargs.pop("rec_types", self.rec_types),
                                  rec_params=kwargs.pop("rec_params",
                                                        self.rec_params),
                                  half_edges=kwargs.get("half_edges", self.half_edges),
                                  node_index=kwargs.get("node_index", self.node_index),
                                  eindex=kwargs.get("eindex", self.eindex),
                                  max_BE=kwargs.get("max_BE", self.max_BE),
                                  base_g=kwargs.get("base_g", self.base_g),
                                  Lrecdx=kwargs.pop("Lrecdx", self.Lrecdx.copy()),
                                  epsilon=kwargs.pop("epsilon",
                                                     self.epsilon.copy()),
                                  allow_empty=kwargs.pop("allow_empty",
                                                         self.allow_empty),
                                  **dmask(kwargs, ["half_edges", "node_index",
                                                   "eindex", "base_g", "drec",
                                                   "max_BE"]))
        if self._coupled_state is not None:
            state._couple_state(state.get_block_state(b=state.get_bclabel(),
                                                      vweight="nonempty",
                                                      copy_bg=False,
                                                      Lrecdx=state.Lrecdx,
                                                      allow_empty=False),
                                self._coupled_state[1])
        return state

    def __getstate__(self):
        state = dict(g=self.g,
                     b=self.b,
                     B=self.B,
                     clabel=array(self.clabel.fa),
                     deg_corr=self.deg_corr,
                     recs=ungroup_vector_property(self.rec,
                                                  range(len(self.recs))),
                     drec=self.drec,
                     rec_types=list(self.rec_types),
                     rec_params=self.rec_params,
                     half_edges=self.half_edges,
                     node_index=self.node_index,
                     eindex=self.eindex,
                     max_BE=self.max_BE,
                     base_g=self.base_g)
        return state

    def __setstate__(self, state):
        conv_pickle_state(state)
        self.__init__(**state)

    def get_E(self):
        r"Returns the total number of edges."
        return self.g.num_edges()

    def get_N(self):
        r"Returns the total number of nodes."
        return self.base_g.num_vertices()

    def get_B(self):
        r"Returns the total number of blocks."
        return self.bg.num_vertices()

    def get_nonempty_B(self):
        r"Returns the total number of nonempty blocks."
        return int((self.wr.a > 0).sum())

    def get_edge_blocks(self):
        r"""Returns an edge property map which contains the block labels pairs for each
        edge."""
        be = self.base_g.new_edge_property("vector<int>")
        self._state.get_be_overlap(self.base_g._Graph__graph,
                                    _prop("e", self.base_g, be))
        return be

    def get_overlap_blocks(self):
        r"""Returns the mixed membership of each vertex.

        Returns
        -------
        bv : :class:`~graph_tool.PropertyMap`
           A vector-valued vertex property map containing the block memberships
           of each node.
        bc_in : :class:`~graph_tool.PropertyMap`
           The labelled in-degrees of each node, i.e. how many in-edges belong
           to each group, in the same order as the ``bv`` property above.
        bc_out : :class:`~graph_tool.PropertyMap`
           The labelled out-degrees of each node, i.e. how many out-edges belong
           to each group, in the same order as the ``bv`` property above.
        bc_total : :class:`~graph_tool.PropertyMap`
           The labelled total degrees of each node, i.e. how many incident edges
           belong to each group, in the same order as the ``bv`` property above.

        """
        bv = self.base_g.new_vertex_property("vector<int>")
        bc_in = self.base_g.new_vertex_property("vector<int>")
        bc_out = self.base_g.new_vertex_property("vector<int>")
        bc_total = self.base_g.new_vertex_property("vector<int>")
        self._state.get_bv_overlap(self.base_g._Graph__graph,
                                    _prop("v", self.base_g, bv),
                                    _prop("v", self.base_g, bc_in),
                                    _prop("v", self.base_g, bc_out),
                                    _prop("v", self.base_g, bc_total))
        return bv, bc_in, bc_out, bc_total

    def get_nonoverlap_blocks(self):
        r"""Returns a scalar-valued vertex property map with the block mixture
        represented as a single number."""

        bv = self.get_overlap_blocks()[0]
        b = self.base_g.new_vertex_property("int")
        self._state.get_overlap_split(self.base_g._Graph__graph,
                                       _prop("v", self.base_g, bv),
                                       _prop("v", self.base_g, b))
        return b

    def get_majority_blocks(self):
        r"""Returns a scalar-valued vertex property map with the majority block
        membership of each node."""

        bv = self.get_overlap_blocks()
        bv, bc = bv[0], bv[-1]
        b = self.base_g.new_vertex_property("int")
        self._state.get_maj_overlap(self.base_g._Graph__graph,
                                     _prop("v", self.base_g, bv),
                                     _prop("v", self.base_g, bc),
                                     _prop("v", self.base_g, b))
        return b

    def entropy(self, adjacency=True, dl=True, partition_dl=True,
                degree_dl=True, degree_dl_kind="distributed", edges_dl=True,
                dense=False, multigraph=True, deg_entropy=True, recs=True,
                recs_dl=True, beta_dl=1., exact=True, **kwargs):
        r"""Calculate the entropy associated with the current block partition.

        Parameters
        ----------
        adjacency : ``bool`` (optional, default: ``True``)
            If ``True``, the adjacency term of the description length will be
            included.
        dl : ``bool`` (optional, default: ``True``)
            If ``True``, the description length for the parameters will be
            included.
        partition_dl : ``bool`` (optional, default: ``True``)
            If ``True``, and ``dl == True`` the partition description length
            will be included.
        degree_dl : ``bool`` (optional, default: ``True``)
            If ``True``, and ``dl == True`` the degree sequence description
            length will be included (for degree-corrected models).
        degree_dl_kind : ``str`` (optional, default: ``"distributed"``)
            This specifies the prior used for the degree sequence. It must be
            one of: ``"uniform"``, ``"distributed"`` (default) or ``"entropy"``.
        edges_dl : ``bool`` (optional, default: ``True``)
            If ``True``, and ``dl == True`` the edge matrix description length
            will be included.
        dense : ``bool`` (optional, default: ``False``)
            If ``True``, the "dense" variant of the entropy will be computed.
        multigraph : ``bool`` (optional, default: ``True``)
            If ``True``, the multigraph entropy will be used.
        deg_entropy : ``bool`` (optional, default: ``True``)
            If ``True``, the degree entropy term that is independent of the
            network partition will be included (for degree-corrected models).
        recs : ``bool`` (optional, default: ``True``)
            If ``True``, the likelihood for real or discrete-valued edge
            covariates is computed.
        recs_dl : ``bool`` (optional, default: ``True``)
            If ``True``, and ``dl == True`` the edge covariate description
            length will be included.
        beta_dl : ``double`` (optional, default: ``1.``)
            Prior inverse temperature.
        exact : ``bool`` (optional, default: ``True``)
            If ``True``, the exact expressions will be used. Otherwise,
            Stirling's factorial approximation will be used for some terms.

        Notes
        -----

        The "entropy" of the state is minus the log-likelihood of the
        microcanonical SBM, that includes the generated graph
        :math:`\boldsymbol{A}` and the model parameters :math:`\boldsymbol{\theta}`,

        .. math::

           \mathcal{S} &= - \ln P(\boldsymbol{A},\boldsymbol{\theta}) \\
                       &= - \ln P(\boldsymbol{A}|\boldsymbol{\theta}) - \ln P(\boldsymbol{\theta}).

        This value is also called the `description length
        <https://en.wikipedia.org/wiki/Minimum_description_length>`_ of the data,
        and it corresponds to the amount of information required to describe it
        (in `nats <https://en.wikipedia.org/wiki/Nat_(unit)>`_).

        For the traditional blockmodel (``deg_corr == False``), the model
        parameters are :math:`\boldsymbol{\theta} = \{\boldsymbol{e},
        \boldsymbol{b}\}`, where :math:`\boldsymbol{e}` is the matrix of edge
        counts between blocks, and :math:`\boldsymbol{b}` is the `overlapping`
        partition of the nodes into blocks. For the degree-corrected blockmodel
        (``deg_corr == True``), we have an additional set of parameters, namely
        the `labelled` degree sequence :math:`\boldsymbol{k}`.

        The model likelihood :math:`P(\boldsymbol{A}|\theta)` is given
        analogously to the non-overlapping case, as described in
        :meth:`graph_tool.inference.blockmodel.BlockState.entropy`.

        If ``dl == True``, the description length :math:`\mathcal{L} = -\ln
        P(\boldsymbol{\theta})` of the model will be returned as well. The
        edge-count prior :math:`P(\boldsymbol{e})` is described in described in
        :func:`~graph_tool.inference.blockmodel.model_entropy`. For the
        overlapping partition :math:`P(\boldsymbol{b})`, we have

        .. math::

           -\ln P(\boldsymbol{b}) = \ln\left(\!\!{D \choose N}\!\!\right) + \sum_d \ln {\left(\!\!{{B\choose d}\choose n_d}\!\!\right)} + \ln N! - \sum_{\vec{b}}\ln n_{\vec{b}}!,

        where :math:`d \equiv |\vec{b}|_1 = \sum_rb_r` is the mixture
        size, :math:`n_d` is the number of nodes in a mixture of size :math:`d`,
        :math:`D` is the maximum value of :math:`d`, :math:`n_{\vec{b}}` is the
        number of nodes in mixture :math:`\vec{b}`.

        For the degree-corrected model we need to specify the prior
        :math:`P(\boldsymbol{k})` for the `labelled` degree sequence as well:

        .. math::

            -\ln P(\boldsymbol{k}) = \sum_r\ln\left(\!\!{m_r \choose e_r}\!\!\right) - \sum_{\vec{b}}\ln P(\boldsymbol{k}|{\vec{b}}),

        where :math:`m_r` is the number of non-empty mixtures which contain type
        :math:`r`, and :math:`P(\boldsymbol{k}|{\vec{b}})` is the likelihood of
        the labelled degree sequence inside mixture :math:`\vec{b}`. For this
        term we have three options:

        1. ``degree_dl_kind == "uniform"``

            .. math::

                P(\boldsymbol{k}|\vec{b}) = \prod_r\left(\!\!{n_{\vec{b}}\choose e^r_{\vec{b}}}\!\!\right)^{-1}.

        2. ``degree_dl_kind == "distributed"``

            .. math::

                P(\boldsymbol{k}|\vec{b}) = \prod_{\vec{b}}\frac{\prod_{\vec{k}}\eta_{\vec{k}}^{\vec{b}}!}{n_{\vec{b}}!} \prod_r q(e_{\vec{b}}^r - n_{\vec{b}}, n_{\vec{b}})

            where :math:`n^{\vec{b}}_{\vec{k}}` is the number of nodes in
            mixture :math:`\vec{b}` with labelled degree :math:`\vec{k}`, and
            :math:`q(n,m)` is the number of `partitions
            <https://en.wikipedia.org/wiki/Partition_(number_theory)>`_ of
            integer :math:`n` into at most :math:`m` parts.

        3. ``degree_dl_kind == "entropy"``

            .. math::

                P(\boldsymbol{k}|\vec{b}) = \prod_{\vec{b}}\exp\left(-n_{\vec{b}}H(\boldsymbol{k}_{\vec{b}})\right)

            where :math:`H(\boldsymbol{k}_{\vec{b}}) =
            -\sum_{\vec{k}}p_{\vec{b}}(\vec{k})\ln p_{\vec{b}}(\vec{k})` is the
            entropy of the labelled degree distribution inside mixture
            :math:`\vec{b}`.

            Note that, differently from the other two choices, this represents
            only an approximation of the description length. It is meant to be
            used only for comparison purposes, and should be avoided in practice.

        For the directed case, the above expressions are duplicated for the in-
        and out-degrees.

        """

        return BlockState.entropy(self, adjacency=adjacency, dl=dl,
                                  partition_dl=partition_dl,
                                  degree_dl=degree_dl,
                                  degree_dl_kind=degree_dl_kind,
                                  edges_dl=edges_dl, dense=dense,
                                  multigraph=multigraph,
                                  deg_entropy=deg_entropy, recs=recs,
                                  recs_dl=recs_dl, beta_dl=beta_dl, exact=exact,
                                  **kwargs)

    def _clear_egroups(self):
        self._state.clear_egroups()

    def _mcmc_sweep_dispatch(self, mcmc_state):
        dS, nattempts, nmoves = \
                    libinference.overlap_mcmc_sweep(mcmc_state, self._state,
                                                    _get_rng())
        if self.__bundled:
            ret = libinference.overlap_mcmc_bundled_sweep(mcmc_state,
                                                          self._state,
                                                          _get_rng())
            dS += ret[0]
            nattempts += ret[1]
            nmoves += ret[2]
        del self.__bundled
        return dS, nattempts, nmoves

    def _mcmc_sweep_parallel_dispatch(states, mcmc_states):
        return libinference.overlap_mcmc_sweep_parallel(mcmc_states,
                                                        [s._state for s in states],
                                                        _get_rng())

    def mcmc_sweep(self, bundled=False, **kwargs):
        r"""Perform sweeps of a Metropolis-Hastings rejection sampling MCMC to sample
        network partitions. If ``bundled == True``, the half-edges incident of
        the same node that belong to the same group are moved together. All
        remaining parameters are passed to
        :meth:`graph_tool.inference.blockmodel.BlockState.mcmc_sweep`."""
        self.__bundled = bundled
        return BlockState.mcmc_sweep(self, **kwargs)

    def _multiflip_mcmc_sweep_dispatch(self, mcmc_state):
        return libinference.overlap_multiflip_mcmc_sweep(mcmc_state,
                                                         self._state,
                                                         _get_rng())

    def _multiflip_mcmc_sweep_parallel_dispatch(states, mcmc_states):
        return libinference.overlap_multiflip_mcmc_sweep_parallel(mcmc_states,
                                                                  [s._state for s in states],
                                                                  _get_rng())

    def _multicanonical_sweep_dispatch(self, multicanonical_state):
        if multicanonical_state.multiflip:
            return libinference.overlap_multicanonical_sweep(multicanonical_state,
                                                             self._state,
                                                             _get_rng())
        else:
            return libinference.overlap_multicanonical_multiflip_sweep(multicanonical_state,
                                                                       self._state,
                                                                       _get_rng())

    def _exhaustive_sweep_dispatch(self, exhaustive_state, callback, hist):
        if callback is not None:
            return libinference.overlap_exhaustive_sweep(exhaustive_state,
                                                         self._state, callback)
        else:
            if hist is None:
                return libinference.overlap_exhaustive_sweep_iter(exhaustive_state,
                                                                  self._state)
            else:
                return libinference.overlap_exhaustive_dens(exhaustive_state,
                                                            self._state,
                                                            hist[0], hist[1],
                                                            hist[2])

    def _gibbs_sweep_dispatch(self, gibbs_state):
        return libinference.gibbs_overlap_sweep(gibbs_state, self._state,
                                                _get_rng())

    def _gibbs_sweep_parallel_dispatch(states, gibbs_states):
        return libinference.overlap_gibbs_sweep_parallel(gibbs_states,
                                                         [s._state for s in states],
                                                         _get_rng())

    def _merge_sweep_dispatch(self, merge_state):
        return libinference.vacate_overlap_sweep(merge_state, self._state,
                                                 _get_rng())

    def shrink(self, B, **kwargs):
        """Reduces the order of current state by progressively merging groups,
        until only ``B`` are left. All remaining keyword arguments are passed to
        :meth:`graph_tool.inference.blockmodel.BlockState.merge_sweep`.

        This function leaves the current state untouched and returns instead a
        copy with the new partition.
        """

        b = self.b.copy()
        continuous_map(b)
        bstate = self.copy(b=b)

        assert self.get_nonempty_B() == bstate.get_nonempty_B(), \
            "Error: inconsistent number of groups after copying (%d, %d)" % \
            (self.get_nonempty_B(), bstate.get_nonempty_B())

        if bstate.get_nonempty_B() < B:
            raise ValueError("cannot shrink state to a larger number" +
                             " of groups: %d -> %d (total: %d)" %
                             (bstate.get_nonempty_B(), B, self.B))

        while bstate.get_nonempty_B() > B:
            bstate.merge_sweep(bstate.get_nonempty_B() - B, **kwargs)

        continuous_map(bstate.b)
        bstate = self.copy(b=bstate.b.a, Lrecdx=bstate.Lrecdx)

        if _bm_test():
            assert bstate.get_nonempty_B() == B, \
                "wrong number of groups after shrink: %d, %d" % \
                (bstate.get_nonempty_B(), B)
            assert bstate.wr.a.min() > 0, "empty group after shrink!"

        return bstate

    def draw(self, **kwargs):
        r"""Convenience wrapper to :func:`~graph_tool.draw.graph_draw` that
        draws the state of the graph as colors on the vertices and edges."""

        bv, bc_in, bc_out, bc_total = self.get_overlap_blocks()
        if self.deg_corr:
            pie_fractions = bc_total.copy("vector<double>")
        else:
            pie_fractions = self.base_g.new_vp("vector<double>",
                                               vals=[ones(len(bv[v])) for v
                                                     in self.base_g.vertices()])

        gradient = kwargs.get("edge_gradient",
                              get_block_edge_gradient(self.base_g,
                                                      self.get_edge_blocks(),
                                                      cmap=kwargs.get("ecmap",
                                                                      None)))
        from graph_tool.draw import graph_draw
        return graph_draw(self.base_g,
                          vertex_shape=kwargs.get("vertex_shape", "pie"),
                          vertex_pie_colors=kwargs.get("vertex_pie_colors", bv),
                          vertex_pie_fractions=kwargs.get("vertex_pie_fractions",
                                                          pie_fractions),
                          edge_gradient=gradient,
                          **dmask(kwargs, ["vertex_shape", "vertex_pie_colors",
                                           "vertex_pie_fractions",
                                           "edge_gradient"]))

def half_edge_graph(g, b=None, B=None, rec=None):
    r"""Generate a half-edge graph, where each half-edge is represented by a node,
    and an edge connects the half-edges like in the original graph."""

    E = g.num_edges()

    b_array = None
    if b is None:
        # if no partition is given, obtain a random one.
        ba = random.randint(0, B, 2 * E)
        ba[:B] = arange(B)        # avoid empty blocks
        if B < len(ba):
            random.shuffle(ba)
        b = ba

    if isinstance(b, numpy.ndarray):
        # if given an array, assume it corresponds to the *final* half-edge
        # partitions
        b_array = b
        b = g.new_vertex_property("int")

    if b.key_type() == "v":
        # If a vertex partition is given, we convert it into a
        # non-overlapping edge partition
        be = g.new_edge_property("vector<int>")
        libinference.get_be_from_b_overlap(g._Graph__graph,
                                           _prop("e", g, be),
                                           _prop("v", g, b))
        b = be
    else:
        # If an half-edge partition is provided, we incorporate it
        b = b.copy(value_type="vector<int32_t>")

    if B is None:
        if b_array is None:
            bs, bt = ungroup_vector_property(b, [0, 1])
            B = int(max(bs.fa.max(), bt.fa.max())) + 1
        else:
            B = b_array.max() + 1

    bs, bt = ungroup_vector_property(b, [0, 1])

    if bs.fa.max() >= B or bt.fa.max() >= B or (b_array is not None and b_array.max() >= B):
        raise ValueError("Maximum value of b is larger or equal to B!")

    eg = Graph(directed=g.is_directed())
    node_index = eg.new_vertex_property("int64_t")
    half_edges = g.new_vertex_property("vector<int64_t>")
    be = eg.new_vertex_property("int")
    eindex = eg.new_edge_property("int64_t")
    erec = eg.new_edge_property("vector<double>")

    if rec is None:
        rec_ = g.new_edge_property("vector<double>")
    else:
        rec_ = g.own_property(rec)

    # create half-edge graph
    libinference.get_eg_overlap(g._Graph__graph,
                                eg._Graph__graph,
                                _prop("e", g, b),
                                _prop("v", eg, be),
                                _prop("v", eg, node_index),
                                _prop("v", g, half_edges),
                                _prop("e", eg, eindex),
                                _prop("e", g, rec_),
                                _prop("e", eg, erec))

    if b_array is not None:
        be.a = b_array

    if rec is None:
        erec = None

    return eg, be, node_index, half_edges, eindex, erec

def augmented_graph(g, b, node_index, eweight=None):
    r"""Generates an augmented graph from the half-edge graph ``g`` partitioned
    according to ``b``, where each half-edge belonging to a different group
    inside each node forms a new node."""

    node_map = g.new_vertex_property("int")
    br_b = libcore.Vector_int32_t()
    br_ni = libcore.Vector_int32_t()
    libinference.get_augmented_overlap(g._Graph__graph,
                                       _prop("v", g, b),
                                       _prop("v", g, node_index),
                                       _prop("v", g, node_map),
                                       br_b, br_ni)

    au, idx, vcount, ecount = condensation_graph(g, node_map,
                                                 eweight=eweight,
                                                 self_loops=True)[:4]
    anidx = idx.copy("int")
    libinference.vector_map(anidx.a, br_ni.a)

    ab = idx.copy("int")
    libinference.vector_map(ab.a, br_b.a)

    return au, ab, anidx, ecount, node_map

def get_block_edge_gradient(g, be, cmap=None):
    r"""Get edge gradients corresponding to the block membership at the endpoints of
    the edges given by the ``be`` edge property map.

    Parameters
    ----------
    g : :class:`~graph_tool.Graph`
        The graph.
    be : :class:`~graph_tool.PropertyMap`
        Vector-valued edge property map with the block membership at each
        endpoint.
    cmap : :class:`matplotlib.colors.Colormap` (optional, default: ``default_cm``)
        Color map used to construct the gradient.

    Returns
    -------
    cp : :class:`~graph_tool.PropertyMap`
       A vector-valued edge property map containing a color gradient.
    """

    if cmap is None:
        from .. draw import default_cm
        cmap = default_cm

    cp = g.new_edge_property("vector<double>")
    rg = [numpy.inf, -numpy.inf]
    for e in g.edges():
        s, t = be[e]
        rg[0] = min(s, rg[0])
        rg[0] = min(t, rg[0])
        rg[1] = max(s, rg[1])
        rg[1] = max(t, rg[1])

    for e in g.edges():
        if int(e.source()) < int(e.target()) or g.is_directed():
            s, t = be[e]
        else:
            t, s = be[e]
        cs = cmap((s - rg[0]) / max(rg[1] - rg[0], 1))
        ct = cmap((t - rg[0]) / max(rg[1] - rg[0], 1))
        cp[e] = [0] + list(cs) + [1] + list(ct)
    return cp


"""An interface between Skyfield and the Python ``sgp4`` library."""

from numpy import array, cross, einsum, zeros_like
from sgp4.earth_gravity import wgs72
from sgp4.io import twoline2rv
from sgp4.propagation import sgp4

from .constants import AU_KM, DAY_S, T0, tau
from .functions import rot_x, rot_y, rot_z
from .positionlib import Apparent, Geocentric, ITRF_to_GCRS
from .timelib import JulianDate, takes_julian_date

# important ones:
# jdsatepoch
# bstar
# inclo - inclination
# nodeo - right ascension of ascending node
# ecco - eccentricity
# argpo - argument of perigee
# mo - mean anomaly
# no - mean motion

_minutes_per_day = 1440.

class EarthSatellite(object):
    """An Earth satellite loaded from a TLE file and propagated with SGP4."""

    def __init__(self, lines, earth):
        sat = twoline2rv(*lines[-2:], whichconst=wgs72)
        self._sgp4_satellite = sat
        self._earth = earth
        self.epoch = JulianDate(utc=(sat.epochyr, 1, sat.epochdays - 1.0))

    def __repr__(self):
        sat = self._sgp4_satellite
        return '<EarthSatellite number={1!r} epoch={0}>'.format(
            self.epoch.utc_iso(), sat.satnum)

    def _position_and_velocity_TEME_km(self, jd):
        """Return the raw true equator mean equinox (TEME) vectors from SGP4.

        Returns a tuple of NumPy arrays ``([x y z], [xdot ydot zdot])``
        expressed in kilometers and kilometers per second.  Note that we
        assume the TLE epoch to be a UTC date, per AIAA 2006-6753.

        """
        sat = self._sgp4_satellite
        epoch = sat.jdsatepoch
        minutes_past_epoch = (jd._utc_float() - epoch) * 1440.
        if getattr(minutes_past_epoch, 'shape', None):
            position = []
            velocity = []
            error = []
            for m in minutes_past_epoch:
                p, v = sgp4(sat, m)
                position.append(p)
                velocity.append(v)
                error.append(sat.error_message)
            return array(position).T, array(velocity).T, error
        else:
            position, velocity = sgp4(sat, minutes_past_epoch)
            return array(position), array(velocity), sat.error_message

    def _compute_GCRS(self, jd):
        """Compute where satellite is in space on a given date."""

        rTEME, vTEME, error = self._position_and_velocity_TEME_km(jd)
        rTEME /= AU_KM
        vTEME /= AU_KM
        vTEME *= DAY_S

        rITRF, vITRF = TEME_to_ITRF(jd.ut1, rTEME, vTEME)
        rGCRS = ITRF_to_GCRS(jd, rITRF)
        vGCRS = zeros_like(rGCRS)  # todo: someday also compute vGCRS?

        return rGCRS, vGCRS, error

    @takes_julian_date
    def gcrs(self, jd):
        """Return a GCRS position for this Earth satellite.

        Uses standard SGP4 theory to predict the satellite location.

        """
        position_au, velociy_au_per_d, error = self._compute_GCRS(jd)
        g = Geocentric(position_au, velociy_au_per_d, jd)
        g.sgp4_error = error
        return g

    def _observe_from_bcrs(self, observer):
        # TODO: what if someone on Mars tries to look at the ISS?

        jd = observer.jd
        rGCRS, vGCRS, error = self._compute_GCRS(jd)
        rGCRS - observer.rGCRS
        vGCRS - observer.vGCRS
        g = Apparent(rGCRS - observer.rGCRS, vGCRS - observer.vGCRS, jd)
        g.sgp4_error = error
        g.observer = observer
        # g.distance = euclidian_distance
        return g

_second = 1.0 / (24.0 * 60.0 * 60.0)

def theta_GMST1982(jd_ut1):
    """Return the angle of Greenwich Mean Standard Time 1982 given the JD.

    This angle defines the difference between the idiosyncratic True
    Equator Mean Equinox (TEME) frame of reference used by SGP4 and the
    more standard Pseudo Earth Fixed (PEF) frame of reference.

    From AIAA 2006-6753 Appendix C.

    """
    t = (jd_ut1 - T0) / 36525.0
    g = 67310.54841 + (8640184.812866 + (0.093104 + (-6.2e-6) * t) * t) * t
    dg = 8640184.812866 + (0.093104 * 2.0 + (-6.2e-6 * 3.0) * t) * t
    theta = (jd_ut1 % 1.0 + g * _second % 1.0) * tau
    theta_dot = (1.0 + dg * _second / 36525.0) * tau
    return theta, theta_dot

def TEME_to_ITRF(jd_ut1, rTEME, vTEME, xp=0.0, yp=0.0):
    """Convert TEME position and velocity into standard ITRS coordinates.

    This converts a position and velocity vector in the idiosyncratic
    True Equator Mean Equinox (TEME) frame of reference used by the SGP4
    theory into vectors into the more standard ITRS frame of reference.
    The velocity should be provided in units per day, not per second.

    From AIAA 2006-6753 Appendix C.

    """
    theta, theta_dot = theta_GMST1982(jd_ut1)
    zero = theta_dot * 0.0
    angular_velocity = array([zero, zero, -theta_dot])
    R = rot_z(-theta)

    if len(rTEME.shape) == 1:
        rPEF = (R).dot(rTEME)
        vPEF = (R).dot(vTEME) + cross(angular_velocity, rPEF)
    else:
        rPEF = einsum('ij...,j...->i...', R, rTEME)
        vPEF = einsum('ij...,j...->i...', R, vTEME) + cross(
            angular_velocity, rPEF, 0, 0).T

    if xp == 0.0 and yp == 0.0:
        rITRF = rPEF
        vITRF = vPEF
    else:
        W = (rot_x(yp)).dot(rot_y(xp))
        rITRF = (W).dot(rPEF)
        vITRF = (W).dot(vPEF)
    return rITRF, vITRF


# -*- coding: utf-8 -*-
"""
Synfit
======

This routine allows the user to automatically fit a spectral line or wavelength
region. Is is allowed any number of parameters.

The best fit is found by minimizing the chi^2. `Synfit` calculates the
synthetic spectrum for all values asked and then select the one with minimum
chi^2 and returns it to the user.

It is also possible to select a subregion of the spectrum by using the
`windows` argument.
"""
import os
import shutil
import re
import json
import shutil
from copy import deepcopy
import numpy as np
import matplotlib.pyplot as plt
from itertools import product
from ..io import specio
from ..synthesis import Synplot
from ..spectools import rvcorr

# Loads the chemical elements and their atomic numbers
PERIODIC = json.load(open(os.getenv('HOME')+'/.s4/extra_resc'+\
                          '/chemical_elements.json'))

REVERSE_PERIODIC = {val:key for key, val in PERIODIC.iteritems()}

def iterator(fit_keys, iter_params):
    """
    Create an array with the values to iterate.

    Parameters
    ----------

    fit_keys: dict;
        A dictionary containing the name of the parameters,

    fit_params: dict;
        A dictionary containing the parameters as keys and the values to be
        fitted.

    Returns
    -------

    iter_values: numpy.ndarray;
        The values to be fitted as a structured array
    """
    # Create the iterator vector.
    args = [iter_params[k] for k in fit_keys]
    iter_values = list(product(*args))

    ## add the iterrating values to self as a numpy array
    data_type = [float] *len(fit_keys)

    return np.array(iter_values, dtype={'names':fit_keys, 'formats':data_type})

class Synfit:
    """
    Fit a spectral line by iterating on user defined parameter
    an returns the best fit by minimizing the $\chi^2$ of the
    synthetic spectrum and the observed spectrum.

    Based on the homonym IDL software created by Dr. Ivan
    Hubeny.

    Parameters
    ----------

    fit_params: dic;
        Parameters to be fitted. It should be a dictionary in
        which each key is the setllar parameter of the chemical
        element to be fitted. The values should be list with three
        elements: the initial and findal values and the step.

    kwargs;
        All Synplot parameters desired to be use, including `teff`,
        `logg`, `synplot_path`, `idl`, `noplot`.

        abund: dic (optional);
            Abundance of chosen chemical elements.

            Should be passsed as a dictionary which the keys are the chemical
            element symbol or atomic number i(can be mixed) and the values are
            the abundance, e.g, for Helium and Oxygen abundance of 10.93 and
            8.69, respectively:

            {2:10.93, 'O':8.69}

            ATTENTION: it does not accept the Synplot format.

        windows: list (optional);
            Spectral region in which the chi-sqaure will be calculated. At
            this point, it only accept a single window, i.e., the list should
            contains only two values: the lower and the upper wavelength.
    """

    def __init__(self, fit_params, **kwargs):
        # Parameters to be fitted
        self.fit_params = deepcopy(fit_params)
        # Fixed parameters
        self.syn_params = deepcopy(kwargs)
        ##########

        # Initalizaze varibales.
        self.iter_params = None
        self.fit_keys = None
        self.no_rot_keys = None
        self.rot_keys = None

        # Check for radial velocity
        if 'rv' in self.syn_params:
            self.rad_vel = self.syn_params['rv']
        else:
            self.rad_vel = 0

        # Check if there is a an observed spectrum.
        # If not quit.
        try:
            # If there is, load it to calculate the weights
            obs_spec = specio.load_spectrum(self.syn_params['observ'])
            # Correct for radial velocity
            obs_spec[:,0] *= rvcorr(self.rad_vel)
        except IOError:
            raise IOError('There is not any observed spectrum.')

        #Prepare kwargs
        if 'synplot_path' in self.syn_params:
            self.synplot_path = self.syn_params.pop('synplot_path')
        else:
            self.synplot_path = os.getenv('HOME')+'/.s4/synthesis/synplot/'

        if 'idl' in self.syn_params:
            self.idl = self.syn_params.pop('idl')
        else:
            self.idl = False  # It will run GDL

        if 'noplot' in self.syn_params and self.syn_params['noplot'] == True:
            self.noplot = True
            del self.syn_params['noplot']
        else:
            self.noplot = False

        if 'windows' in self.syn_params:
            self.weights = np.zeros_like(obs_spec[:, 0])

            for window in np.reshape(self.syn_params['windows'], (-1, 2)):
                index_lower = obs_spec[:, 0] > window[0]
                index_upper = obs_spec[:, 0] < window[1]
                self.weights[index_lower & index_upper] = 1

            self.windows = self.syn_params.pop('windows')
        else:
            self.weights = np.ones_like(obs_spec[:,0])
            self.windows = None

        # Obtain teff and logg if set on syn_params
        if 'teff' in self.syn_params:
            self.teff = self.syn_params.pop('teff')

        if 'logg' in self.syn_params:
            self.logg = self.syn_params.pop('logg')

        # Initialize variable to store the best fit values
        self.best_fit = {}

    def sample_params(self):
        """
        Creates the values to fit for each parameter. It also merge
        the chemical elements to a parameter accepted by `Synplot`.
        """
        def values_to_fit(key, val):
            """
            Creates the values to fit from a three-value list.

            Parameters
            ----------

            key: string
                Parameter to be fitted.

            val: list;
                List with initial value, final value and the step.
            """
            try:
            # Test if the parameters were inserted correctly.
                assert len(val) == 3

                min_value, max_value, step = [float(i) for i in val]

                n_values = np.rint((max_value - min_value)/step + 1)
                vector = np.linspace(min_value, max_value, n_values)

                return vector
            except :
                raise Exception("Value of '{}' must be a ".format(key)+\
                                "list with three values.")

        self.iter_params = {key:values_to_fit(key, val)
                            for key, val in self.fit_params.iteritems()}

        # Get the name of the parameters to be fitted
        self.fit_keys = self.iter_params.keys()

        # Segregate paramates in convolution related and unrelated.
        ## Get the keys convolution unrelated (i.e. not in ['vrot', 'vmac_rt'])
        self.no_rot_keys = [key
                            for key in self.fit_keys
                            if key not in ['vrot', 'vmac_rt']]
        ## Get the keys convolution related (i.e. in ['vrot', 'vmac_rt'])
        self.rot_keys = [key
                         for key in self.fit_keys
                         if key in ['vrot', 'vmac_rt']]

    def fit(self):
        r"""
        Fit a spectral line by iterating on user defined parameter
        an returns the best fit by minimizing the $\chi^2$ of the
        synthetic spectrum and the observed spectrum.

        Based on the homonym IDL software created by Dr. Ivan
        Hubeny.
        """
        # Creates the values in which each parameter will be fitted
        self.sample_params()

        # Create iterator.
        iter_values = iterator(self.fit_keys, self.iter_params)

        #Obtain the number of varying params
        n_params = len(self.fit_keys)

        ######
        # Array to store the values of each parameter and the chisquare

        ## Create an array of NaN
        chisquare = np.empty([len(iter_values), 1])
        chisquare.fill(np.nan)

        ## Create the array with the iteration avlues + NaN for the chisquare

        ### Remove the data type of the iter_values array.
        ### this is necessary in order to add the chisquare array
        #### Removes data type
        tmp_array = iter_values.view((float, n_params))
        #### Guarantee that the format will be correct for any number of
        #### parameters
        tmp_array = tmp_array.reshape(len(iter_values), -1)

        ### Join the arrays
        self.chisq_values = np.hstack((tmp_array, chisquare))

        ### Set the data type for the chisq_values array
        data_type = self.fit_keys + ['chisquare']
        self.chisq_values.dtype = {'names':data_type,
                                   'formats':[float]*len(data_type)}
        ######

        # Create a library of unconvolved spectra
        self.build_library()

        # Loop it!
        for n, it in enumerate(iter_values):
            self.iteration(n, it)

        # Find the best value
        self.find_best_fit()

    def build_library(self):
        """Build spectra library of unconvolved spectra"""
        if len(self.no_rot_keys) > 0 and len(self.rot_keys) > 0:
            # Build library for non rotation parameters with vsini=vmac_rt=0
            no_rot_values = iterator(self.no_rot_keys, self.iter_params)
            for n, it in enumerate(no_rot_values):
                ## Creates a dic with the parameters and values to be fitted
                ## in this loop
                params = {key:val for key, val in zip(it.dtype.names, it)}

                ## Set vsini=vmac_rt=0
                params['vrot'] = 0
                params['vmac_rt'] = 0

                ## Check if teff and logg were selected to be fitted.
                ## If yes, set a variable to them.
                if 'teff' in params:
                    self.teff = params.pop('teff')

                if 'logg' in params:
                    self.logg = params.pop('logg')

                ## Join the abundances

                ### Gets all chemical elements asked to be fit
                abund = {key:it[key]
                         for key in it.dtype.names
                         if (key in PERIODIC) or (key in REVERSE_PERIODIC)}
                if abund:
                    ### delete the chemical elements parameters from the
                    ### dictionary
                    for key in abund:
                        del params[key]

                ## Set parameters for synplot
                synplot_params = deepcopy(self.syn_params)
                synplot_params.update(params)

                ## Deal with fixed and varying abundances
                self.merge_abundances(abund, synplot_params)

                ## Synthesize spectrum
                self.synthesis = Synplot(self.teff, self.logg,
                                         self.synplot_path, self.idl,
                                         **synplot_params)
                self.synthesis.run()

                ## Backup fort.7 and fort.17
                spec_name = '_'.join(['{}_{}'.format(key, val)
                                      for key, val in zip(it.dtype.names, it)])
                shutil.move('{}fort.7'.format(self.synplot_path),
                            '/tmp/synfit_{}.7'.format(spec_name))
                shutil.move('{}fort.17'.format(self.synplot_path),
                            '/tmp/synfit_{}.17'.format(spec_name))

        elif len(self.no_rot_keys) == 0 and len(self.rot_keys) > 0:
            # There is only 'vrot' or/and 'vmac_rt'. All iteration fits will
            # be just convolutions. It creates only one spectrum

            ## Set parameters for synplot
            synplot_params = deepcopy(self.syn_params)

            ## Set values of rotation to 0
            synplot_params['vrot'] = 0
            synplot_params['vmac_rt'] = 0

            ## Synthesize spectrum
            synthesis = Synplot(self.teff, self.logg, self.synplot_path,
                                self.idl, **synplot_params)
            synthesis.run()

            ## Backup fort.7 and fort.17
            shutil.move('{}fort.7'.format(self.synplot_path),
                        '/tmp/synfit.7')
            shutil.move('{}fort.17'.format(self.synplot_path),
                        '/tmp/synfit.17')

        elif len(self.no_rot_keys) > 0 and len(self.rot_keys) == 0:
            # No rotational parameters
            # Do not build library.
            # There is no need since Synspec will have to run every time.
            pass
        else:
            # There is no parameters. Something wen wrong?
            raise RuntimeError("There is no parameters or it was not " + \
                               "classified as rotational or non rotational. " +\
                               "It seems that something went wrong.")

    def iteration(self, n, it):
        """Code to be iterated on a loop."""

        # Creates a dic with the parameters and values to be fitted
        #in this loop
        params = {key:val for key, val in zip(it.dtype.names, it)}

        #make plot title before removing teff and logg
        if self.noplot == False:
            plot_title = ', '.join(['{}={}'.format(key, val)
                                    for key, val in params.iteritems()])

        # Check if teff and logg were selected to be fitted.
        # If yes, set a variable to them.
        if 'teff' in params:
            self.teff = params.pop('teff')

        if 'logg' in params:
            self.logg = params.pop('logg')

        # Join the abundances

        ## Gets all chemical elements asked to be fit
        abund = {key:it[key]
                 for key in it.dtype.names
                 if (key in PERIODIC) or (key in REVERSE_PERIODIC)}
        if abund:
            ## delete the chemical elements parameters from the dictionary
            for key in abund:
                del params[key]

        # Set parameters for synplot
        synplot_params = deepcopy(self.syn_params)
        synplot_params.update(params)

        # Deal with fixed and varying abundances
        self.merge_abundances(abund, synplot_params)

        # Copy not convolved spectrum to Synplot folder
        if len(self.no_rot_keys) > 0 and len(self.rot_keys) > 0:
            # Set to not calculate spectrum, just convolve
            ## I tried to set the parameter 'ispec' to -1 but it didn't work.
            ## So it will use the parameter 'norun'.
            synplot_params['norun'] = 1

            ## There are rotational and non rotational parameters
            spec_name = '_'.join(['{}_{}'.format(key, val)
                                  for key, val in zip(it.dtype.names, it)
                                  if key not in ['vrot', 'vmac_rt']])
            shutil.copy('/tmp/synfit_{}.7'.format(spec_name),
                        '{}fort.7'.format(self.synplot_path))
            shutil.copy('/tmp/synfit_{}.17'.format(spec_name),
                        '{}fort.17'.format(self.synplot_path))

        elif len(self.no_rot_keys) == 0 and len(self.rot_keys) > 0:
            # Set to not calculate spectrum, just convolve
            ## I tried to set the parameter 'ispec' to -1 but it didn't work.
            ## So it will use the parameter 'norun'.
            synplot_params['norun'] = 1

            ## There is only 'vrot' or/and 'vmac_rt'.
            shutil.copy('/tmp/synfit.7',
                        '{}fort.7'.format(self.synplot_path))
            shutil.copy('/tmp/synfit.17',
                        '{}fort.17'.format(self.synplot_path))
        elif len(self.no_rot_keys) > 0 and len(self.rot_keys) == 0:
            # No rotational parameters
            pass
        else:
            # There is no parameters. Something wen wrong?
            raise RuntimeError("There is no parameters or it was not " + \
                               "classified as rotational or non rotational. " +\
                               "It seems that something went wrong.")

        # Synthesize spectrum
        self.synthesis = Synplot(self.teff, self.logg, self.synplot_path,
                                 self.idl, **synplot_params)
        self.synthesis.run()

        # Apply scale and radial velocity if needed
        if 'scale' in self.synthesis.parameters:
            self.synthesis.apply_scale()

        self.synthesis.observation[:, 0] *= rvcorr(self.rad_vel)

        #Do an interpolation

        flm = np.interp(self.synthesis.observation[:,0],
                        self.synthesis.spectrum[:,0],
                        self.synthesis.spectrum[:,1])
                        #/max(syn.observation[:,1])

        #Some kind of normalization on the observed flux?
        fobm = self.synthesis.observation[:,1]#/max(syn.observation[:,1])

        # Calculate the chi**2
        chisq = np.sum(((fobm - flm)**2/flm) * self.weights)
        #chisq = chisq * max(fobs)                 #????

        # store the values of the parameters
        self.chisq_values['chisquare'][n] = chisq

        # Plot, if desired
        if self.noplot == False:
            # The synthetic spectrum was corrected by scale and the
            #observed one by radial velocity. The plot function in Synplot
            #also does that, so we need to set those parameters to 1 and 0,
            # respectively.
            self.synthesis.parameters['scale'] = 1
            self.synthesis.parameters['rv'] = 0

            # Adds the value of chisquare to the title
            plot_title += r'$\chi^2$='+'{:.06f}'.format(chisq)
            self.synthesis.plot(title=plot_title, windows=self.windows)

    @staticmethod
    def merge_abundances(abund, synplot_params):
        """
        Merge varying and fixed parameters. It change the `abund` and
        `synplot_params` in place.
        """
        # The abundance should be merged individually because it could be
        # a mix of fixed and varying abundances.
        if abund and 'abund' in synplot_params:
            ## Check for overlapping elements.
            ### Transform all elements to its symbol
            for key, val in deepcopy(abund).iteritems():
                try:
                    abund[REVERSE_PERIODIC[key]] = val
                    del abund[key]
                except KeyError:
                    #### The chemical element is already as a symbol
                    pass

            for key, val in deepcopy(synplot_params['abund']).iteritems():
                try:
                    synplot_params['abund'][REVERSE_PERIODIC[key]] = val
                    del synplot_params['abund'][key]
                except KeyError:
                    #### The chemical element is already as a symbol
                    pass

            ## Get the fixed abundances
            try:
                abund.update({k:v
                              for k, v in synplot_params['abund'].iteritems()
                              if k not in deepcopy(abund)})
            except ValueError:
                # There no fixed abundance
                pass

            ## Adds to the synplot_params dictionary
            synplot_params['abund'] = abund
        if abund and 'abund' not in synplot_params:
            synplot_params['abund'] = abund

    def find_best_fit(self):
        """
        Obtain the fitted parameters for the chosen parameters and the
        value of the chi^2.
        """
        fitted_vals = self.chisq_values[np.argmin(
                                        self.chisq_values['chisquare'])]

        self.best_fit = {param:fitted_value
                           for param, fitted_value
                           in zip(self.iter_params.keys(),
                                  fitted_vals.view(float))}
        self.best_fit['chisquare'] = fitted_vals['chisquare'][0]

    def plot_best_fit(self, title=None):
        """
        Plot the observed spectrum and the synthetic using the best values
        found.

        Parameters
        ----------

        title: str (optional);
            A title for the plot. If the string 'default' is passed, it will
            contain the parameters fitted with the best values found.
        """

        # Set parameters for synplot
        synplot_params = deepcopy(self.syn_params)
        best_fit = deepcopy(self.best_fit)
        chisq = best_fit.pop('chisquare')

        #make plot title before removing teff and logg
        if title == 'default':
            title = r'$\chi^2$=' + '{:.4f}: '.format(chisq)
            title += ', '.join(['{}={}'.format(key, val)
                                    for key, val in best_fit.iteritems()])

        ## Replace parameters for best value
        if 'teff' in best_fit:
            self.teff = best_fit.pop('teff')

        if 'logg' in best_fit:
            self.logg = best_fit.pop('logg')

        # Join the abundances

        ## Gets all chemical elements asked to be fit
        abund = {key:val
                 for key, val in self.best_fit.iteritems()
                 if (key in PERIODIC) or (key in REVERSE_PERIODIC)}
        if abund:
            for key in abund:
                del best_fit[key]

        # Deal with fixed and varying abundances
        self.merge_abundances(abund, synplot_params)

        # Add the best values found to the Synplot parameters
        for key, value in best_fit.iteritems():
            synplot_params[key] = value

        # Synthesize spectrum
        synthesis = Synplot(self.teff, self.logg, self.synplot_path,
                            self.idl, **synplot_params)
        synthesis.plot(title=title, windows=self.windows)

    def plot_chisquare(self, interpolation='linear', logscale=False, **kwargs):
        """
        Plot the distribution of chi-square for one given parameter.
        It obly works if the number of parameters to be fitted are one or two.

        Parameters
        ----------

        interpolation: str (optional);
            Type of interpolation for the color plot when fitting two
            parameters. One of ['nearest', 'linear', 'cubic']. More explanation
            here:

            http://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.griddata.html

        scale: boolean (optional):
            If set to True, it will plot the logarithmic of the chi-square.

        kwargs;
            Matplotlib.pyplot.plot kwargs.
        """
        # Number of parameters fitted.
        number_params = len(self.fit_keys)

        fig = plt.figure()
        ax = fig.add_subplot(111)

        if number_params == 1:
            if logscale:
                y_axis = np.log(self.chisq_values['chisquare'])
            else:
                y_axis = self.chisq_values['chisquare']

            ax.plot(self.iter_params[self.fit_keys[0]],
                    y_axis, zorder=5, **kwargs)

            xlabel = self.fit_keys[0]

            ax.set_xlabel(xlabel)
            ax.set_xticks(self.iter_params[xlabel])
            if xlabel in PERIODIC:
                ax.set_xticklabels(ax.get_xticks(), rotation=-45)

            ax.set_ylabel(r'$\chi^2$')

            ax.grid(axis='x', zorder=0)
        elif number_params == 2:
            from scipy.interpolate import griddata

            # Transform the chisquare array in a proper array
            #and not an array of tuples
            chisq_values = deepcopy(self.chisq_values)

            ## Check for abundance
            if 'abund' in self.chisq_values.dtype.names:
                elem = [param for param in self.fit_params.keys()
                        if param in PERIODIC][0]

                chisq_values['abund'] = [abund
                                         for abund in chisq_values['abund']]

            ## Removes data type
            chisquare_arr = chisq_values.view((float, 3))
            ## Guarantee that the format will be correct for any number of
            ## parameters
            chisquare_arr = chisquare_arr.reshape(len(chisq_values), -1)

            edges = np.hstack([(min(param_vector), max(param_vector))
                               for param_vector in chisquare_arr.T[:-1]])

            # Points in wich the grid will be calculated
            grid_points = [np.linspace(edges[i], edges[i+1], 200)
                           for i in np.arange(0, 2*number_params, 2)]

            # Makes a mesh grid to griddata
            grid_params = np.meshgrid(*grid_points)
            grid_params = tuple([i for i in grid_params])

            # Grid the data
            Z = griddata(chisquare_arr[:,:number_params], chisquare_arr[:,-1],
                         grid_params, method=interpolation)

            if logscale:
                # Get the log to increase the contrast between limits
                from matplotlib.colors import LogNorm
                kwargs['norm'] = LogNorm()

            # Plot
            cs = ax.contour(Z, aspect='auto', extent=edges,
                            origin='lower', zorder=5, **kwargs)
            plt.clabel(cs, inline=1)

            # Set labels
            xlabel = self.fit_keys[0]
            ylabel = self.fit_keys[1]

            ax.set_xlabel(xlabel)
            ax.set_xticks(self.iter_params[xlabel])
            if xlabel in PERIODIC:
                ax.set_xticklabels(ax.get_xticks(), rotation=-45)

            ax.set_ylabel(ylabel)
            ax.set_yticks(self.iter_params[ylabel])
            if ylabel in PERIODIC:
                ax.set_yticklabels(ax.get_yticks(), rotation=-45)

            ax.grid(zorder=0)
        else:
            raise ValueError('The number of parameters is greater than 2.')


#!/usr/bin/env python
"""
Calculate Fisher matrix and P(k) constraints for all redshift bins for a given
experiment.
"""
import numpy as np
import pylab as P
import radiofisher as rf
from radiofisher import experiments
from radiofisher.units import *
from mpi4py import MPI
import sys

comm = MPI.COMM_WORLD
myid = comm.Get_rank()
size = comm.Get_size()

################################################################################
# Set-up experiment parameters
################################################################################

# Load cosmology and experimental settings
e = experiments
cosmo = experiments.cosmo

# Label experiments with different settings
#EXPT_LABEL = "_2yr_3pbwedge"
EXPT_LABEL = "_1bin_4000hr" #"_2yr"
#EXPT_LABEL = "_2yr_horizwedge"

#cosmo['A_xi'] = 0.01
#cosmo['logkmg'] = np.log10(0.005)

# A_xi: 0.01 0.1
# logkmg: 0.05, 0.01, 0.005, 0.001

expt_list = [
    ( 'exptS',            e.exptS ),        # 0
    ( 'aexptM',           e.exptM ),        # 1
    ( 'exptL',            e.exptL ),        # 2
    ( 'iexptL',           e.exptL ),        # 3
    ( 'cexptL',           e.exptL ),        # 4
    ( 'GBT',              e.GBT ),          # 5
    ( 'Parkes',           e.Parkes ),       # 6
    ( 'GMRT',             e.GMRT ),         # 7
    ( 'WSRT',             e.WSRT ),         # 8
    ( 'APERTIF',          e.APERTIF ),      # 9
    ( 'VLBA',             e.VLBA ),         # 10
    ( 'JVLA',             e.JVLA ),         # 11
    ( 'iJVLA',            e.JVLA ),         # 12
    ( 'BINGO',            e.BINGO ),        # 13
    ( 'iBAOBAB32',        e.BAOBAB32 ),     # 14
    ( 'iBAOBAB128',       e.BAOBAB128 ),    # 15
    ( 'yCHIME',           e.CHIME ),        # 16
    ( 'iAERA3',           e.AERA3 ),        # 17
    ( 'iMFAA',            e.MFAA ),         # 18
    ( 'yTIANLAIpath',     e.TIANLAIpath ),  # 19
    ( 'yTIANLAI',         e.TIANLAI ),      # 20
    ( 'yTIANLAIband2',    e.TIANLAIband2 ), # 21
    ( 'FAST',             e.FAST ),         # 22
    ( 'KAT7',             e.KAT7 ),         # 23
    ( 'iKAT7',            e.KAT7 ),         # 24
    ( 'cKAT7',            e.KAT7 ),         # 25
    ( 'MeerKATb1',        e.MeerKATb1 ),    # 26
    ( 'iMeerKATb1',       e.MeerKATb1 ),    # 27
    ( 'cMeerKATb1',       e.MeerKATb1 ),    # 28
    ( 'MeerKATb2',        e.MeerKATb2 ),    # 29
    ( 'iMeerKATb2',       e.MeerKATb2 ),    # 30
    ( 'cMeerKATb2',       e.MeerKATb2 ),    # 31
    ( 'ASKAP',            e.ASKAP ),        # 32
    ( 'SKA1MIDbase1',     e.SKA1MIDbase1 ), # 33
    ( 'iSKA1MIDbase1',    e.SKA1MIDbase1 ), # 34
    ( 'cSKA1MIDbase1',    e.SKA1MIDbase1 ), # 35
    ( 'SKA1MIDbase2',     e.SKA1MIDbase2 ), # 36
    ( 'iSKA1MIDbase2',    e.SKA1MIDbase2 ), # 37
    ( 'cSKA1MIDbase2',    e.SKA1MIDbase2 ), # 38
    ( 'SKA1MIDfull1',     e.SKA1MIDfull1 ), # 39
    ( 'iSKA1MIDfull1',    e.SKA1MIDfull1 ), # 40
    ( 'cSKA1MIDfull1',    e.SKA1MIDfull1 ), # 41
    ( 'SKA1MIDfull2',     e.SKA1MIDfull2 ), # 42
    ( 'iSKA1MIDfull2',    e.SKA1MIDfull2 ), # 43
    ( 'cSKA1MIDfull2',    e.SKA1MIDfull2 ), # 44
    ( 'fSKA1SURbase1',    e.SKA1SURbase1 ), # 45
    ( 'fSKA1SURbase2',    e.SKA1SURbase2 ), # 46
    ( 'fSKA1SURfull1',    e.SKA1SURfull1 ), # 47
    ( 'fSKA1SURfull2',    e.SKA1SURfull2 ), # 48
    ( 'exptCV',           e.exptCV ),       # 49
    ( 'GBTHIM',           e.GBTHIM ),       # 50
    ( 'SKA0MID',          e.SKA0MID ),      # 51
    ( 'fSKA0SUR',         e.SKA0SUR ),      # 52
    ( 'SKA1MID900',       e.SKA1MID900 ),   # 53
    ( 'SKA1MID350',       e.SKA1MID350 ),   # 54
    ( 'iSKA1MID900',      e.SKA1MID900 ),   # 55
    ( 'iSKA1MID350',      e.SKA1MID350 ),   # 56
    ( 'fSKA1SUR650',      e.SKA1SUR650 ),   # 57
    ( 'fSKA1SUR350',      e.SKA1SUR350 ),   # 58
    ( 'aSKA1LOW',         e.SKA1LOW ),      # 59
    ( 'SKAMID_PLUS',      e.SKAMID_PLUS ),  # 60
    ( 'SKAMID_PLUS2',     e.SKAMID_PLUS2 ), # 61
    ( 'yCHIME_nocut',     e.CHIME_nocut ),  # 62
    ( 'yCHIME_avglow',    e.CHIME_avglow ), # 63
    ( 'MID_B1_Base',      e.MID_B1_Base ),  # 64
    ( 'MID_B1_Alt',       e.MID_B1_Alt ),   # 65
    ( 'MID_B2_Base',      e.MID_B2_Base ),  # 66
    ( 'MID_B2_Upd',       e.MID_B2_Upd ),   # 67
    ( 'MID_B2_Alt',       e.MID_B2_Alt ),   # 68
    ( 'aLOW_Base',        e.LOW_Base ),     # 69
    ( 'aLOW_Upd',         e.LOW_Upd ),      # 70
    ( 'aLOW_Alt',         e.LOW_Alt ),      # 71
    ( 'MID_B2_Alt2',      e.MID_B2_Alt2 ),  # 72
    ( 'iMID_B1_Base',     e.MID_B1_Base ),  # 73
    ( 'iMID_B1_Alt',      e.MID_B1_Alt ),   # 74
    ( 'iMID_B2_Base',     e.MID_B2_Base ),  # 75
    ( 'hMID_B1_Rebase',   e.MID_B1_Rebase), # 76
    ( 'hMID_B1_Octave',   e.MID_B1_Octave), # 77
    ( 'hMID_B2_Rebase',   e.MID_B2_Rebase), # 78
    ( 'hMID_B2_Octave',   e.MID_B2_Octave), # 79
    ( 'iCVTEST1',         e.CVlimited_z0to3), # 80
    ( 'iCVTEST2',         e.CVlimited_z2to5), # 81
    ( 'iHIRAX',           e.HIRAX),         # 82
    ( 'iCosVis32x32',     e.CosVis32x32),   # 83
    ( 'iCosVis32x32_dmin10m', e.CosVis32x32_dmin10m), # 84
    ( 'iCosVis256x256',   e.CosVis256x256),   # 85
    ( 'MID_B1_RedBook',   e.MID_B1_RedBook), # 86
    ( 'MID_B2_RedBook',   e.MID_B2_RedBook), # 87
    ( 'MID_B1_SKAonly_RedBook', e.MID_B1_SKAonly_RedBook), # 88
    ( 'MID_B1_MK_RedBook', e.MID_B1_MK_RedBook), # 89
    ( 'MID_B2_MK_RedBook', e.MID_B2_MK_RedBook), # 90
    ( 'iHIRAX_highz',      e.HIRAX_highz),  # 91
    ( 'MeerKATL',      e.MeerKAT_Lband),    # 92
    ( 'MeerKATUHF',        e.MeerKAT_UHF),  # 93
]
names, expts = zip(*expt_list)
names = list(names); expts = list(expts)

################################################################################

# Take command-line argument for which survey to calculate, or set manually
if len(sys.argv) > 1:
    k = int(sys.argv[1])
    try:
        Sarea = float(sys.argv[2])
    except:
        Sarea = None
        pass
else:
    raise IndexError("Need to specify ID for experiment.")

names[k] += EXPT_LABEL
if myid == 0:
    print("="*50)
    print("Survey:", names[k])
    print("="*50)

# Tweak settings depending on chosen experiment
cv_limited = False
expts[k]['mode'] = "dish"
if names[k][0] == "i": expts[k]['mode'] = "idish"
if names[k][0] == "c": expts[k]['mode'] = "combined"
if names[k][0] == "y": expts[k]['mode'] = "icyl"
if names[k][0] == "f": expts[k]['mode'] = "paf"
if names[k][0] == "t": expts[k]['mode'] = "ipaf"
if names[k][0] == "a": expts[k]['mode'] = "iaa"
if names[k][0] == "h": expts[k]['mode'] = "hybrid"

expt = expts[k]
if Sarea is None:
    survey_name = names[k]
    root = "output/" + survey_name
else:
    expt['Sarea'] = Sarea * (D2RAD)**2.
    survey_name = names[k] + "_" + str(int(Sarea))
    root = "output/" + survey_name

# Define redshift bins
expt_zbins = rf.overlapping_expts(expt)
zs, zc = rf.zbins_equal_spaced(expt_zbins, dz=0.1)
#zs, zc = rf.zbins_equal_spaced(expt_zbins, dz=0.25) # 0.2
#zs, zc =  rf.zbins_const_dnu(expt_zbins, cosmo, dnu=20.)
#zs, zc = rf.zbins_const_dr(expt_zbins, cosmo, bins=14)
#zs, zc = rf.zbins_const_dnu(expt_zbins, cosmo, dnu=60.)
#zs, zc = rf.zbins_const_dnu(expt_zbins, cosmo, dnu=30.)
#zs = rf.zbins_fixed(expt_zbins, dz=0.1)

# FIXME
#print("FIXME! zbins set to manual")
#zs = np.array([0.25, 0.48])
#zc = 0.5*(zs[1:] + zs[:-1])

# Define kbins (used for output)
kbins = np.logspace(np.log10(0.001), np.log10(1.), 61)
#cosmo['f0_kbins'] = np.array([1e-4, 1e-2, 1e-1, 1e1])

expt['epsilon_fg'] = 1e-14
#expt['ttot'] *= 17520. / 1e4 #8765. / 1e4 # 2 calendar years on-sky
expt['ttot'] *= 4000. / 1e4 #8765. / 1e4
expt['k_nl0'] = 0.2 # = 0.3 h/Mpc
expt['wedge'] = False #'horizon' #'3pb' #False

# Neutrino mass
cosmo['mnu'] = 0.

# Precompute cosmological functions, P(k), massive neutrinos, and T(k) for f_NL
cosmo_fns = rf.background_evolution_splines(cosmo)
if cosmo['mnu'] != 0.:
    # Massive neutrinos
    mnu_str = "mnu%03d" % (cosmo['mnu']*100.)
    fname_pk = "cache_pk_%s.dat" % mnu_str
    fname_nu = "cache_%s" % mnu_str
    survey_name += mnu_str; root += mnu_str
    cosmo = rf.load_power_spectrum(cosmo, fname_pk, comm=comm)
    mnu_fn = rf.deriv_neutrinos(cosmo, fname_nu, mnu=cosmo['mnu'], comm=comm)
else:
    # Normal operation (no massive neutrinos or non-Gaussianity)
    cosmo = rf.load_power_spectrum(cosmo, "cache_pk.dat", comm=comm)
    mnu_fn = None

# Non-Gaussianity
#transfer_fn = rf.deriv_transfer(cosmo, "cache_transfer.dat", comm=comm)
transfer_fn = None

# Effective no. neutrinos, N_eff
#Neff_fn =  rf.deriv_neutrinos(cosmo, "cache_Neff", Neff=cosmo['N_eff'], comm=comm)
Neff_fn = None

# Optional additional parameters
switches = []
#switches = ['mg', ] #'sdbias']

# Scale-dependent growth
#cosmo['fs8_kbins'] = [0., 1e-2, 1e-1, 1e0, 1e2]

H, r, D, f = cosmo_fns

################################################################################
# Store cosmological functions
################################################################################

# Store values of cosmological functions
if myid == 0:
    # Calculate cosmo fns. at redshift bin centroids and save
    _H = H(zc)
    _dA = r(zc) / (1. + np.array(zc))
    _D = D(zc)
    _f = f(zc)
    np.savetxt(root+"-cosmofns-zc.dat", np.column_stack((zc, _H, _dA, _D, _f)))

    # Calculate cosmo fns. as smooth fns. of z and save
    zz = np.linspace(0., 1.05*np.max(zc), 1000)
    _H = H(zz)
    _dA = r(zz) / (1. + zz)
    _D = D(zz)
    _f = f(zz)
    np.savetxt(root+"-cosmofns-smooth.dat", np.column_stack((zz, _H, _dA, _D, _f)) )

# Precompute derivs for all processes
eos_derivs = rf.eos_fisher_matrix_derivs(cosmo, cosmo_fns, fsigma8=True)

"""
# Output all cosmo/instrumental parameters
print "*"*50
for key in cosmo.keys():
    print "%20s: %s" % (key, cosmo[key])
print "*"*50
for key in expt.keys():
    print "%20s: %s" % (key, expt[key])
print "*"*50
"""

################################################################################
# Loop through redshift bins, assigning them to each process
################################################################################

for i in range(zs.size-1):
    if i % size != myid:
      continue
    print(">>> %2d working on redshift bin %d / %d -- z = %3.3f" \
          % (myid, i, zs.size, zc[i]))

    # Calculate effective experimental params. in the case of overlapping expts.
    Sarea_rad = Sarea*(D2RAD)**2. if Sarea is not None else None
    expt_eff = rf.overlapping_expts(expt, zs[i], zs[i+1], Sarea=Sarea_rad)

    # Calculate basic Fisher matrix
    # (A, bHI, Tb, sigma_NL, sigma8, n_s, f, aperp, apar, [Mnu], [fNL], [pk]*Nkbins)
    F_pk, kc, binning_info, paramnames = rf.fisher(
                                         zs[i], zs[i+1], cosmo, expt_eff,
                                         cosmo_fns=cosmo_fns,
                                         transfer_fn=transfer_fn,
                                         massive_nu_fn=mnu_fn,
                                         Neff_fn=Neff_fn,
                                         return_pk=True,
                                         cv_limited=cv_limited,
                                         switches=switches,
                                         kbins=kbins )

    # Expand Fisher matrix with EOS parameters
    ##F_eos =  rf.fisher_with_excluded_params(F, [10, 11, 12]) # Exclude P(k)
    F_eos, paramnames = rf.expand_fisher_matrix(zc[i], eos_derivs, F_pk,
                                                names=paramnames, exclude=[],
                                                fsigma8=True)

    # Expand Fisher matrix for H(z), dA(z)
    # Replace aperp with dA(zi), using product rule. aperp(z) = dA(fid,z) / dA(z)
    # (And convert dA to Gpc, to help with the numerics)
    paramnames[paramnames.index('aperp')] = 'DA'
    da = r(zc[i]) / (1. + zc[i]) / 1000. # Gpc
    F_eos[7,:] *= -1. / da
    F_eos[:,7] *= -1. / da

    # Replace apar with H(zi)/100, using product rule. apar(z) = H(z) / H(fid,z)
    paramnames[paramnames.index('apar')] = 'H'
    F_eos[8,:] *= 1. / H(zc[i]) * 100.
    F_eos[:,8] *= 1. / H(zc[i]) * 100.

    # Save Fisher matrix and k bins
    np.savetxt(root+"-fisher-full-%d.dat" % i, F_eos, header=" ".join(paramnames))
    if myid == 0: np.savetxt(root+"-fisher-kc.dat", kc)

    # Save P(k) rebinning info
    np.savetxt(root+"-rebin-Fbase-%d.dat" % i, np.array(binning_info['F_base']) )
    np.savetxt(root+"-rebin-cumul-%d.dat" % i, np.array(binning_info['cumul']) )
    np.savetxt(root+"-rebin-kgrid-%d.dat" % i, np.array(binning_info['kgrid']) )
    np.savetxt(root+"-rebin-Vfac-%d.dat" % i, np.array([binning_info['Vfac'],]) )

comm.barrier()
if myid == 0: print("Finished.")


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
.. py:currentmodule:: xray.mac.models.chantler2005
.. moduleauthor:: Hendrix Demers <hendrix.demers@mail.mcgill.ca>

Chantler 2005 MAC model.
"""

###############################################################################
# Copyright 2021 Hendrix Demers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
###############################################################################

# Standard library modules.
import csv
import logging

# Third party modules.
from scipy.interpolate import interp1d

# Local modules.
from xray.mac import get_current_module_path

# Globals and constants variables.
SECTION_NAME = "Chantler2005"
OPTION_PATHNAME = "pathname"
OPTION_FILENAME = "filename"
OPTION_ENERGY_UNIT = "energyUnit"

ENERGIES_eV = "energies_eV"
MAC_cm2_g = "mac_cm2_g"

ENERGY_UNIT_eV = "eV"
ENERGY_UNIT_keV = "keV"

class Chantler2005:
    def __init__(self):
        self.minimum_energy_eV = 0.0
        self.minimum_mac_cm2_g = 0.0

        self.mac_data = {}
        self.edge_energies_eV = {}

        self.experimental_data = {}

    def read_mac_data(self, file_path=None, energy_unit=ENERGY_UNIT_keV):
        if file_path is None:
            file_path = get_current_module_path(__file__, "../../../data/chantler2005/FFastMAC.csv")
        self.reset_data()
        input_file = csv.reader(open(file_path))

        for items in input_file:
            self.extract_data_from_items_line(items, energy_unit)

    def reset_data(self):
        self.experimental_data = {}

    def extract_data_from_items_line(self, items, energy_unit):
        index = 0
        maximum_index = len(items)
        atomic_number = 1

        while index < maximum_index:
            try:
                if items[index] != '':
                    if energy_unit == ENERGY_UNIT_keV:
                        energy_keV = float(items[index])  # noqa
                        energy_eV = energy_keV * 1.0e3  # noqa
                    else:
                        energy_eV = float(items[index])  # noqa

                    mac_cm2_g = float(items[index+1])

                    if energy_eV > 0.0:
                        self.experimental_data.setdefault(atomic_number, {})
                        self.experimental_data[atomic_number].setdefault(ENERGIES_eV, []).append(energy_eV)
                        self.experimental_data[atomic_number].setdefault(MAC_cm2_g, []).append(mac_cm2_g)
            except ValueError as status:
                logging.error(status)
                logging.info(items)

            atomic_number += 1
            index += 2

    def compute_mac_cm2_g(self, energy_emitter_eV, atomic_number_absorber):  # noqa
        return self._compute_mac_cm2_g(energy_emitter_eV, atomic_number_absorber)

    def _compute_mac_cm2_g(self, energyEmitter_eV, atomic_number_absorber):  # noqa
        if atomic_number_absorber not in self.mac_data:
            self.read_mac_data()

            energies_eV = self.experimental_data[atomic_number_absorber][ENERGIES_eV]  # noqa
            macs_cm2_g = self.experimental_data[atomic_number_absorber][MAC_cm2_g]

            if len(energies_eV) > 0:
                self.mac_data.setdefault(atomic_number_absorber, {})

                self.minimum_energy_eV = energies_eV[0]

                self.minimum_mac_cm2_g = macs_cm2_g[0]

                self.maximum_energy_eV = energies_eV[-1]

                self.maximum_mac_cm2_g = macs_cm2_g[-1]

                self.mac_data[atomic_number_absorber] = interp1d(energies_eV, macs_cm2_g)

            else:
                logging.error("No mac for %i and %0.1f", atomic_number_absorber, energyEmitter_eV)
                return 0.0

        if energyEmitter_eV <= self.minimum_energy_eV:
            return self.minimum_mac_cm2_g

        if energyEmitter_eV >= self.maximum_energy_eV:
            return self.maximum_mac_cm2_g

        if atomic_number_absorber in self.mac_data:
            try:
                mac_value = self.mac_data[atomic_number_absorber](energyEmitter_eV)
            except ValueError:
                print(atomic_number_absorber, energyEmitter_eV)
                mac_value = self.minimum_mac_cm2_g
            return mac_value
        else:
            logging.error("No mac for %i and %0.1f", atomic_number_absorber, energyEmitter_eV)
            return 0.0

def compare_all_versions():
    import matplotlib.pyplot as plt
    from xray.mac.models.ionization_energies import IonizationEnergies, SUBSHELLS

    mac = Chantler2005()

    atomic_numbers = [1, 6, 22, 92]
    filenames = {"Default": "FFastMAC.csv", "NISTMonte2": "FFastMAC_nistMonte2.csv", "DTSA2": "FFastMAC_DTSA2.csv"}

    for atomic_number in atomic_numbers:
        plt.figure()
        for filename_key in filenames:
            filename = filenames[filename_key]

            if filename_key == "NISTMonte2":
                file_path = get_current_module_path(__file__, "../../data/chantler2005/%s" % filename)
                mac.read_mac_data(file_path, ENERGY_UNIT_eV)
            else:
                file_path = get_current_module_path(__file__, "../../data/chantler2005/%s" % filename)
                mac.read_mac_data(file_path)

            energies_eV = mac.experimental_data[atomic_number][ENERGIES_eV]  # noqa
            mac_cm2_g = mac.experimental_data[atomic_number][MAC_cm2_g]

            if filename_key == "Default":
                plt.loglog(energies_eV, mac_cm2_g, '.', label=filename_key)
            else:
                plt.loglog(energies_eV, mac_cm2_g, label=filename_key)

        plt.legend()
        plt.title(atomic_number)

        ionization_energies = IonizationEnergies()

        for subshell in SUBSHELLS:
            edge_energy_eV = ionization_energies.ionization_energy_eV(atomic_number, subshell)  # noqa
            if edge_energy_eV > 0.0:
                plt.axvline(edge_energy_eV, zorder=-10)
    plt.show()

def create_hdf5_file():
    import h5py
    import numpy as np

    filename = "chantler2005.hdf5"
    file_path = get_current_module_path(__file__, "../../data/chantler2005/%s" % filename)

    with h5py.File(file_path, "w") as hdf5_file:
        mac = Chantler2005()
        mac.read_mac_data()

        group_elements = hdf5_file.require_group("elements")

        for atomic_number in sorted(mac.experimental_data.keys()):
            group_name = "{:02d}".format(atomic_number)
            group_atomic_number = group_elements.require_group(group_name)
            energies_eV = np.array(mac.experimental_data[atomic_number][ENERGIES_eV])  # noqa
            macs_cm2_g = np.array(mac.experimental_data[atomic_number][MAC_cm2_g])

            group_atomic_number.create_dataset(ENERGIES_eV, data=energies_eV)
            group_atomic_number.create_dataset(MAC_cm2_g, data=macs_cm2_g)