#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Mar 4 09:30:07 2026
@author: ava
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.constants import *
import pandas as pd
from pathlib import Path
import sys

sys.path.append(str(Path(__file__).resolve().parents[1]))
import plotstyle  # Dein plotstyle!

MODE = "paper"
plotstyle.set_style(MODE)

e0 = epsilon_0
me = m_e

def J_to_MeV(J):
    return J / (e * 1e6)

def MeV_to_J(MeV):
    return MeV * (e * 1e6)

Na = 6.022140857e23
u = 1e-3 / Na


particles_BB = {
    'Proton': {
        'E0_min': 30,   # MeV
        'E0_max': 10000,  # MeV
        'm0': m_p,
        'z': 1
    },
    'Helium': {
        'E0_min': 30,   # MeV
        'E0_max': 10000,  # MeV
        'm0': 6.644657e-27,
        'z': 2
    },
    'Muon': {
        'E0_min': 30,   # MeV
        'E0_max': 10000,  # MeV
        'm0': 1.883531627e-28,
        'z': 1
    }
}

materials = {
    'BGO': {
        # Elektronendichte in BGO (gewichteter Durchschnitt)
        'ne': (4 * (9.78e3 / (208.98e-3)) * 83 * N_A + 3 * (5.323e3 / (72.63e-3)) * 32 * N_A + 12 * (1.429e3 / (16.00e-3)) * 8 * N_A) / (4 + 3 + 12),
        # Ionisierungsenergie in BGO (gewichteter Durchschnitt)
        'I': (4 * 250 * e + 3 * 290 * e + 12 * 150 * e) / (4 + 3 + 12),
        # Dicke in m
        'd': 2e-2
    }
}

def rel_speed(Ekin, m0):
    return c * np.sqrt(1 - 1 / ((Ekin / m0 / c / c + 1) ** 2))

def gamma(v):
    return 1 / np.sqrt(1 - v * v / c / c)

def bethebloch(Ekin, m0, z, ne, I):
    v = rel_speed(J_to_MeV(Ekin) if isinstance(Ekin, np.ndarray) else Ekin, m0)
    lorentz_gamma = gamma(v)
    C = ne * z**2 * e**4 / (4 * np.pi * me * v**2 * e0**2)
    D = np.log(2 * lorentz_gamma**2 * me * v**2 / I)
    beta = v**2 / c**2
    return C * (D - beta)

def Eloss(E0, m0, z, ne, I, dx, distance):
    steps = round(distance / dx)
    E = MeV_to_J(E0)
    for i in range(steps):
        dE = bethebloch(E, m0, z, ne, I) * dx
        E -= dE
    Ediff = MeV_to_J(E0) - E
    return J_to_MeV(Ediff)


# -----------------------------
# Geant4‑Konfiguration (ohne Elektron)
# -----------------------------

# Nur Proton, Helium, Mu‑ (kein e-)
PARTICLES_BASE_G4 = {
    'proton':  'G4outfiles/proton_70energies',
    'alpha':   'G4outfiles/helium_70energies',
    'mu-':     'G4outfiles/muon_70energies'
}

colors_G4 = {
    'proton': 'red',
    'alpha':  'green',
    'mu-':    (0.5, 0.3, 0.8)
}

LABELS_G4 = {
    'proton': 'Proton',
    'alpha':  'Helium',
    'mu-':    'Muon'
}

LINSTYLES = {
    2: '-',   # 2 cm
    4: '--',  # 4 cm
    6: ':',   # 6 cm
}

THICKNESSES = [2, 4, 6]

# -----------------------------
# Funktion: GESAMTER Energieverlust (ALLE Zeilen!) – nur für G4
# -----------------------------

def load_particle_data(filename, part_name, label):
    try:
        df = pd.read_csv(filename, sep='\t')
        df['z_cm'] = df['z'] / 10.0
        df_part = df[df['part'] == part_name].copy()

        if len(df_part) == 0:
            print(f"     → {label}: KEINE {part_name}-Hits gefunden!")
            return None

        total_eloss = df_part.groupby(['primaryE', 'event'])['edep'].sum().reset_index()
        total_eloss.columns = ['E0', 'event', 'dE_total']

        means = total_eloss.groupby('E0')['dE_total'].agg(['mean','std']).reset_index()
        means.columns = ['E0', 'dE_mean', 'dE_std']
        return means

    except FileNotFoundError:
        print(f"     → Datei '{filename}' nicht gefunden!")
        return None
    except Exception as e:
        print(f"     → Fehler bei {label}: {e}")
        return None


# -----------------------------
# 9 Dateien laden: Proton/Helium/Mu‑, 2/4/6 cm (ohne e-)
# -----------------------------

all_data = {}  # (part, thickness) -> DataFrame

