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


"""
 cs: 
 

 en: 

"""

import numpy as np
from scipy import stats
import copy


class Ingot:
    """
    Prazdná třida pro "nevypalené" vzorky, tj. bez přiřazeného rozdělení
    """
    def __init__(self, data, attr='R'):
        # data? takový neslaný nazev...
        # data suppose to be pandas compatible, i.e. 
        # nsim, nvar = data.shape
        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)
        #if attr in ('R', 'Rn', 'GK', 'G', 'P', 'U') and attr==self.space:
        #    return f._data
        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):
        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)
    
    def __repr__(self):
        return "%s(%s)"%('SNorm', self.__nvar)
        
    def __str__(f):
        return str(f.R)
        
    def __call__(f, ns=0):
        f_copy = eval(f.__repr__())
        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 = eval(f.__repr__())
        # robim kopiu z _data, nebo ne?
        f_copy._data = np.atleast_2d(f._data[slice,:]) #.copy()
        return f_copy
        
        
        # slajsy. Nejsem jist, zda mám robit kopiu, nebo ne. Takže bacha!
    def __getattr__(f, attr):
        if attr in ('pdf_R', 'pdf_Rn', 'pdf_GK', 'pdf_G'):
            return f._data[:,-1]
        elif attr in ('R', 'Rn', 'GK', 'G'):
            return f._data[:,:f.__nvar]
        elif attr in ('P', 'U'):
            return f._data[:,f.__nvar:2*f.__nvar]
        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)])
            
        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'):
            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
        elif space in ('R', 'Rn', 'GK', 'G'):
            # does sample is another f_model object? 
            try:
                sample_R = getattr(sample, space)
            except:
                # 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:
                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 = eval(f.__repr__())
        if sample is not None:
            f_copy.add_sample(sample, space)
        return f_copy
    
#    # drobná pomucka
#    def new_random_sample(f, ns=1):  
#        f_copy = eval(f.__repr__())
#        f_copy.add_sample(np.random.random((ns, f.nvar)), 'U')
#        return f_copy
    
    




class UnCorD: # nic moc nazev, ale je přece lepší nez CommonJointDistribution
    """
    Takes tuple of scipy stats distribution objects
    """
    def __init__(self, marginals):
        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)
        
    def __repr__(self):
        return "%s(%s)"%('UnCorD', repr(self.__marginals))
        
    def __str__(f):
        return 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 == 'pdf_R':
            return f._data[:,-2]
        elif attr in ('pdf_GK', 'pdf_G'):
            return f._data[:,-1]
        elif 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)])
        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'):
            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
        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:
            # 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
    
#    # drobná pomucka
#    def new_random_sample(f, ns=1):  
#        f_copy = f()
#        f_copy.add_sample(np.random.random((ns, f.nvar)), 'U')
#        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 Rn2R(f, sample):
#        sample_Rn = np.array(sample).reshape(-1, f.nvar)
#        sample_R = np.empty_like(sample_Rn)
#        for i in range(f.nvar):
#            sample_R[:, i] = sample_Rn[:, i]*f.marginals[i].std() + f.marginals[i].mean()
#        return sample_R
#        
#    def R2P(f, sample):
#        sample_R = np.array(sample).reshape(-1, f.nvar)
#        sample_P = np.empty_like(sample_R)
#        for i in range(f.nvar):
#            sample_P[:, i] = sample_Rn[:, i]*f.marginals[i].std() + f.marginals[i].mean()
#        return sample_P