import numpy as np
import matplotlib.pyplot as plt

x_pb = np.array([2, 5, 7, 10, 12, 15, 17, 20])
y_pb = np.array([5350, 3728, 3079, 2191, 1745, 1313, 1045, 717])

x_cu = np.array([2, 5, 7, 10, 12, 15, 17, 20])
y_cu = np.array([5725, 4953, 4410, 3689, 3257, 2620, 2433, 2132])

x_al = np.array([1, 2, 3, 5, 7, 10, 12, 15, 17, 20])
y_al = np.array([6437, 6080, 5832, 5841, 5515, 5278, 5096, 4771, 4524, 4394])

# Logarytm naturalny liczby zliczeń
log_y_pb = np.log(y_pb)
log_y_cu = np.log(y_cu)
log_y_al = np.log(y_al)

# Funkcja liniowa do dopasowania
def linear_model(x, a, b):
    return a * x + b

params_pb = [-0.1096, 8.793]
params_cu = [-0.0576, 8.783]
params_al = [-0.01924, 8.7597]

# Przygotowanie danych do wykresu eksponencjalnego
x_fit = np.linspace(1, 20, 400)
y_fit_pb = np.exp(linear_model(x_fit, *params_pb))
y_fit_cu = np.exp(linear_model(x_fit, *params_cu))
y_fit_al = np.exp(linear_model(x_fit, *params_al))

# Tworzenie wykresu
errors_pb = np.sqrt(y_pb)
errors_cu = np.sqrt(y_cu)
errors_al = np.sqrt(y_al)

# Tworzenie wykresu logarytmiczno-liniowego
plt.figure(figsize=(12, 6))
plt.errorbar(x_pb, log_y_pb, yerr=errors_pb / y_pb, fmt='o', markersize=5, label='Ołów', color='blue')
plt.errorbar(x_cu, log_y_cu, yerr=errors_cu / y_cu, fmt='o', markersize=5, label='Miedź', color='red')
plt.errorbar(x_al, log_y_al, yerr=errors_al / y_al, fmt='o', markersize=5, label='Aluminium', color='green')
plt.plot(x_fit, linear_model(x_fit, *params_pb), '-', color='blue', label=f'Dopasowanie dla ołowiu')
plt.plot(x_fit, linear_model(x_fit, *params_cu), '-', color='red', label=f'Dopasowanie dla miedzi')
plt.plot(x_fit, linear_model(x_fit, *params_al), '-', color='green', label=f'Dopasowanie dla aluminium')

plt.xlabel('Grubość absorbentu [mm]')
plt.ylabel('$\ln(N(d))$')
plt.title('Logarytmiczno-liniowy wykres zależności liczby kwantów $\gamma$ zarejestrowanych przy przejściu przez absorbent od grubości absorbentu')
plt.legend()
plt.grid(True)
plt.savefig('log_linear_plot.png')
plt.clf()

# Tworzenie wykresu liniowo-liniowego
plt.figure(figsize=(12, 6))
plt.errorbar(x_pb, y_pb, yerr=errors_pb, fmt='o', markersize=5, label='Ołów', color='blue')
plt.errorbar(x_cu, y_cu, yerr=errors_cu, fmt='o', markersize=5, label='Miedź', color='red')
plt.errorbar(x_al, y_al, yerr=errors_al, fmt='o', markersize=5, label='Aluminium', color='green')
plt.plot(x_fit, y_fit_pb, '-', color='blue', label=f'Dopasowanie dla ołowiu')
plt.plot(x_fit, y_fit_cu, '-', color='red', label=f'Dopasowanie dla miedzi')
plt.plot(x_fit, y_fit_al, '-', color='green', label=f'Dopasowanie dla aluminium')

plt.xlabel('Grubość absorbentu [mm]')
plt.ylabel('$N(d)$')
plt.title('Liniowo-liniowy wykres zależności liczby kwantów $\gamma$ zarejestrowanych przy przejściu przez absorbent od grubości absorbentu')
plt.legend()
plt.grid(True)
plt.savefig('linear_linear_plot.png')
plt.clf()