import json
import csv
import ipaddress
import os
import re

import openpyxl
from openpyxl.styles import Font, PatternFill, Alignment
from openpyxl.styles.colors import Color
from openpyxl.worksheet.dimensions import RowDimension
from openpyxl.utils.cell import column_index_from_string, get_column_letter

from bluecat import util
from bluecat.entity import Entity
from bluecat.api_exception import BAMException, PortalException

from flask import g

import config.default_config as config

# Workaround for DHCPV4ClientOption, DHCPServiceOption
def get_deployment_options(entity, option_types, server_id):
    options = []
    try:
        options = entity._api_client.getDeploymentOptions(entity.get_id(), option_types, server_id)
    except GeneralError as e:
        print('Exception at get_deployment_options(%s)' % util.safe_str(e))
        raise BAMException(safe_str(e))
    return options

def get_value_from_properties(properties, key):
    props = re.split(r'(?<!\\)\|', properties)
    for prop in props:
        value = prop.split('=')
        if value[0] == key:
            return value[1]
    return ''

col_config = {}
CONFIG_FILE = '%s/config_%s.json'

def module_path():
    return os.path.dirname(os.path.abspath(__file__))

def get_resource_text():
    return util.get_text(module_path(), config.language)

r_text = get_resource_text()
SHEET_TITLE_FOR_CONFIGURATION = r_text['sheet_title_for_configuration']
SHEET_TITLE_FOR_BLOCK = r_text['sheet_title_for_block']
SHEET_TITLE_FOR_NETWORK = r_text['sheet_title_for_network']

NETWORKID_DISPLAY_NAME = r_text['network_id_display_name']
BROADCAST_DISPLAY_NAME = r_text['broadcast_display_name']

BOLD_FONT = Font(bold=True)
LINK_FONT = Font(underline='single', color=Color(theme=10, type="theme"))
BOLD_LINK_FONT = Font(bold=True, underline='single', color=Color(theme=10, type="theme"))
TITLE_FILL_COLOR = PatternFill('solid', fgColor='9BC2E6')

def load_config(module_dir, nodes):
    global col_config
    col_config = json.load(open(CONFIG_FILE % (module_dir, config.language)))
    for col in col_config['props']:
        if col['gwidth'] == 0:
            continue

        node = {}
        node['label'] = col['title']
        node['name'] = col['id']
        node['index'] = col['id']
        node['width'] = col['gwidth']
        node['sortable'] = False
        nodes.append(node)

def get_order(entity):
    order = 0
    etype = entity.get_type()
    if etype in [Entity.IP4Block, Entity.IP4Network]:
        range = entity.get_property('CIDR')
        if range is not None:
            val = range.split('/')
            order = util.ip42int(val[0])
        else:
            order = util.ip42int(entity.get_property('start'))
    elif etype == Entity.DHCP4Range:
        order = util.ip42int(entity.get_property('start'))
    elif etype == Entity.IP4Address:
        order = util.ip42int(entity.get_address())
    return order

def get_range(entity):
    range = '-'
    etype = entity.get_type()
    if etype == Entity.Configuration:
        range = entity.get_name()
    elif etype in [Entity.IP4Block, Entity.IP4Network]:
        range = entity.get_property('CIDR')
        if range is None:
            range = entity.get_property('start') + '-' + entity.get_property('end')
    elif etype == Entity.DHCP4Range:
        range = entity.get_property('start') + '-' + entity.get_property('end')
    elif etype == Entity.IP4Address:
        range = entity.get_address()
    return range

def get_sorted_children(parent):
    nodes = []
    for ip4_block in parent.get_children_of_type(Entity.IP4Block):
        node = {}
        node['entity'] = ip4_block
        node['order'] = get_order(ip4_block)
        nodes.append(node)

    for ip4_network in parent.get_children_of_type(Entity.IP4Network):
        node = {}
        node['entity'] = ip4_network
        node['order'] = get_order(ip4_network)
        nodes.append(node)

    for dhcp4_range in parent.get_children_of_type(Entity.DHCP4Range):
        node = {}
        node['entity'] = dhcp4_range
        node['order'] = get_order(dhcp4_range)
        nodes.append(node)

    for ip4_address in parent.get_children_of_type(Entity.IP4Address):
        node = {}
        if ip4_address.get_state() == 'DHCP_RESERVED':
            node['entity'] = ip4_address
            node['order'] = get_order(ip4_address)
            nodes.append(node)

    nodes.sort(key = lambda node: node['order'])
    return nodes

