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

"""
Zde leží whiteboxy
Zatimco BlackBox pěčlivě ukladá věškerá data,
věškeré sady vzorků, průběžné odhady a tak,
WhiteBox musí obsahovat jen pf_exact, které buď už předem zná,
nebo jej spočítá a není vůbec/nachren nutný cokoliv ukladat.

en:
f_model + g_model = pf
WhiteBox = [f_model, g_model, pf_exact]
pf_exact is most important part of WhiteBox 
Knowledge of pf is the only reason to create WhiteBox

whitebox actually IS g_model PLUS:
.f f_model
.pf_exact
.pf_exact_method
"""
import numpy as np
from scipy import stats
from scipy import special # for S_ball
from scipy import integrate # for S_ball

from .samplebox import SampleBox

from . import f_models
import copy
from . import g_models

from . import plot


class WhiteBox:
    """
    Bazová třida pro dědictví
    
    úkolem whiteboxu je spočítat pf-ko
    .pf_exact
    .pf_exact_method
    """
    pf_exact_method = 'None'
    
    def __init__(wt,  f_model, g_model):
        wt.gm = g_model
        wt.f = f_model
        # na začatku nemáme vzorky - pouze rozdělení a podpís
        wt.sample_box = SampleBox(wt.f(), gm_signature=wt.gm_signature)
        
        try: # no to jsme líný. Možná samotný g_model tuší jak se spočte pf-ko?
            wt.pf_exact, wt.pf_exact_method = g_model.pf_expression(f_model)
        except AttributeError:
            pass
        
    def __str__(wt):
        #ёč хощется přidat nějaký popísek k testkejsům
        #č chce se odlišovat úlohy, co už získaly slavu 
        #č v časopísech od těch, co ještě nie,
        #č od těch, se kterými ještě hrajeme
        #č od těch, co jsme před chvilkou vymysleli
        if 'description' in wt.__dict__:
            return  wt.description
        else:
            return  wt.__class__.__name__ + ' of' + str(wt.gm)
        
    def __repr__(wt):
        return  'WhiteBox(%s, %s)' %(repr(wt.f), repr(wt.gm))
        
    def __len__(wt):
        return len(wt.sample_box)
        
        
    def __call__(wt, sample=None):
        if sample is None:
            sample = wt.sample_box()
        
        # zamykame se do sebe
        result = wt.gm(sample)
        wt.sample_box.add_sample(result)
        return result
    
        
    def __getitem__(self, slice):
        self_copy = copy.copy(self)
        self_copy.sample_box = self.sample_box[slice]
        return self_copy
        
    def __getattr__(wt, attr):
        if attr == 'whitebox':
            return wt
        # co patři g_modelovi?
        elif attr == 'get_2D_R_boundary':
            return wt.gm.get_2D_R_boundary
        elif attr == 'gm_signature':
            try: # byla volná funkce?
                return wt.gm.__name__
                # asi to byla trida?
            except AttributeError:
                return repr(wt.gm)
        
        
        # co mělo být definováno ve WhiteBoxu? Ale teda není?
        elif attr == 'pf_exact_method':
            raise AttributeError
            
          # пытка-непытка
        elif attr == 'pf_exact': 
            if 'beta_exact' in wt.__dict__:
                return stats.norm.cdf(-wt.beta_exact)
            else:
                # мы нищего не знаем
                raise AttributeError("Nothing known about failure probability")
                
        elif attr == 'beta_exact':
            if 'pf_exact' in wt.__dict__:
                return -stats.norm.ppf(wt.pf_exact)
            else:
                raise AttributeError("Nothing known about failure probability")
        
        
        # branime sa rekurzii
        # defend against recursion
        elif attr == 'sample_box':
            raise AttributeError
            
        # zbytek teda nievím
        else:
            return getattr(wt.sample_box, attr)
    
    # just plot, green points, red points...
    plot2D = plot.plot2D
    plot3D = plot.plot3D
    show2D = plot.show2D
    show3D = plot.show3D
            
        
    def get_2D_boundary(wt, nrod=100, viewport_sample=None, viewport_space='R'):
        """
        Fence off!
        nrod - number of rods in fencing
        viewport_sample - limit points of viewport
        (function will get xlim and ylim from there, 
        assuming there will be at least two points)
        """
        
        if (viewport_sample is not None): # and viewport_sample.nsim: #>0
            if viewport_space=='R':
                viewport_sample_R = viewport_sample
            else:
                viewport_sample_R = wt.f_model.new_sample(viewport_sample, viewport_space, extend=True).R
        else:
            viewport_sample_R = wt.f_model.new_sample([[-7,-7],[7,7]], 'G', extend=True).R
            
            # should I tolerate nD?
        viewport_sample_R = np.vstack((viewport_sample_R[:,:2], wt.sample_box.R[:,:2]))
        xmin = np.ma.masked_invalid(viewport_sample_R[:,0]).min()
        xmax = np.ma.masked_invalid(viewport_sample_R[:,0]).max()
        ymin = np.ma.masked_invalid(viewport_sample_R[:,1]).min()
        ymax = np.ma.masked_invalid(viewport_sample_R[:,1]).max()
        
        # získám seznam polí
        bounds_R = wt.get_2D_R_boundary(nrod, (xmin, xmax), (ymin, ymax))
        # transformuji na seznam vzorků
        return [wt.f_model.new_sample(bounds_R[i], extend=True)  for i in range(len(bounds_R))]
        
        


    # Monte Carlo, n-krátá realizace
    def MC(wt, Nsim=int(1e6)):
        
        
        # tohlensto může bejt dost těžkým
        result = wt.gm(wt.f(Nsim))
        # should I stay or should I go?
        wt.sample_box.add_sample(result)
        
        # je tu jakoby že g_model vždy vrací nějakej sample_box
        pf_exact = np.count_nonzero(result.failsi)/Nsim
        
        # šlo by to?
        if wt.pf_exact_method == 'None' or wt.pf_exact_method == 'MC' and wt.Nsim <= Nsim:
            wt.pf_exact = pf_exact
            wt.Nsim = Nsim
            wt.pf_exact_method = 'MC'
        
        print('Monte Carlo estimation of pf is %s (%s simulations)'%(pf_exact, Nsim))
        return result
        


    # IS, (n-2)-krátá realizace, n>2
    def IS(wt, Nsim=int(1e4), h_model=None, IS_mode='G'):
        """
        IS_mode - v jakých souřadnicích robím merge a jaká PDF použiváme?
        může být 'R' nebo 'G'
        jinde # čo jinde?
        """
        
        if h_model is not None:
            wt.h = h_model 
            wt.IS_mode = IS_mode
        elif 'h' not in wt.__dict__:
            wt.h = f_models.UnCorD([stats.norm(0,2.5) for i in range(wt.f.nvar)])
            wt.IS_mode = 'G'
            
            
        #
        # jdeme na to, koťě!
        #
        
        # zgenerujeme vzorky
        # nic zajimavýho
        wt.h = wt.h(Nsim)
        
        # a teď bacha!
        if wt.IS_mode == 'R':
            # jestli máme to právé vzorkovácí rozdělení - tak nemáme čo robiť
            to_sample = wt.f.new_sample(wt.h) # smerdží se to po R
            # w like weights
            wt.w = to_sample.pdf_R / wt.h.pdf_R
        elif wt.IS_mode == 'G':
            # tady musíme provést jeden trik
            to_sample = wt.f.new_sample(wt.h.R, 'G') # R-ko smerdžíme ako G-čko
            wt.w = to_sample.pdf_G / wt.h.pdf_R # snad je to správně
        else:
            # шо-то тут не то...
            # čo blbnéš, kámo?
            # What's going on with my IS_mode?
            raise ValueError("IS_mode should be either 'R' or 'G'")
        
        # vahy máme, jedeme dál
        # sample_box jíž není prázdnej
        result = wt.gm(to_sample)
        wt.sample_box.add_sample(result)
        
        # hodilo by sa to?
        pf_exact = np.sum(wt.w[result.failsi])/Nsim
        
        if wt.pf_exact_method in ('None', 'IS_norm', 'IS'):
            wt.Nsim = Nsim
            if pf_exact < 1:
                wt.pf_exact = pf_exact
                wt.pf_exact_method = 'IS'
            else:
                # ať mně nerobí ostudu
                wt.pf_exact = np.sum(wt.w[result.failsi]) / np.sum(wt.w)
                wt.pf_exact_method = 'IS_norm'
        
        print('Importance Sampling pure estimation of pf is %s (%s simulations)'%(pf_exact, Nsim))
        return result
       


