import os
import numpy as np
from multiprocessing import Pool
import astropy.io.fits as pyfits

from noisy import add_noise_X, add_noise_S


def write_fits_dep(event_id, energy, en_primary, theta, phi, type, en_dep_noisy, en_dep, name_id, pos_id, x_pixel, y_pixel):  # ! è qui che decido quali informazioni sulla simulazione voglio nei file fits_dep
    
    x1_pixel = (pos_id // 89) +1
    y1_pixel = (pos_id % 89) +1
    
    X_Y = np.vstack([x1_pixel, y1_pixel]).T
    
    # Write FITS file
    # "Null" primary array
    prhdu = pyfits.PrimaryHDU()
    
    # Extension
    tbhdu = pyfits.BinTableHDU.from_columns([
        pyfits.Column(name='EventID', format='1J', array = event_id),
        pyfits.Column(name='En_Primary', format='1D', unit='keV', array = en_primary),
        pyfits.Column(name='En_dep', format='1D', unit='keV', array = en_dep),
        pyfits.Column(name='En_dep_Noise', format='1D', unit='keV', array = en_dep_noisy),
        #pyfits.Column(name=name_id, format='1I', array = pos_id),
        pyfits.Column(name='X_Y', format='2I', array = X_Y), 
        pyfits.Column(name='X_Pixel_cm', format='1D', unit='cm', array = x_pixel),
        pyfits.Column(name='Y_Pixel_cm', format='1D', unit='cm', array = y_pixel)
        ])
    
    tbhdu.header['EXTNAME'] = 'EVENTS'
    tbhdu.header['TELESCOP'] = 'THESEUS'
    tbhdu.header['INSTRUME'] = 'XGIS'
    tbhdu.header['COMMENT'] = "Simulated GRB spectrum, Band function alpha1 = -1, alpha2 = -3, Ebreak = 100 keV"
    tbhdu.header['TUNIT1'] = "keV" 
    tbhdu.header['TYPE'] = type
    #tbhdu.header['ENERGY'] = energy
    tbhdu.header['THETA'] = theta
    tbhdu.header['PHI'] = phi
    tbhdu.header['COMMENT'] = "Alfonso Pisapia"
    
    hdulist = pyfits.HDUList([prhdu, tbhdu])
    
    fits_dep_dir = "/home/alfonso/Scrivania/THESEUS/xgis_m7-main/fits_dep"
    if not os.path.exists(fits_dep_dir):
        os.makedirs(fits_dep_dir)
        
    file_name = f"{type}_grb_{theta}_{phi}.fits"
    
    output_file = os.path.join(fits_dep_dir, file_name)
    hdulist.writeto(output_file, overwrite=True)
    
    return file_name

# Funzione che prende in input un singolo file FITS da processare.
def get_data(files):
    # Estrazione di energia, theta e phi dal nome del file
    base_name = os.path.basename(files)  # Ottiene il nome del file senza il percorso
    name_parts = base_name.replace('.fits', '').split('_')  # Rimuove l'estensione e suddivide il nome del file in una lista di stringhe usando _ come delimitatore.
    # Energia, theta e phi sono in posizioni fisse nel nome
    #energy = float(name_parts[1])
    theta = float(name_parts[2])
    phi = float(name_parts[3])
    
    # Apertura del file fits e lettura dei dati                                        # ! in base alle info che voglio nel fits_dep qui devo richiamare le info
    with pyfits.open(files) as hdul:
        data = hdul[1].data
        tot_events = len(data)
        
        pixel_id_col = data['Pixel_ID']
        pixel_id = pixel_id_col[pixel_id_col != -1000]
        
        scint_id_col = data['Scint_ID']
        scint_id = scint_id_col[scint_id_col != -1000]
        
        # mask_X = (data['Scint_ID'] == -1000)   # condizione              //TOdO questo è un altro modo di gestire i dati in modo piu compatto, non è concluso
        # eventi_X = data[mask_X]
        # mask_S = (data['Scint_ID'] != -1000)   # condizione
        # eventi_S = data[mask_S]
        # columns_for_X = ['EventID','En_Dep','Pixel_ID','X_Pixel','Y_Pixel','En_Primary']   #  qui scelgo le info
        # columns_for_S = ['EventID','En_Dep','Scint_ID','X_Pixel','Y_Pixel','En_Primary'] 
        # eventi_X_reduced = eventi_X[columns_for_X]
        # eventi_S_reduced = eventi_S[columns_for_S]
        
        event_id = data['EventID']
        X_event_id = event_id[scint_id_col == -1000] 
        S_event_id = event_id[scint_id_col != -1000]
        
        energy = data['En_Primary'][0]
        
        en_primary = data['En_Primary']
        X_en_primary = en_primary[scint_id_col == -1000] 
        S_en_primary = en_primary[scint_id_col != -1000]
        
        en_dep = data['En_dep']
        X_en_dep = en_dep[scint_id_col == -1000]
        S_en_dep = en_dep[scint_id_col != -1000]
        
        x_pixel = data['X_Pixel']
        X_x_pixel = x_pixel[scint_id_col == -1000]
        S_x_pixel = x_pixel[scint_id_col != -1000]
        
        y_pixel = data['Y_Pixel']
        X_y_pixel = y_pixel[scint_id_col == -1000]
        S_y_pixel = y_pixel[scint_id_col != -1000]
        
        num_eventi_X = (scint_id_col == -1000).sum()
        num_eventi_S = (scint_id_col != -1000).sum()
        ratio_X = num_eventi_X/tot_events
        ratio_S = num_eventi_S/tot_events
        if num_eventi_X + num_eventi_S != tot_events:
            raise ValueError("Errore: Eventi X + Eventi S DIVERSO da Eventi Totali!")

        X_en_dep_noisy = [add_noise_X(X_dep) for X_dep in X_en_dep] # list comprehension di Python, una struttura compatta che consente di applicare una funzione a ciascun elemento di un array o di una lista, creando direttamente un nuovo array con i risultati.
        S_en_dep_noisy = [add_noise_S(S_dep) for S_dep in S_en_dep]

    return X_event_id, S_event_id, energy, X_en_primary, S_en_primary, ratio_X, ratio_S, X_x_pixel, S_x_pixel, X_y_pixel, S_y_pixel, theta, phi, X_en_dep_noisy, S_en_dep_noisy, X_en_dep, S_en_dep, pixel_id, scint_id

# scrive sia il file di testo che i file fits
def write_results(file_dat, results):

    with open(file_dat, "w") as f:
        for X_event_id, S_event_id, energy, X_en_primary, S_en_primary, ratio_X, ratio_S, X_x_pixel, S_x_pixel, X_y_pixel, S_y_pixel, theta, phi, X_en_dep_noisy, S_en_dep_noisy, X_en_dep, S_en_dep, pixel_id, scint_id in results:
            
            X_fits = write_fits_dep(X_event_id, energy, X_en_primary, theta, phi, 'X', X_en_dep_noisy, X_en_dep, 'Pixel_ID', pixel_id, X_x_pixel, X_y_pixel)
            S_fits = write_fits_dep(S_event_id, energy, S_en_primary, theta, phi, 'S', S_en_dep_noisy, S_en_dep, 'Scint_ID', scint_id, S_x_pixel, S_y_pixel)
            
            #f.write(f"{energy:.1f}\t{ratio_X:.3f}\t{ratio_S:.3f}\t{theta}\t{phi}\t{X_fits}\t{S_fits}\n")




# Estrae i dati, processa in parallelo tutti i file .fits presenti nella folder, usa le due funzioni già definite
def data_estraction(file_dat):
    # Crea una lista di file FITS trovati nella folder. Sorted per leggerli in ordine
    folder = "/home/alfonso/Scrivania/THESEUS/xgis_m7-main/fits"

    files = sorted([os.path.join(folder, file) for file in os.listdir(folder) if file.endswith(".fits")], key=os.path.getmtime)
    
    # Crea un pool di processi che permette di eseguire funzioni in parallelo.
    # Usa il metodo map() del pool per eseguire la funzione get_data() su ciascun file FITS in parallelo.
    # I risultati vengono raccolti in una lista chiamata results.
    with Pool() as pool:
        results = pool.map(get_data, files) # qui files è l'argomento che passi a get_data
    
    write_results(file_dat, results)         # per tutti i file fits che trova, in parallelo, ottiene i dati  e li scrive in file_dat