def is_inherited(option):
    result = False
    inherited = get_value_from_properties(option['properties'], 'inherited')
    return True if inherited == 'true' else False

mac_pools = {}
dhcp_classes = {}

def clear_pool_class_cache():
    global mac_pools, dhcp_classes
    mac_pools = {}
    dhcp_classes = {}

def get_mac_pool_name(properties):
    api = api = g.user.get_api()
    global mac_pools
    id = get_value_from_properties(properties, 'macPool')

    pool_name = None
    if id in mac_pools.keys():
        pool_name = mac_pools[id]
    else:
        entity = api.get_entity_by_id(id)
        if entity is not None:
            pool_name = entity.get_name()
            mac_pools[id] = pool_name
    return pool_name

def get_dhcp_class_name(properties):
    api = api = g.user.get_api()
    global dhcp_classes
    id = get_value_from_properties(properties, 'dhcpMatchClass')

    class_name = None
    if id in dhcp_classes.keys():
        class_name = mac_pools[id]
    else:
        entity = api.get_entity_by_id(id)
        if entity is not None:
            class_name = entity.get_name()
            dhcp_classes[id] = class_name
    return class_name

def collect_options(entity):
    options = []
    try:
        for dop in get_deployment_options(entity, 'DHCPV4ClientOption', -1):
            if not is_inherited(dop):
                options.append(dop['name'] + '=' + dop['value'])
        for dop in get_deployment_options(entity, 'DHCPServiceOption', -1):
            if not is_inherited(dop):
                if '-pool' in dop['name']:
                    pool_name = get_mac_pool_name(dop['properties'])
                    options.append(dop['name'] + '=' + pool_name)
                elif '-dhcp-class-' in dop['name']:
                    print(dop)
                    class_name = get_dhcp_class_name(dop['properties'])
                    options.append(dop['name'] + '=' + class_name)
                else:
                    options.append(dop['name'] + '=' + dop['value'])
    except PortalException as pe:
        print(pe)

    return options

def set_common_props(node, parent, level, entity):
    node['id'] = entity.get_id()
    node['order'] = get_order(entity)
    node['parent'] = parent
    node['level'] = level + 1
    node['loaded'] = True

def set_props(node, entity, gui_flag = False):
    props = col_config['props']
    type_names = col_config['type_names']
    for prop in props:
        if gui_flag and prop['gwidth'] == 0:
            continue

        id = prop['id']
        if id == 'range':
            node[id] = get_range(entity)
        elif id == 'name':
            node[id] = entity.get_name()
        elif id == 'type':
            node[id] = type_names[entity.get_type()]
        else:
            node[id] = entity.get_property(id)

def construct_ip4_block_node(parent, level, entity):
    node = {}
    set_common_props(node, parent, level, entity)
    set_props(node, entity, True)
    node['isLeaf'] = False
    return node

def construct_ip4_network_node(parent, level, entity):
    node = {}
    set_common_props(node, parent, level, entity)
    set_props(node, entity, True)
    node['isLeaf'] = True
    return node

def construct_structure_row(indent, entity):
    row = []
    props = col_config['props']
    type_names = col_config['type_names']
    for prop in props:
        if prop['id'] == 'range':
            nw_range = ''
            for index in range(indent):
                nw_range += '\t'
            nw_range += get_range(entity)
            row.append(nw_range)
        elif prop['id'] == 'name':
            row.append(entity.get_name())
        elif prop['id'] == 'type':
            row.append(type_names[entity.get_type()])
        elif prop['id'] == 'options':
            row.append('\n'.join(collect_options(entity)))
        else:
            row.append(entity.get_property(prop['id']))
    return row

