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

import numpy as np
#from scipy import stats
from scipy import special # for S_ball
from scipy import integrate # for S_ball

# нельзя просто так взять и написать Ньютонову методу
# fails on nvar = 501, fails on Sball(500).get_r(0), fails on Sball(800).get_r(0.999)

class Sball:
    def __init__(self, nvar):
        self.nvar = nvar
        if nvar != 2:
            self.C = 2**(1-nvar/2) / special.gamma(nvar/2)
            self.logC = (1-nvar/2)*np.log(2) - special.gammaln(nvar/2)
            self.flex = self.current_r = np.sqrt(self.nvar-1)
            self.flex_pf = self.current_pf = self.get_pf(self.flex)
            
    def get_pf(self, r):
        """
        returns pf, i.e. complementary part of multidimensional Gaussian distribution
        """
        if self.nvar == 1:
            #return 1 - 2**(1-nvar/2) / special.gamma(nvar/2)    *    (np.sqrt(np.pi)*special.erf(r/np.sqrt(2)))/np.sqrt(2)
            return 1 - special.erf(r/1.4142135623730951)
        elif self.nvar == 2:
            return np.exp(-r**2/2)
        elif self.nvar == 3:
            #return 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)
            return 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 self.nvar == 4:
            return (r**2/2+1)*np.exp(-r**2/2)
        elif self.nvar == 6:
            return (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
            # ale ve výsokých dimenzích aspoň počítá
        elif self.nvar % 2 == 0: # sudé
            poly = [1]
            for i in range(self.nvar-2, 0, -2):
                poly.append(0)
                poly.append(i*poly[-2])
            return np.polyval(np.array(poly) / poly[-1], r) * np.exp(-r**2/2) 
            
        else:
            try:
                pf = self.C * integrate.quad(lambda x: np.exp(-(x**2)/2)*x**(self.nvar-1), r, np.inf)[0] 
            except OverflowError:
                pf = 1 - self.C * integrate.quad(lambda x: np.exp(-(x**2)/2)*x**(self.nvar-1), 0, r)[0] 
                    
            return pf
            
    def get_r(self, desired_pf):
        """
        sball_inversion
        returns r
        """
        if self.nvar == 2:
            return np.sqrt(-2*np.log(desired_pf))
        elif self.flex_pf == desired_pf:
            return self.flex
        else:
            # je to jistější
            self.current_r = self.flex 
            self.current_pf = previous_pf = self.flex_pf
            
            self.__do_iter(desired_pf)
            self.current_pf = self.get_pf(self.current_r)
            # hrůza
            # pokračujeme, dokud to nezkonverguje, přenejmenším pokud to konvergue a neosciluje.
            while self.current_pf != previous_pf and self.current_pf != desired_pf\
                     and (self.current_pf > desired_pf or previous_pf < desired_pf):
                     
                previous_pf = self.current_pf
                self.__do_iter(desired_pf)
                self.current_pf = self.get_pf(self.current_r)
                
            return self.current_r
            
    def __do_iter(self, desired_pf):
        r = self.current_r
        denominator =  (self.C * np.exp(-(r**2)/2)*r**(self.nvar-1))
        if denominator != 0 and not np.isnan(denominator):
            self.current_r += (self.current_pf - desired_pf) / denominator
        else:
            # zkombinujeme C a r^nvar, ale stejně nikoho to nezahraní
            log_delta = np.log(abs(self.current_pf - desired_pf)) + (r**2)/2 - (np.log(r)*(self.nvar-1) + self.logC)
            self.current_r += np.exp(log_delta)*np.sign(self.current_pf - desired_pf)
            
        if self.current_r < 0:
            self.current_r = r/2
                
            
    def get_r_iteration(self, desired_pf):
        """
        Same as get_r, but do just one iteration
        """
        
        
        if self.nvar == 2:
            return np.sqrt(-2*np.log(desired_pf)), desired_pf
            
         # logaritmus je na nulu citelný
        elif self.current_pf - desired_pf != 0:    
            
            # hrůza, nečitelný
            # pokud je současné r-ko v jiné straně od chtěného r-ka, tak se vrátíme do inflexního bodu
            if (self.flex_pf > self.current_pf) is (self.flex_pf < desired_pf):
                # vstupní kontrola
                self.current_r = self.flex 
                self.current_pf = self.flex_pf
                
               
            r = self.current_r # pro výstupní kontrolu 
            
            self.__do_iter(desired_pf)
            
            
            # vystupní kontrola
            if (self.flex > self.current_r) is (self.flex < r):
                # preskočili jsme inflexní bod
                self.current_r = self.flex 
                self.current_pf = self.flex_pf
                # ještě jednou
                self.__do_iter(desired_pf)
            
            self.current_pf = self.get_pf(self.current_r) # ne že bychom pf potrebovali v tomto kroce, ale...
        return self.current_r, self.current_pf