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 =false

if FIRST_RUN
    for pkg in ["Statistics", "FITSIO", "Images", "OffsetArrays", "SpecialFunctions", "NativeFileDialog", "Gtk"]
        Pkg.add(pkg)
    end
end

using Statistics
using Printf

# fits file  : https://juliaastro.org/FITSIO.jl/stable/
using FITSIO

# package pour filtrage des images (imfilter) et tableaux à indices décalés (OffsetVector)
using Images
using OffsetArrays: OffsetVector

# fonction d'erreur erf(), utilisée pour le seuillage doux des coefficients d'ondelette
using SpecialFunctions: erf

# package pour sélection interactive du répertoire et du fichier (fenêtres de dialogue)
using NativeFileDialog

# package pour la petite boîte de dialogue Oui/Non (choix du débruitage) : évite de dépendre de
# readline()/stdin, qui peut se bloquer de façon intempestive selon la façon dont le script est
# lancé depuis VS Code (extension Julia). ask_dialog() ouvre une fenêtre GTK native avec 2 boutons.
using Gtk: ask_dialog

# v3 : le paramètre h_wow est maintenant écrit sur 4 chiffres dans le nom du fichier de sortie
#   (h_wow*1000 arrondi, au lieu de h_wow*100) afin de distinguer des valeurs proches de 1,
#   ex. h_wow=0.995 -> "h0995" (au lieu de "h100" avec l'ancienne précision au centième)

##################################################################################################################
# Fonctions du programme :
#   Portage Julia de l'algorithme WOW! (Wavelets Optimized Whitening), développé par Frédéric
#   Auchère, disponible dans le package Python `watroo` (PyPI : https://pypi.org/project/watroo/,
#   dépôt : https://github.com/frederic-auchere/wavelets), et intégré au package affilié SunPy
#   `sunkit-image` (sunkit_image.enhance.wow) :
#   https://docs.sunpy.org/projects/sunkit-image/en/stable/api/sunkit_image.enhance.wow.html
#
#   Référence :
#   Auchère, F., Soubrié, E., Pelouze, G., Buchlin, É. 2023, "Image Enhancement With Wavelets
#   Optimized Whitening", Astronomy & Astrophysics, 670, A66. doi:10.1051/0004-6361/202245345
#   Pour la transformée en ondelettes "à trous" (spline B3) et la propagation du bruit à travers
#   les échelles :
#   Starck, J.-L. & Murtagh, F. 2002, Handbook of Astronomical Data Analysis, Springer-Verlag.
#
# NOTE SUR LA FIDÉLITÉ DU PORTAGE :
#   Ce script est basé sur la lecture directe du code source réel de `watroo` (version 0.0.4,
#   fichiers wavelets.py et utils.py, package téléchargé depuis PyPI), et pas seulement sur la
#   description de l'article ou de la documentation sunkit-image. La logique de la transformée en
#   ondelettes à trous, les coefficients de propagation du bruit (sigma_e), le mode de débruitage
#   (pondération continue par la fonction d'erreur erf, et non un seuillage doux classique à la
#   Donoho), le blanchiment local, et le mélange gamma optionnel (qui porte en réalité sur la
#   reconstruction débruitée de l'image, pas sur l'image brute) reproduisent fidèlement cette
#   logique.
#   Seule la variante "bilatérale" (edge-aware, paramètre `bilateral` de watroo, qui évite les
#   halos autour des forts gradients en pondérant la convolution par la ressemblance locale des
#   pixels) N'EST PAS implémentée ici, pour rester raisonnable en complexité. Elle peut être
#   ajoutée dans une version ultérieure si besoin.
#
# Principe (algorithme WOW! non bilatéral) :
#   1. Transformée en ondelettes "à trous" (spline B3), sur n_scales échelles :
#        à l'échelle j (j = 1 : la plus fine), l'image lissée c_j est obtenue par convolution de
#        c_(j-1) avec un noyau spline B3 "dilaté" (pas de dilatation = 2^(j-1), c'est-à-dire j-1
#        "trous" - des zéros - insérés entre les coefficients du noyau de base [1,4,6,4,1]/16) ;
#        le coefficient d'ondelette est W_j = c_(j-1) - c_j. c_0 = image d'origine ; le résidu
#        final c_n (n = n_scales) est l'image la plus lissée, jamais différenciée.
#   2. Débruitage (optionnel), scinde chaque coefficient W_j par sa "significativité" :
#        seuil_j = coefficient_debruitage_j * sigma_bruit * sigma_e[j]
#        où sigma_bruit est estimé une fois pour toutes, de façon robuste (MAD), sur la 1re
#        échelle (dominée par le bruit), et sigma_e[j] est un facteur standard (tabulé) de
#        propagation du bruit à l'échelle j pour la transformée en ondelettes spline B3 (2D).
#        En seuillage doux (par défaut) : W_j *= erf(|W_j| / seuil_j)  (pondération continue,
#        proche de 0 pour les coefficients dominés par le bruit, proche de 1 au-delà du seuil).
#        En seuillage dur : W_j *= (|W_j| > seuil_j)  (masque binaire).
#   3. Blanchiment ("whitening", optionnel mais activé par défaut) : chaque coefficient W_j est
#      normalisé par une estimation LOCALE de son écart-type (obtenue en lissant W_j² avec le
#      même noyau dilaté que celui utilisé à cette échelle), ce qui égalise la variance locale à
#      toutes les échelles et en tout point de l'image : les structures fines sont ainsi
#      renforcées partout, y compris dans les zones peu contrastées.
#   4. Synthèse : l'image de sortie est la somme pondérée des coefficients W_j (débruités et
#      blanchis) sur toutes les échelles, plus le résidu c_n (lui aussi normalisé par son écart-
#      type global si le blanchiment est actif).
#   5. Mélange optionnel (paramètre h, désactivé par défaut, h=0) : le résultat de l'étape 4 est
#      mélangé avec une version étirée en gamma de la reconstruction DÉBRUITÉE (mais non blanchie)
#      de l'image d'origine, c'est-à-dire la somme des W_j débruités (étape 2) + le résidu brut,
#      AVANT normalisation par le blanchiment.
###################################################################################################################