def is_in_dhcp_range(dhcp_ranges, address):
    result = False
    for dhcp_range in dhcp_ranges:
        start_ip = util.ip42int(dhcp_range.get_property('start'))
        end_ip = util.ip42int(dhcp_range.get_property('end'))
        ip4_address = util.int2ip4(address)
        if start_ip <= address and address <= end_ip:
            result = True
            break
    return result

def has_dhcp_range_or_reserve(entity):
    result = False
    if Entity.IP4Block == entity.get_type():
        for ip4_block in entity.get_ip4_blocks():
            result = has_dhcp_range_or_reserve(ip4_block)
            if result:
                break

        if result == False:
            for ip4_network in entity.get_ip4_networks():
                result = has_dhcp_range_or_reserve(ip4_network)
                if result:
                    break

    elif Entity.IP4Network == entity.get_type():
        dhcp_ranges = entity.get_children_of_type(Entity.DHCP4Range)
        if 0 < len(list(dhcp_ranges)):
            result = True
        else:
            for ip4_address in entity.get_children_of_type(Entity.IP4Address):
                if ip4_address.get_state() == 'DHCP_RESERVED':
                    result = True
                    break
    return result

# --------------------------------------------
# Network Structure CSV writing functions.
# --------------------------------------------
def write_structure_header(writer):
    header = []
    props = col_config['props']
    for prop in props:
        header.append(prop['title'])

    writer.writerow(header)

def write_structure(writer, level, entity, full):
    etype = entity.get_type()
    if etype in [Entity.DHCP4Range, Entity.IP4Address]:
        row = construct_structure_row(level, entity)
        writer.writerow(row)

    elif has_dhcp_range_or_reserve(entity):
        if etype == Entity.IP4Network or full:
            row = construct_structure_row(level, entity)
            writer.writerow(row)
            level += 1

        nodes = get_sorted_children(entity)
        for node in nodes:
            write_structure(writer, level, node['entity'], full)

# --------------------------------------------
# Top Writing function for CSV format.
# --------------------------------------------
def export_as_csv(api, dirname, filename, id, full):
    clear_pool_class_cache()
    entity = api.get_entity_by_id(id)

    encoding = str(col_config['encoding'])
    f = open(dirname + '/' + filename, 'w', encoding=encoding)
    writer = csv.writer(f, lineterminator = '\n')

    write_structure_header(writer)
    if entity.get_type() == Entity.Configuration:
        for ip4_block in entity.get_children_of_type(Entity.IP4Block):
            write_structure(writer, 0, ip4_block, full)
    else:
        write_structure(writer, 0, entity, full)

    f.close()

# -------------------------------------
# IP Address Excel writing functions.
# -------------------------------------
def get_sheet_name(entity):
    return get_range(entity).replace('/', ' MASK ')

def write_row_for_excel(sheet, start_column, index, row):
    start_column_idx = column_index_from_string(start_column)
    for column in range(len(row)):
        cell_value = row[column]
        cell = sheet.cell(row=index, column=start_column_idx + column, value=cell_value)
        if (cell_value is not None) and (1 < len(cell_value.split('\n'))):
            cell.alignment = Alignment(wrapText=True)

def set_fill_color(sheet, start_column, index, row):
    start_column_idx = column_index_from_string(start_column)

    for column in range(len(row)):
        cell = sheet.cell(row=index, column=start_column_idx + column)
        cell.fill = TITLE_FILL_COLOR
        cell.font = BOLD_FONT

def set_column_width(sheet, start_column, widthes):
    start_column_idx = column_index_from_string(start_column)

    for index in range(len(widthes)):
        column_letter = get_column_letter(start_column_idx + index)
        sheet.column_dimensions[column_letter].width = widthes[index]

def add_auto_filters(sheet, start_column, start_row, length):
    start_column_idx = column_index_from_string(start_column)
    end_column  = get_column_letter(start_column_idx + length - 1)
    filter_range = start_column + str(start_row) +  ':' + end_column + str(start_row)
    sheet.auto_filter.ref = filter_range

