''' Module to provide the Layer - Stack - Sample classes to build a sample for reflectivity modelling.

Classes:
ReflFunction - A function class that can be used as a parameter in the other classes.
is_reflfunction - Funtion that checks if an object belongs to the class ReflFunction
ReflBase - Base class for all the physical classes.
LayerBase - Base Layer class.
StackBase - Base Stack class.
SampleBase - Base Sample class.
InstrumentBase - Base Instrument class.
ModelFactory - A factory class that creates specific instances of the Base classes for a certain set of parameters.
MakeClasses - Function that creates the classes (backward comparability).
'''

import numpy as np
from genx.gui_logging import iprint


class ReflFunction:
    def __init__(self, function, validation_args, validation_kwargs, id=None):
        """ Creates the Refl Function given function. The arguments validation_args and
        validation_kwargs will be used to the validate the returned type of the function by passing
        them to function. The variable id should be a unique string to identify the type ReflFunction.
        """
        self.__func__ = function
        self.validation_args = validation_args
        self.validation_kwargs = validation_kwargs
        self.id = id

    def __call__(self, *args, **kwargs):
        return self.__func__(*args, **kwargs)

    def validate(self):
        ''' Function to test that the function returns the anticipated type
        '''
        return self.__call__(*self.validation_args, **self.validation_kwargs)

    def _check_obj(self, other):
        """ Checks the object other so that it fulfills the demands for arithmetic operations.
        """
        supported_types = [int, float, int, complex, np.float64, np.float32]
        if is_reflfunction(other):
            if self.id != other.id:
                raise TypeError("Two ReflFunction objects must have identical id's to conduct arithmetic operations")
        elif not type(other) in supported_types:
            raise TypeError("%s is not supported for arithmetic operations "% (repr(type(other))) +
                            "of a ReflFunction. It has to int, float, long or complex" )

    def __mul__(self, other):
        self._check_obj(other)
        if is_reflfunction(other):
            def new_func(*args, **kwargs):
                return self(*args, **kwargs)*other(*args, **kwargs)
        else:
            def new_func(*args, **kwargs):
                return self(*args, **kwargs)*other
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)

    def __rmul__(self, other):
        self._check_obj(other)
        if is_reflfunction(other):
            def new_func(*args, **kwargs):
                return other(*args, **kwargs)*self(*args, **kwargs)
        else:
            def new_func(*args, **kwargs):
                return other*self(*args, **kwargs)
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)

    def __add__(self, other):
        self._check_obj(other)
        if is_reflfunction(other):
            def new_func(*args, **kwargs):
                return self(*args, **kwargs) + other(*args, **kwargs)
        else:
            def new_func(*args, **kwargs):
                return self(*args, **kwargs) + other
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)

    def __radd__(self, other):
        self._check_obj(other)
        if is_reflfunction(other):
            def new_func(*args, **kwargs):
                return other(*args, **kwargs) + self(*args, **kwargs)
        else:
            def new_func(*args, **kwargs):
                return other + self(*args, **kwargs)
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)

    def __sub__(self, other):
        self._check_obj(other)
        if is_reflfunction(other):
            def new_func(*args, **kwargs):
                return self(*args, **kwargs) - other(*args, **kwargs)
        else:
            def new_func(*args, **kwargs):
                return self(*args, **kwargs) - other
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)

    def __rsub__(self, other):
        self._check_obj(other)
        if is_reflfunction(other):
            def new_func(*args, **kwargs):
                return other(*args, **kwargs) - self(*args, **kwargs)
        else:
            def new_func(*args, **kwargs):
                return other - self(*args, **kwargs)
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)

    def __div__(self, other):
        self._check_obj(other)
        if is_reflfunction(other):
            def new_func(*args, **kwargs):
                return self(*args, **kwargs)/other(*args, **kwargs)
        else:
            def new_func(*args, **kwargs):
                return self(*args, **kwargs)/other
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)

    def __rdiv__(self, other):
        self._check_obj(other)
        if is_reflfunction(other):
            def new_func(*args, **kwargs):
                return other(*args, **kwargs)/self(*args, **kwargs)
        else:
            def new_func(*args, **kwargs):
                return other/self(*args, **kwargs)
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)

    def __neg__(self):
        def new_func(*args, **kwargs):
            return -self(*args, **kwargs)
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)

    def __pos__(self):
        def new_func(*args, **kwargs):
            return self(*args, **kwargs)
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)

    def __pow__(self, other):
        self._check_obj(other)
        if is_reflfunction(other):
            def new_func(*args, **kwargs):
                return self(*args, **kwargs)**other(*args, **kwargs)
        else:
            def new_func(*args, **kwargs):
                return self(*args, **kwargs)**other
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)

    def __rpow__(self, other):
        self._check_obj(other)
        if is_reflfunction(other):
            def new_func(*args, **kwargs):
                return other(*args, **kwargs)**self(*args, **kwargs)
        else:
            def new_func(*args, **kwargs):
                return other**self(*args, **kwargs)
        return ReflFunction(new_func, self.validation_args, self.validation_kwargs, self.id)