class HyperPlane(WhiteBox): # куда ж без него...
    def __init__(self, betas=(1,2,3)):
        """
        Class takes for inicialization tuple of betas 
        Betas are coeffitients in sense of Regression Analysis (well, not really)
        g= a*X1 + b*X2 + c
        betas=(a,b,c)
        """
        self._betas = betas
        self.gm = g_models.Linear_nD(betas)
        self.f = f_models.SNorm(len(betas)-1)
        # na začatku nemáme vzorky - pouze rozdělení a podpís
        self.sample_box = SampleBox(self.f(), gm_signature=self.gm_signature)
        
        # tady už je to ta, "náše" beta )
        # beta = c/np.sqrt(a**2 + b**2)
        self.beta_exact = betas[-1]/np.sqrt(np.sum(np.array(betas[:-1])**2)) 
        self.pf_exact = stats.norm.cdf(-self.beta_exact)
        self.pf_exact_method = 'FORM (exact solution)' # Ang, Tang and Pythagoras
        
        
    def __str__(wt):
        return  'HyperPlaneBox%sD'%(len(wt._betas)-1)
        
    def __repr__(wt):
        return  'HyperPlane(%s)' % repr(wt._betas)



class Weibull_Z_min(WhiteBox):
    def __init__(self, wb_scales=(1.,1.), shape=5, **kwargs):
        """
        parametry pravdepodobnostniho rozdeleni pro Z_min s Weib. velicinami
        wb_scales=(1,1) - tuple of Weibull scale parameters, len(wb_scales)==nvar
        shape = 5 
        je třeba zadat buď pf_exact, nebo konštantu u funkce minima Z_min
        """
        self.wb_scales = wb_scales
        self.shape = 5
        self.f = f_models.UnCorD([stats.weibull_min(shape, scale=sc_i) for sc_i in wb_scales])
        
        # scale parametr minima z nvar Weibullovskych
        # tohle by platilo pro stejná rozdělení
        #sn = scale * nvar ** (-1.0 / shape)
        # pro nás musí to být něco takovýho
        sn = np.sum(np.power(wb_scales, -shape)) ** (-1.0 / shape)
        self.rvweibmin = stats.weibull_min(shape, scale=sn)
        
        # je třeba zadat buď pf_exact, nebo konštantu u funkce minima Z_min
        self.pf_exact_method = 'exact solution'
        if 'pf_exact' in kwargs:
            self.pf_exact = kwargs['pf_exact']
            self.const = -self.rvweibmin.ppf(self.pf_exact)
        elif 'beta_exact' in kwargs:
            self.beta_exact = kwargs['beta_exact']
            self.const = -self.rvweibmin.ppf(self.pf_exact)
        elif 'const' in kwargs:
            self.const = kwargs['const']
            self.pf_exact = self.rvweibmin.cdf(-self.const) # asi
        else:
            # no to teda uživatele пошли
            self.pf_exact = 1e-4
            self.const = -self.rvweibmin.ppf(self.pf_exact)
        
        self.gm = g_models.Z_min(self.const)
        # na začatku nemáme vzorky - pouze rozdělení a podpís
        self.sample_box = SampleBox(self.f(), gm_signature=self.gm_signature)

    def __str__(self):
        return  'Weibull_Z_min%sD'%(len(self.wb_scales))
        
    def __repr__(self):
        return  'Weibull_Z_min(%s, %s, pf_exact=%s)' % (repr(self.wb_scales), repr(self.shape), repr(self.pf_exact))




