# -*- coding: utf-8 -*-

# standard imports
import numpy as np
import cpnest
import nessai



class NessaiWrapper(nessai.model.Model):

    """
    wrapper class for nessai
    takes an already set up cpnest.model.Model object and wraps it as a nessai.model.Model object
    subsiquent model class can then be passed to nessai.flowsampler.FlowSampler
    handles data type conversions: names, bounds, and live points for methods
    """

    def __init__(self, signal_model_obj, uniform_prior=True):

        super().__init__()
        self.signal_model_obj = signal_model_obj # existing cpnest model object
        # atributes: cpnest --> nessai
        self.names = signal_model_obj.names
        self.bounds = dict(zip(signal_model_obj.names,signal_model_obj.bounds))
        self.cpnest_output_is_scalar=False


    def convert_live_point(self,x_nessai):
        # convert x: nessai --> cpnest
        # returns: signple cpnest live point, or a list of cpnest live points (depending on input shape)

        names = self.names
        if(x_nessai["it"].size ==1):
            x_cpnest = cpnest.parameter.LivePoint(names=names)
            for n in names:
                x_cpnest[n] = x_nessai[n]
            x_cpnest.logL = x_nessai["logL"]
            x_cpnest.logP = x_nessai["logP"]
        else:
            x_cpnest=[]
            for row in range(0,x_nessai["it"].size):
                current_x_cpnest = cpnest.parameter.LivePoint(names=names)
                for n in names:
                    current_x_cpnest[n] = x_nessai[n][row]
                current_x_cpnest.logL = x_nessai["logL"][row]
                current_x_cpnest.logP = x_nessai["logP"][row]
                x_cpnest.append(current_x_cpnest)
                del(current_x_cpnest)

        return x_cpnest


    def feed_into_cpnest(self,method, x_cpnest):
        # pass cpnest live points to cpnest methods
        # add extra safty

        if(type(x_cpnest) == cpnest.parameter.LivePoint):
            cpnest_output=method(x_cpnest)
        else:
            self.cpnest_output_is_scalar=False
            cpnest_output=[self.scalar_safe(method,x) for x in x_cpnest]
        nessai_output=np.array(cpnest_output)
        return nessai_output


    def scalar_safe(self, method, x):
        # safty check for boxed scalars

        if(self.cpnest_output_is_scalar == True):
            output = method(x)
        else:
            try:
                output=method(x)[0]
            except:
                output = method(x)
                self.cpnest_output_is_scalar = True

        return output

    # deligate log_proir and log_likelihood to existing cpnest model object
    # use live point converter and feeder function to convert formatting
    def log_prior(self,x_nessai):
        x_cpnest = self.convert_live_point(x_nessai)
        output_nessai = self.feed_into_cpnest(self.signal_model_obj.log_prior, x_cpnest)
        return output_nessai

    def log_likelihood(self,x_nessai):
        x_cpnest = self.convert_live_point(x_nessai)
        output_nessai = self.feed_into_cpnest(self.signal_model_obj.log_likelihood, x_cpnest)
        return output_nessai