# -----------------------------------------------------------------------------------------------
# noyau spline B3 "à trous" 1D, dilaté à l'échelle "scale" (scale = 1 : échelle la plus fine, pas
# de dilatation ; scale = j : pas de dilatation = 2^(j-1), c'est-à-dire j-1 "trous" (zéros) insérés
# entre deux coefficients consécutifs du noyau de base [1,4,6,4,1]/16)
# -----------------------------------------------------------------------------------------------
function b3_kernel_1d(scale::Int)
    h = (1 / 16, 1 / 4, 3 / 8, 1 / 4, 1 / 16)
    step = 2^(scale - 1)
    r = 2 * step
    hk = OffsetVector(zeros(Float64, 4 * step + 1), -r:r)
    hk[-2*step] = h[1]
    hk[-1*step] = h[2]
    hk[0] = h[3]
    hk[1*step] = h[4]
    hk[2*step] = h[5]
    return hk
end

# noyau 2D séparable (paire de noyaux 1D) utilisable directement par imfilter ; le noyau spline B3
# 2D du package watroo est exactement le produit extérieur (outer product) du noyau 1D avec
# lui-même, donc la convolution séparable ci-dessous est mathématiquement identique à la
# convolution 2D pleine utilisée dans watroo (cv2.filter2D avec le noyau 2D complet).
# IMPORTANT : il faut envelopper le tuple de noyaux 1D avec KernelFactors.kernelfactors (comme le
# fait Kernel.gaussian en interne) pour qu'imfilter les traite comme une convolution 2D séparable ;
# un tuple brut (k, k) est mal interprété par imfilter comme un noyau à une seule dimension.
function atrous_kernel(scale::Int)
    k = b3_kernel_1d(scale)
    return KernelFactors.kernelfactors((k, k))
end