# já teda tridu zababachnu, ale že to teoreticky platí jsem neodvozoval
# UPD: testy neprochází, logicky hodnoty nesedí
# vůbec netuším jak by se to mohlo odvozovat 
class Gaussian_Z_sumexp(WhiteBox):
    def __init__(self, nvar=2, **kwargs):
        """
        je třeba zadat buď pf_exact, nebo konštantu u funkce Z_sumexp
        """
        
        # je tam předpoklad SNormu?
        self.f = f_models.SNorm(nvar)
        
        self.C1 = np.sqrt(np.sqrt(5) / 3. - 1. / 3.)
        self.C2 = np.sqrt(3.) / 3.
        # je třeba zadat buď pf_exact, nebo konštantu u funkce Z_sumexp
        self.pf_exact_method = 'tešt solution'
        if 'pf_exact' in kwargs:
            self.pf_exact = kwargs['pf_exact']
            self.C = self.beta_exact * self.C1 * np.sqrt(nvar) - self.C2 * nvar
        elif 'beta_exact' in kwargs:
            self.beta_exact = kwargs['beta_exact']
            self.C = self.beta_exact * self.C1 * np.sqrt(nvar) - self.C2 * nvar
        elif 'const' in kwargs:
            self.const = kwargs['const']
            self.C = self.const
            self.beta_exact = (self.C + self.C2 * nvar) / self.C1 / np.sqrt(nvar)
        elif 'C' in kwargs:
            self.C = kwargs['C']
            self.beta_exact = (self.C + self.C2 * nvar) / self.C1 / np.sqrt(nvar)
        else:
            # no to teda uživatele пошли
            self.pf_exact = 1e-4
            self.C = self.beta_exact * self.C1 * np.sqrt(nvar) - self.C2 * nvar
        
        
        self.const = self.C
        self.gm = g_models.Z_sumexp(self.const)
        # na začatku nemáme vzorky - pouze rozdělení a podpís
        self.sample_box = SampleBox(self.f(), gm_signature=self.gm_signature)
     
    def __str__(self):
        return  'Gaussian_Z_sumexp%sD'%(self.nvar)
        
    def __repr__(self):
        return  'Gaussian_Z_sumexp(%s, pf_exact=%s)' % (repr(self.nvar), repr(self.pf_exact))   



