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


"""
 cs: 
 

 en: 

"""

import numpy as np
from scipy import stats
import copy

# alpha atribut beztak veřejný, ale tato funkce ještě provadí normalizaci a kontrolu
def set_alpha(f_model, input_alpha):
    alpha = np.atleast_1d(input_alpha).flatten()
    if len(alpha) != f_model.nvar:
        raise ValueError
    else:
        f_model.alpha = alpha / stats.gmean(alpha)

class Ingot:
    """
    Prazdná třida pro "nevypalené" vzorky, tj. bez přiřazeného rozdělení
    """
    def __init__(self, data, attr='R'):
        self.attr = attr
        try:
            self.data = np.atleast_2d(getattr(data, attr))
        except AttributeError:
            self.data = np.atleast_2d(data)
        
        
    def __repr__(self):
        return "%s(np.%s, '%s')"%('Ingot', repr(self.data), self.attr)
        
    def __str__(self):
        return str(self.data)
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, slice):
        self_copy = copy.copy(self)
        # robim kopiu z _data, nebo ne?
        self_copy.data = np.atleast_2d(self.data[slice,:]) #.copy()
        return self_copy
        
        
    def __getattr__(self, attr):
        # Рекурсилы пезьдэт!
        if attr in ('attr', 'data'):
            raise AttributeError(attr)
        elif attr == 'nvar':
            nsim, nvar = self.data.shape
            return nvar
        elif attr == 'nsim':
            return len(self.data)
            
        # гулять так гулять
        elif attr == self.attr:
            return self.data
            
        raise AttributeError(attr)
      
        
        # vstupné vzorky jsou implicitně R, 
        # nikoliv z prostoru tohoto krámu
    def add_sample(self, sample, space='R'):
        # first of all, are you one of us?
        if space == self.attr:
            # does sample is another sample object? 
            # zde nechcu, aby spadlo
            try:
                self.data = np.vstack((self.data, getattr(sample, self.attr)))
            except AttributeError:
                self.data = np.vstack((self.data, sample))
            
        # no, actually
        else:
            # does sample is another sample object? 
            # self.data = np.vstack((self.data, getattr(sample, self.attr)))
            # ale zde chcu 
            # (aby spadlo)
            raise ValueError
            
        
    # drobná pomucka
    def new_sample(f, sample=None, space='R'):  
        return Ingot(sample, space) # too easy


