using Pkg

# -----------------------------------------------------------------------------------------------
# Installation des packages : mettre à false une fois les packages installés une première fois.
# Pkg.add() vérifie le registre à chaque exécution, ce qui coûte plusieurs secondes/minutes
# à chaque lancement du script pour rien si les packages sont déjà présents.
# -----------------------------------------------------------------------------------------------
#const FIRST_RUN = true
const FIRST_RUN = false

if FIRST_RUN
    for pkg in ["NativeFileDialog", "FITSIO"]
        Pkg.add(pkg)
    end
end

using Printf
using NativeFileDialog
using FITSIO

# -----------------------------------------------------------------------------------------------
# Choix du type de recentrage utilisé pour les images en entrée : détermine uniquement le nom
# du fichier HDR de sortie (le calcul du stack HDR lui-même est identique dans les deux cas).
#   "recentrage_soleil" : nom de sortie inchangé ("stack_HDR_...")
#   "recentrage_lune"   : nom de sortie avec l'indicateur "moon" ("stack_HDR_moon_...")
# -----------------------------------------------------------------------------------------------
const recentrage = "recentrage_soleil"
# const recentrage = "recentrage_lune"

##################################################################################################################
# Program Functions:
#   read a series of images with various exposures (images already registered, ordre indifférent)
#   no sensor response curve applied : the sensor is supposed to be linear
#   rescale images (factor = 1/T)
#   stack rescaled images with weighting function (three cases depending on exposure time : 1s, 1/250s et autres)
#   write HDR resulting FIT file
#
# Inputs :
#      sélection interactive du répertoire et des fichiers FIT à stacker
#
# Outputs :
#       HDR image with generic name : stack_HDR_xxxx-nn, with xxxx = numéro de la 1er image stackée, nn = nombre d'images stackées
#       the HDR image is normalized. Meaning that the max ADU is identical whether 9 images are stacked or n cycles of 9 images.
#
# 22/08/2026 : mise à jour pour TSE2026 (noms fichiers et répertoire, fonction poids idem pour 1/250s et 1/125s)
#
# v15 : refonte pour la performance
#   - toute la logique de calcul est déplacée dans une fonction (stack_hdr) au lieu du scope global :
#     en Julia, le code exécuté hors fonction n'est pas type-stable, ce qui ralentit énormément
#     les boucles (gain potentiel de 10x à 50x rien que sur ce point)
#   - Pkg.add() désactivable via FIRST_RUN pour ne pas re-vérifier le registre à chaque lancement
#   - utilisation de vues (@view) au lieu de copies de tranches d'image
#   - la fonction de pondération (FW / FW_1s / FW_250s) est choisie une seule fois par image
#     (et non recalculée à chaque canal de couleur)
#   - accumulation en place (.+=) pour éviter des allocations de tableaux temporaires
#   - fermeture systématique des fichiers FITS (évite un leak de handles sur de grandes séries)
#
# v16 : ajout du paramètre `recentrage` ("recentrage_soleil" / "recentrage_lune") : si "recentrage_lune",
#   le nom du fichier HDR de sortie est préfixé par "stack_HDR_moon_" au lieu de "stack_HDR_"
#
# v17 parallel : parallelisation pour réduire temps CPU
###################################################################################################################

function FW(x)
    # Fonction de pondération générique appliquée à une intensité de pixel normée 16 bits
    ADU_seuil_max = 64000

    # symmetrical sinus version
    # il faut passer à 0 au-delà de ADU_seuil_max, sinon on récupère du bruit dans les hautes valeurs
    weight_coeff = sin(x * pi / ADU_seuil_max)
    # suppression des valeurs négatives ou nulles
    return max(weight_coeff, 0.0001)
end

#..........................................................................................................................
function FW_1s(x)
    # Fonction de pondération spécifique au temps de pose 1 s
    ADU_seuil_max = 64000
    ADU_mid = 2^15

    if x < ADU_mid
        return 1.0
    else
        weight_coeff = 1 - (x - ADU_mid) / (ADU_seuil_max - ADU_mid)
        return max(weight_coeff, 0.0001)
    end
end

#..........................................................................................................................
function FW_250s(x)
    # Fonction de pondération spécifique au temps de pose 1/250 s (et 1/125s)
    ADU_max = (2^16) - 1
    ADU_mid = 2^15

    if x < ADU_mid
        weight_coeff = sin(x * pi / ADU_max)
        return max(weight_coeff, 0.0001)
    else
        return 1.0
    end
end

#############################################################################################
#
## begining of main program
#
###############################################################################################