def cast_to_array(list_of_obj, *args, **kwargs):
    ''' Casts an list_of_obj, can be a number or an ReflFunction, into an list of evaluated values'''
    id = ''
    shape = False
    ret_list = []
    for obj in list_of_obj:
        if is_reflfunction(obj):
            if id == '':
                id = obj.id
                ret_list.append(obj(*args, **kwargs))
                # Check if we have got an array
                if not np.isscalar(ret_list[-1]):
                    shape = ret_list[-1].shape
            elif id == obj.id:
                ret_list.append(obj(*args, **kwargs))
            else:
                TypeError("Two ReflFunction objects must have identical id's in order to merge them into an array")
        else:
            # We assume that this is an object that can be transformed into an array later on.
            ret_list.append(obj)
    # if we have an array make sure that all the objects have the same shape
    if shape:
        base_array = np.ones(shape)
        nret_list = []
        for item in ret_list:
            if np.isscalar(item):
                nret_list.append(item*base_array)
            else:
                nret_list.append(item)
        ret_list = nret_list

    return np.array(ret_list)

def harm_sizes(ar, shape, dtype = np.float64):
    '''Utility function to add an additional axis if needed to fulfill the size in shape'''
    ar = np.array(ar, dtype = dtype)
    if shape is None:
        return ar
    elif len(ar.shape) < len(shape):
        return np.array(ar[..., np.newaxis]*np.ones(shape), dtype = dtype)
    elif ar.shape == shape:
        return ar
    else:
        raise TypeError('The size of the array, %s, can not be changed to shape %s'%(ar.shape, shape))



def is_reflfunction(obj):
    ''' Convenience function to determine whether obj belongs to the ReflFunction class.
    Return boolean.
    '''
    return obj.__class__.__name__ == 'ReflFunction'

def resolve_par(obj, par_name, max_recur=100):
    """ Resolves a layer parameter and takes into account a coupling to another object.
    If obj.par_name is of type obj the relation will be recursive, obj.par_name.par_name.
    A maximum of 100 levels are conducted before an error is thrown.

    :param obj: An object
    :param par_name: A name of a member
    :param max_recur: Maximum number of recursions.
    :return: a value
    """
    val = getattr(obj, par_name)
    i = 0
    while type(val) is type(obj) and i < max_recur and not is_reflfunction(val):
        val = getattr(val, par_name)
        i += 1

    return val