# --------------------------------------------
# Network Structure Excel writing functions.
# --------------------------------------------
def write_structure_title_for_excel(sheet, entity):
    title_cell_index = col_config['excel']['title_cell_index']
    title_string = ''
    if entity.get_type() == Entity.Configuration:
        title_string = SHEET_TITLE_FOR_CONFIGURATION + entity.get_name()
    else:
        if entity.get_type() == Entity.IP4Block:
            title_string = SHEET_TITLE_FOR_BLOCK
        else:
            title_string = SHEET_TITLE_FOR_NETWORK
        title_string += get_range(entity)
        if entity.get_name() != None:
            title_string += ' - ' + entity.get_name()

    sheet[title_cell_index].value = title_string
    sheet[title_cell_index].font = BOLD_FONT

def write_structure_header_for_excel(start_column, start_row, sheet):
    header = []
    widthes = []
    props = col_config['props']
    for prop in props:
        header.append(prop['title'])
        widthes.append(prop['width'])

    write_row_for_excel(sheet, start_column, start_row, header)
    set_fill_color(sheet, start_column, start_row, header)
    set_column_width(sheet, start_column, widthes)
    sheet.freeze_panes = 'A' + str(start_row + 1)

def write_structure_node_for_excel(sheet, level, start_column, index, entity):
    row = construct_structure_row(0, entity)
    write_row_for_excel(sheet, start_column, index, row)
    cell = sheet.cell(row=index, column=column_index_from_string(start_column))
    cell.alignment = Alignment(indent=level, vertical='center')

def write_structure_for_excel(sheet, level, start_column, index, entity, full):
    etype = entity.get_type()
    if etype in [Entity.DHCP4Range, Entity.IP4Address]:
        write_structure_node_for_excel(sheet, level, start_column, index, entity)
        if 0 < level:
            sheet.row_dimensions[index].outline_level = level
        index += 1
    elif has_dhcp_range_or_reserve(entity):
        if etype == Entity.IP4Network or full:
            write_structure_node_for_excel(sheet, level, start_column, index, entity)
            if 0 < level:
                sheet.row_dimensions[index].outline_level = level
            index += 1
            level += 1

        nodes = get_sorted_children(entity)
        for node in nodes:
            index = write_structure_for_excel(sheet, level , \
                                              start_column, index, node['entity'], full)
    return index

# --------------------------------------------
# Top Writing function for Excel format.
# --------------------------------------------
def export_as_excel(api, dirname, filename, id, full):
    clear_pool_class_cache()
    excel_config = col_config['excel']
    start_column = excel_config['start_column']
    start_row = excel_config['start_row']

    workbook = openpyxl.load_workbook(dirname + '/templates/Default.xlsx')
    sheet = workbook.active
    entity = api.get_entity_by_id(id)
    sheet.title = get_sheet_name(entity)
    write_structure_title_for_excel(sheet, entity)
    write_structure_header_for_excel(start_column, start_row, sheet)
    add_auto_filters(sheet, start_column, start_row, len(col_config['props']))
    index =  start_row + 1
    if entity.get_type() == Entity.Configuration:
        for ip4_block in entity.get_children_of_type(Entity.IP4Block):
            index = write_structure_for_excel(sheet, 0, \
                                              start_column, index, ip4_block, full)
    else:
        write_structure_for_excel(sheet, 0, start_column, index, entity, full)
    workbook.save(dirname + '/' + filename)


import json
from openpype_modules.ftrack.lib import BaseAction, statics_icon

