import numpy as np
import matplotlib.pyplot as plt
from numba import jit
import scipy.optimize as scopt
import scipy.integrate as scint
import time


real_value = 2/np.sqrt(3*np.pi) # hodnota, ktera je spoctena analyticky


#%% Náhled funkce
plt.rcParams.update({"font.size":20})


"""

Tož se podívejme, jaký budou cca hranice integrace, páč nemůžeme integrovat přes celej vesmír, že
"""


def y(x):
    return np.abs(x*np.exp(-x**2/2))

def z(x,y):
    return np.abs((x-y)*np.exp(-(x**2+y**2)/2))


x = np.linspace(-7.5,7.5,2500)
x2 = np.linspace(-5,5,2500)
X, Y = np.meshgrid(x2,x2)

fig = plt.figure(figsize=(13,6))
ax = fig.add_subplot(121)
ax.set_title("$y = |x| \cdot e^{-x^2/2}$")
ax.plot(x,y(x))
ax.set_xlabel("x")

ax = fig.add_subplot(122, projection="3d")
ax.set_title("$z = |x-y| \cdot e^{-(x^2 + y^2 )/2}$")
ax.plot_surface(X,Y, z(X,Y), cmap="inferno", alpha=1)
ax.set_xlim(-5,5)
ax.set_ylim(-5,5)
ax.set_xlabel("x")
ax.set_ylabel("y")

plt.tight_layout(h_pad=3)


#%% Různá N
"""Vidímě, že nám bohatě bude stačit integrační interval $[-5,5]$

"""

@jit(nopython=True)
def sumup(a,b,N): # a a b sou hranice, mezi kterymi integrujeme, N je pocet vycislovanych hodnot
    V = (b-a)**6 # objem podporstoru, kde itnegrujeme

    def fce(random_vals): # funkce argumentu integralu / sumy
        return np.sqrt((random_vals[0]-random_vals[1])**2 + (random_vals[2]-random_vals[3])**2 + (random_vals[4]-random_vals[5])**2) * np.exp(-(random_vals[0]**2+random_vals[1]**2+random_vals[2]**2+random_vals[3]**2+random_vals[4]**2+random_vals[5]**2)/2)

    random = np.random.uniform(low = a, high = b, size=(N,6)) #vektor sesti random uniform cisel z intervalu (a,b)
    sum = 0 # vysledek sumy / integralu bez prefaktoru (pro vypocet erroru f_n)
    f = 0 # suma kvadrátů (pro vypocet erroru (f_n)^2)

    for i in range(0,N):
        sum = sum + fce(random[i])
        f = f + fce(random[i])**2
    I = 1/(2*np.sqrt(3)) * (1/(2*np.pi))**3 * V/N * sum # vysledka hodnota integralu

    err = 1/(2*np.sqrt(3)) * (1/(2*np.pi))**3 * V * np.sqrt((f/N - (sum/N)**2)/N) # vypocet erroru
    return I, err

Ns = [10000,20000,30000, 50000,100000,200000,300000,400000,1000000,2000000,3000000,4000000, 5000000,10000000,100000000] # soubor poctu "vystrelu", pro ktere budeme vyhodnocovat
vals = []
errs = []

start = time.time()
for i in range(0,len(Ns)):
    val, err = sumup(-5,5,Ns[i]) # tu si volame funkci pro jednotliva N.
    vals.append(val)
    errs.append(err)
end = time.time()

takes = round(end - start,1)

print("-------------------------------------------------------") 
print("Výpočet pro různá N:")
print("-------------------------------------------------------") 
print(f"Výpočet trval: {takes} s")
print("sigma pro N = 10e5:" + str(errs[4]))


def fit(x,A,B): # zkusme si fitnout zavisost chyby na N pomocí odmocniny. 
    return A/np.sqrt(x) +B

popt, pcov = scopt.curve_fit(fit, Ns, errs, p0=[20,0.5], bounds = ([0,0], [50,1])) # fit odmocninou

