#!/usr/bin/env python
"""
Tunes the "zinf" parameter for each atomic line in atomic lines file

The "zinf" parameter is a distance in angstrom from the centre of an atomic
line. It specifies the calculation range for the line:
[centre-zinf, centre+zinf].

This script runs pfant for each atomic line to determine the width of each
atomic line and thus zinf.

Note: if input files main.dat and abonds.dat are not present in current directory, will copy the
      Sun data.

Note: if modeles.mod is not present in current directory, will run innewmarcs to generate this file.

Note: the precision in the zinf found depends on the calculation step ("pas")
specified in main.dat. A higher "pas" means lower precision and a tendency to
get higher zinf's. This is really not critical. pas=0.02 or pas=0.04 should do.

"""

import argparse
import logging
import copy
import numpy as np
import os.path
import a99
import pyfant


a99.logging_level = logging.INFO
a99.flag_log_file = True
a99.flag_log_console = True

# Fraction of peak to be considered its "settlement"
EPSILON = 1e-4

DEFOUT = "<made-up filename>"


def _get_zinf(lambda_centre, norm):
    """Returns zinf given *normalized* spectrum, or 0 if line if flat.

    Args:
      lambda_centre: centre of line
      norm: Spectrum instance

    Spectrum is expected to contain one line only.
    """

    y = norm.y  # normalized spectrum: values between 0 and 1
    if np.all(np.diff(y) == 0):
        # flat line
        return 0
    ymax = 1
    ymin = min(y)
    i_zinf = np.argmax(norm.y < ymax-(ymax-ymin)*EPSILON)
    zinf = lambda_centre-norm.x[i_zinf]
    return zinf


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
     description=__doc__,
     formatter_class=a99.SmartFormatter
     )
    parser.add_argument('--min', type=float, nargs='?', default=.1,
     help='minimum zinf. If zinf found for a particular line is smaller than '
          'this value, this value will be used instead')
    parser.add_argument('--max', type=float, nargs='?', default=50.,
     help='maximum zinf. If zinf found for a particular line is greater than '
          'this value, this value will be used instead')
    parser.add_argument('--inflate', type=float, nargs='?', default=1.1,
     help='Multiplicative constant to apply a "safety margin". '
          'Each zinf found will be multiplied by this value. For example '
          'a value of INFLATE=1.1 means that all the zinf\'s saved will be 10 '
          'percent larger than those calculated')
    parser.add_argument('--ge_current', action="store_true",
     help='"Greater or Equal to current": If this option is set, the current '
          'zinf in the atomic lines file is used as a lower boundary.')
    parser.add_argument('fn_input', type=str, help='input file name', nargs=1)
    parser.add_argument('fn_output', type=str, help='output file name',
     nargs="?", default=DEFOUT)
    parser.add_argument('--no_clean', action="store_true",
    help='If set, will not remove the session directories.')

    args = parser.parse_args()
    logger = a99.get_python_logger()

    fn_out = args.fn_output
    if fn_out == DEFOUT:
        temp = args.fn_input[0].replace("atoms-untuned-", "")
        fn_out = "atoms-"+temp
        logger.info("Made up output filename: '%s'" % fn_out)

    logger.info("Using minimum zinf = %g" % args.min)
    logger.info("Using MAXIMUM zinf = %g" % args.max)
    logger.info("Using inflate = %g" % args.inflate)
    logger.info("Using ge_current = %s" % args.ge_current)

    file_atoms = pyfant.FileAtoms()
    file_atoms.load(args.fn_input[0])

    logger.info("Number of lines in file '%s': %d" % \
     (args.fn_input[0], file_atoms.num_lines))

    # # Tries to get input files together
    pyfant.setup_inputs(h=False, atoms=False, molecules=False, opa=False)


    # ## Runs innewmarcs if modeles.mod not present
    
    if not os.path.isfile('modeles.mod'):
        logger.info("'modeles.mod' does not exist, running innewmarcs...")
        inn = pyfant.Innewmarcs()
        inn.conf.opt.opa = False
        inn.run()



    # ## As suggested by M. Trevisan, adds 0.5 to all abundances, except EXCEPTIONS
    K_ADD = 0.5
    EXCEPTIONS = ["FE", "HE"]
    logger.info("Preparing abundances (Adding %g to all abundances, " % K_ADD)
    logger.info("except those of %s)..." % str(EXCEPTIONS))
    fa = pyfant.FileAbonds()
    fa.load("abonds.dat")
    cnt = 0
    for i in range(len(fa)):
       if fa.ele[i] not in EXCEPTIONS:
           fa.abol[i] += K_ADD
           cnt += 1
    logger.info("%d abundance(s) changed" % cnt)


    # This is old code, with a minimum value for the abundances but I don't think it is very reallistic
    # MIN_ABOND = 3
    # logger.info("Preparing abundances file with all abundances >= %g ..." % MIN_ABOND)
    # fa = FileAbonds()
    # fa.load("abonds.dat")
    # cnt = 0
    # for i in range(len(fa)):
    #     if fa.abol[i] < MIN_ABOND:
    #         fa.abol[i] = MIN_ABOND
    #         cnt += 1
    # logger.info("%d abundance(s) changed" % cnt)


    # ## Creates the Runnable instances
    PAS = 0.04
    logger.info("Calculation step will be set to %g" % PAS)
    logger.info("Preparing pfant's...")
    combos, ll, i = [], [], 0
    ii = 0
    n = file_atoms.num_lines
    for atom in file_atoms.atoms:
        for line in atom.lines:
            a = copy.copy(atom)  # shallow copy
            a.lines = [line]
            f = pyfant.FileAtoms()
            f.atoms = [a]
            combo = pyfant.Combo([pyfant.FOR_PFANT])
            # Fortran messages will not be displayed in terminal
            combo.conf.flag_log_console = False
            combo.conf.sid.flag_split_dirs = True
            # Forces pfant to run one single iteration (ikeytot will be =1)
            combo.conf.opt.aint = args.max+10
            combo.conf.file_atoms = f
            combo.conf.file_abonds = fa
            combo.conf.opt.logging_level = "warning"
            combo.conf.opt.pas = PAS
            combo.conf.flag_output_to_dir = True
            combo.conf.opt.zinf = args.max
            combo.conf.opt.no_molecules = True
            combo.conf.opt.opa = False
            combo.conf.opt.no_h = True
            # Note that half of the line (needs to)/(will be) calculated
            combo.conf.opt.llzero = line.lambda_-args.max
            combo.conf.opt.llfin = line.lambda_
            combos.append(combo)
            ll.append(a)

            i += 1
            ii += 1
            if ii == 1000:
                print(a99.format_progress(i, n))
                ii = 0

    # # Runs pfant
    logger.info("Running pfant's...")






    # FLAG_CLEAN = True
    # rm = RunnableManager(flag_auto_clean=FLAG_CLEAN)






    rm = pyfant.RunnableManager()
    rm.add_runnables(combos)
    app = a99.get_QApplication([])
    form = pyfant.XRunnableManager(None, rm)
    form.show()
    # it is good to start the manager as late as possible, otherwise
    # the program will hang if, for example, the form fails to be created.
    rm.start()
    app.exec_()
    rm.exit()

    # # Saves log
    LOG_FILENAME = "tune-zinf-status.log"
    with open(LOG_FILENAME, "w") as h:
      h.write(str(rm))
    logger.info("Final status saved to file '%s'" % LOG_FILENAME)

    # # Finds zinf's and creates output file
    flag_success, flag_clean = False, False
    if not rm.flag_finished:
        logger.info("Not finished")
    elif rm.num_failed > 0:
        logger.info("Number of failed tasks: %d" % rm.num_failed)
        logger.info("Please check log files inside directories of sessions that failed")
        logger.info("(session directories will not be removed).")
    else:
        logger.info("Finding zinf's, please wait...")
        # # Calculates zinf and save new atomic lines file
        n = len(combos)
        X = np.zeros((n, 3))  # [algf, kiex, zinf], ...]
        ii = 0
        print(a99.format_progress(0, n))
        cnt_min = 0
        cnt_max = 0
        for i, (a, combo) in enumerate(zip(ll, combos)):
            combo.pfant.load_result()


            # if not FLAG_CLEAN:
            #     combo.pfant.load_result()
            # else:
            #     pass  # result will be already loaded



            norm = combo.pfant.norm  # normalized spectrum
            line = a.lines[0]  # a is single-line FileAtoms object
            zinf = _get_zinf(line.lambda_, norm)
            zinf *= args.inflate
            if zinf == 0:
                # Note that some lines don't appear at all, so there is no way
                # to determine zinf
                logger.warning("zinf is ZERO: (%s) (%s) (%s)" %
                 (combo.conf.sid.dir, a.one_liner_str(), line.one_liner_str()))
            if zinf < args.min:
                zinf = args.min
                cnt_min += 1
            if zinf > args.max:
                # print " %gOLHOLHO (%s) (%s) (%s)" % \
                #  (zinf, combo.conf.session_dir, a.one_liner_str(), line.one_liner_str())
                zinf = args.max
                cnt_max += 1
            if args.ge_current and line.zinf > zinf:
                zinf = line.zinf
            line.zinf = zinf

            i += 1
            ii += 1
            if ii == 100:
                print(a99.format_progress(i, n))
                ii = 0
        print(a99.format_progress(n, n))
        logger.info("zinf's clipped to minimum: %d" % cnt_min)
        logger.info("zinf's clipped to maximum: %d" % cnt_max)

        logger.info("Creating output file...")
        file_atoms.save_as(fn_out)
        flag_clean = not args.no_clean
        flag_success = True


    # # Removes session-* directories
    if flag_clean:
        # this part removes sessionxxxx/sessionyyyy
        logger.info("Cleaning...")

        for i, combo in enumerate(combos):
            try:
                combo.sid.clean()
            except Exception as E:
                logger.info("Error cleaning session for %s: %s" % (combo.name, str(E)))
        # this part removes sessionxxxxx
        # but also removes previous sessions
        # not gonna do it dd = glob.glob("session*")

    if flag_success:
        logger.info("Successfully created file '%s' !!" % fn_out)