class SNorm:
    """
    Standard Gauss distribution
    """
    def __init__(self, nvar, alpha=None):
        self.__nvar = nvar
        # nvar_R + nvar_U + pdf_R
        rowsize = nvar*2 + 1
        # data? takový neslaný nazev...
        # data suppose to be cKDTree compatible, i.e. 
        # nsim, nvar = data.shape
        self._data = np.empty((0, rowsize), dtype=float)
        if alpha is None:
            self.alpha = np.ones(nvar)
        else:
            self.alpha = self.set_alpha(alpha)
        
     # nemusím duplikovat a dědit 
    set_alpha = set_alpha
    
    def __repr__(self):
        return "%s(%s, %s)"%('SNorm', self.__nvar, repr(self.alpha))
        
    def __str__(f):
        return "SNorm sample: " + str(f.R)
        
    def __call__(f, ns=0):
        f_copy = copy.copy(f)
        if ns:
            sample_P = np.random.random((ns, f_copy.nvar))
            sample_R = stats.norm.ppf(sample_P)
                
            pdfs_R = stats.norm.pdf(sample_R)
            pdf_R = np.prod(pdfs_R, axis=1).reshape(-1, 1)
            f_copy._data = np.hstack((sample_R, sample_P, pdf_R))
        return f_copy
        
    def __len__(f):
        return len(f._data)
    
    def __getitem__(f, slice):
        f_copy = copy.copy(f)
        # robim kopiu z _data, nebo ne?
        f_copy._data = np.atleast_2d(f._data[slice,:]) #.copy()
        return f_copy
        
        
    def __getattr__(f, attr):
        if attr in ('R', 'Rn', 'GK', 'G'):
            return f._data[:,:f.__nvar]
        elif attr in ('aR', 'aRn', 'aGK', 'aG'):
            return f._data[:,:f.__nvar] * f.alpha
        elif attr in ('P', 'U'):
            return f._data[:,f.__nvar:2*f.__nvar]
        elif attr in ('aP', 'aU'):
            return f._data[:,f.__nvar:2*f.__nvar] * f.alpha
        elif attr == 'nvar':
            return f.__nvar
        elif attr == 'nsim':
            return len(f._data)
        elif attr == 'marginals':
            return [stats.norm for __ in range(f.__nvar)]
        elif attr == 'cor':
            return np.diag([1 for __ in range(f.__nvar)])
            
        # hustoty
        # I'm considering to deprecate attribute access
        elif attr[:4] == 'pdf_':
            return f.pdf(attr[4:])
            
        raise AttributeError(attr)
        
        
       # pro určitou konzistenci. Ne že bych chtěl zamykat dveře v Pythonu
    def __setattr__(f, attr, value):
        if attr in ('_SNorm__nvar','_data', 'alpha'):
            f.__dict__[attr] = value
        else:
            #raise AttributeError('Čo tu robíš?')
            #raise AttributeError('Враг не пройдёт!')
            #raise AttributeError('Иди отсюда!')
            #raise AttributeError('Аслыкъёсы воштыны уг луи!')
            raise AttributeError('Atribute %s of %s object is not writable' % (attr, f.__class__.__name__))
        
        
        
    def add_sample(f, sample, space='R'):
        # does sample is exactly me?
        if f.__repr__() == sample.__repr__():
            newdata = sample._data
            
            # new piece of code "alpha"-related
        elif space in ('aR', 'aRn', 'aGK', 'aG', 'aP', 'aU'):
            # does sample is another f_model object?
            try:
                sample = getattr(sample, space) / f.alpha
            except AttributeError:
                sample = sample / f.alpha
            space = space[1:]
            
        if space in ('R', 'Rn', 'GK', 'G'):
            # does sample is another f_model object? 
            try:
                sample_R = getattr(sample, space)
            except AttributeError:
                # no
                sample_R = sample
            sample_P = stats.norm.cdf(sample_R)
            
            pdfs_R = stats.norm.pdf(sample_R)
            pdf_R = np.prod(pdfs_R, axis=pdfs_R.ndim-1)
            if pdfs_R.ndim == 2:
                newdata = np.hstack((sample_R, sample_P, pdf_R.reshape(len(pdf_R), 1)))
            else:
                newdata = np.hstack((sample_R, sample_P, pdf_R))
            
        elif space in ('P', 'U'):
            try:
                sample_P = getattr(sample, space)
            except AttributeError:
                sample_P = sample
            sample_R = stats.norm.ppf(sample_P)
                
            pdfs_R = stats.norm.pdf(sample_R)
            pdf_R = np.prod(pdfs_R, axis=pdfs_R.ndim-1)
            if pdfs_R.ndim == 2:
                newdata = np.hstack((sample_R, sample_P, pdf_R.reshape(len(pdf_R), 1)))
            else:
                newdata = np.hstack((sample_R, sample_P, pdf_R))
            
        f._data = np.vstack((f._data, newdata))
            
        
    # drobná pomucka
    def new_sample(f, sample=None, space='R'):  
        f_copy = copy.copy(f)
        if sample is not None:
            f_copy.add_sample(sample, space)
        return f_copy
    
    
    def pdf(f, space='R'):
        """
        Returns own probability densities in the given space
        """
        if space in ('R', 'Rn', 'GK', 'G'):
            return f._data[:,-1]
        elif space in ('P', 'U'):
            return f.sample_pdf(f, space)
            
        elif space in ('aR', 'aRn', 'aGK', 'aG', 'aP', 'aU'):
            return f.pdf(space[1:]) / np.prod(f.alpha)
            
        else:   
            raise ValueError('Unknown space %s' % space)
            
        
        
    def sample_pdf(f, input_sample, space='R'):
        """
        Calculates probability density for the given external sample in the given space.
        Function intended for the case no transformation needed. 
        Otherwise new f_sample should be instanciated.
        """
        
        # does sample is another f_model object? 
        # кинлы со zase кулэ?
        try:
            sample = getattr(input_sample, space)
        except AttributeError:
            # no
            sample = np.atleast_2d(input_sample)
            
        
        if space in ('R', 'Rn', 'GK', 'G'):
            pdfs_R = stats.norm.pdf(sample)
            return np.prod(pdfs_R, axis=pdfs_R.ndim-1)
            
        elif space in ('P', 'U'):
            return np.where(np.all((sample >= 0) & (sample <= 1), axis=1), 1, 0)
            
        # new piece of code "alpha"-related
        elif space in ('aR', 'aRn', 'aGK', 'aG', 'aP', 'aU'):
            return f.sample_pdf(sample / f.alpha, space=space[1:]) / np.prod(f.alpha)
            
        else:   
            raise ValueError('Unknown space %s' % space)
        
            
    