class ReflBase:
    # _parameters is a dict of parameter names with their default values
    _parameters = {}

    def __init__(self, **kargs):
        # Setup all the parameters in the class
        for par in self._parameters:
            setattr(self, par, self._parameters[par])
            iscomplex = type(self._parameters[par]) is complex
            # Making the set function
            self._make_set_func(par, iscomplex)
            # Creating the get function
            self._make_get_func(par, iscomplex)

        # Set all parameters given as keyword arguments
        for k in kargs:
            if not k in self._parameters and not k in dir(self):
                continue
                raise ValueError('%s is not an parameter in %s' %
                                 (k, self.__class__))
            else:
                setattr(self, k, kargs[k])

    def _make_get_func(self, par, iscomplex=False):
        ''' Creates a get function for parameter par and binds it to the object
        '''

        def get_func():
            return getattr(self, par)

        get_func.__name__ = 'get' + par.capitalize()
        setattr(self, get_func.__name__, get_func)

        if iscomplex:
            def get_real_func():
                return getattr(self, par).real

            get_real_func.__name__ = get_func.__name__ + 'real'
            setattr(self, get_real_func.__name__, get_real_func)

            def get_imag_func():
                return getattr(self, par).imag

            get_imag_func.__name__ = get_func.__name__ + 'imag'
            setattr(self, get_imag_func.__name__, get_imag_func)

    def _make_set_func(self, par, iscomplex=False):
        ''' Creates a set function for parameter par and binds it to the object
            '''

        def set_func(val):
            setattr(self, par, val)

        set_func.__name__ = 'set' + par.capitalize()
        setattr(self, set_func.__name__, set_func)

        if iscomplex:
            def set_real_func(val):
                setattr(self, par, val + getattr(self, par).imag * 1J)

            set_real_func.__name__ = set_func.__name__ + 'real'
            setattr(self, set_real_func.__name__, set_real_func)

            def set_imag_func(val):
                setattr(self, par, val * 1J + getattr(self, par).real)

            set_imag_func.__name__ = set_func.__name__ + 'imag'
            setattr(self, set_imag_func.__name__, set_imag_func)


    def get_parameters(self):
        ''' Returns all the parameters of the current object '''

        return self.parameters


class InstrumentBase(ReflBase):
    _parameters = {}

    def __repr__(self):
        s = 'Instrument('

        for k in self._parameters:
            # if the type is a string...
            if type(getattr(self, k)) == type(''):
                stemp = "%s='%s', " % (k, str(getattr(self, k)))
            else:
                stemp = '%s=%s, ' % (k, str(getattr(self, k)))
            s = s + stemp
        return s[:-1] + ')'


class LayerBase(ReflBase):
    _parameters = {}

    def __repr__(self):
        s = 'Layer('
        for k in self._parameters:
            s += '%s=%s, ' % (k, str(getattr(self, k)))
        return s[:-2] + ')'


class StackBase(ReflBase):
    _parameters = {'Repetitions': 1}
    Layers = []
    Repetitions = 1

    def __init__(self, **kargs):
        ReflBase.__init__(self, **kargs)
        if 'Layers' in kargs:
            self.Layers = kargs['Layers']
            kargs.pop('Layers')
        else:
            self.Layers = []


    def __repr__(self):
        s = 'Stack: '
        for k in self._parameters:
            s += '%s=%s, ' % (k, str(getattr(self, k)))
        s = s[:-2] + '\n'
        it = len(self.Layers)
        for lay in range(it - 1, -1, -1):
            s += '\t' + repr(self.Layers[lay]) + '\n'
        return s

    def resolveLayerParameter(self, parameter):
        par = [resolve_par(lay, parameter)
               for lay in self.Layers] * int(self.Repetitions)
        return par


# Class Sample

class SampleBase(ReflBase):
    _parameters = {}
    Ambient = None
    Substrate = None
    Stacks = []

    def __init__(self, **kargs):
        ReflBase.__init__(self, **kargs)
        special_objects = ['Ambient', 'Substrate', 'Stacks']
        for p in special_objects:
            if p in kargs:
                setattr(self, p, kargs[p])
            else:
                raise ValueError('%s has to be defined to create a Sample object' % p)
        # Create a new dictonary for the new arguments
        nkargs = {}
        for p in kargs:
            if not p in special_objects:
                nkargs[p] = kargs[p]


    def __repr__(self):
        Add = 'Sample: '
        for k in self._parameters:
            if type(getattr(self, k)) == type(''):
                Add += "%s='%s', " % (k, str(getattr(self, k)))
            else:
                Add += '%s=%s, ' % (k, str(getattr(self, k)))
            #Add += '%s=%s, ' % (k, getattr(self, k).__repr__())
        if len(self._parameters) > 0:
            Add = Add[:-2] + '\n'
        else:
            Add = Add[:-1] + '\n'
        temp = [repr(x) for x in self.Stacks]
        temp.reverse()
        return (Add + 'Ambient:\n\t' + repr(self.Ambient) + '\n' + ''.join(temp)
                + 'Substrate:\n\t' + repr(self.Substrate))

    def resolveLayerParameters(self):
        par = self.Substrate._parameters.copy()
        for k in par:
            par[k] = [resolve_par(self.Substrate, k)]
        for k in self.Substrate._parameters:
            for stack in self.Stacks:
                par[k] = par[k] + stack.resolveLayerParameter(k)
            par[k] = par[k] + [resolve_par(self.Ambient, k)]
        return par


