"""
.. module:: cosmo_hammer
    :synopsis: Invoke the Cosmo Hammer
.. moduleauthor:: Benjamin Audren <benjamin.audren@epfl.ch>

This module interface Monte Python with the CosmoHammer, available
`publicly <http://www.astro.ethz.ch/refregier/research/Software/cosmohammer>`_
and developped by Joel Akeret and Sebastian Seehars, who helped me a lot
creating this interface.

The link between the two codes is that some functions have been added to the
classes defined in Monte Python, so that they can become CoreModules and
LikelihoodModules of CosmoHammer. Mostly, all the important classes in Monte
Python (:class:`Data` <data.Data>, Class and :class:`Likelihood`
<likelihood_class.Likelihood>)

"""
import numpy as np
import os
import warnings
import logging

import io_mp
import sampler
from cosmoHammer import MpiCosmoHammerSampler
from cosmoHammer import LikelihoodComputationChain
#from cosmoHammer.likelihood.chain.LikelihoodComputationChain import (
#    LikelihoodComputationChain)
#from cosmoHammer.sampler.CosmoHammerSampler import CosmoHammerSampler
from cosmoHammer.util.SampleFileUtil import SampleFileUtil

# Cosmo Hammer subfolder and name separator
CH_subfolder = 'CH'
CH_separator = '-'
# Cosmo Hammer file names ending, after the defined 'base_name'
name_arguments = '.arguments'
name_chain = 'chain_CH__sampling.txt'
name_error_chain =  'chain_CH__sampling-error_log.txt'

# Cosmo Hammer option prefix
CH_prefix = 'CH_'
# User-defined arguments for the Hammer
CH_user_arguments = {
    'walkersRatio':
    {'help': 'Number of walkers',
     'type': int},
    'burninIterations':
    {'help': 'Number of burnin phase iterations',
     'type': int},
    'sampleIterations':
    {'help': 'Number of sample iterations',
     'type': int}}

# List of strings that will contain the derived parameters names.
# Filled in run using the MontePython class data
# Used in DerivedUtil.persistValues
CH_derived_names = []

def run(cosmo, data, command_line):
    """
    Sample with the CosmoHammer

    """
    # Store the parameters inside the format expected by CosmoHammer
    # TODO: about the derived params?
    parameter_names = data.get_mcmc_parameters(["varying"])
    # Fill now CH_derived_names to be used in DerivedUtil.persistValues
    global CH_derived_names
    CH_derived_names= data.get_mcmc_parameters(["derived"])

    # Ensure that their prior is bound and flat
    is_flat, is_bound = sampler.check_flat_bound_priors(
        data.mcmc_parameters, parameter_names)
    if not is_flat:
        raise io_mp.ConfigurationError(
            'The Cosmo Hammer is only available with flat ' +
            'priors. Sorry!')
    if not is_bound:
        raise io_mp.ConfigurationError(
            'The Cosmo Hammer is only available for bound ' +
            'parameters. Set reasonable bounds for them in the ".param"' +
            'file.')

    params = []
    for parameter in parameter_names:
        params.append(data.mcmc_parameters[parameter]['initial'])
    params = np.array(params)

    # If absent, create the sub-folder CH
    CH_folder = os.path.join(command_line.folder, CH_subfolder)
    if not os.path.exists(CH_folder):
        os.makedirs(CH_folder)

    # Initialize a chain object (Beware, the concept is quite different than
    # the chain of the module :mod:`mcmc`)
    chain = LikelihoodComputationChain(
        min=params[:, 1],
        max=params[:, 2])

    # Add data and cosmo as two core modules. Note that the order is important
    # here, since data must be called before cosmo.
    chain.addCoreModule(data)
    chain.addCoreModule(cosmo)
    chain.addCoreModule(store_cosmo_derived)

    # Add each likelihood class as a LikelihoodModule
    for likelihood in data.lkl.itervalues():
        chain.addLikelihoodModule(likelihood)

    # Define the file prefix
    chain_name = [a for a in command_line.folder.split(os.path.sep) if a][-1]
    file_prefix = os.path.join(command_line.folder, CH_subfolder, chain_name)

    # Recover the User options
    data.CH_arguments = {'walkersRatio': 50,
                         'burninIterations': 10,
                         'sampleIterations': 30}
    for arg in CH_user_arguments:
        value = getattr(command_line, CH_prefix+arg)
        if value != -1:
            data.CH_arguments[arg] = value
        # else, do not define them, and leave the default Cosmo Hammer ones.

    # Write the CosmoHammer arguments
    with open(file_prefix+name_arguments, 'w') as arg_file:
        for arg in data.CH_arguments:
            arg_file.write(
                ' = '.join([str(arg), str(data.CH_arguments[arg])]) + '\n')

    # Create an extension to the SampleFileUtil from cosmoHammer
    derived_util = DerivedUtil(file_prefix)

    try:
        num_threads = int(os.environ['OMP_NUM_THREADS'])
    except KeyError:
        warnings.warn(
            "The environment variable OMP_NUM_THREADS is not set. "
            "To run the Cosmo Hammer meaningfully, you should better "
            "set it to something! Defaulting to 1 for now.")
        num_threads = 1

    # Create the Sampler object
    sampler_hammer = MpiCosmoHammerSampler(
        params=params,
        likelihoodComputationChain=chain,
        filePrefix=file_prefix,
        storageUtil=derived_util,
        threadCount=num_threads,
        **data.CH_arguments)

    # create console handler and set level to debug (does not seem to appear)
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.DEBUG)
    logging.getLogger().addHandler(console_handler)

    sampler_hammer.startSampling()

    #print sampler_hammer
    #for elem in dir(sampler_hammer):
        #print elem
    #_, data_from_context = chain()
    #print data_from_context