class SNorm_Z_sumsq(WhiteBox):
    def __init__(self, nvar=2, **kwargs):
        """
        je třeba zadat buď pf_exact, nebo konštantu u funkce Z_sumsq
        """
        
        # je tu předpoklad SNormu, to vím
        self.f = f_models.SNorm(nvar)
        
        self.rvchisq = stats.chi2(nvar)
        
        # je třeba zadat buď pf_exact, nebo konštantu u funkce Z_sumsq
        self.pf_exact_method = 'exact solution'
        if 'pf_exact' in kwargs:
            self.pf_exact = kwargs['pf_exact']
            self.C = self.rvchisq.ppf(self.pf_exact)
        elif 'beta_exact' in kwargs:
            self.beta_exact = kwargs['beta_exact']
            self.C = self.rvchisq.ppf(self.pf_exact)
        elif 'const' in kwargs:
            self.const = kwargs['const']
            self.C = self.const
            self.pf_exact = self.rvchisq.cdf(self.C)
        elif 'C' in kwargs:
            self.C = kwargs['C']
            self.pf_exact = self.rvchisq.cdf(self.C)
        else:
            # no to teda uživatele пошли
            self.pf_exact = 1e-4
            self.C = self.rvchisq.ppf(self.pf_exact)
        
        
        self.const = self.C
        self.gm = g_models.Z_sumsq(self.C)
        # na začatku nemáme vzorky - pouze rozdělení a podpís
        self.sample_box = SampleBox(self.f(), gm_signature=self.gm_signature)
     
    def __str__(self):
        return  'SNorm_Z_sumsq%sD'%(self.nvar)
        
    def __repr__(self):
        return  'SNorm_Z_sumsq(%s, pf_exact=%s)' % (repr(self.nvar), repr(self.pf_exact))   






