#!/usr/bin/env python3
"""
Script to create a comparison plot of total beta spectra for edge neutrino masses (1.5 keV and 2.5 keV)
"""

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# Input files
mass_1_5_csv = "/home/unzhakov/workspace/massive_nu_pr/pr144_nu_mass_1_50keV.csv"
mass_2_45_csv = "/home/unzhakov/workspace/massive_nu_pr/pr144_nu_mass_2_45keV.csv"

# Output file
output_png = (
    "/home/unzhakov/workspace/massive_nu_pr/pr144_comparison_1_5_vs_2_45_keV.png"
)

print(f"Reading data from {mass_1_5_csv} and {mass_2_45_csv}...")


# Read CSV files, skipping metadata headers
# The metadata headers start with '#'
def read_csv_with_metadata(filepath):
    with open(filepath, "r") as f:
        lines = f.readlines()

    # Find the line where data starts (not starting with '#')
    data_start_line = 0
    for i, line in enumerate(lines):
        if not line.startswith("#"):
            data_start_line = i
            break

    # Read the data
    df = pd.read_csv(filepath, skiprows=data_start_line)
    return df


# Read the data
df_1_5 = read_csv_with_metadata(mass_1_5_csv)
df_2_45 = read_csv_with_metadata(mass_2_45_csv)

print("Data loaded successfully.")

# Extract energy and spectrum data
energies_1_5 = df_1_5["energy_MeV"].values
spectrum_1_5 = df_1_5["spectrum"].values

energies_2_45 = df_2_45["energy_MeV"].values
spectrum_2_45 = df_2_45["spectrum"].values

print(
    f"Energy range: {energies_1_5[0] * 1000:.1f} eV to {energies_1_5[-1] * 1000:.1f} eV"
)

# Create comparison plot
fig, ax = plt.subplots(figsize=(12, 8))

# Plot both spectra
ax.plot(energies_1_5 * 1000, spectrum_1_5, "b-", lw=2, label="m_ν = 1.5 keV")
ax.plot(energies_2_45 * 1000, spectrum_2_45, "r--", lw=2, label="m_ν = 2.45 keV")

# Set labels and title
ax.set_xlabel(r"Electron kinetic energy $E$ [keV]", fontsize=12)
ax.set_ylabel("Normalized Spectrum", fontsize=12)
ax.set_title(
    "Pr144 -> Nd144 Beta Decay: Comparison of Neutrino Mass Effects\n"
    "(1.5 keV vs 2.45 keV)",
    fontsize=14,
)

# Set log scale for y-axis to see the endpoint region clearly
ax.set_yscale("log")

# Add grid
ax.grid(True, alpha=0.3)

# Add legend
ax.legend(fontsize=12)

# Focus on the endpoint region (last part of the spectrum)
# Find the energy range near the endpoint
max_energy_keV = max(energies_1_5) * 1000
# Show last 10 keV of the spectrum
start_energy_keV = max_energy_keV - 10.0
ax.set_xlim(start_energy_keV, max_energy_keV)

# Save the plot
plt.tight_layout()
plt.savefig(output_png, dpi=150, bbox_inches="tight")
print(f"Comparison plot saved to {output_png}")

# Also create a linear scale plot for the endpoint region
fig_linear, ax_linear = plt.subplots(figsize=(12, 6))

# Plot both spectra in linear scale
ax_linear.plot(energies_1_5 * 1000, spectrum_1_5, "b-", lw=2, label="m_ν = 1.5 keV")
ax_linear.plot(energies_2_45 * 1000, spectrum_2_45, "r--", lw=2, label="m_ν = 2.45 keV")

# Set labels and title
ax_linear.set_xlabel(r"Electron kinetic energy $E$ [keV]", fontsize=12)
ax_linear.set_ylabel("Normalized Spectrum", fontsize=12)
ax_linear.set_title(
    "Pr144 -> Nd144 Beta Decay: Comparison of Neutrino Mass Effects (Endpoint Region)\n"
    "(1.5 keV vs 2.45 keV) - Linear Scale",
    fontsize=14,
)

# Set linear scale for y-axis
ax_linear.set_yscale("linear")

# Set y-axis limits to show the variation in the spectrum
# Find the min and max values in the endpoint region to set appropriate limits
min_spec_1_5 = np.min(spectrum_1_5[-100:])  # Last 100 points
max_spec_1_5 = np.max(spectrum_1_5[-100:])
min_spec_2_45 = np.min(spectrum_2_45[-100:])
max_spec_2_45 = np.max(spectrum_2_45[-100:])

overall_min = (
    min(min_spec_1_5, min_spec_2_45) * 0.9
    if min(min_spec_1_5, min_spec_2_45) > 0
    else 0
)
overall_max = max(max_spec_1_5, max_spec_2_45) * 1.1

ax_linear.set_ylim(overall_min, overall_max)

# Add grid
ax_linear.grid(True, alpha=0.3)

# Add legend
ax_linear.legend(fontsize=12)

# Focus on the endpoint region (last part of the spectrum)
ax_linear.set_xlim(start_energy_keV, max_energy_keV)

# Save the linear plot
plt.tight_layout()
linear_output_png = (
    "/home/unzhakov/workspace/massive_nu_pr/pr144_comparison_1_5_vs_2_45_keV_linear.png"
)
plt.savefig(linear_output_png, dpi=150, bbox_inches="tight")
print(f"Linear comparison plot saved to {linear_output_png}")