# Factory needs to:
#            add members and their names in the _parameter list
#            Add set and get methods for each parameter
class ModelFactory:
    ''' Class that creates new classes to implement certain parameters 
        dynamically.
    '''



    def __init__(self):
        class Layer(LayerBase): _parameters = LayerBase._parameters.copy()
        self.Layer = Layer

        class Stack(StackBase): _parameters = StackBase._parameters.copy()
        self.Stack = Stack

        class Sample(SampleBase): _parameters = SampleBase._parameters.copy()
        self.Sample = Sample

        class Instrument(InstrumentBase): _parameters = InstrumentBase._parameters.copy()
        self.Instrument = Instrument

    @staticmethod
    def _add_parameter(cls, par_name, def_val):
        cls._parameters.update({par_name: def_val})

    @staticmethod
    def _add_sim_method(cls, name, func):
        def method(self, *args):
            nargs = args[:-1] + (self, ) + (args[-1], )
            return func(*nargs)

        method.__name__ = 'Sim' + name
        setattr(cls, method.__name__, method)

    def set_layer_parameters(self, parameters):
        '''Adds extra parameters (dict with name: value entries) to the
            Layer class'''
        [self._add_parameter(self.Layer, k, parameters[k]) for k in parameters]

    def set_stack_parameters(self, parameters):
        '''Adds extra parameters (dict with name: value entries) to the
            Stack class'''
        # The if not is a hook to so that everything works with the old implementation
        [self._add_parameter(self.Stack, k, parameters[k]) for k in parameters
         if not k in ['Layers', 'Repetitions']]

    def set_sample_parameters(self, parameters):
        '''Adds extra parameters (dict with name: value entries) to the
            Sample class'''
        # The if not is a hook to so that everything works with the old implementation
        [self._add_parameter(self.Sample, k, parameters[k]) for k in parameters
         if not k in ['Ambient', 'Substrate', 'Stacks']]

    def set_simulation_functions(self, functions):
        ''' Adds the simulation functions to the sample class
        '''
        [self._add_sim_method(self.Sample, k, functions[k]) for k in functions]

    def set_instrument_parameters(self, parameters):
        '''Adds extra parameters (dict with name: value entries) to the
            Instrument class'''
        [self._add_parameter(self.Instrument, k, parameters[k]) for k in parameters]

    def get_layer_class(self):
        return self.Layer

    def get_stack_class(self):
        return self.Stack

    def get_sample_class(self):
        return self.Sample

    def get_instrument_class(self):
        return self.Instrument


# Function! to create classes with the model given by the input parameters
def MakeClasses(InstrumentParameters={'Wavelength': 1.54, 'Coordinates': 1}, \
                LayerParameters={'d': 0.0, 'sigma': 0.0, 'n': 0.0 + 0.0j}, \
                StackParameters={}, \
                SampleParameters={}, \
                SimulationFunctions={'Specular': lambda x: x, 'OffSpecular': lambda x: x}, \
                ModelID='Standard'):
    factory = ModelFactory()
    factory.set_layer_parameters(LayerParameters)
    factory.set_stack_parameters(StackParameters)
    factory.set_sample_parameters(SampleParameters)
    factory.set_simulation_functions(SimulationFunctions)
    factory.set_instrument_parameters(InstrumentParameters)
    Layer = factory.get_layer_class()
    Stack = factory.get_stack_class()
    Sample = factory.get_sample_class()
    Instrument = factory.get_instrument_class()
    return (Instrument, Layer, Stack, Sample)


