import math
import numpy as np
from matplotlib import pyplot as plt


def trap_proton():
    def pdf(E, C1=559.76377, C2=-1.53424, C3=2.40873, sum=1.51565):
        return C1 * pow(E - C2, -C3) / sum
    pdf = np.vectorize(pdf)

    def pdf_int(E, C1=559.76377, C2=-1.53424, C3=2.40873, sum=1.5156509):
        return C1 * math.pow(E - C2, 1 - C3) / (1 - C3) / sum

    def pdf_inv(y, C1=559.76377, C2=-1.53424, C3=2.40873, sum=1.5156509):
        return pow(sum * (pdf_int(50) + y) * (1 - C3) / C1, 1 / (1 - C3)) + C2

    with open('debug.txt', 'r') as f:
        data = f.readlines()
    data = np.array([float(s) for s in data[0].split(', ')[:-1]])
    x = np.linspace(50, 1000, 1000)

    plt.plot(x, pdf(x), label="f(x)")
    plt.hist(data, bins=np.linspace(50, 1000, 500), density=True, histtype='step', label="Simulation")
    plt.legend()
    plt.ylabel('Frequency')
    plt.xlabel('Energy/(MeV)')
    plt.savefig('docs/simu-tp.png', bbox_inches='tight')


def gcr_proton():
    M = 931.5

    def get_pc(E):
        return math.sqrt(2 * M * E + E * E)

    def get_beta(E):
        pc = get_pc(E)
        return pc / math.sqrt(M * M + pc * pc)

    def pdf(E, C1=125.89223, alpha=44.00673, C2=-0.000460574, C3=112356, sum=361.781665965026):
        pc = get_pc(E)
        beta = get_beta(E)
        return C1 * math.pow(beta, alpha - 1) * math.pow(pc / (pc + C2), C3) / pc / sum
    pdf = np.vectorize(pdf)

    with open('debug.txt', 'r') as f:
        data = f.readlines()

    data = np.array([float(s) for s in data[0].split(', ')[:-1]])
    x = np.linspace(220, 1e5, 10000)

    plt.plot(x, pdf(x), label="f(x)")
    plt.hist(data, bins=np.linspace(220, 1e5, 100), density=True, histtype='step', label="Simulation")
    plt.legend()
    plt.ylabel('Frequency')
    plt.xlabel('Energy/(MeV)')
    plt.savefig('docs/simu-gcrp.png', bbox_inches='tight')


def gcr_ion():
    def pdf(E, C1=0.01815, C2=0.000297726, C3=0.00184, C4=0.49271, sum=48.4143048):
        return C1 * math.exp(-C2 * E) * (1 - math.exp(-C3 * E + C4)) / sum
    pdf = np.vectorize(pdf)

    with open('debug.txt', 'r') as f:
        data = f.readlines()
    data = np.array([float(s) for s in data[0].split(', ')[:-1]])
    x = np.linspace(220, 1e5, 10000)

    plt.plot(x, pdf(x), label="f(x)")
    plt.hist(data, bins=np.linspace(220, 1e5, 300), density=True, histtype='step', label="Simulation")
    plt.legend()
    plt.ylabel('Frequency')
    plt.xlabel('Energy/(MeV)')
    plt.savefig('docs/simu-gcri.png', bbox_inches='tight')


trap_proton()
gcr_proton()
gcr_ion()