#!/usr/bin/env python3
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
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)


# -----------------------------
# MODUS‑EINSTELLUNG GANZ OBEN
# -----------------------------

# Beides (proton, helium, mu-, e-) → True
WITH_ELECTRON = False

# Nur Elektron plotten? → True
PLOT_ELECTRON_ONLY = False


# Partikel-Konfiguration (Basis ohne Dicke)
PARTICLES_BASE = {
    'proton':  'G4outfiles/proton_70energies',
    'alpha':   'G4outfiles/helium_70energies',
    'mu-':     'G4outfiles/muon_70energies',
    'e-':      'G4outfiles/electron_70energies'
}


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


LABELS = {
    'proton': 'Proton',
    'alpha': 'Helium',
    'mu-': 'Muon',
    'e-': 'Electron'
}


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


THICKNESSES = [2, 4, 6]



# -----------------------------
# Funktionen
# -----------------------------

# Funktion: GESAMTER Energieverlust (ALLE Zeilen!)
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 als COPY, um SettingWithCopyWarning zu vermeiden
        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



# -----------------------------
# 12 Dateien laden (mit Filter)
# -----------------------------

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

for part_g4 in PARTICLES_BASE.keys():

    # Modus‑Filter
    if part_g4 == 'e-' and not WITH_ELECTRON:
        continue
    if PLOT_ELECTRON_ONLY and part_g4 != 'e-':
        continue

    label = LABELS[part_g4]
    base = PARTICLES_BASE[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
# -----------------------------

plt.figure()

# Sammle tatsächlich geplottete Teilchen
plotted_particles = set()

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.get(part_g4, 'black')
    linestyle = LINSTYLES[thickness]

    plt.errorbar(
        df.loc[mask, 'E0'],
        df.loc[mask, 'dE_mean'],
        yerr=df.loc[mask, 'dE_std'],
        color=color,
        linestyle=linestyle,
        marker='None',
        linewidth=2,
        label=f"{LABELS[part_g4]}"  # Label nur für Teilchen-Farbe
    )

    plotted_particles.add(part_g4)



# Referenzlinie: IDEAL falls ΔE = E0
x = np.linspace(30, 5000, 100)
plt.plot(x, x, color="black", linewidth=1)


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)



# Legenden wie gewünscht
from matplotlib.lines import Line2D

# Teilchen‑Legende: nur die Teilchen, die wirklich geplottet werden
legend_particle_handles = []
for part in ['proton', 'alpha', 'mu-', 'e-']:
    if part not in plotted_particles:
        continue
    color = colors[part]
    label = LABELS[part]
    legend_particle_handles.append(
        Line2D([0], [0], color=color, lw=2, label=label)
    )


# BGO‑Thickness‑Legende bleibt fix (2/4/6 cm, schwarz)
legend_thickness_handles = [
    Line2D([0], [0], color='black', linestyle='-',  lw=2, label='2 cm'),
    Line2D([0], [0], color='black', linestyle='--', lw=2, label='4 cm'),
    Line2D([0], [0], color='black', linestyle=':',  lw=2, label='6 cm')
]


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


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