fit_names = ["$A =$ ", "$B =$ "] # vypsani parametru fitnute odmocniny

print("-------------------------------------------------------") 
print(r"hodnoty fitu 1/sqrt(N):")

for i in range(0,len(popt)):
    print(fit_names[i] + str(popt[i]) + " +/- " + str(np.sqrt(pcov[i][i])))
    
print("-------------------------------------------------------")    
print("-------------------------------------------------------")      
fig = plt.figure(figsize=(13,6))
ax = fig.add_subplot(121)
ax.plot(Ns, errs, "o", label = "error", color = "dodgerblue")
ax.plot(Ns, fit(Ns, *popt) - popt[1], color = "crimson", label = r"fit úměrný $\frac{1}{\sqrt{N}}$")
ax.set_xscale("log")
ax.set_xlabel("Počet \"výstřelů\" N")
ax.set_ylabel("error")
ax.legend(frameon=False)
ax = fig.add_subplot(122)
ax.errorbar(Ns, vals, errs, color = "dodgerblue", ecolor = "crimson", label = "spočtená hodnota \ns chybou")
ax.axhline(real_value, color = "black", label = "přesná hodnota")
ax.set_xscale("log")
ax.set_xlabel("Počet \"výstřelů\" N")
ax.set_ylabel("hodnota integrálu")
ax.legend(frameon=False, loc="best")

plt.tight_layout(h_pad=3)

#%% Více vyčíslení


"""Ověření, že funguje centrální limitní věta, vypocet funguje velice obdobne
"""

N = 100000 #pocet "vystrelu" pouzitych pri jednom vypoctu
N_rep = 2000 # pocet opakovani vypoctu
a = -5 # spodni hranice integrace
b = 5 # horni hranice integrace
@jit(nopython=True)
def sumup(a,b,N, N_rep): #a a b sou hranice, mezi kterymi integrujeme
    V = (b-a)**6
    def fce(random_vals):
        return np.sqrt((random_vals[0]-random_vals[1])**2 + (random_vals[2]-random_vals[3])**2 + (random_vals[4]-random_vals[5])**2) * np.exp(-(random_vals[0]**2+random_vals[1]**2+random_vals[2]**2+random_vals[3]**2+random_vals[4]**2+random_vals[5]**2)/2)
    random = np.random.uniform(low = a, high = b, size=(N_rep,N,6))
    values = []
    for k in range(0,N_rep):
        sum = 0
        for i in range(0,N):
            sum = sum + fce(random[k][i])
        I = 1/(2*np.sqrt(3)) * (1/(2*np.pi))**3 * V/N * sum
        values.append(I)
    return values

start = time.time()

data = sumup(a,b,N,N_rep) # tu si volame vypocet datasetu, ktery se pak analyzuje. Pri techto hodnotach N, N_rep mi to trvalo cca  10 s
counted_value = np.mean(data) #stredni hodnota datasetu

end = time.time()

takes = round(end - start,1)

print("Výpočet pro více vyčíslení:")
print("-------------------------------------------------------") 

print(f"Výpočet trval: {takes} s")


print("counted value = " + str(counted_value))
print("real value = " + str(real_value))

plt.figure(figsize=(11,8))
hist_values = plt.hist(data, bins = int(N_rep/50), color = "lightgrey", label = "počet výsledků v intervalu") # vyrobime histogram

# print(hist_values[0], hist_values[1])
def Gaussian(x, A, c, s): # gaussovka pro fit histogramu
    return A * np.exp(-(x-c)**2/(2*s**2))

cents = [] # stredy x-hodnot histogramu
for i in range(1, len(hist_values[1])):
    cents.append((hist_values[1][i-1] + hist_values[1][i])/2)

popt, pcov = scopt.curve_fit(Gaussian, cents, hist_values[0], p0 = [N_rep/6.7, counted_value, np.min(cents + (np.max(cents) - np.min(cents))/2) ]) # fit gaussovkou. Muze se stat, ze se ten fit nechytne, pak by melo stacit to jen pustit znovu