# -----------------------------------------------------------------------------------------------
# facteurs EXACTS de propagation de l'écart-type d'un bruit blanc gaussien à travers les échelles
# de la transformée en ondelettes "à trous" spline B3 (2D), extraits du code source de watroo
# (classe B3spline, attribut sigma_e_2d, wavelets.py) :
#   sigma_bruit(échelle j) = sigma_bruit(pixel) * B3_SIGMA_E[j]
# (watroo ne fournit ces facteurs tabulés que jusqu'à 11 échelles ; au-delà, ce script utilise une
# extrapolation approximative par division par 2 à chaque échelle supplémentaire)
# -----------------------------------------------------------------------------------------------
const B3_SIGMA_E = [0.8907, 0.20072, 0.085551, 0.041261, 0.020470,
    0.010232, 0.0051435, 0.0026008, 0.0013161, 0.00067359, 0.00040040]

function sigma_e(scale::Int)
    if scale <= length(B3_SIGMA_E)
        return B3_SIGMA_E[scale]
    else
        return B3_SIGMA_E[end] * 0.5^(scale - length(B3_SIGMA_E))
    end
end

# "significativité" d'un coefficient d'ondelette : pondération dans [0,1] (seuillage doux, via la
# fonction d'erreur erf) ou masque binaire (seuillage dur), utilisée pour le débruitage
function significance(W::AbstractMatrix{T}, threshold::T; soft::Bool=true) where {T<:AbstractFloat}
    if soft
        return T.(erf.(abs.(W) ./ threshold))
    else
        return T.(abs.(W) .> threshold)
    end
end

# estimation locale de l'écart-type d'un coefficient d'ondelette (pour le blanchiment), obtenue en
# lissant W² avec le noyau dilaté propre à l'échelle considérée
function local_sigma(W::AbstractMatrix{T}, scale::Int) where {T<:AbstractFloat}
    kf = atrous_kernel(scale)
    var_local = imfilter(W .^ 2, kf, "symmetric")
    return T.(sqrt.(max.(var_local, T(1e-15))))
end

