# NOTES 
#
# - installer miniconda 
# - ajouter conda au "PATH"
#  > conda init
#  > conda create -n feelgrib python=3.13 -y
#  > conda activate feelgrib
#  > conda install -c conda-forge pygrib requests
#  > conda install conda-forge::wgrib2
#  > conda install -c conda-forge cdo


# WSL wsl --install 
# dans terminal wsl
# apt-get install python3-grib
# apt-get install cdo


# telecharger https://www.ftp.cpc.ncep.noaa.gov/wd51we/wgrib2/wgrib2.tgz.v3.1.3



import os
import requests
import bz2
import subprocess
import pygrib
import re
from datetime import date,timedelta

def build_url(base_url, basename,previousday=0):
    # today = datetime.datetime.utcnow().strftime("%d%m%y")
    today = date.today()
    target = (today - timedelta(previousday)).strftime("%y%m%d")
    filename = f"{basename}_{target}-12.grb.bz2"
    grib_filename = filename.replace(".bz2", "")
    return base_url + filename, filename, grib_filename

def download_file(url, local_filename):
    print(f"Téléchargement de : {url}")
    response = requests.get(url, stream=True)
    if response.status_code != 200:
        raise Exception(f"Erreur de téléchargement : HTTP {response.status_code}")
    with open(local_filename, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)
    print(f"Fichier téléchargé : {local_filename}")

def decompress_bz2(input_file, output_file):
    print(f"Décompression de : {input_file}")
    with bz2.open(input_file, 'rb') as f_in:
        with open(output_file, 'wb') as f_out:
            f_out.write(f_in.read())
    print(f"Fichier décompressé : {output_file}")
    
def merge_grib_spatially_cdo(file1, file2, output_file, grid_file="target_grid.txt"):
    """
    Fusionne deux fichiers GRIB1 spatialement avec CDO.
    file1 : chemin du premier fichier GRIB (grille cible)
    file2 : second fichier GRIB (sera interpolé sur la même grille)
    output_file : nom du fichier GRIB de sortie fusionné
    grid_file : nom temporaire pour stocker la définition de grille
    """
    try:
        # 1. Extraire la grille du fichier 1
        print(f"[CDO] Extraction de la grille de référence depuis {file1}")
        subprocess.run(["cdo", "griddes", file1], check=True, stdout=open(grid_file, "w"))

        # Calcul des coordonnées
        output = subprocess.check_output(['cdo', 'griddes', file1], text=True)
        a_xfirst = float(re.search(r'xfirst\s*=\s*([-\d.]+)', output).group(1))
        a_xinc   = float(re.search(r'xinc\s*=\s*([-\d.]+)', output).group(1))
        a_xsize  = int(re.search(r'xsize\s*=\s*(\d+)', output).group(1))
        a_yfirst = float(re.search(r'yfirst\s*=\s*([-\d.]+)', output).group(1))
        a_yinc   = float(re.search(r'yinc\s*=\s*([-\d.]+)', output).group(1))
        a_ysize  = int(re.search(r'ysize\s*=\s*(\d+)', output).group(1))
        
        output = subprocess.check_output(['cdo', 'griddes', file2], text=True)
        b_xfirst = float(re.search(r'xfirst\s*=\s*([-\d.]+)', output).group(1))
        b_xinc   = float(re.search(r'xinc\s*=\s*([-\d.]+)', output).group(1))
        b_xsize  = int(re.search(r'xsize\s*=\s*(\d+)', output).group(1))
        b_yfirst = float(re.search(r'yfirst\s*=\s*([-\d.]+)', output).group(1))
        b_yinc   = float(re.search(r'yinc\s*=\s*([-\d.]+)', output).group(1))
        b_ysize  = int(re.search(r'ysize\s*=\s*(\d+)', output).group(1))
        
        xfirst = a_xfirst if a_xfirst < b_xfirst else b_xfirst
        xend   = a_xfirst+(a_xinc*a_xsize) if a_xfirst+(a_xinc*a_xsize) > b_xfirst+(b_xinc*b_xsize) else b_xfirst+(b_xinc*b_xsize)
        xsize  = round((xend - xfirst)/a_xinc)
        print(f"[CDO] Calcul first a:{a_xfirst} b:{b_xfirst} ==> {xfirst}")
        print(f"[CDO] Calcul end a:{a_xfirst+(a_xinc*a_xsize)} b:{b_xfirst+(b_xinc*b_xsize)} ==> {xend}   ==> size : {xsize}")
        
        yfirst = a_yfirst if a_yfirst < b_yfirst else b_yfirst
        yend   = a_yfirst+(a_yinc*a_ysize) if a_yfirst+(a_yinc*a_ysize) > b_yfirst+(b_yinc*b_ysize) else b_yfirst+(b_yinc*b_ysize)
        ysize  = round((yend - yfirst)/a_yinc)
        print(f"[CDO] Calcul first a:{a_yfirst} b:{b_yfirst} ==> {yfirst}")
        print(f"[CDO] Calcul end a:{a_yfirst+(a_yinc*a_ysize)} b:{b_yfirst+(b_yinc*b_ysize)} ==> {yend}   ==> size : {ysize}")
        
        with open(grid_file, "w") as f:
            f.write(f"gridtype  = lonlat\n")
            f.write(f"gridsize  = {xsize*ysize}\n")
            f.write(f"xsize     = {xsize}\n")
            f.write(f"ysize     = {ysize}\n")
            f.write(f"xname     = lon\n")
            f.write(f"xlongname = \"longitude\"\n")
            f.write(f"xunits    = \"degrees_east\"\n")
            f.write(f"yname     = lat\n")
            f.write(f"ylongname = \"latitude\"\n")
            f.write(f"yunits    = \"degrees_north\"\n")
            f.write(f"xfirst    = {xfirst}\n")
            f.write(f"xinc      = {a_xinc}\n")
            f.write(f"yfirst    = {yfirst}\n")
            f.write(f"yinc      = {a_yinc}\n")
            f.write(f"scanningMode = 64\n")
        
        # 2. Remappage bilinéaire des deux fichiers vers la même grille
        file1_remap = file1.replace(".grb", "_remap.grb")
        file2_remap = file2.replace(".grb", "_remap.grb")    

        print(f"[CDO] Remappage de {file1} → {file1_remap}")
        subprocess.run(["cdo", "remapbil," + grid_file, file1, file1_remap], check=True)

        print(f"[CDO] Remappage de {file2} → {file2_remap}")
        subprocess.run(["cdo", "remapbil," + grid_file, file2, file2_remap], check=True)

        # 3. Fusion des deux fichiers remappés
        print(f"[CDO] Fusion des deux fichiers remappés → {output_file}")
        subprocess.run(["cdo", "mergegrid", file1_remap, file2_remap, output_file], check=True)

        print(f"✅ Fichier fusionné créé : {output_file}")

        # Nettoyage temporaire
        os.remove(grid_file)
        os.remove(file1_remap)
        os.remove(file2_remap)

    except subprocess.CalledProcessError as e:
        print(f"❌ Erreur lors de l'exécution de CDO : {e}")
    except Exception as ex:
        print(f"❌ Erreur inattendue : {ex}")
  