if __name__ == '__main__':
    InstrumentParameters = {'probe': 'x-ray', 'wavelength': 1.54, 'coords': 'tth',
                            'I0': 1.0, 'res': 0.001,
                            'restype': 'no conv', 'respoints': 5, 'resintrange': 2, 'beamw': 0.01,
                            'footype': 'no corr', 'samplelen': 10.0, 'incangle': 0.0, 'pol': 'uu',
                            'Ibkg': 0.0, 'tthoff': 0.0}

    LayerParameters = {'sigma': 0.0, 'dens': 1.0, 'd': 0.0, 'f': (1.0 + 1.0j) * 1e-20,
                       'b': 0.0 + 0.0J, 'xs_ai': 0.0, 'magn': 0.0, 'magn_ang': 0.0}
    StackParameters = {'Layers': [], 'Repetitions': 1}
    SampleParameters = {'Stacks': [], 'Ambient': None, 'Substrate': None}

    (Instrument, Layer, Stack, Sample) = MakeClasses(InstrumentParameters, \
                                                     LayerParameters, StackParameters, SampleParameters,
                                                     {'Specualr': lambda x: x}, \
                                                     'test')
    inst = Instrument(footype='gauss beam', probe='x-ray', beamw=0.04, resintrange=2, pol='uu', wavelength=1.54,
                      respoints=5, Ibkg=0.0, I0=2, samplelen=10.0, restype='fast conv', coords='2θ', res=0.001,
                      incangle=0.0)

    # BEGIN Sample DO NOT CHANGE
    Amb = Layer(b=0, d=0.0, f=(1e-20 + 1e-20j), dens=1.0, magn_ang=0.0, sigma=0.0, xs_ai=0.0, magn=0.0)
    topPt = Layer(b=0, d=11.0, f=58, dens=4 / 3.92 ** 3, magn_ang=0.0, sigma=3.0, xs_ai=0.0, magn=0.0)
    TopFe = Layer(b=0, d=11.0, f=26, dens=2 / 2.866 ** 3, magn_ang=0.0, sigma=2.0, xs_ai=0.0, magn=0.0)
    Pt = Layer(b=0, d=11.0, f=58, dens=4 / 3.92 ** 3, magn_ang=0.0, sigma=2.0, xs_ai=0.0, magn=0.0)
    Fe = Layer(b=0, d=11, f=26, dens=2 / 2.866 ** 3, magn_ang=0.0, sigma=2.0, xs_ai=0.0, magn=0.0)
    bufPt = Layer(b=0, d=45, f=58, dens=4 / 3.92 ** 3, magn_ang=0.0, sigma=2, xs_ai=0.0, magn=0.0)
    bufFe = Layer(b=0, d=2, f=26, dens=2 / 2.866 ** 3, magn_ang=0.0, sigma=2, xs_ai=0.0, magn=0.0)
    Sub = Layer(b=0, d=0.0, f=12 + 16, dens=2 / 4.2 ** 3, magn_ang=0.0, sigma=4.0, xs_ai=0.0, magn=0.0)

    top = Stack(Layers=[TopFe, topPt], Repetitions=1)
    ML = Stack(Layers=[Fe, Pt], Repetitions=19)
    buffer = Stack(Layers=[bufFe, bufPt], Repetitions=1)

    sample = Sample(Stacks=[buffer, ML, top], Ambient=Amb, Substrate=Sub)

    # Test case for ReflFunction class
    from scipy import interpolate

    def create_dispersion_func(name='Fe'):
        path = '../databases/f1f2_nist/'
        e, f1, f2 = np.loadtxt(path + '%s.nff' % name.lower(), skiprows=1, \
                               unpack=True)
        f1interp = interpolate.interp1d(e, f1, kind='linear')
        f2interp = interpolate.interp1d(e, f2, kind='linear')

        def f(energy):
            return f1interp(energy) - 1.0J * f2interp(energy)

        return f


    fFe = ReflFunction(create_dispersion_func('Fe'), (1000,), {}, id = 'f(E)')
    fCo = ReflFunction(create_dispersion_func('Co'), (1000,), {}, id = 'f(E)')
    iprint(fFe.validate())
    iprint('Cast to array tests:')
    ltest = [fFe, 3.0, 3.0]
    iprint('Test single value: ', cast_to_array(ltest, 1000))
    iprint('Test array: ', cast_to_array(ltest, np.arange(1000, 1010)))
    iprint('Cast to array finished.')
