"""Functions to provide default ADAS files for Aurora modelling, including capabilities to fetch
derived data files remotely from the OPEN-ADAS website.
"""
# MIT License
#
# Copyright (c) 2021 Francesco Sciortino
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import shutil, os
import requests

# location of the "adas_data" directory relative to this script:
adas_data_dir = (
    os.path.dirname(os.path.realpath(__file__)) + os.sep + "adas_data" + os.sep
)


def get_adas_file_loc(filename, filetype="adf11"):
    """Find location of requested atomic data file for the indicated ion.
    Accepts all ADAS "derived" data file types 
    (adf11, adf12, adf13, adf15, adf21, adf22).
    The search proceeds with the following attempts, in this order:

    #. If the file is available in Aurora/adas_data/*filetype*, with filetype given by the user, 
       always use this data.

    #. If the input filename is actually a full path to the file on the local system, use this file 
       and copy it to Aurora/aurora/adas_data/actual_filename, where actual_filename is the file
       name rather than the full path. 

    #. If the environmental variable "AURORA_ADAS_DIR" is defined, attempt to find the file there, 
       under adf11/ and adf15/ directories. AURORA_ADAS_DIR may for example be the path to a central
       repository of ADAS files for Aurora on a cluster where not everyone may have write-permissions.
       For this option, files are not copied at all.

    #. Attempt to fetch the file remotely via open.adas.ac.uk and save it in 
       Aurora/aurora/adas_data/*filetype*/. 

    Parameters
    ----------
    filename : str
        Name of the ADAS file of interest, e.g. 'plt89_ar.dat'.
    filetype : str
        ADAS file type. Only derived data types are allowed, i.e. one of 
        "adf11", "adf12", "adf13", "adf15", "adf21", "adf22"
    Returns
    -------
    file_loc : str
        Full path to the requested file. 
    """

    def fetch_file(filename, filetype, loc):
        # make sure that aurora/adas_data/adf** directory exists
        if not os.path.isdir(os.path.dirname(os.path.dirname(loc))):
            os.makedirs(os.path.dirname(os.path.dirname(loc)))
        if not os.path.isdir(os.path.dirname(loc)):
            os.makedirs(os.path.dirname(loc))

        url = "https://open.adas.ac.uk/download/" + filetype + os.sep
        if filetype == "adf11":
            filename_mod = filename.split("_")[0] + os.sep + filename
            
        elif filetype in ["adf12", "adf13", "adf21", "adf22"]:
            filename_mod = filename.replace("#","][").split("_")[0] +\
                           os.sep + filename.replace("#", "][")

        elif filetype == "adf15":
            
            if filename.startswith("pec"):
                # more standard format, the following patterns should work fine:
                num = filename[3:5]
                spec = filename.split("#")[1].split("_")[0]
                filename_mod = filename.replace("#", "][").replace("_", f"/pec{num}][{spec}_")
            elif filename.startswith("transport"):
                # different link format for files with "transport" name:
                filename_mod = "transport/" + filename.replace("#", "][")
            else:
                # patterns may be different, attempt simple guess:
                filename_mod = filename.split("_")[0] + os.sep + filename.replace("#", "][")
                
        else:
            raise ValueError(
                "ADAS file type/format not recognized.\n"+\
                "Could not find it or download it automatically!"
            )

        r = requests.get(url + os.sep + filename_mod)
            
        if len(r.text) < 1000:
            # OPEN-ADAS reports short URL error text rather than an error code
            raise ValueError(f"Could not fetch file {filename} from ADAS!")
        
        with open(loc, "wb") as f:
            f.write(r.content) 
            
    if filename == 'none':
        # user doesn't want to load this file
        return 
    
    elif os.path.exists(adas_data_dir+filetype+os.sep+filename):
        # file is available in adas_data:
        return adas_data_dir+filetype+os.sep+filename

    elif os.path.exists(filename):
        # user gave a complete filepath. Don't copy the file from the original location
        return filename

    elif "AURORA_ADAS_DIR" in os.environ:
        loc = os.environ["AURORA_ADAS_DIR"] + os.sep + filetype + os.sep + filename
        if os.path.exists(loc):
            pass
        else:
            # File could not be found. Download it and save it in AURORA_ADAS_DIR:
            fetch_file(filename, filetype, loc=loc)
        return loc

    else:
        # File could not be found. Download it and save it in adas_data:
        if not os.path.isdir(adas_data_dir):
            # make sure that aurora/adas_data exists
            os.makedirs(adas_data_dir)

        loc = adas_data_dir + filetype + os.sep + filename
        fetch_file(filename, filetype, loc=loc)
        return loc



