#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import matplotlib.pyplot as plt
import argparse
import sys
import numpy as np
import pandas as pd
import os.path
from scipy.optimize import curve_fit
from fit_functions import lan

# ---------------------------
# Argumente
# ---------------------------
parser = argparse.ArgumentParser()
parser.add_argument("filenumber", type=int)
parser.add_argument("-s", "--sumfit", action='store_true')
parser.add_argument("-e", "--export", action='store_true')
args = parser.parse_args()

# ---------------------------
# File mapping
# ---------------------------
if args.filenumber == 5:
    file = "2026-03-10-ahbgoa-test-5"
elif args.filenumber == 7:
    file = "2026-04-13-ahbgoa-test-7"
else:
    print("Only filenumbers 5 and 7 are supported.")
    sys.exit()

filename = "/net/asterix/home/asterix/ava/NMAHEPAM/hists/" + file + "_V"

# ---------------------------
# Layout & Positionen
# ---------------------------
pos = [3,2,1,6,5,4,11,10,9,8,7,14,13,12,17,16,15]
pos = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17]

#alle Positionen verwenden
fpos = list(range(1,18))

# ---------------------------
# MAIN
# ---------------------------
def main():

    fig = plt.figure(figsize=(19,11))
    axes = fig.subplot_mosaic(
    """
    ..CCBBAA..
    .FF.EE.DD.
    KKJJIIHHGG
    .NN.MM.LL.
    ..QQPPOO..
    """
    )

    axs = list(axes.values())

    resV = 0.838214 / 2
    energy = 'value'

    a = 6*[40] + [10]
    b = 6*[110] + [310]

    startparams = 6*[[200,75,6,0.5,100]] + [[200,75,6,0.5,100]]

    fitparams = {'P': [], 'mpv': [], 'err_mpv':[]}

    for i in reversed(range(len(axs))):

        axs[i].text(10, 1, pos[i], fontsize=14, color="grey")
        axs[i].set_xlim([3,150])
        axs[i].set_ylim([0.09,15])
        axs[i].set_yscale('log')
        axs[i].grid(True)

        if pos[i] in fpos:
            fitparams['P'].append(pos[i])
            data = pd.read_csv(filename + "%02d" % pos[i] + ".Bhist", sep='\s+')
            Itime = read_Itime(filename)

            channels = list(data.columns)[1:]

            for c in range(len(channels)):

                y = data[channels[c]] * 3600 / resV / Itime

                # ---------------------------
                # FALL 1: kein sumfit → alles normal plotten
                # ---------------------------
                if not args.sumfit:
                    axs[i].step(data[energy], y, label=channels[c])

                # ---------------------------
                # FALL 2: sumfit → nur "sum" fitten
                # ---------------------------
                elif args.sumfit and channels[c] == 'SUM':

                    sig = np.sqrt(data[channels[c]] + 1)

                    par, cov = curve_fit(
                        lan,
                        data[energy][a[c]:b[c]],
                        data[channels[c]][a[c]:b[c]],
                        startparams[c],
                        sigma=sig[a[c]:b[c]],
                        absolute_sigma=True
                    )

                    print(str(channels[c])+': ',par)
                    fitparams['mpv'].append(par[1])
                    perr = np.sqrt(np.diag(cov))
                    fitparams['err_mpv'].append(perr[1])

                    axs[i].step(
                        data[energy],
                        y,
                        label=f"sum ; MPV={par[1]:.2f} ± {perr[1]:.2f}"
                    )

                    axs[i].plot(
                        data[energy][a[c]:b[c]],
                        lan(data[energy], *par)[a[c]:b[c]] * 3600 / resV / Itime,
                        'k--'
                    )

                    frame = pd.DataFrame(fitparams)
                    frame.to_csv('fitparameters/'+file+'_mpv_sums.txt', sep = ' ', index = False)

                # ---------------------------
                # FALL 3: sumfit aktiv → andere Kanäle nur plotten
                # ---------------------------
                else:
                    axs[i].step(data[energy], y)

            axs[i].legend(fontsize=7)

        # keine x-Ticks außer unten
    no_x = [0,1,2,3,4,5,7,8,9,12]
    for i in no_x:
        plt.setp(axs[i].get_xticklabels(), visible=False)

    # keine y-Ticks außer links
    no_y = [1,2,4,5,7,8,9,10,12,13,15,16]
    for i in no_y:
        plt.setp(axs[i].get_yticklabels(), visible=False)

    # ---------------------------
    # Detektor-Labels (C1–C6 im Uhrzeigersinn)
    # ---------------------------
    # oben
    plt.figtext(0.51, 0.9, "  B6  ", fontsize=18, ha='center', va='bottom',
            bbox=dict(boxstyle="square", facecolor='grey', alpha=0.5))

    # unten
    plt.figtext(0.51, 0.07, "  B3  ", fontsize=18, ha='center', va='top',
            bbox=dict(boxstyle="square", facecolor='grey', alpha=0.5))

    # links oben
    plt.figtext(0.2, 0.75, "  B5  ", fontsize=18, ha='center', va='bottom',
            bbox=dict(boxstyle="square", facecolor='grey', alpha=0.5),
            rotation=40)

    # links unten
    plt.figtext(0.2, 0.25, "  B4  ", fontsize=18, ha='center', va='top',
            bbox=dict(boxstyle="square", facecolor='grey', alpha=0.5),
            rotation=-40)

    # rechts oben
    plt.figtext(0.82, 0.75, "  B1  ", fontsize=18, ha='center', va='bottom',
            bbox=dict(boxstyle="square", facecolor='grey', alpha=0.5),
            rotation=-40)

    # rechts unten
    plt.figtext(0.82, 0.25, "  B2  ", fontsize=18, ha='center', va='top',
            bbox=dict(boxstyle="square", facecolor='grey', alpha=0.5),
            rotation=40)

    fig.suptitle(f"Overview position dependency '{file}'", fontsize=18)
    fig.supxlabel("signal [mV]")
    fig.supylabel("counts [1/h/mV]")

    if args.export:
        name = f"plots/overview-{args.filenumber}"
        plt.savefig(name + ".png")
    else:
        plt.show()

# ---------------------------
# Itime
# ---------------------------
def read_Itime(file):
    f = file + '.Itime'
    while len(f) > 0:
        if os.path.isfile(f + '.Itime'):
            return int(pd.read_csv(f + '.Itime', sep='\s+', names=['time'])['time'][0])
        f = f[:-1]
    return 1

# ---------------------------
main()