class SNorm_S_ball(WhiteBox):
    #r = 5.256521 # pf 1.00000404635e-06
    def __init__(self, nvar=2, r=5.256521):
        
        
        # SNorm
        self.f = f_models.SNorm(nvar)
        
        
        #
        #   pf, jen tak, hračka
        #
        self.pf_exact_method = 'precise solution'
        if nvar == 1:
            #self.pf_exact = 1 - 2**(1-nvar/2) / special.gamma(nvar/2)    *    (np.sqrt(np.pi)*special.erf(r/np.sqrt(2)))/np.sqrt(2)
            self.pf_exact = 1 - special.erf(r/1.4142135623730951)
        elif nvar == 2:
            self.pf_exact = np.exp(-r**2/2)
        elif nvar == 3:
            #self.pf_exact = 1 - 2**(1-nvar/2) / special.gamma(nvar/2)    *    (np.exp(-r**2/2)*(np.sqrt(np.pi)*np.exp(r**2/2)*special.erf(r/np.sqrt(2))-np.sqrt(2)*r))/np.sqrt(2)
            self.pf_exact = 1 - 0.5641895835477564 * (np.exp(-r**2/2)*(np.sqrt(np.pi)*np.exp(r**2/2)*special.erf(r/np.sqrt(2))-np.sqrt(2)*r))
        elif nvar == 4:
            self.pf_exact = (r**2/2+1)*np.exp(-r**2/2)
        elif nvar == 6:
            self.pf_exact = (r**4+4*r**2+8)*np.exp(-r**2/2)/8
            
            # nvar=8:  (48-(r^6+6*r^4+24*r^2+48)*e^(-r^2/2)  / 2**(nvar/2))/48
            
            # hračička ve hračce
            # nemám žádnou jistotu, že tohle počítá přesněji
        elif nvar % 2 == 0: # sudé
            poly = [1]
            for i in range(nvar-2, 0, -2):
                poly.append(0)
                poly.append(i*poly[-2])
            self.pf_exact = np.polyval(np.array(poly) / poly[-1], r) * np.exp(-r**2/2) 
            
        else:
            self.pf_exact = 1 - 2**(1-nvar/2) / special.gamma(nvar/2)    *    integrate.quad(lambda x: np.exp(-(x**2)/2)*x**(nvar-1), 0, r)[0] 
        
        
        self.r = r
        self.gm = g_models.S_ball(r)
        # na začatku nemáme vzorky - pouze rozdělení a podpís
        self.sample_box = SampleBox(self.f(), gm_signature=self.gm_signature)
     
    def __str__(self):
        return  'SNorm_S_ball%sD'%(self.nvar)
        
    def __repr__(self):
        return  'SNorm_S_ball(nvar=%s, r=%s)' % (repr(self.nvar), repr(self.r))   


#
#corner_values = np.full(2**nvar, 1, np.int8) # -1 means inertní, 0 failure, 1 success.
#alpha = [1, 1]
#
#if fce_por == 'Sin2D':
#    corner_values[0] = 1
#    corner_values[3] = 0
#    pf_exact = 4.1508e-4
#elif fce_por == 'S_ball':
#    corner_values = np.full(2**nvar, 0, np.int8)
#    pf_exact = 1 - 2**(1-nvar/2) / gamma(nvar/2) * integrate.quad(lambda x: np.exp(-(x**2)/2)*x**(nvar-1), 0, r)[0] 
#elif fce_por == 'S_ballSin2D':
#    corner_values = np.full(2**nvar, 0, np.int8)
#    pf_exact = 0 # I don't know, really
#elif fce_por == 'Prod_FourBetas':
#    corner_values = np.full(2**nvar, 0, np.int8) # dont know in any corner
#    beta = 2.0
#    pf_exact = 2*np.sqrt(2)*stats.norm.cdf(-beta) # true only for beta -> infinity!!
#elif fce_por == 'BlackSwan2D':
#    corner_values[0] = 1
#    corner_values[3] = 0
#    #corner_values = np.full(2**nvar, 1, np.int8) # Success in all corners
#    #corner_values[0] = 0 #failure somwehere in top right corner
#    a = 2.0
#    b = 5.0
#    p = stats.norm.cdf(-a)*stats.norm.cdf(-b)
#    # a<b
#    pf_exact = p
#    # a>b
#    # pf_exact = stats.norm.cdf(a) - stats.norm.cdf(b) + p 
#elif fce_por == 'Metaballs2D':
#    corner_values = np.full(2**nvar, 0, np.int8) # dont know in any corner




