#!/usr/bin/env python3
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import sys
from matplotlib.lines import Line2D

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

MODE = "paper"
plotstyle.set_style(MODE)

WITH_ELECTRON = False
PLOT_ELECTRON_ONLY = False

PARTICLES_BASE_NEW = {
    'proton':  'outBGO/threeBGO_prob_proton.0.hits',
    'alpha':   'outBGO/threeBGO_prob_helium.0.hits',
    'mu-':     'outBGO/threeBGO_prob_muon.0.hits',
    'e-':      'outBGO/threeBGO_electron1000.0.hits'
}

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: '-',
    4: '--',
    6: ':'
}

THICKNESSES = [2, 4, 6]
DETECTOR_MAP = {
    2: ['B1'],
    4: ['B1', 'B2'],
    6: ['B1', 'B2', 'B3']
}


def load_stopprob_data(filename, part_name, label):
    try:
        columns = ['event', 'track', 'det', 'primaryE', 'edep', 'nioni']
        df = pd.read_csv(filename, sep=r'\s+', comment='#', names=columns)
        df = df[df['track'] == 1].copy()

        if len(df) == 0:
            print(f"     -> {label}: KEINE Primärteilchen gefunden!")
            return None

        all_stats = []

        for thickness, detectors in DETECTOR_MAP.items():
            df_sel = df[df['det'].isin(detectors)].copy()

            if len(df_sel) == 0:
                continue

            df_sel['edep_abs'] = df_sel['edep'].abs()

            df_event = (
                df_sel
                .groupby(['event', 'primaryE'])['edep_abs']
                .sum()
                .reset_index()
            )

            df_event['stopped'] = (df_event['edep_abs'] >= 0.999 * df_event['primaryE']).astype(int)

            hit_stats = (
                df_event
                .groupby('primaryE')['stopped']
                .agg(['mean', 'count'])
                .reset_index()
            )

            hit_stats.columns = ['E0', 'stop_prob', 'N']
            hit_stats['stop_prob'] *= 100.0
            hit_stats['thickness'] = thickness
            all_stats.append(hit_stats)

        if len(all_stats) == 0:
            return None

        return pd.concat(all_stats, ignore_index=True)

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


all_stats = {}

for part_g4 in PARTICLES_BASE_NEW.keys():
    if part_g4 == 'e-' and not WITH_ELECTRON:
        continue

    if PLOT_ELECTRON_ONLY and part_g4 != 'e-':
        continue

    label = LABELS[part_g4]
    filename = PARTICLES_BASE_NEW[part_g4]

    print(f"\n=== {label} stopping probability ({part_g4}) ===")
    print(f"     Datei: {filename}")

    df = load_stopprob_data(filename, part_g4, label)

    if df is None:
        print(f"     -> {label}: KEINE DATEN")
        continue

    print(f"     -> {label}: {len(df)} Energien")

    for thickness in THICKNESSES:
        df_thick = df[df['thickness'] == thickness].copy()
        if len(df_thick) > 0:
            all_stats[(part_g4, thickness)] = df_thick
        else:
            print(f"     -> {label} / {thickness} cm: KEINE DATEN")
            all_stats[(part_g4, thickness)] = None


print("\n=== Stopping probability [%] pro Startenergie ===")
for (part_g4, thickness), df in all_stats.items():
    label = LABELS[part_g4]
    thickness_label = f"{thickness} cm"
    if df is not None:
        print(f"\n{label} ({thickness_label}):")
        print(df.round(3))
    else:
        print(f"\n{label} ({thickness_label}): KEINE DATEN")


plt.figure()
plotted_particles = set()

for (part_g4, thickness), stats in all_stats.items():
    if stats is None or len(stats) == 0:
        continue

    color = colors.get(part_g4, 'black')
    linestyle = LINSTYLES[thickness]

    x = stats['E0']
    y = stats['stop_prob']
    n = stats['N']
    yerr = np.sqrt((y / 100.0) * (1.0 - y / 100.0) / n) * 100.0 *3

    # Energie bestimmen, bei der die 50%-Schwelle unterschritten wird
    if np.any(y <= 50):

        idx = np.where(y <= 50)[0][0]

        if idx > 0:
            # lineare Interpolation zwischen den beiden Punkten um E50 genauer zu bestimmen
            x1, x2 = x.iloc[idx - 1], x.iloc[idx]
            y1, y2 = y.iloc[idx - 1], y.iloc[idx]

            E50 = x1 + (50.0 - y1) * (x2 - x1) / (y2 - y1)
        else:
            E50 = x.iloc[0]

        plt.axvline(
            E50,
            color=color,
            linestyle=linestyle,
            linewidth=1,
            alpha=0.7
        )

        plt.text(
            E50 * 0.9, 
            45,                  # y-ausrichtung
            f"{E50:.0f} MeV",
            rotation=90,
            color=color,
            fontsize=22,     
            ha='left',
            va='center'
        )

    # plt.errorbar(
    #     x, y, yerr=yerr,
    #     color=color,
    #     linestyle=linestyle,
    #     marker='None',
    #     linewidth=2,
    #     capsize=2,
    #     label=f"{LABELS[part_g4]} ({thickness} cm)"
    # )

    plt.plot(
    x, y,
    color=color,
    linestyle=linestyle,
    linewidth=2,
    label=f"{LABELS[part_g4]} ({thickness} cm)"
    )

    plt.fill_between(
        x,
        y - yerr,
        y + yerr,
        color=color,
        alpha=0.2
    )

    plotted_particles.add(part_g4)

plt.xscale('log')
plt.ylim(-2, 102)
plt.xlim(29, 3000)
plt.xlabel('Start energy $E_0$ in MeV')
#plt.ylabel("Stopping probability in %")
plt.ylabel(r"$\mathrm{Stopping\ probability\ in\ \%}$")
plt.grid(True, which='both', ls='--', lw=0.5)

legend_particle_handles = []
for part in ['proton', 'alpha', 'mu-', 'e-']:
    if part not in plotted_particles:
        continue
    legend_particle_handles.append(
        Line2D([0], [0], color=colors[part], lw=2, label=LABELS[part])
    )

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='lower right', title='BGO thickness')
plt.gca().add_artist(legend1)

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