def from_CH_output_to_chains(folder):
    """
    Translate the output of the Cosmo Hammer into Monte Python chains

    This function will be called by the module :mod:`analyze`.
    """
    chainsClean, lklClean = [], []  # Arrays to store the valid steps.
    error = []  # Array to store the step did not could not be computed.

    chain_name = [a for a in folder.split(os.path.sep) if a][-2]
    base_name = os.path.join(folder, chain_name)
    # Recover the points in parameter space
    with open(base_name+'.out', 'r') as param_values_file:
        chains = np.loadtxt(param_values_file)
    # Recover the associated likelihood (it is the -log likelihood)
    with open(base_name+'prob.out', 'r') as lkl_values_file:
        lkl = np.loadtxt(lkl_values_file)

    # Glue them together, with an additional column of ones, for the
    # multiplicity. This does not mean that the acceptance rate is one, but
    # points where the code stayed are duplicated in the file.

    # First, separe errors from accepted points:
    for x, p in zip(chains, lkl):
        if p != -np.inf:
            chainsClean.append(x)
            lklClean.append([-p])  # MP stores -loglkl in its files.
        else:
            error.append([val for val in x if not np.isnan(val)])  # Remove nan's

    ## Create the array of ones
    ones = np.array([[1] for _ in range(len(lklClean))])


    ## Concatenate everything and save to file
    final = np.concatenate((ones, lklClean, chainsClean), axis=1)
    output_folder = os.path.join(folder, '..')
    output_chain_path = os.path.join(output_folder, name_chain)
    output_error_chain_path = os.path.join(output_folder, name_error_chain)
    np.savetxt(output_chain_path, final)
    np.savetxt(output_error_chain_path, error)


class DerivedUtil(SampleFileUtil):
    """
    Extends the writing class from CosmoHammer to include derived parameters.
    """
    def persistValues(self, posFile, probFile, pos, prob, data):
        """
        Writes the walker positions and the likelihood to the disk
        """
        # extend the pos array to also contain the value of the derived
        # parameters

        derived_not_computed = [np.NaN] * len(CH_derived_names)

        derived = np.array(
            #[[a for a in elem.itervalues()] if elem else derived_not_computed for elem in data])
            [[elem[key] for key in sorted(elem.iterkeys())] if elem else derived_not_computed for elem in data])
        final = np.concatenate((pos, derived), axis=1)

        posFile.write("\n".join(
            [" ".join([str(q) for q in p]) for p in final]))
        posFile.write("\n")
        posFile.flush()

        probFile.write("\n".join([str(p) for p in prob]))
        probFile.write("\n")
        probFile.flush()


def store_cosmo_derived(ctx):
    """
    Store derived parameters. Based on sampler.compute_lkl
    """
    from classy import CosmoSevereError

    cosmo = ctx.get("cosmo")
    data = ctx.get("data")

    if data.get_mcmc_parameters(['derived']) != []:
        try:
            derived = cosmo.get_current_derived_parameters(
                data.get_mcmc_parameters(['derived']))
        except AttributeError:
            # This happens if the classy wrapper is still using the old
            # convention, expecting data as the input parameter
            cosmo.get_current_derived_parameters(data)
        except CosmoSevereError:
            raise io_mp.CosmologicalModuleError(
                "Could not write the current derived parameters")

        ctx.add('key_data', derived)

        # If we want the derived parameters stored rescaled
        # derived_scaled = {elem: value/data.mcmc_parameters[elem]['scale'] for elem, value in derived.iteritems()}
        # ctx.add('key_data', derived_scaled)
