# Download exomol opacities
# author: Aurelien Falco

import numpy as np
import os
import wget
from chemistry import species_name_to_common_isotopologue_name

### output formats available on exomol.com
output_formats = [
"R15000_0.3-50mu.xsec.TauREx.h5",
"R1000_0.3-50mu.ktable.petitRADTRANS.h5",
"R1000_0.3-50mu.ktable.ARCiS.fits",
"R1000_0.3-50mu.ktable.NEMESIS.kta",
]

def download_exomol(molecule=None, isotopologue=None, line_list=None, output_dir="tmp", output_format="taurex"):
    """
    Downloads ExoMol data for a given molecule, isotopologue, and line list in the specified output format.
    Parameters:
    molecule (str): The name of the molecule.
    isotopologue (str): The isotopologue of the molecule.
    line_list (str): The line list to download. If you don't know, leave empty
    output_dir (str): The directory where the downloaded file will be saved.
    output_format (str): The desired output format for the downloaded file. Among [ taurex, petitradtrans, arcis, nemesis ].
    Example:
    download_exomol("H2O", "1H2-16O", "BT2", "/path/to/output", "taurex")
    """
    if molecule is not None:
        isotopologue = species_name_to_common_isotopologue_name(molecule)

    exomol_url = "https://exomol.com/db/"
    exfile = "Exomol_species.dat"

    #check if Exomol file exists and create it otherwise
    exists = os.path.isfile(exfile)
    if(exists == 0):
        # raise FileNotFoundError(f"Exomol file {exfile} not found. Run exomol2.py to create it. https://raw.githubusercontent.com/exoclime/HELIOS-K/refs/heads/master/exomol2.py")
        wget.download("https://raw.githubusercontent.com/exoclime/HELIOS-K/refs/heads/master/exomol2.py")
        from exomol2 import main as exomol2_main
        exomol2_main()

    isotopologue_ex, name_linelist_ex, path_ex = np.loadtxt(exfile, usecols=(1,2,3), unpack=True, dtype=str)
    #name_linelist = "1H2-16O__BT2"

    name_linelist0 = name_linelist_ex[isotopologue_ex == isotopologue]
    path0 = path_ex[isotopologue_ex == isotopologue]

    if line_list is None:
         print("Available linelists for", molecule, "are:")
         print(name_linelist0)
         return
    elif line_list == "first-found":
        ### try to get first line list that exists
        for name_linelist in name_linelist0:
            if name_linelist == "first-found": continue # avoid infinite loop
            try:
                download_exomol(molecule, isotopologue, name_linelist, output_dir, output_format)
                break
            except Exception as e:
                print(f"For {name_linelist}: {e}")
                pass
        return
    else:
        name_linelist = [i for i in name_linelist0 if line_list in i][0]
        path = [p for i,p in enumerate(path0) if line_list in name_linelist0[i]][0]

    ### get desired output format from user input
    matching_formats = [fmt for fmt in output_formats if output_format.lower() in fmt.lower()]
    if not matching_formats:
        raise ValueError(f"Invalid output format: {output_format}. Valid formats are: {', '.join(output_formats)}")
    output_format = matching_formats[0]

    ### download archive and extract it
    filename = name_linelist + "." + output_format
    url = exomol_url + path + "/" + filename
    print(f"URL: {url}")
    filename = os.path.join(output_dir, filename)
    os.makedirs(output_dir, exist_ok=True)

    output_name = filename
    if molecule is not None:
        # if molecule is set, do not use isotopologue in output name
        output_name = molecule + "__" + name_linelist.split("__")[1] + "." + output_format

    wget.download(url, output_name)