gauss_names = ["amplituda: $A =$ ", "střed: $x_0 =$ ", "sigma: $\sigma = $ "] # vypsani parametru fitnute gaussovky

print("-------------------------------------------------------")
print("hodnoty fitu Gaussem: ")
for i in range(0,len(popt)):
    print(gauss_names[i] + str(popt[i]) + " +/- " + str(np.sqrt(pcov[i][i])))

print("-------------------------------------------------------")    
print("-------------------------------------------------------") 

plt.plot(cents, Gaussian(cents, *popt), color = "dodgerblue", label = "Gaussian fit")
plt.axvline(real_value, 0, np.max(hist_values[0]) + 20, color = "black", label = "přesná hodnota")
plt.axvline(counted_value, 0, np.max(hist_values[0]) + 20, color = "crimson", label = "spočtená hodnota")
plt.xlabel("Hodnota určená jedním opakováním metody M-C")
plt.ylabel("Četnost výskytu")
#plt.ylim(0, np.max(hist_values[0])+10)
plt.xlim(0.5,0.9)
plt.legend(frameon=False, loc = "best")


#%% Trapezoid

#definice parametru. a,b okraje integrace, N je pocet dilku
a = -5
b = 5
N = 10
dx = (b - a) / (N - 1)

real_value = 2/np.sqrt(3*np.pi) # hodnota, ktera je spoctena analyticky
factor = 1/(2*np.sqrt(3)) * (1/(2*np.pi))**3 # prefaktor integralu

def fce(argument): # funkce argumentu integralu / sumy
        return np.sqrt((argument[0]-argument[1])**2 + (argument[2]-argument[3])**2 + (argument[4]-argument[5])**2) * np.exp(-(argument[0]**2+argument[1]**2+argument[2]**2+argument[3]**2+argument[4]**2+argument[5]**2)/2)

points = [] # tu mame linspacei pro jednotlive dimenze
for i in range(0,6):
    points.append(np.linspace(a, b, N))

U,V,W,U_,V_,W_ = np.meshgrid(points[0],points[1],points[2],points[3],points[4],points[5]) # vytvorime si adekvatni meshgrid

values = fce([U,V,W,U_,V_,W_]) #spocteme hodnoty fce v jednotlivych bodech

#vypocet pomoci scipy
sctime1 = time.time()
integral_value = scint.trapezoid(values, dx=dx, axis = 5) #pocatecni integral podel prvni dimenze
for i in range(0,5): #spocteni integralu podel zbylych dimenzi
    integral_value = scint.trapezoid(integral_value, dx=dx, axis = np.abs(i-4))
sctime2 = time.time()

#vypocet pomoci numpy
nptime1 = time.time()
integral_value_np = np.trapz(values, dx=dx, axis = 5) #pocatecni integral podel prvni dimenze
for i in range(0,5): #spocteni integralu podel zbylych dimenzi
    integral_value_np = np.trapz(integral_value_np, dx=dx, axis = np.abs(i-4))
nptime2 = time.time()

sctime = sctime2 - sctime1
nptime = nptime2 - nptime1

print("Výpočet lichoběžníkovým pravidlem:")
print("-------------------------------------------------------") 
print("Men value:"+ str(counted_value))
print("real value:" + str(real_value))
print("-------------------------------------------------------")  
print("scipy.trapezoid value:" + str(factor*integral_value))
print("scipy.trapezoid time:" + str(sctime))
print("-------------------------------------------------------")  
print("numpy.trapz value:" + str(factor*integral_value_np))
print("numpy.trapz time:" + str(nptime))
print("-------------------------------------------------------")    
print("-------------------------------------------------------") 
#%% Porovnání

machine_error = abs(factor*integral_value -real_value)
men_error = abs(counted_value - real_value)

delta = machine_error - men_error 

if machine_error < men_error:
    print("Machines WIN!!! by" + str(delta))
else:
    print("Men WIN!!! by " + str(delta))