import numpy as np
import astropy.io.fits as pyfits
import matplotlib.pyplot as plt
import os
import re
np.set_printoptions(threshold=np.inf)

def image_fits():
    
    os.chdir('/home/alfonso/Scrivania/THESEUS/xgis_m7-main/fits_dep')
    
    fits_imag_dir = "/home/alfonso/Scrivania/THESEUS/xgis_m7-main/fits_matrix"
    if not os.path.exists(fits_imag_dir):
        os.makedirs(fits_imag_dir)
    
    try:
        for file in os.listdir():
            base_name = os.path.basename(file)
            name_parts = base_name.replace('.fits', '').split('_')
            
            type = name_parts[0]
            energy = float(name_parts[3])
            theta = float(name_parts[4])
            phi = float(name_parts[5])

            with pyfits.open(file) as hdul:
                data = hdul[1].data
            
            if type == 'X':
                pixel_id = data['Pixel_ID']
                events = pixel_id

            elif type == 'S':
                scint_id = data['Scint_ID']
                events = scint_id


            ######################### MATRICI ##########################################
            
            # matrice di zeri 89x89
            counts_matrix = np.zeros((89, 89), dtype=int)

            # per ogni eventoe aggiorna il conteggio nel pixel corrispondente
            for event in events:
                # dall'array trova le coo del pixel
                row = event // 89   # divisione intera per riga
                col = event % 89    # resto per colonna
                counts_matrix[row, col] += 1

            #print(counts_matrix)

            dead_matrix = np.ones((89, 89), dtype=int)

            # Imposta a zero ogni 9 righe e ogni 9 colonne
            dead_matrix[8::9, :] = 0  # ogni nona riga
            dead_matrix[:, 8::9] = 0  # ogni nona colonna

            IMMAGINE = counts_matrix*dead_matrix

            #print(IMMAGINE)

            elements_sum = np.sum(IMMAGINE)
            print("Somma totale dei valori nella matrice è:", elements_sum)


            ######################### Write FITS file ##################################
            
            # "Null" primary array
            prhdu = pyfits.PrimaryHDU()
            imhdu = pyfits.ImageHDU(data=IMMAGINE)

            imhdu.header['EXTNAME'] = 'MATRIX'
            imhdu.header['TELESCOP'] = 'THESEUS'
            imhdu.header['INSTRUME'] = 'XGIS'
            imhdu.header['COMMENT'] = "Simulazione in asse"
            imhdu.header['COMMENT'] = f"Direzione theta:{theta} - phi{phi}"
            imhdu.header['COMMENT'] = f"{energy} kev"
            imhdu.header['COMMENT'] = f"eventi {type}"
            imhdu.header['COMMENT'] = "Alfonso Pisapia"

            hdulist = pyfits.HDUList([prhdu, imhdu])

            fits_imag= re.sub(r'en_dep', 'imag', file)

            output_file = os.path.join(fits_imag_dir, fits_imag)
            hdulist.writeto(output_file, overwrite=True)


            # plt.imshow(IMMAGINE, cmap='viridis', interpolation='nearest')
            # plt.colorbar(label='Conteggi')
            # plt.title('Matrice dei Conteggi degli Eventi')
            # plt.xlabel('Colonne')
            # plt.ylabel('Righe')
            # plt.show()
            
        print("Immagini.fits creati correttamente.")
    except Exception as e:
        print(f"Errore durante la creazione di Immagini.fits: {e}")
        
image_fits()