for part_g4 in PARTICLES_BASE_G4.keys():
    label = LABELS_G4[part_g4]
    base = PARTICLES_BASE_G4[part_g4]

    for thickness in THICKNESSES:
        thickness_label = f"{thickness} cm"
        filename = f"{base}_{thickness}cm_0.hits"
        print(f"\n=== {label} ({part_g4}), {thickness_label} ===")
        print(f"     Datei: {filename}")

        df = load_particle_data(filename, part_g4, label)
        if df is None:
            print(f"     → {label} / {thickness_label}: KEINE DATEN")
        else:
            print(f"     → {label} / {thickness_label}: {len(df)} Energiepunkte")
        all_data[(part_g4, thickness)] = df


# -----------------------------
# PLOT: Bethe‑Bloch + Geant4 in einem Plot
# -----------------------------

plt.figure()

material = materials['BGO']
dx = 1e-5

distances = {
    '2 cm': 0.02,
    '4 cm': 0.04,
    '6 cm': 0.06
}

# Farben und Linestyles für Bethe‑Bloch (nur Proton/Helium/Mu‑, wie in BB‑Skript)
colors_BB = {
    'Proton': 'red',
    'Helium': 'green',
    'Muon':   (0.5, 0.3, 0.8)
}

linestyles_BB = {
    '2 cm': '-',
    '4 cm': '--',
    '6 cm': ':'
}


# 1) Bethe‑Bloch‑Kurven
for name, particle_params in particles_BB.items():
    E_values = np.logspace(
        np.log10(particle_params['E0_min']),
        np.log10(particle_params['E0_max']),
        700
    )

    for label, distance in distances.items():
        Eloss_values = np.array([
            Eloss(
                E0,
                particle_params['m0'],
                particle_params['z'],
                material['ne'],
                material['I'],
                dx,
                distance
            ) for E0 in E_values
        ])

        plt.plot(
            E_values,
            Eloss_values,
            color=colors_BB[name],
            linestyle=linestyles_BB[label],
            linewidth=2,
            label=f"{name} (Bethe‑Bloch, {label})"
        )


# 2) Geant4‑Kurven (Errorbars)
for (part_g4, thickness), df in all_data.items():
    if df is None or len(df) == 0:
        continue

    mask = df['dE_mean'] > 0
    if mask.sum() == 0:
        continue

    color = colors_G4[part_g4]
    linestyle = LINSTYLES[thickness]

    x = df.loc[mask, 'E0']
    y = df.loc[mask, 'dE_mean']
    yerr = df.loc[mask, 'dE_std']

    # Mittlere Linie für Geant4
    plt.plot(
        x, y,
        color=color,
        linestyle=linestyle,
        linewidth=1,
        label=f"{LABELS_G4[part_g4]} (G4, {thickness} cm)"
    )

    # Fehlerschlauch (50% sichtbar, leicht durchsichtig)
    plt.fill_between(
        x,
        y - yerr,
        y + yerr,
        color=color,
        linestyle=linestyle,
        alpha=0.4
    )

# Referenzlinie: IDEAL ΔE = E0
x = np.linspace(30, 10000, 100)
plt.plot(x, x, color="black", linewidth=1, linestyle='-', label=r"$\Delta E = E_{\text{kin}}$")


plt.xscale('log')
plt.yscale('log')
plt.ylim(10, 1000)
plt.xlim(29, 10000)
plt.xlabel('Kinetic energy $E_{kin}$ in MeV')
plt.ylabel('Energy loss $\\Delta E$ in MeV')
plt.grid(True, which="both", ls="--", lw=0.5)


# -----------------------------
# ZWEI LEGENDEN
# 1) Particles (mit BB‑ und G4‑Labels, aber nur 3 Teilchen)
from matplotlib.lines import Line2D

legend_particle_handles = []

for name in ['Proton', 'Helium', 'Muon']:
    color = colors_BB[name]
    legend_particle_handles.append(
        Line2D([0], [0], color=color, lw=2, label=name)
    )

# 2) BGO thickness (2/4/6 cm, Linestyles, schwarz)
legend_thickness_handles = []
for label in distances.keys():
    legend_thickness_handles.append(
        Line2D([0], [0], color='black', linestyle=linestyles_BB[label], lw=2, label=label)
    )


legend_particles = plt.legend(
    handles=legend_particle_handles,
    loc='upper right',
    title="Particles"
)
legend_thickness = plt.legend(
    handles=legend_thickness_handles,
    loc='upper left',
    title="BGO thickness"
)
plt.gca().add_artist(legend_particles)

plt.tight_layout()
plotstyle.savefig("BGO_Eloss_BetheBloch_vs_G4_246cm", category="BB")
plt.show()