#def get_quadrant_probability(center_point=(0,0), quadrant='I'):
        """
       sebemenší pomocná funkce pri hranici L tvaru
        quadrants also сэрегъёс-compatible
        #### CORNERS 2D #####
        # print(сэрегъёс)
        # numbering:
        #    2  |  3
        #  -----|-----
        #    0  |  1
        """

#class Quadrant2D(WhiteBox):
#   def __new__(cls, g_model, f_model=f_models.SNorm(2)):
#       """
#       Trik je v tom, že já odhaduji pf na základě hranici poruchy,
#       prohlašenou g_modelem. 
#       g_model MUST have .mins or .maxs setted up!
#       """
#       if f_model.nvar != 2:
#           raise ValueError("Reliability problem is supposed to be 2D")
#           
#       else:
#           wt = super(Quadrant2D, cls).__new__(cls)
#           wt.gm = g_model
#           wt.f = f_model
#           # na začatku nemáme vzorky - pouze rozdělení a podpís
#           wt.sample_box = SampleBox(wt.f(), gm_signature=wt.gm_signature)
#           
#           f1, f2 = f_model.marginals
#           try:
#               (c1, k1), (c2, k2) = g_model.mins
#               if k1 > 0:
#                   a = f1.cdf(-c1/k1)
#               else:
#                   a = f1.sf(-c1/k1)
#               if k2 > 0:
#                   b = f2.cdf(-c2/k2)
#               else:
#                   b = f2.sf(-c2/k2)
#           
#           quadrant = g_model.get_2D_R_boundary.quadrant
#           xc, yc = g_model.get_2D_R_boundary.center_point
#           
#           #quadrants also сэрегъёс-compatible
#           #### CORNERS 2D #####
#           # print(сэрегъёс)
#           # numbering:
#           #    2  |  3
#           #  -----|-----
#           #    0  |  1
#           
#           wt.pf_exact_method = 'exact solution' # under condition of right g_model
#               # mně nic hezčího prostě nenapadá(
#           if self.quadrant in ('I', 3):
#               xbound = np.append(np.full(nrod, xc), np.linspace(xc, xmax, nrod, endpoint=True))
#               ybound = np.append(np.linspace(ymax, yc, nrod, endpoint=True), np.full(nrod, yc))
#           elif self.quadrant in ('II', 2):
#               xbound = np.append(np.linspace(xmin, xc, nrod, endpoint=True), np.full(nrod, xc))
#               ybound = np.append(np.full(nrod, yc), np.linspace(yc, ymax, nrod, endpoint=True))
#           elif self.quadrant in ('III', 0):
#               xbound = np.append(np.linspace(xmin, xc, nrod, endpoint=True), np.full(nrod, xc))
#               ybound = np.append(np.full(nrod, yc), np.linspace(yc, ymin, nrod, endpoint=True))
#           else: # self.quadrant in ('IV', 1):
#               xbound = np.append(np.full(nrod, xc), np.linspace(xc, xmax, nrod, endpoint=True))
#               ybound = np.append(np.linspace(ymin, yc, nrod, endpoint=True), np.full(nrod, yc))
#        
#elif fce_por == 'Logistic2D':
#    a = 4.
#    b = 5.  #one line, second line and subtract the twice calculated intersection
#    pf_exact = lambda a=4, b=5: stats.norm.cdf(-a) + stats.norm.cdf(-b) - stats.norm.cdf(-a)*stats.norm.cdf(-b) 
       



# no pf information - no reason to create such wbox

#class UniformBranin2D(WhiteBox):
#    def __init__(self):
#        
#        # Uniform-uniform
#        self.f = f_models.UnCorD((stats.uniform, stats.uniform))
#       
#        self.gm = g_models.branin_2D
#        # na začatku nemáme vzorky - pouze rozdělení a podpís
#        self.sample_box = SampleBox(self.f, gm_signature=self.gm_signature)
#     
#    def __str__(self):
#        return  'UniformBranin2D'
#        
#    def __repr__(self):
#        return  'UniformBranin2D()'