#!/usr/bin/env python3

import argparse
import atexit
import contextlib
import os
import os.path

import importlib_resources
import memoir
import memoir.check_line

parser = argparse.ArgumentParser(
    description="MEMOIR: Molecular EMissiOn IdentifieR. Searches for lines in a spectrum and assignes IDs by comparing their frequencies with those of known lines. Returns a file text with the ID and the peak values (frequency, velocity, intensity/flux) of the detected lines.",
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
    "-i",
    "--fits_image",
    type=str,
    nargs="+",
    help="Fits image to take spectrum on mask <<fits_mask>>.",
    default=[""],
)
parser.add_argument(
    "-m",
    "--fits_mask",
    type=str,
    nargs="+",
    help="Mask used to take spectrum in image <<fits_image>>.",
    default=[""],
)
parser.add_argument(
    "-t",
    "--frequency_tolerance",
    type=float,
    help="Frequency tolerance (in the frequency units of the input spectrum). If a detected line has a known line with a frequency separation less than <<frequency_tolerance>>, it will assume its ID, if not, will remain unidentified (U).",
    default=0.01,
)
parser.add_argument(
    "-w",
    "--line_width",
    type=float,
    help="Line width (in the velocity units of the input spectrum). Local peaks closer than <<line_width>> will be considered as belonging to the same line, and only the one with higher flux will be returned.",
    default=20,
)
parser.add_argument(
    "--snr",
    type=float,
    help="Signal-to-noise ratio. Only peaks with flux higher than <<SNR>> will be returned",
    default=5,
)
parser.add_argument(
    "-s",
    "--spectrum_file_name",
    type=str,
    nargs="+",
    help="Name of the file (including extension) with the spectrum to analyse. If not set (default behaviour) will use all the files in the working directory end-named 'spectrum.txt'.",
    default=[""],
)

# The default file used by --known_lines_file_name; included in the package.
KNOWN_LINES_DEFAULT_FILE = "data/allmols_combined_transitions.csv"

parser.add_argument(
    "-l",
    "--known_lines_file_name",
    type=str,
    help="Name of the file (including extension) with the known lines. If not specified, use memoir's built-in one.",
    default=None,
)
parser.add_argument(
    "-o",
    "--output",
    type=str,
    help="Name of the output folder.",
    default="output",
)
parser.add_argument(
    "--vlsr",
    type=float,
    help="Source radial velocity (local standard of rest) in km/s.",
    default=0,
)
parser.add_argument(
    "-v",
    "--version",
    action="version",
    version=memoir.__version__,
)

if __name__ == "__main__":
    args = parser.parse_args()

    output = memoir.check_line.output_folder(args.output)

    for image_name, mask_name in zip(args.fits_image, args.fits_mask):
        if image_name != "":
            print(f"For image {image_name} and mask {mask_name}:")
            if mask_name != "":
                region_pix, mask = memoir.check_line.make_union_mask(mask_name)
                if region_pix != False:
                    image = memoir.check_line.Image(image_name, mask)
                    (
                        chan,
                        pix,
                        freq,
                        velocity,
                        spectrum,
                        beam_area,
                        beam_pix,
                    ) = image.get_spectrum(region_pix)
                    image.write_spectrum(
                        chan, pix, freq, velocity, spectrum, beam_area, beam_pix
                    )
            else:
                print("You need to provide a fits for the mask.")

    if args.spectrum_file_name == [""]:
        spectrum_file_names = memoir.check_line.spectrum_exist()
    else:
        spectrum_file_names = args.spectrum_file_name

    # Fall back to the built-in known lines file if the user didn't specify any.
    # https://importlib-resources.readthedocs.io/en/latest/migration.html#pkg-resources-resource-filename.
    if args.known_lines_file_name is None:
        file_manager = contextlib.ExitStack()
        atexit.register(file_manager.close)
        ref = importlib_resources.files(memoir) / KNOWN_LINES_DEFAULT_FILE
        known_lines_path = file_manager.enter_context(importlib_resources.as_file(ref))

    else:
        known_lines_path = args.known_lines_file_name

    if spectrum_file_names:
        for file_name in spectrum_file_names:
            print(f"For spectrum {file_name}:")
            spectrum = memoir.check_line.Spectrum(file_name)
            lines = spectrum.potential_lines(known_lines_path, args.vlsr)
            peak_frequencies, peak_velocities, peak_fluxes = spectrum.find_lines(
                args.snr, args.line_width
            )
            actual_lines = memoir.check_line.match_lines(
                lines, peak_frequencies, args.frequency_tolerance
            )
            spectrum.write_parameters(
                actual_lines,
                peak_frequencies,
                peak_velocities,
                peak_fluxes,
                output,
                file_name,
            )
            spectrum.make_plot(
                file_name, actual_lines, peak_frequencies, peak_fluxes, output
            )