function wow(data::AbstractMatrix{T};
    n_scales::Union{Nothing,Int}=nothing,
    weights::Union{Nothing,Vector{<:Real}}=nothing,
    whitening::Bool=true,
    denoise_coefficients::Union{Nothing,Vector{<:Real}}=nothing,
    noise::Union{Nothing,Real}=nothing,
    soft_threshold::Bool=true,
    gamma::Real=3.2,
    gamma_min::Union{Nothing,Real}=nothing,
    gamma_max::Union{Nothing,Real}=nothing,
    h::Real=0.0) where {T<:AbstractFloat}

    # Fonction :
    #   applique l'algorithme WOW! (Wavelets Optimized Whitening) sur une image 2D (un seul canal)
    #
    # Entrée :
    #   data                 : image 2D (Float32 ou Float64)
    #   n_scales             : nombre d'échelles de la transformée en ondelettes à trous ; si
    #                          `nothing` (par défaut), calculé automatiquement à partir de la
    #                          taille de l'image, comme dans watroo/sunkit-image
    #   weights              : poids relatif de chaque échelle (y compris le résidu final) dans la
    #                          synthèse ; par défaut : 1 pour toutes les échelles
    #   whitening            : applique le blanchiment (étape 3) si true (comportement WOW! standard)
    #   denoise_coefficients : seuils (en unités d'écart-type de bruit local) par échelle, du plus
    #                          fin au plus grossier ; ex. [5, 2, 1] ne débruite que les 3 échelles
    #                          les plus fines (les plus sensibles au bruit). `nothing` ou `[]`
    #                          désactive le débruitage.
    #   noise                : écart-type du bruit (par pixel), si connu à l'avance ; sinon estimé
    #                          automatiquement, de façon robuste, sur la 1re échelle
    #   soft_threshold       : seuillage doux par erf (true, recommandé) ou dur (false)
    #   gamma, gamma_min, gamma_max, h : mélange optionnel avec la reconstruction débruitée (mais
    #                          non blanchie) de l'image, étirée en gamma (h = poids de cette
    #                          partie ; h = 0, par défaut, désactive le mélange)
    #
    # Sortie :
    #   image transformée (mêmes dimensions que data)

    dcoefs = denoise_coefficients === nothing ? Float64[] : Float64.(denoise_coefficients)

    # nombre d'échelles par défaut : la même formule que watroo (log2 de la plus petite dimension
    # de l'image, moins log2 de la taille du noyau spline B3 de base, soit log2(5))
    size_based_max = round(Int, log2(minimum(size(data))) - log2(5))
    ns = if n_scales === nothing
        h < 1 ? size_based_max : length(dcoefs)
    else
        min(n_scales, size_based_max)
    end
    ns = clamp(ns, 1, length(B3_SIGMA_E))   # borné à la table de facteurs de bruit disponible

    # poids de synthèse : ns échelles d'ondelettes + 1 pour le résidu final, tous égaux à 1 par défaut
    wgt = ones(Float64, ns + 1)
    if weights !== nothing
        for i in 1:min(length(weights), ns + 1)
            wgt[i] = weights[i]
        end
    end

    # seuils de débruitage : ns échelles d'ondelettes, 0 (pas de débruitage) au-delà de ce qui est fourni
    dns = zeros(Float64, ns)
    for i in 1:min(length(dcoefs), ns)
        dns[i] = dcoefs[i]
    end

    # 1. transformée en ondelettes à trous : calcul de toutes les échelles W_1..W_ns + résidu c
    c = copy(data)
    W_scales = Vector{Matrix{T}}(undef, ns)
    for s in 1:ns
        kf = atrous_kernel(s)
        c_next = T.(imfilter(c, kf, "symmetric"))
        W_scales[s] = c .- c_next
        c = c_next
    end

    # estimation robuste (MAD) de l'écart-type du bruit, sur la 1re échelle (dominée par le bruit)
    sigma_noise = noise === nothing ? T(median(abs.(W_scales[1])) / 0.6745 / sigma_e(1)) : T(noise)

    gamma_scaled = h > 0 ? zeros(T, size(data)) : nothing

    # 2. débruitage + 3. blanchiment, échelle par échelle
    for s in 1:ns
        W = W_scales[s]

        if dns[s] != 0 && sigma_noise > 0
            threshold = T(dns[s]) * sigma_noise * T(sigma_e(s))
            W .*= significance(W, threshold; soft=soft_threshold)
        end

        if h > 0
            gamma_scaled .+= W
        end

        if whitening && h < 1
            sigma_loc = local_sigma(W, s)
            W ./= sigma_loc
        end
        W .*= T(wgt[s])
    end

    # blanchiment du résidu final (normalisation par son écart-type GLOBAL, pas local)
    if h > 0
        gamma_scaled .+= c
    end
    if whitening && h < 1
        residual_sigma = T(std(c))
        residual_sigma = residual_sigma > 0 ? residual_sigma : T(1e-15)
        c = c ./ residual_sigma
    end
    c = c .* T(wgt[ns+1])

    # 4. synthèse : somme des échelles débruitées/blanchies + résidu
    image = zeros(T, size(data))
    for s in 1:ns
        image .+= W_scales[s]
    end
    image .+= c

    # 5. mélange optionnel avec la reconstruction débruitée (non blanchie), étirée en gamma
    if h > 0
        gmin = gamma_min === nothing ? minimum(gamma_scaled) : T(gamma_min)
        gmax = gamma_max === nothing ? maximum(gamma_scaled) : T(gamma_max)
        gamma_scaled .-= gmin
        if (gmax - gmin) != 0
            gamma_scaled ./= T(gmax - gmin)
        end
        gamma_scaled .= clamp.(gamma_scaled, zero(T), one(T))
        gamma_scaled .^= T(1 / gamma)
        image = T(1 - h) .* image .+ T(h) .* gamma_scaled
    end

    return image
end

###############################################################################################
#
## début du programme principal
#
###############################################################################################