class UnCorD: # nic moc nazev, ale je přece lepší nez CommonJointDistribution
    """
    Takes tuple of scipy stats distribution objects
    """
    def __init__(self, marginals, alpha=None):
        self.__marginals = marginals
        # nvar_Rn + nvar_R + nvar_P + nvar_G + pdf_R + pdf_G
        rowsize = len(marginals)*4 + 2
        # data? takový neslaný nazev...
        # data suppose to be cKDTree compatible, i.e. 
        # nsim, nvar** = data.shape
        self._data = np.empty((0, rowsize), dtype=float)
        if alpha is None:
            self.alpha = np.ones(self.nvar)
        else:
            self.alpha = self.set_alpha(alpha)
        
     # nemusím duplikovat a dědit 
    set_alpha = set_alpha
        
    def __repr__(self):
        return "%s(%s, %s)"%('UnCorD', repr(self.__marginals), repr(self.alpha))
        
    def __str__(f):
        return "UnCorD: " + str(f.R)
    
    def __call__(f, ns=0):  
        f_copy = copy.copy(f) # nebo deep?
        f_copy._data = np.empty((0, f_copy._data.shape[1]), dtype=float)
        
        if ns:
            sample_dict = {'P':np.random.random((ns, f_copy.nvar))}
            f._chain(sample_dict)
            pdf_G = np.prod(stats.norm.pdf(sample_dict['G']), axis=1).reshape(-1, 1)
            pdfs_R = [f.marginals[i].pdf(sample_dict['R'][:, i]) for i in range(f.nvar)]
            # je tu fakt axis=0. Dochazí totíž v iterátoru k převracení
            pdf_R = np.prod(pdfs_R, axis=0).reshape(-1, 1)
            # nvar_Rn + nvar_R + nvar_P + nvar_G + pdf_R + pdf_G
            f_copy._data = np.hstack((sample_dict['Rn'], sample_dict['R'], sample_dict['P'], sample_dict['G'], pdf_R, pdf_G))
        return f_copy
    
    def __len__(f):
        return len(f._data)
    
    def __getitem__(f, slice):
        f_copy = copy.copy(f) # nebo deep?
        f_copy._data = np.atleast_2d(f._data[slice,:]) #.copy()
        return f_copy
        
        # dúfám, že tyhle slajsy sa vyplatí
    def __getattr__(f, attr):
        if attr == 'Rn':
            return f.__frame(0)
        elif attr == 'R':
            return f.__frame(1)
        elif attr in ('P', 'U'):
            return f.__frame(2)
        elif attr in ('GK', 'G'):
            return f.__frame(3)
            
        elif attr == 'nvar':
            return len(f.__marginals)
        elif attr == 'nsim':
            return len(f._data)
        elif attr == 'marginals':
            return f.__marginals
        elif attr == 'cor':
            return np.diag([1 for __ in range(f.nvar)])
            
        elif attr in ('aR', 'aRn', 'aGK', 'aG', 'aP', 'aU'):
            return getattr(f, attr[1:]) * f.alpha
            
        # hustoty
        # I'm considering to deprecate attribute access
        elif attr[:4] == 'pdf_':
            return f.pdf(attr[4:])
            
        raise AttributeError(attr)
        
    # pro určitou konzistenci. Ne že bych chtěl zamykat dveře v Pythonu
    def __setattr__(f, attr, value):
        if attr in ('_UnCorD__marginals','_data', 'alpha'):
            f.__dict__[attr] = value
        else:
            #raise AttributeError('Аслыкъёсы воштыны уг луи!')
            raise AttributeError('Atribute %s of %s object is not writable' % (attr, f.__class__.__name__))
        
    def __frame(f, i):
        nvar = f.nvar
        sl = slice(i*nvar, (i+1)*nvar)
        return f._data[:,sl]
        
    def add_sample(f, sample, space='R'):
        # isinstance, ne?
        if f.__class__.__name__ == sample.__class__.__name__:
            if f.marginals == sample.marginals:
                f._data = np.vstack((f._data, sample._data))
                return f
        
        # new piece of code "alpha"-related
        elif space in ('aR', 'aRn', 'aGK', 'aG', 'aP', 'aU'):
            # does sample is another f_model object?
            try:
                sample = getattr(sample, space) / f.alpha
            except AttributeError:
                sample = sample / f.alpha
            space = space[1:]        
                
        elif space not in ('R', 'Rn', 'P', 'GK', 'G', 'U'):
            # co jako, mám gettext sem tahnout?!
            raise ValueError('Zadaný prostor %s mi není znám' % space)
            raise ValueError('Unknown space %s' % space)
            
        # does sample is another f_model object? 
        try:
            sample_ = getattr(sample, space)
        except AttributeError:
            # no
            sample_ = sample
            
        if space=='GK':
            space='G'
        elif space=='U':
            space='P'
            
        sample_dict = {space:np.array(sample_, dtype=float).reshape(-1, f.nvar)}
        f._chain(sample_dict)
        pdf_G = np.prod(stats.norm.pdf(sample_dict['G']), axis=1).reshape(-1, 1)
        pdfs_R = [f.marginals[i].pdf(sample_dict['R'][:, i]) for i in range(f.nvar)]
        # je tu fakt axis=0. Dochazí totíž v iterátoru k převracení
        pdf_R = np.prod(pdfs_R, axis=0).reshape(-1, 1)
        # nvar_Rn + nvar_R + nvar_P + nvar_G + pdf_R + pdf_G
        newdata = np.hstack((sample_dict['Rn'], sample_dict['R'], sample_dict['P'], sample_dict['G'], pdf_R, pdf_G))
            
        f._data = np.vstack((f._data, newdata))
            
        
    # drobná pomucka
    def new_sample(f, sample=None, space='R'):  
        f_copy = f()
        if sample is not None:
            f_copy.add_sample(sample, space)
        return f_copy
    
     
    def _chain(f, sample_dict):
        # chain tam
        # чаль татысь
        if 'R' not in sample_dict and 'Rn' in sample_dict:
            sample_dict['R'] = np.empty_like(sample_dict['Rn'])
            for i in range(f.nvar):
                sample_dict['R'][:, i] = sample_dict['Rn'][:, i]*f.marginals[i].std() + f.marginals[i].mean()
                
        if 'P' not in sample_dict and 'R' in sample_dict:
            sample_dict['P'] = np.empty_like(sample_dict['R'])
            for i in range(f.nvar):
                sample_dict['P'][:, i] = f.marginals[i].cdf(sample_dict['R'][:, i])
                
        if 'G' not in sample_dict and 'P' in sample_dict:
            sample_dict['G'] = stats.norm.ppf(sample_dict['P'])
        
        # chain sem
        # чаль татчи
        elif 'P' not in sample_dict and 'G' in sample_dict:
            sample_dict['P'] = stats.norm.cdf(sample_dict['G'])
            
        if 'R' not in sample_dict and 'P' in sample_dict:
            sample_dict['R'] = np.empty_like(sample_dict['P'])
            for i in range(f.nvar):
                sample_dict['R'][:, i] = f.marginals[i].ppf(sample_dict['P'][:, i])
                
        if 'Rn' not in sample_dict and 'R' in sample_dict:
            sample_dict['Rn'] = np.empty_like(sample_dict['R'])
            for i in range(f.nvar):
                sample_dict['Rn'][:, i] = (sample_dict['R'][:, i] - f.marginals[i].mean())/f.marginals[i].std()
     
     
    def pdf(f, space='R'):
        """
        Returns own probability densities in the given space
        """
        if space == 'R':
            return f._data[:,-2]
        elif space == 'Rn':
            return f._data[:,-2] * np.prod(list(f.marginals[i].std() for i in range(f.nvar)))
        elif space in ('GK', 'G'):
            return f._data[:,-1]
        elif space in ('P', 'U'):
            return f.sample_pdf(f, space)
            
        elif space in ('aR', 'aRn', 'aGK', 'aG', 'aP', 'aU'):
            return f.pdf(space[1:]) / np.prod(f.alpha)
        
        else:   
            raise ValueError('Unknown space %s' % space)
            
        
        
    def sample_pdf(f, input_sample, space='R'):
        """
        Calculates probability density for the given external sample in the given space.
        Function intended for the case no transformation needed. 
        Otherwise new f_sample should be instanciated.
        """
        
        # does sample is another f_model object? 
        # кинлы со zase кулэ?
        try:
            sample = getattr(input_sample, space)
        except AttributeError:
            # no
            sample = np.atleast_2d(input_sample)
            
        
        if space == 'R':
            pdfs_R = [f.marginals[i].pdf(sample[:, i]) for i in range(f.nvar)]
            # je tu fakt axis=0. Dochazí totíž v iterátoru k převracení
            return np.prod(pdfs_R, axis=0)
                
        elif space == 'Rn':
            pdf_R = f.sample_pdf(input_sample, space='R')
            return pdf_R * np.prod(list(f.marginals[i].std() for i in range(f.nvar)))
        
        elif space in ('GK', 'G'):
            pdfs_R = stats.norm.pdf(sample)
            return np.prod(pdfs_R, axis=pdfs_R.ndim-1)
            
        elif space in ('P', 'U'):
            return np.where(np.all((sample >= 0) & (sample <= 1), axis=1), 1, 0)
            
        # new piece of code "alpha"-related
        elif space in ('aR', 'aRn', 'aGK', 'aG', 'aP', 'aU'):
            return f.sample_pdf(sample / f.alpha, space=space[1:]) / np.prod(f.alpha)
            
        else:   
            raise ValueError('Unknown space %s' % space)