def extract_surface_current(grib_file, output_file):
    grbs = pygrib.open(grib_file)
    selected = []
    print(f"Extraction des informations de courrant de surface depuis {grib_file}")

    for grb in grbs:
        if "component of current" in grb.parameterName:
            # param_id = getattr(grb, 'parameterNumber', 'N/A')
            #level = getattr(grb, 'level', 'N/A')
            #units = getattr(grb, 'units', 'N/A')
            #print(f"ParamId: {param_id}, Niveau: {level}, Unités: {units}")
            #print(f"OK  Niveau: {grb.level}, Unités: {grb.units}")
            selected.append(grb)

    if not selected:
        raise Exception("Aucun champ 'unknown' trouvé.")

    with open(output_file, 'wb') as f_out:
        for grb in selected:
            f_out.write(grb.tostring())

    print(f"Extraction terminée, fichier sauvegardé : {output_file}")

def main():
    output_filename = "surface_currents.grb"
    merged_filename = "merged.grb"
    
    url, bz2_filename, grib_a_filename = build_url("https://openskiron.org/gribs_wrf_4km/","Dunkirk_4km_WRF_WAM")
    download_file(url, bz2_filename)
    decompress_bz2(bz2_filename, grib_a_filename)

    url, bz2_filename, grib_b_filename = build_url("https://openskiron.org/gribs_wrf_4km/","Hastings_4km_WRF_WAM")
    download_file(url, bz2_filename)
    decompress_bz2(bz2_filename, grib_b_filename)

    extract_surface_current(grib_a_filename,"surfacea.grb")
    extract_surface_current(grib_b_filename,"surfaceb.grb")


    merge_grib_spatially_cdo("surfacea.grb","surfaceb.grb",merged_filename);
    
    extract_surface_current(merged_filename, output_filename)

if __name__ == "__main__":
    main()