function main()

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

    println("Nombre de threads Julia disponibles : ", Threads.nthreads())
    if Threads.nthreads() == 1
        println("ATTENTION : un seul thread disponible - le calcul des 3 canaux R,G,B restera")
        println("            séquentiel. Relancez avec 'julia -t 3 ...' (ou -t auto) pour paralléliser.")
    end

    # sélection interactive du répertoire de travail
    work_dir = pick_folder(pwd())
    isempty(work_dir) && error("Aucun répertoire sélectionné - programme arrêté")
    cd(work_dir)

    println(pwd())

    # sélection interactive du fichier fit à traiter (par ex. le résultat HDR/_color d'un script précédent)
    selected_file = pick_file(work_dir; filterlist="fit,fits")
    isempty(selected_file) && error("Aucun fichier sélectionné - programme arrêté")
    file_name_in = basename(selected_file)

    # lecture de l'image d'entrée (fit, 3 couches RGB)
    f_in = FITS(file_name_in)
    header = read_header(f_in[1])
    image_in = read(f_in[1])
    close(f_in)

    dim_x = header["NAXIS1"]
    dim_y = header["NAXIS2"]
    println("dim_x = ", dim_x)
    println("dim_y = ", dim_y)

    # conversion en Float32 (au cas où l'image d'entrée serait dans un autre type)
    image_in = Float32.(image_in)

    # -----------------------------------------------------------------------------------------------
    # choix interactif des options de débruitage (denoise) et de blanchiment (whitening), via de
    # petites boîtes de dialogue à 2 boutons (Oui/Non) plutôt que par le terminal (readline() s'est
    # révélé peu fiable selon la façon dont le script est lancé depuis VS Code)
    # -----------------------------------------------------------------------------------------------
    denoise_on = ask_dialog("Appliquer le débruitage (denoise) des coefficients d'ondelettes ?",
        "Non", "Oui")
    println(denoise_on ? "-> débruitage activé" : "-> débruitage désactivé")
    flush(stdout)

    whitening_on = ask_dialog("Appliquer le blanchiment (whitening) ? (Non = juste une synthèse\n" *
                               "multi-échelle pondérée, sans normalisation locale - bien plus rapide)",
        "Non", "Oui")
    println(whitening_on ? "-> blanchiment activé" : "-> blanchiment désactivé")
    flush(stdout)

    # -----------------------------------------------------------------------------------------------
    # paramètres de la WOW : à ajuster selon le rendu souhaité (voir docstring de wow() ci-dessus)
    #
    # n_scales_wow : nombre d'échelles de la transformée en ondelettes à trous. ATTENTION : la
    # taille du noyau DOUBLE à chaque échelle supplémentaire (dilatation 2^(j-1)), donc le temps de
    # calcul explose avec n_scales - les 2-3 dernières échelles dominent largement le coût total.
    # Laisser à `nothing` calcule automatiquement le nombre maximal compatible avec la taille de
    # l'image (round(log2(min(dim_x,dim_y)) - log2(5))), ce qui peut être TRÈS long sur une grande
    # image (ex. environ 10 échelles, donc un noyau d'environ 2000 pixels de large, pour une image
    # 8000x5000) sans forcément apporter de structures utiles en plus.
    #
    # Correspondance approximative avec un sigma gaussien (comme utilisé en MGN) : l'échelle j de
    # la transformée en ondelettes à trous correspond à peu près à un lissage gaussien d'écart-type
    # sigma ≈ 2^(j-1) * 1.12 pixels. Par exemple, si les sigmas MGN utiles pour cette image vont
    # jusqu'à environ 40 px (structures plus grandes jugées non pertinentes), l'échelle 6 (sigma
    # effectif ≈ 36 px) ou 7 (≈ 72 px, marge de sécurité) suffit largement - inutile de monter à 10.
    # -----------------------------------------------------------------------------------------------
    n_scales_wow = 6
    whitening_wow = whitening_on
    # gamma_wow, h_wow : réglages du mélange optionnel avec la version étirée en gamma (voir
    # l'étape 5 du principe de l'algorithme, en tête de fichier). h_wow = 0 (désactivé) à 1 (tout
    # gamma, plus d'effet WOW) ; gamma_wow est l'exposant de l'étirement (entre 2.5 et 4, typiquement)
    gamma_wow = 4.0
    h_wow = 0.995
    soft_threshold_wow = true
    # seuils de débruitage (en unités d'écart-type de bruit local), du plus fin au plus grossier ;
    # valeurs de l'exemple officiel sunkit-image/watroo (3 échelles les plus fines, les plus sensibles au bruit)
    denoise_coefficients_wow = [5.0, 2.0, 1.0]

    image_wow = similar(image_in)

    println("Démarrage du calcul WOW (n_scales, débruitage, etc. voir paramètres ci-dessus)...")
    flush(stdout)

    # -----------------------------------------------------------------------------------------------
    # traitement indépendant de chaque canal couleur (R,G,B) : les 3 canaux ne communiquent jamais
    # entre eux dans wow(), donc on peut les calculer en parallèle sur des threads Julia distincts.
    # Chaque itération écrit dans une tranche disjointe de image_wow (pas de risque de conflit mémoire).
    #
    # IMPORTANT : pour que ceci accélère réellement le calcul, Julia doit être démarré avec plusieurs
    # threads, par ex. :
    #     julia -t 3 TSE2026-4-eclipse-WOW-v2.jl
    # ou en définissant la variable d'environnement avant de lancer Julia :
    #     export JULIA_NUM_THREADS=3      (Linux/macOS)
    #     set JULIA_NUM_THREADS=3         (Windows)
    # "auto" fonctionne aussi (JULIA_NUM_THREADS=auto) et laisse Julia choisir selon les coeurs dispo.
    # Avec 1 seul thread disponible, Threads.@threads ne fait aucun mal mais n'accélère rien non plus.
    # -----------------------------------------------------------------------------------------------
    Threads.@threads for i_ch in [1, 2, 3]
        println("calcul WOW, canal ", i_ch, " (thread ", Threads.threadid(), ")")
        flush(stdout)   # sans ce flush, les messages des threads peuvent rester bloqués dans le
                         # buffer de sortie jusqu'à la fin du programme (voir échange précédent)
        image_wow[:, :, i_ch] = Base.invokelatest(wow, image_in[:, :, i_ch];
            n_scales=n_scales_wow,
            whitening=whitening_wow,
            denoise_coefficients=denoise_on ? denoise_coefficients_wow : nothing,
            soft_threshold=soft_threshold_wow,
            gamma=gamma_wow,
            h=h_wow)
        println("-> canal ", i_ch, " terminé")
        flush(stdout)
    end

    # nom du fichier de sortie : ajout de "_WOW" juste avant l'extension ".fit"
    # (l'extension du fichier reste .fit, obligatoire pour un fichier FITS valide ; "WOW" apparaît
    # dans le nom du fichier, comme "MGN" pour le script MGN)
    # chaîne "_sX_hxxxx_gyy_" juste après "WOW" : X = n_scales_wow (1 chiffre) ; xxxx et yy = valeurs
    # de h_wow et gamma_wow, sur respectivement 4 et 2 chiffres (point supprimé) - h_wow*1000 arrondi
    # (ex. h_wow=0.995 -> "h0995", précision au millième pour distinguer des valeurs proches de 1),
    # gamma_wow*10 arrondi (ex. gamma_wow=3.2 -> "g32")
    params_str = @sprintf("%s_", denoise_on ? "denoise" : "raw")
    s_str = @sprintf("s%d_", n_scales_wow)
    hg_str = @sprintf("h%04d_g%02d_", round(Int, h_wow * 1000), round(Int, gamma_wow * 10))
    file_name_out = chop(file_name_in, tail=4) * "_" * params_str * "WOW_" * s_str * hg_str * ".fit"

    f_out = FITS(file_name_out, "w")
    write(f_out, image_wow[:, :, :]; header=header)
    close(f_out)

    println("fichier écrit : ", file_name_out)

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

end # fin de la fonction main()

# invokelatest en défense contre les avertissements "world age" de Julia 1.12 (peuvent apparaître
# selon la façon dont le script est exécuté, ex. VS Code plutôt qu'un terminal classique)
Base.invokelatest(main)