function stack_hdr()

    # chronométrage du temps total de calcul
    t_start = time()

    println("Nombre de threads Julia disponibles : ", Threads.nthreads())
    if Threads.nthreads() == 1
        println("(pour utiliser plusieurs threads, relancer avec : julia --threads=4 nom_du_script.jl)")
    end

    # -----------------------------------------------------------------------------------------------
    # Sélection interactive du répertoire de travail et des fichiers FIT à stacker
    # -----------------------------------------------------------------------------------------------

    work_dir = pick_folder(pwd())
    isempty(work_dir) && error("Aucun répertoire sélectionné - programme arrêté")
    cd(work_dir)
    println("Répertoire de travail : ", pwd())

    selection = pick_multi_file(work_dir; filterlist="fit,fits")
    isempty(selection) && error("Aucun fichier sélectionné - programme arrêté")

    selected_paths = isa(selection, AbstractString) ? split(selection, "|") : selection
    file_name_in = sort(String[basename(f) for f in selected_paths])

    nb_images = length(file_name_in)
    println("Nombre d'images sélectionnées pour le stack HDR : ", nb_images)
    for f in file_name_in
        println("  ", f)
    end

    # choix d'écrire ou non les fichiers filtrés par la fonction de pondération
    choix_print_W = "NOK"

    file_name_out = Array{String}(undef, nb_images)
    for i_image in 1:nb_images
        file_name_out[i_image] = file_name_in[i_image][1:22] * "_W.fit"
    end

    # préfixe du nom du fichier HDR de sortie : "stack_HDR_" (recentrage sur le Soleil), ou
    # "stack_HDR_moon_" si les images en entrée sont recentrées sur la Lune
    prefix_HDR = recentrage == "recentrage_lune" ? "stack_HDR_moon_" : "stack_HDR_"
    file_name_HDR = prefix_HDR * file_name_in[1][5:8] * "_" * @sprintf("%02i", nb_images) * ".fit"

    local HDR_image::Array{Int32,3}
    local W_map_C_cumul::Array{Float64,3}

    # loop on images to be stacked
    for i_image in 1:nb_images

        println("indice de l'image courante = ", i_image)

        f_in = FITS(file_name_in[i_image])
        header = read_header(f_in[1])
        println("T = ", header["EXPOSURE"])
        image = read(f_in[1])
        close(f_in)

        dim_x = header["NAXIS1"]
        dim_y = header["NAXIS2"]
        T = header["EXPOSURE"]

        if i_image == 1
            HDR_image = zeros(Int32, dim_x, dim_y, 3)
            W_map_C_cumul = zeros(Float64, dim_x, dim_y, 3)
        end

        # sélection de la fonction de pondération adaptée au temps de pose,
        # une seule fois par image (au lieu d'un test répété pour chaque canal)
        wfun = T > 0.95 ? FW_1s : (T < 0.009 ? FW_250s : FW)

        # image_W ne sert que si on écrit les fichiers filtrés par la pondération (choix_print_W == "OK") ;
        # on évite sinon d'allouer un tableau inutile
        if choix_print_W == "OK"
            image_W = similar(image)
        else
            image_W = nothing
        end

        # ---------------------------------------------------------------------------------------
        # Accumulation pondérée de cette image dans le stack HDR (et dans la carte des poids
        # cumulés), canal par canal.
        #
        # Parallélisation : la boucle est réécrite en une seule boucle FUSIONNÉE sur les pixels
        # (lignes j, en parallèle avec Threads.@threads) et les 3 canaux, au lieu de la boucle
        # "for k in 1:3" d'origine (qui ne pouvait exploiter au mieux que 3 threads, un par
        # canal, et allouait un tableau temporaire Wk = wfun.(layer) à chaque canal). Ici :
        #   - chaque thread traite un sous-ensemble de lignes (j) qui lui est propre : pour un
        #     (i, j) donné, seule une thread écrit dans HDR_image[i,j,k] et W_map_C_cumul[i,j,k],
        #     donc aucune donnée n'est partagée en écriture entre threads (pas de race condition,
        #     pas besoin de verrou) ;
        #   - le calcul est fusionné en une seule passe par pixel (poids + accumulation, sans
        #     tableau intermédiaire Wk), ce qui réduit aussi les allocations et les accès mémoire ;
        #   - le gain dépend du nombre de threads Julia disponibles (voir message ci-dessus) :
        #     sans threads (--threads=1), le comportement et le temps de calcul restent identiques
        #     à la version séquentielle d'origine.
        # ---------------------------------------------------------------------------------------
        Threads.@threads for j in 1:dim_y
            @inbounds for i in 1:dim_x, k in 1:3
                pixel = image[i, j, k]
                w = wfun(pixel)

                # accumulation dans l'image HDR (couche filtrée, divisée par le temps de pose, et passage en entiers 32 bits)
                HDR_image[i, j, k] += round(Int32, (w * pixel) / T)

                # coefficients de pondération cumulés
                W_map_C_cumul[i, j, k] += w

                if choix_print_W == "OK"
                    image_W[i, j, k] = round(Int, w * pixel)
                end
            end
        end

        if choix_print_W == "OK"
            f_out = FITS(file_name_out[i_image], "w")
            write(f_out, image_W; header=header)
            close(f_out)
        end
    end

    # calculation of the renormalized HDR image using the cumulative weighting coefficients
    for k in 1:3
        @views HDR_image[:, :, k] .= round.(Int32, HDR_image[:, :, k] ./ W_map_C_cumul[:, :, k])
    end

    # header du fichier HDR = header de la première image de la série
    f_init = FITS(file_name_in[1])
    header_init = read_header(f_init[1])
    close(f_init)

    f_out = FITS(file_name_HDR, "w")
    write(f_out, HDR_image; header=header_init)
    close(f_out)

    println("Fichier HDR écrit : ", file_name_HDR)

    # affichage du temps de calcul total
    elapsed = time() - t_start
    println(@sprintf("Temps de calcul total : %.1f s (%.1f min)", elapsed, elapsed/60))

    return file_name_HDR
end

stack_hdr()