class ThumbToChildren(BaseAction):
    '''Custom action.'''

    # Action identifier
    identifier = 'thumb.to.children'
    # Action label
    label = 'Thumbnail'
    # Action variant
    variant = " to Children"
    # Action icon
    icon = statics_icon("ftrack", "action_icons", "Thumbnail.svg")

    def discover(self, session, entities, event):
        """Show only on project."""
        if (len(entities) != 1 or entities[0].entity_type in ["Project"]):
            return False
        return True

    def launch(self, session, entities, event):
        '''Callback method for action.'''

        userId = event['source']['user']['id']
        user = session.query('User where id is ' + userId).one()

        job = session.create('Job', {
            'user': user,
            'status': 'running',
            'data': json.dumps({
                'description': 'Push thumbnails to Childrens'
            })
        })
        session.commit()
        try:
            for entity in entities:
                thumbid = entity['thumbnail_id']
                if thumbid:
                    for child in entity['children']:
                        child['thumbnail_id'] = thumbid

            # inform the user that the job is done
            job['status'] = 'done'
        except Exception as exc:
            session.rollback()
            # fail the job if something goes wrong
            job['status'] = 'failed'
            raise exc
        finally:
            session.commit()

        return {
            'success': True,
            'message': 'Created job for updating thumbnails!'
        }

def register(session):
    '''Register action. Called when used as an event plugin.'''

    ThumbToChildren(session).register()


"""
Routines that implement fits directly to samples of masses (or other values
drawn from a TGGD).

For fits to binned data, see :module:`mrpy.fit_curve`. For the
definition of the likelihood involved in the fits within this module see
:class:`mrpy.likelihoods.PerObjLike`.

This module also provides pre-defined prior functions, specifically, the ``normal_prior``.
"""

import numpy as np
import scipy.optimize as opt
from scipy.stats import truncnorm

import mrpy.extra.likelihoods as lk

try:
    import emcee
except ImportError:
    print("Warning: emcee not installed, some routines won't work.")

def _retarg(ll, jac, ret_jac):
    if ret_jac:
        return ll, jac
    else:
        return ll

# Define the likelihood function
def _lnl(p, m, nm, mmin, V,  bounds, prior_func=None, prior_kwargs=None,debug=0, ret_jac=False):
    ## Note m,nm, mmin are always interpreted as being a list of arrays/scalars

    if prior_kwargs is None:
        prior_kwargs = {}

    # Some absolute bounds
    if p[2] < 0 or p[0] < 0 :
        if debug > 0:
            print("OUT OF BOUNDS: ", p)
        return _retarg(-np.inf, np.inf, ret_jac)

    # Enforced bounds
    for i,pp in enumerate(p):
        if not bounds[i][0] <= pp <= bounds[i][1]:
            if debug > 0:
                print("parameter out of bounds: ", p, bounds)
            return _retarg(-np.inf, np.inf, ret_jac)

    # Priors
    if prior_func is None:  # default uniform prior.
        ll, jac = 0, np.zeros_like(p)
    else:
        ll, jac = prior_func(p,**prior_kwargs)

    # Likelihood
    for mi, nmi, mmini,Vi in zip(m, nm, mmin,V):
        _mod = lk.SampleLikeWeights(weights=nmi, logm=np.log10(mi),
                                    log_mmin=np.log10(mmini),
                                    logHs=p[0], alpha=p[1], beta=p[2], lnA = p[3] + np.log(Vi))
        ll += _mod.lnL
        if ret_jac:
            jac += _mod.jacobian

    if debug > 1:
        print("pars, ll, jac: ", p, ll, jac)

    if np.isnan(ll):
        return _retarg(-np.inf, jac, ret_jac)
    else:
        return _retarg(ll, jac, ret_jac)

def _objfunc(*args):
    out = _lnl(*args)
    if np.isscalar(out):
        return -out
    else:
        return -out[0], -out[1]

def normal_prior(p,mean,sd):
    """
    A normal prior on each parameter.

    Parameters
    ----------
    p : list
        Values of the parameters

    mean : list
        Values of the prior mean

    sd : list
        Values of the prior standard deviations. Set to inf
        if uniform prior required on a single parameter.

    Returns
    -------
    ll :
        The likelihood of the parameters given the priors

    jac :
        The jacobian of the likelihood.

    """
    ll = 0
    jac = []

    for v,mu,s in zip(p,mean,sd):
        ll += -((v-mu)/(2*s))**2
        jac.append(-(v-mu)/s)

    return ll, jac