def adas_files_dict():
    """Selections for ADAS files for Aurora runs and radiation calculations.
    This function can be called to fetch a set of default files, which can then be modified (e.g. to 
    use a new file for a specific SXR filter) before running a calculation. 

    Returns
    -------
    files : dict
        Dictionary with keys equal to the atomic symbols of many of the most common ions of
        interest in fusion research. For each ion, a sub-dictionary contains recommended file 
        names for the relevant ADAS data types. Not all files types are available for all ions. 
        Files types are usually a subset of 
        'acd','scd','prb','plt','ccd','prc','pls','prs','fis','brs','pbs',prc'
        Refer to :py:func:`~aurora.atomic.get_adas_file_types` for a description of the meaning of 
        each file.
    """

    files = {}
    files["H"] = {}  # 1
    files["H"]["acd"] = "acd12_h.dat"  # "acd96_h.dat"
    files["H"]["scd"] = "scd12_h.dat"  # "scd96_h.dat"
    files["H"]["prb"] = "prb12_h.dat"  # "prb96_h.dat"
    files["H"]["plt"] = "plt12_h.dat"  # "plt96_h.dat"
    files["H"]["ccd"] = "ccd96_h.dat"
    files["H"]["prc"] = "prc96_h.dat"
    files["H"]["pls"] = "pls_H_14.dat"
    files["H"]["prs"] = "prs_H_14.dat"
    files["H"]["fis"] = "sxrfil14.dat"
    files["H"]["brs"] = "brs05360.dat"
    files["H"]["pbs"] = "pbsx7_h.dat"
    files["H"]["prc"] = "prc89_h.dat"
    files["He"] = {}  # 2
    files["He"]["acd"] = "acd96_he.dat"
    files["He"]["scd"] = "scd96_he.dat"
    files["He"]["prb"] = "prb96_he.dat"
    files["He"]["plt"] = "plt96_he.dat"
    files["He"]["ccd"] = "ccd96_he.dat"
    files["He"]["pls"] = "pls_He_14.dat"
    files["He"]["prs"] = "prs_He_14.dat"
    files["He"]["fis"] = "sxrfil14.dat"
    files["He"]["brs"] = "brs05360.dat"
    files["He"]["pbs"] = "pbsx5_he.dat"
    files["He"]["prc"] = "prc96_he.dat"
    files["Li"] = {}  # 3
    files["Li"]["acd"] = "acd96_li.dat"
    files["Li"]["scd"] = "scd96_li.dat"
    files["Li"]["ccd"] = "ccd89_li.dat"
    files["Li"]["prb"] = "prb96_li.dat"
    files["Li"]["plt"] = "plt96_li.dat"
    files["Li"]["prc"] = "prc89_li.dat"
    files["Li"]["pls"] = "pls89_li.dat"
    files["Li"]["pbs"] = ""
    files["Li"]["prc"] = "prc89_li.dat"
    files["Be"] = {}  # 4
    files["Be"]["acd"] = "acd96_be.dat"
    files["Be"]["scd"] = "scd96_be.dat"
    files["Be"]["prb"] = "prb96_be.dat"
    files["Be"]["plt"] = "plt96_be.dat"
    files["Be"]["ccd"] = "ccd89_be.dat"
    files["Be"]["prc"] = "prc89_be.dat"
    files["Be"]["pls"] = "plsx5_be.dat"
    files["Be"]["prs"] = "prsx5_be.dat"
    files["Be"]["pbs"] = "pbsx5_be.dat"
    files["Be"]["prc"] = "prc89_be.dat"
    files["B"] = {}  # 5
    files["B"]["acd"] = "acd89_b.dat"
    files["B"]["scd"] = "scd89_b.dat"
    files["B"]["ccd"] = "ccd89_b.dat"
    files["B"]["prb"] = "prb89_b.dat"
    files["B"]["plt"] = "plt89_b.dat"
    files["B"]["prc"] = "prc89_b.dat"
    files["B"]["pls"] = "plsx5_b.dat"
    files["B"]["prs"] = "prsx5_b.dat"
    files["B"]["pbs"] = "pbsx5_b.dat"
    files["B"]["prc"] = "prc89_b.dat"
    files["C"] = {}  # 6
    files["C"]["acd"] = "acd96_c.dat"
    files["C"]["scd"] = "scd96_c.dat"
    files["C"]["prb"] = "prb96_c.dat"
    files["C"]["plt"] = "plt96_c.dat"
    files["C"]["ccd"] = "ccd96_c.dat"
    files["C"]["pls"] = "pls_C_14.dat"
    files["C"]["prs"] = "prs_C_14.dat"
    files["C"]["fis"] = "sxrfil14.dat"
    files["C"]["brs"] = "brs05360.dat"
    files["C"]["pbs"] = "pbsx5_c.dat"
    files["C"]["prc"] = "prc96_c.dat"
    files["N"] = {}  # 7
    files["N"]["acd"] = "acd96_n.dat"
    files["N"]["scd"] = "scd96_n.dat"
    files["N"]["ccd"] = "ccd89_n.dat"
    files["N"]["prb"] = "prb96_n.dat"
    files["N"]["plt"] = "plt96_n.dat"
    files["N"]["pls"] = "plsx8_n.dat"
    files["N"]["prs"] = "prsx8_n.dat"
    files["N"]["fis"] = "sxrfilD1.dat"
    files["N"]["brs"] = "brs05360.dat"
    files["N"]["ccd"] = "ccd96_n.dat"
    files["N"]["pbs"] = "pbsx5_n.dat"
    files["N"]["prc"] = "prc96_n.dat"
    files["O"] = {}  # 8
    files["O"]["acd"] = "acd96_o.dat"
    files["O"]["scd"] = "scd96_o.dat"
    files["O"]["ccd"] = "ccd89_o.dat"
    files["O"]["prb"] = "prb96_o.dat"
    files["O"]["plt"] = "plt96_o.dat"
    files["O"]["pls"] = "plsx5_o.dat"
    files["O"]["prs"] = "prsx5_o.dat"
    files["O"]["pbs"] = "pbsx5_o.dat"
    files["O"]["prc"] = "prc89_o.dat"
    files["F"] = {}  # 9
    files["F"]["acd"] = "acd89_f.dat"
    files["F"]["scd"] = "scd89_f.dat"
    files["F"]["ccd"] = "ccd89_f.dat"
    files["F"]["prb"] = "prb89_f.dat"
    files["F"]["plt"] = "plt89_f.dat"
    files["F"]["fis"] = "sxrfil14.dat"
    files["F"]["brs"] = "brs05360.dat"
    files["F"]["pls"] = "pls_F_14.dat"
    files["F"]["prs"] = "prs_F_14.dat"
    files["F"]["pbs"] = "pbsx5_f.dat"
    files["F"]["prc"] = "prc89_f.dat"
    files["Ne"] = {}  # 10
    files["Ne"]["acd"] = "acd96_ne.dat"
    files["Ne"]["scd"] = "scd96_ne.dat"
    files["Ne"]["prb"] = "prb96_ne.dat"
    files["Ne"]["plt"] = "plt96_ne.dat"
    files["Ne"]["ccd"] = "ccd89_ne.dat"
    files["Ne"]["pls"] = "plsx8_ne.dat"
    files["Ne"]["prs"] = "prsx8_ne.dat"
    files["Ne"]["fis"] = "sxrfilD1.dat"
    files["Ne"]["brs"] = "brs05360.dat"
    files["Ne"]["pbs"] = "pbsx5_ne.dat"
    files["Ne"]["prc"] = "prc89_ne.dat"
    files["Na"] = {}  # 10
    files["Na"]["acd"] = "acd85_ne.dat"
    files["Na"]["scd"] = "scd85_ne.dat"
    # files["Na"]['prb'] = ""
    # files["Na"]['plt'] = ""
    # files["Na"]['ccd'] = ""
    # files["Na"]['pls'] = ""
    # files["Na"]['prs'] = ""
    # files["Na"]['fis'] = ""
    # files["Na"]['brs'] = ""
    # files["Na"]["pbs"] = ""
    # files["Na"]['prc'] = ""
    files["Mg"] = {}  # 12
    files["Mg"]["acd"] = "acd85_mg.dat"
    files["Mg"]["scd"] = "scd85_mg.dat"
    # files["Mg"]['prb'] = ""
    # files["Mg"]['plt'] = ""
    # files["Mg"]['ccd'] = ""
    # files["Mg"]['pls'] = ""
    # files["Mg"]['prs'] = ""
    # files["Mg"]['fis'] = ""
    # files["Mg"]['brs'] = ""
    # files["Mg"]["pbs"] = ""
    # files["Mg"]['prc'] = ""
    files["Al"] = {}  # 13
    files["Al"]["acd"] = "acd89_al.dat"
    files["Al"]["scd"] = "scd89_al.dat"
    files["Al"]["prb"] = "prb89_al.dat"
    files["Al"]["plt"] = "plt89_al.dat"
    files["Al"]["ccd"] = "ccd89_al.dat"
    files["Al"]["pls"] = "pls_Al_14.dat"
    files["Al"]["prs"] = "prs_Al_14.dat"
    files["Al"]["fis"] = "sxrfil14.dat"
    files["Al"]["brs"] = "brs05360.dat"
    files["Al"]["pbs"] = "pbsx5_al.dat"
    files["Al"]["prc"] = "prc89_al.dat"
    files["Si"] = {}  # 14
    files["Si"]["acd"] = "acd96_si.dat"
    files["Si"]["scd"] = "scd96_si.dat"
    files["Si"]["prb"] = "prb96_si.dat"
    files["Si"]["plt"] = "plt96_si.dat"
    files["Si"]["pls"] = "pls_Si_14.dat"
    files["Si"]["prs"] = "prs_Si_14.dat"
    files["Si"]["fis"] = "sxrfil14.dat"
    files["Si"]["brs"] = "brs05360.dat"
    files["Si"]["ccd"] = "ccd89_si.dat"
    files["Si"]["pbs"] = "pbsx5_si.dat"
    files["Si"]["prc"] = "prc89_si.dat"
    files["S"] = {}  # 16
    files["S"]["acd"] = "acd89_s.dat"
    files["S"]["scd"] = "scd89_s.dat"
    files["S"]["prb"] = "prb89_s.dat"
    files["S"]["plt"] = "plt89_s.dat"
    files["S"]["pls"] = "pls_S_14.dat"
    files["S"]["prs"] = "prs_S_14.dat"
    files["S"]["fis"] = "sxrfil14.dat"
    files["S"]["brs"] = "brs05360.dat"
    files["S"]["ccd"] = "ccd89_s.dat"
    files["S"]["pbs"] = "pbsx5_s.dat"
    files["S"]["prc"] = "prc89_s.dat"
    files["Cl"] = {}  # 17
    files["Cl"]["acd"] = "acd89_cl.dat"
    files["Cl"]["scd"] = "scd89_cl.dat"
    files["Cl"]["prb"] = "prb89_cl.dat"
    files["Cl"]["plt"] = "plt89_cl.dat"
    files["Cl"]["pls"] = "pls_Cl_14.dat"
    files["Cl"]["prs"] = "prs_Cl_14.dat"
    files["Cl"]["fis"] = "sxrfil14.dat"
    files["Cl"]["brs"] = "brs05360.dat"
    files["Cl"]["ccd"] = "ccd89_cl.dat"
    files["Cl"]["pbs"] = "pbsx5_cl.dat"
    files["Cl"]["prc"] = "prc89_cl.dat"
    files["Ar"] = {}  # 18
    files["Ar"]["acd"] = "acd89_ar.dat"
    files["Ar"]["scd"] = "scd89_ar.dat"
    files["Ar"]["prb"] = "prb89_ar.dat"
    files["Ar"]["plt"] = "plt41_ar.dat"  # Henderson data
    files["Ar"]["ccd"] = "ccd89_ar.dat"
    files["Ar"]["prc"] = "prc89_ar.dat"
    files["Ar"]["pls"] = "pls_Ar_14.dat"
    files["Ar"]["prs"] = "prs_Ar_14.dat"
    files["Ar"]["fis"] = "sxrfil14.dat"
    files["Ar"]["brs"] = "brs05360.dat"
    files["Ar"]["pbs"] = "pbsx5_ar.dat"
    files["Ar"]["prc"] = "prc89_ar.dat"
    files["Ca"] = {}  # 20
    files["Ca"]["acd"] = "acd85_ca.dat"
    files["Ca"]["scd"] = "scd85_ca.dat"
    files["Ca"][
        "ccd"
    ] = "ccd89_w.dat"  # Ca CCD file not available, use first 20 ion stages using Foster scaling
    files["Ca"]["prb"] = "prb85_ca.dat"  # not public on OPEN-ADAS, must request
    files["Ca"]["plt"] = "plt85_ca.dat"  # not public on OPEN-ADAS, must request
    files["Ca"]["pls"] = "pls_Ca_14.dat"
    files["Ca"]["prs"] = "prs_Ca_14.dat"
    files["Ca"]["fis"] = "sxrfil14.dat"
    files["Ca"]["brs"] = "brs05360.dat"
    files["Ca"]["pbs"] = ""
    files["Ca"]["prc"] = ""
    files["Fe"] = {}  # 26
    files["Fe"]["acd"] = "acd89_fe.dat"
    files["Fe"]["scd"] = "scd89_fe.dat"
    files["Fe"]["prb"] = "prb89_fe.dat"
    files["Fe"]["plt"] = "plt41_fe.dat"  # Henderson data
    files["Fe"]["pls"] = "pls_Fe_14.dat"
    files["Fe"]["prs"] = "prs_Fe_14.dat"
    files["Fe"]["fis"] = "sxrfil14.dat"
    files["Fe"]["brs"] = "brs05360.dat"
    files["Fe"]["ccd"] = "ccd89_fe.dat"
    files["Fe"]["pbs"] = "pbsx5_fe.dat"
    files["Fe"]["prc"] = "prc89_fe.dat"
    files["Ni"] = {}  # 28
    files["Ni"]["acd"] = "acd85_ni.dat"
    files["Ni"]["scd"] = "scd85_ni.dat"
    files["Ni"]["prb"] = "prb89_ni.dat"
    files["Ni"]["plt"] = "plt89_ni.dat"
    files["Ni"]["pls"] = "pls_Ni_14.dat"
    files["Ni"]["prs"] = "prs_Ni_14.dat"
    files["Ni"]["fis"] = "sxrfil14.dat"
    files["Ni"]["brs"] = "brs05360.dat"
    files["Ni"]["ccd"] = "ccd89_ni.dat"
    files["Ni"]["pbs"] = "pbsx5_ni.dat"
    files["Ni"]["prc"] = "prc89_ni.dat"
    files["Kr"] = {}  # 36
    files["Kr"]["acd"] = "acd89_kr.dat"
    files["Kr"]["scd"] = "scd89_kr.dat"
    files["Kr"]["ccd"] = "ccd89_kr.dat"
    files["Kr"]["prb"] = "prb89_kr.dat"
    files["Kr"]["plt"] = "plt41_kr.dat"  # Henderson data
    files["Kr"]["pls"] = "plsx5_kr.dat"
    files["Kr"]["prs"] = "prsx5_kr.dat"
    files["Kr"]["pbs"] = "pbsx5_kr.dat"
    files["Kr"]["prc"] = "prc89_kr.dat"
    files["Mo"] = {}  # 42
    files["Mo"]["acd"] = "acd89_mo.dat"
    files["Mo"]["scd"] = "scd89_mo.dat"
    files["Mo"]["ccd"] = "ccd89_mo.dat"
    files["Mo"]["plt"] = "plt89_mo.dat"
    files["Mo"]["prb"] = "prb89_mo.dat"
    files["Mo"]["pls"] = "plsx5_mo.dat"
    files["Mo"]["prs"] = "prsx5_mo.dat"
    files["Mo"]["pbs"] = "pbsx5_mo.dat"
    files["Mo"]["prc"] = "prc89_mo.dat"
    files["Xe"] = {}  # 56
    files["Xe"]["acd"] = "acd89_xe.dat"
    files["Xe"]["scd"] = "scd89_xe.dat"
    files["Xe"]["ccd"] = "ccd89_xe.dat"
    files["Xe"]["plt"] = "plt41_xe.dat"  # Henderson data
    files["Xe"]["prb"] = "prb89_xe.dat"
    files["Xe"]["prs"] = "prsx5_xe.dat"
    files["Xe"]["pls"] = "prsx5_xe.dat"
    files["Xe"]["pbs"] = "pbsx5_xe.dat"
    files["Xe"]["prc"] = "prc89_xe.dat"
    files["W"] = {}  # 74
    files["W"]["acd"] = "acd89_w.dat"
    files["W"]["scd"] = "scd89_w.dat"
    files["W"]["prb"] = "prb89_w.dat"
    files["W"]["plt"] = "plt41_w.dat"
    files["W"]["fis"] = "sxrfil14.dat"
    files["W"]["brs"] = "brs05360.dat"
    files["W"]["pls"] = "pls_W_14.dat"
    files["W"]["prs"] = "prs_W_14.dat"
    files["W"]["ccd"] = "ccd89_w.dat"
    files["W"]["pbs"] = "pbsx5_w.dat"
    files["W"]["prc"] = "prc89_w.dat"

    return files