class SimFit(object):
    """
    Per-object fits.

    Parameters
    ----------
    m : array or list of arrays
        Masses. Either an array or a list of arrays, each of which is a sample *to be
        analysed simultaneously*. In the latter case the samples should have the same
        underlying distribution, but may have differing truncation scales.

    nm : array, optional
        Specifies the number of occurrences of each variate in `m` (which should then
        ideally be unique). If not passed, each variate is assumed to occur once. This
        is useful for speeding up fits on quantized simulations. If `m` is a list of
        arrays, this should be also.

    mmin : array_like, optional
        The truncation mass of the sample. By default takes the lowest value of `m`.
        If `m` is a list of arrays, this should be a list.

    V : array, optional
        The volume of each subsample

    hs_bounds, alpha_bounds, beta_bounds : 2-tuple
        2-tuples specifying minimum and maximum values for each bound.

    prior_func : function, optional
        A function to calculate the likelihood arising from priors on the parameters.
        By default, uniform priors are assumed, which add nothing to the likelihood.
        The first parameter taken by the function should be the vector of parameter
        values, ``[logHs,alpha,beta]``, after which arbitrary values are passed via
        `prior_kwargs`. It should return a tuple, with the first value
        being a float and constituting the likelihood arising from the prior, and the second
        being a 3-vector constituting the modification to the jacobian. This can be zero
        if no jacobian is desired.

    prior_kwargs : dict
        Arguments sent to the `prior_func`.

    Notes
    -----
    Use as stringent bounds as possible, since the algorithm explores the
    edges, which can induce numerical error if values far from the solution are chosen.
    """

    def __init__(self, m, nm=None, mmin=None,V=1.0,
                 hs_bounds=(10, 16), alpha_bounds=(-1.99, -1.3),
                 beta_bounds=(0.1, 2.0), lnA_bounds = (-40,-10),
                 prior_func=None,prior_kwargs=None):

        if prior_kwargs is None:
            prior_kwargs = {}

        self._determine_suite(m, nm, mmin,V)

        # Make sure all masses are above mmin
        for i, (m, mmin) in enumerate(zip(self.m, self.mmin)):
            self.m[i] = m[m >= mmin]
            self.nm[i] = self.nm[i][m >= mmin]

        self.hs_bounds = hs_bounds
        self.alpha_bounds = alpha_bounds
        self.beta_bounds = beta_bounds
        self.lnA_bounds = lnA_bounds
        self.prior_func = prior_func
        self.prior_kwargs = prior_kwargs

    def _determine_suite(self, m, nm, mmin,V):
        ## Determine whether there is a suite of simulations.
        if np.isscalar(m[0]):
            self.m = [m]
            self.logm = [np.log10(m)]

            if nm is None:
                self.nm = [np.ones_like(m)]
            else:
                self.nm = [nm]

            if mmin is None:
                self.mmin = [m.min()]
            else:
                self.mmin = [mmin]

            self.log_mmin = [np.log10(self.mmin[0])]

            self.V = np.array([V]).flatten()

        else:
            self.m = m
            self.logm = [np.log10(x) for x in self.m ]

            if nm is None:
                self.nm = [np.ones_like(x) for x in m]
            else:
                self.nm = nm

            if mmin is None:
                self.mmin = [x.min() for x in m]
            else:
                self.mmin = mmin

            self.log_mmin = [np.log10(x.min()) for x in self.mmin]

            if np.isscalar(V):
                print("WARNING: V is a scalar, but there are multiple datasets")
                self.V = [V]*len(m)
            else:
                self.V = V

    def run_downhill(self, hs0=14.5, alpha0=-1.9, beta0=0.8, lnA0=-40.0,
                     debug=0, jac=True, **minimize_kw):
        """
        Downhill-gradient optimization.

        Parameters
        ----------
        hs0, alpha0, beta0, lnA0: float, optional
            Initial guess for each of the MRP parameters.

        debug : int, optional
            Set the level of info printed out throughout the function. Highest current
            level that is useful is 2.

        jac : bool, optional
            Whether to use analytic jacobian (usually a good idea)

        minimize_kw : dict
            Any other parameters to :func:`scipy.optimize.minimize`.

        Returns
        -------
        downhill_res : `OptimizeResult`
            The optimization result represented as a ``OptimizeResult`` object
            (see scipy documentation).
            Important attributes are: ``x`` the solution array, ``success`` a Boolean flag
            indicating if the optimizer exited successfully and ``message`` which describes
            the cause of the termination.

            The parameters are ordered by `logHs`, `alpha`, `beta`, `[lnA]`.

        downhill_obj : :class:`mrpy.likelihoods.PerObjLikeWeights` object
            An object containing the solution parameters and methods to access
            relevant quantities, such as the mass function, or jacobian and
            hessian at the solution.

        Examples
        --------
        The most obvious example is to generate a sample of variates from given parameters:

        >>> from mrpy.base.stats import TGGD
        >>> r = TGGD(scale=1e14,a=-1.8,b=1.0,xmin=1e12).rvs(1e5)

        Then find the best-fit parameters for the resulting data:

        >>> from mrpy.fitting.fit_sample import SampleFit
        >>> FitObj = SampleFit(r)
        >>> res,obj = FitObj.run_downhill()
        >>> print res.x

        We can also use the ``obj`` object to explore some of the qualities of the fit

        >>> print obj.hessian
        >>> print obj.cov

        The results are also stored in the class as `downhill_obj` and `downhill_res`.

        >>> from matplotlib.pyplot import plot
        >>> plot(FitObj.downhill_obj.logm,FitObj.downhill_obj.dndm(log=True))
        >>> print obj.stats.mean, r.mean()
        """
        p0 =[hs0, alpha0, beta0,lnA0]
        bounds = [self.hs_bounds, self.alpha_bounds, self.beta_bounds,self.lnA_bounds]

        self.downhill_res = opt.minimize(_objfunc, p0, args=(self.m, self.nm, self.mmin, self.V, bounds,
                                                             self.prior_func,self.prior_kwargs,
                                                             debug, jac),
                                         bounds=bounds, jac=jac, **minimize_kw)

        self.downhill_obj = [lk.SampleLikeWeights(logm=np.log10(mi), weights=nmi,
                                                  logHs=self.downhill_res.x[0], alpha=self.downhill_res.x[1],
                                                  beta=self.downhill_res.x[2], lnA = self.downhill_res.x[3] - np.log(V),
                                                  log_mmin=np.log10(mmini)) for mi, nmi, mmini,V in
                             zip(self.m, self.nm, self.mmin,self.V)]

        return self.downhill_res, self.downhill_obj

    # =========================================================================================
    # EMCEE BASED ROUTINES
    # =========================================================================================
    @staticmethod
    def _get_cuts(bound, mu, sigma):
        return (bound - mu)/sigma

    def _get_initial_ball(self, guess, bounds, chains):
        s = 0.05
        if bounds is not None:
            a = np.array([c[0] for c in bounds])
            b = np.array([c[1] for c in bounds])

            aa = self._get_cuts(a, guess, np.abs(s*(b-a)))
            bb = self._get_cuts(b, guess, np.abs(s*(b-a)))
        else:
            aa = [-np.inf]*len(guess)
            bb = [np.inf]*len(guess)

        stacked_val = np.empty((chains, len(guess)))

        for i, (g, A, B) in enumerate(zip(guess, aa, bb)):
            stacked_val[:, i] = truncnorm(A, B, loc=g, scale=np.abs(s*(b[i]-a[i]))).rvs(chains)

        return stacked_val

    def run_mcmc(self, nchains=50, warmup=1000, iterations=1000,
                 hs0=14.5, alpha0=-1.9, beta0=0.8, lnA0=-26.0, logm0 = None, debug=0,
                 opt_init=False, opt_kw=None, chainfile="chain.dat", save_latent = True,
                 **kwargs):
        """
        Per-object MCMC fit for masses `m`, using the `emcee` package.

        This creates an :class:`emcee.EnsembleSampler` object with a correct
        model, and runs warmup and specified iterations. The entire :class:`emcee.EnsembleSampler`
        object is returned, and stored in the instance of this class as :attr:`~mcmc_res`. This affords
        greater flexibility, with the ability to run no warmup and 0 iterations, and run the iterations
        oneself.

        Parameters
        ----------
        nchains : int, optional
            Number of chains to use in the AIES MCMC algorithm

        warmup : int, optional
            Number (discarded) warmup iterations.

        iterations : int, optional
            Number of iterations to keep in the chain.

        hs0, alpha0, beta0: float, optional
            Initial guess for each of the MRP parameters.

        debug : int, optional
            Set the level of info printed out throughout the function. Highest current
            level that is useful is 2.

        opt_init : bool, optional
            Whether to run a downhill optimization routine to get the best
            starting point for the MCMC.

        opt_kw : dict, optional
            Any arguments to pass to the downhill run.

        kwargs :
            Any other parameters to :class:`emcee.EnsembleSampler`.

        Returns
        -------
        mcmc_res : :class:`emcee.EnsembleSampler` object
            This object contains the stored chains, and other attributes.

        Examples
        --------
        The most obvious example is to generate a sample of variates from given parameters:

        >>> from mrpy.base.stats import TGGD
        >>> r = TGGD(scale=1e14,a=-1.8,b=1.0,xmin=1e12).rvs(1e5)

        Then find the best-fit parameters for the resulting data:

        >>> from mrpy.fitting.fit_sample import SampleFit
        >>> FitObj = SampleFit(r)
        >>> mcmc_res = FitObj.run_mcmc(nchains=10,warmup=100,iterations=100)
        >>> print mcmc_res.flatchain.mean(axis=0)
        """
        if opt_kw is None:
            opt_kw = {}

        # First, set guess, either by optimization or passed values
        guess = np.array([hs0,alpha0,beta0,lnA0])
        bounds = [self.hs_bounds, self.alpha_bounds, self.beta_bounds,self.lnA_bounds]

        if opt_init:
            if not hasattr(self, "downhill_res"):
                self.run_downhill(hs0, alpha0, beta0, lnA0, debug, **opt_kw)
            res = self.downhill_res

            if res.success:
                guess = res.x
                if debug:
                    print("Optimization result (used as guess): ", guess)
            else:
                print("WARNING: Optimization failed. Falling back on given initial parameters.")

        initial = self._get_initial_ball(guess, bounds, nchains)

        self.mcmc_res = emcee.EnsembleSampler(nchains, initial.shape[1], _lnl,
                                              args=(self.m, self.nm, self.mmin, self.V, bounds,
                                                    self.prior_func, self.prior_kwargs,
                                                    debug, False), **kwargs)
        if warmup:
            initial, _, _ = self.mcmc_res.run_mcmc(initial, warmup, storechain=False)
            self.mcmc_res.reset()

        self.mcmc_res.run_mcmc(initial, iterations)

        return self.mcmc_res

    def lnL(self, p, ret_jac=False, debug=0):
        """
        Return the log-likelihood of the current model at the parameters `p`.

        Parameters
        ----------
        p : array
            The values of the parameters, ``[logHs, alpha, beta]``.

        ret_jac : bool
            Whether to return the jacobian as the second arg.

        debug : int, optional
            Set the level of info printed out throughout the function. Highest current
            level that is useful is 2.

        Returns
        -------
        ll : float
            The log-likelihood at the parameters. This is exactly the same as used in the
            fitting process.

        jac : length-3 array, optional
            Returned only if `ret_jac` is `True`. The jacobian at the current parameter vector.

        """
        bounds = [self.hs_bounds, self.alpha_bounds, self.beta_bounds,self.lnA_bounds]
        return _lnl(p, self.m, self.nm, self.mmin, self.V,bounds, self.prior_func,self.prior_kwargs,debug, ret_jac)