import contextlib
import posixpath
from glob import glob
from os import getpid, path, rename
from os import remove as remove_file
from time import sleep
from traceback import format_exc

import numpy as np
from nabu.io.writer import h5_write_object, merge_hdf5_files
from silx.io.dictdump import dicttoh5, h5todict
from tomoscan.io import HDF5File

from . import version as integrator_version
from .slurm import submit_bash_script
from .utils import (
    ai_units,
    files_relative_paths,
    get_identifier,
    make_azims,
    retry_n_times,
)


def collect_results_from_files(out_files):
    intensities = []
    errors = []
    for out_file in sorted(out_files):
        results = h5todict(out_file, path="azimuthal integration/results")
        intensities.append(results["intensities"])
        if "errors" in results:
            errors.append(results["errors"])
    intensities = np.vstack(intensities)
    if errors != []:
        errors = np.vstack(errors)
    else:
        errors = None
    return intensities, errors


def create_output_file_id15a(
    scan,
    partial_output_files,
    ai_config,
    output_file,
    delete_partial_files=True,
    use_virtual_dataset=False,
    extra_metadata=None,
):
    """
    Create the final output file, with ESRF ID15A layout.

    Parameters
    ----------
    scan: HDF5Dataset object
        Data structure with information on scan.
    partial_output_files: list of str
        List of paths of partial output files
    ai_config: AIConfiguration
        Data structure with azimuthal integration configuration
    output_file: str or dict
        Path to the final output file. Will be overwritten if already exists.
        If a dict is provided, the output file is output_file[scan.dataset_hdf5_url.path()]
    """
    # Load some processing information from a partial output file
    processing_cfg = h5todict(
        partial_output_files[0],
        path="azimuthal integration/configuration",
    )

    #
    # "info"
    #
    # doing it here to be sure it won't import the whole world and mess with OpenCL/CPU affinity
    from pyFAI import version as pyfai_version

    info_group = {
        "calibration_file": path.abspath(ai_config.poni_file),
        "datafiles": sorted(scan.get_virtual_sources().keys()),
        "end_time": scan.end_time,
        "exposure time": scan.exposure_time,
        ("exposure time", "unit"): "second",
        "geometric_corrections": np.zeros(
            1,
            dtype="f",
        ),  # to be supported in the future
        "mask_file": ai_config.mask_file,
        # "pixel1" is used for both pixel dims as in hdintrange
        "pixel size (mm)": (
            processing_cfg["pixel_size"][0] * 1e3,
            processing_cfg["pixel_size"][0] * 1e3,
        ),
        ("pixel size (mm)", "unit"): "millimeter",
        "pyfai_method": processing_cfg["pyfai_method"],
        "S-D distance (mm)": processing_cfg["distance"] * 1e3,
        ("S-D distance (mm)", "unit"): "millimeter",
        "start_time": scan.start_time,
        "title": scan.title,
        "wavelength (A)": processing_cfg["wavelength"] * 1e10,
        ("wavelength (A)", "unit"): "angstrom",
        "pyFAI_version": pyfai_version,
    }

    #
    # "positions"
    #
    positions_group = {"exposure_time": scan.exposure_time}
    for field_name, field_val in scan.metadata.items():
        if field_name.startswith(("measurement", "instrument/positioners")):
            title = posixpath.basename(field_name)
            positions_group[title] = field_val

    #
    # Integration results (eg. azim_01)
    #
    azimuthal_ranges = make_azims(*ai_config.azimuthal_range)
    integration_results_group = {}
    n_azim_slices = int(ai_config.azimuthal_range[-1])
    # no error propagation when doing median (?)
    do_error_propagation = (
        ai_config.error_model is not None and ai_config.trim_method != "median"
    )
    if not (use_virtual_dataset):
        intensities, errors = collect_results_from_files(partial_output_files)
        for i in range(n_azim_slices):
            azim_results = {
                "intensities": intensities[i::n_azim_slices].reshape(
                    -1,
                    ai_config.n_pts,
                ),
                "azimuthal_range": azimuthal_ranges[i],
            }
            if do_error_propagation:
                azim_results["error"] = errors[i::n_azim_slices].reshape(
                    -1,
                    ai_config.n_pts,
                )
            integration_results_group["azim_%02d" % (i + 1)] = azim_results

    #
    # Finalize structure to dump + Unit (eg. r_mm)
    #
    dict_to_dump = {
        "info": info_group,
        "positions": positions_group,
        ai_config.unit: processing_cfg["radial_positions"],
        (ai_config.unit, "radial"): 1,
        **integration_results_group,
    }

    if isinstance(output_file, dict):
        output_file = output_file[scan.dataset_hdf5_url.path()]

    dicttoh5(dict_to_dump, output_file)

    if use_virtual_dataset:
        entry = "/"
        process_name = "azimuthal integration"
        for azim_idx in range(n_azim_slices):
            local_files = files_relative_paths(partial_output_files, output_file)
            results_types = ["intensities"]
            if do_error_propagation:
                results_types.append("error")
            h5_paths = [
                posixpath.join(
                    entry,
                    process_name,
                    "results",
                    str("azim_%02d" % (azim_idx + 1)),
                    what,
                )
                for what in results_types
            ]
            for h5_path in h5_paths:
                merge_hdf5_files(
                    local_files,
                    h5_path,
                    output_file,
                    "/",
                    output_entry=entry,
                    config=None,
                    base_dir=path.dirname(output_file),
                    overwrite=True,
                    # Not very elegant. 'data_name' expects something like "intensities",
                    # but we pass something like 'azim_01/intensities'.
                    # This works (h5py seems to create groups recursively) but might be brittle.
                    data_name=posixpath.join(
                        "/",
                        *h5_path.split(posixpath.sep)[-2:],
                    ),  # /azim_01/intensities
                )
            with HDF5File(output_file, "a") as f:
                f[str("azim_%02d" % (azim_idx + 1)) + "/azimuthal_range"] = (
                    azimuthal_ranges[azim_idx]
                )
                for what in ["results", "sequence_index"]:
                    if what in f:
                        del f[what]
                h5_write_object(f, "program", "integrator", overwrite=True)
                h5_write_object(f, "version", integrator_version, overwrite=True)
    #
    # Extra metadata
    #
    extra_metadata = extra_metadata or {}
    if len(extra_metadata) > 0:
        with HDF5File(output_file, "a") as f:
            f.create_group("configuration")
            dicttoh5(extra_metadata, f["configuration"], update_mode="add")

    #
    # Clean-up partial files
    #

    if delete_partial_files and not (use_virtual_dataset):
        for fname in partial_output_files:
            try:
                remove_file(fname)
            except PermissionError:
                contextlib.suppress(PermissionError)  # TODO log this somewhere

    return output_file


def create_output_file(
    scan,
    partial_output_files,
    ai_config,
    output_file,
    layout,
    extra_metadata=None,
):
    # Note: will edit for id11 based on feedback from initial testing
    if layout.lower() == "id15a" or layout.lower() == "id11":
        return create_output_file_id15a(
            scan,
            partial_output_files,
            ai_config,
            output_file,
            delete_partial_files=False,
            use_virtual_dataset=True,
            extra_metadata=extra_metadata,
        )
    raise NotImplementedError


def watch_scans_and_create_output_files(
    scans_info,
    ai_config,
    layout,
    period=0.5,
    logger=None,
    extra_metadata=None,
    repack_output=False,
    repack_output_options=None,
):

    def log_message(msg):
        if logger is not None:
            logger.info(msg)
        else:
            print(msg)

    output_created = dict.fromkeys(scans_info.keys(), False)
    try:
        while not (all(output_created.values())):
            for scan_url, scan_info in scans_info.items():
                partial_out_files = scan_info["partial_output_files"]
                output_exists = [path.isfile(fname) for fname in partial_out_files]
                if all(output_exists) and not (output_created.get(scan_url, False)):
                    log_message(
                        "[%d] Creating output file for %s" % (getpid(), scan_url),
                    )
                    out_file = retry_n_times(
                        3,
                        logger,
                        create_output_file,
                        scan_info["scan_structure"],
                        partial_out_files,
                        ai_config,
                        scan_info["output_file"],
                        layout,
                        extra_metadata=extra_metadata,
                    )
                    log_message("[%d] Created %s" % (getpid(), out_file))
                    output_created[scan_url] = True

                    if repack_output:
                        submit_repack_output(
                            out_file,
                            average_xy=ai_config.average_xy,
                            **(repack_output_options or {}),
                        )

            sleep(period)
    except Exception:
        log_message("Something went wrong! Aborting final output files creation")
        log_message(format_exc())
    finally:
        pass


def submit_repack_output(
    out_file,
    average_xy=None,
    python_env=None,
    slurm_resources=None,
    log_dir=None,
    workspace_dir=None,
):
    if python_env is None or slurm_resources is None:
        raise ValueError

    script_fname = path.join(
        workspace_dir,
        "%s_%s_repack.sh" % (get_identifier(use_date=False), path.basename(out_file)),
    )
    content = f"""\
#!/bin/bash -l

{python_env}

python -c "from integrator.writer import repack_output_file; repack_output_file('{out_file}', average_xy='{average_xy}')"
"""
    with open(script_fname, "w") as fdesc:
        fdesc.write(content)

    submit_bash_script(
        slurm_resources,
        script_fname,
        log_dir=log_dir,
        log_basename="repack",
    )


def _reduce_integrated_stack(out_file_content, average_xy, out_file_name):
    if average_xy in [None, "none", "None"]:
        return
    reduction_function = np.median if average_xy.lower() == "median" else np.mean

    # Get the "x axis" (units)
    # ai_unit = out_file_content["configuration"]["configuration"]["azimuthal integration"]["unit"][()] # [()] because asarray=True in h5todict()
    # x_axis = out_file_content.get(ai_unit, None)
    x_axis = next(
        (
            out_file_content.get(u, None)
            for u in set(ai_units.values())
            if out_file_content.get(u, None) is not None
        ),
        None,
    )
    if x_axis is None:
        raise ValueError("Could not get units")

    n_azim_slices = len(
        [
            key
            for key in out_file_content
            if not (isinstance(key, tuple)) and key.startswith("azim_")
        ],
    )

    # for each "azim_xx"
    i_azim_slice = 1
    for key, val in out_file_content.items():
        if isinstance(key, tuple) or not (key.startswith("azim_")):
            continue
        print("Performing %s_xy on %s" % (average_xy, key))
        red_int_xy = reduction_function(val["intensities"], axis=0)
        val["intensities_%s_xy" % average_xy] = red_int_xy
        to_save = [x_axis, red_int_xy]
        if val.get("error", None) is not None:
            red_err_xy = reduction_function(val["error"], axis=0)
            val["error_%s_xy" % average_xy] = red_err_xy
            to_save.append(red_err_xy)
        txtfile_name = out_file_name.replace(".h5", ".xye")
        if n_azim_slices > 1:
            txtfile_name = txtfile_name.replace("azint", "azint%02d" % i_azim_slice)
        np.savetxt(txtfile_name, np.transpose(to_save), fmt="%.5f")
        i_azim_slice += 1


def repack_output_file(output_file, average_xy=None):
    """
    Transform a output file containing virtual dataset, to a output file containing all the data.
    This was written for ID15A where using too many files (partial results) was problematic for data archival.
    """
    output_file_repacked = output_file + ".repacked"
    print("Repacking file %s ..." % (path.basename(output_file)), end="", flush=True)
    # Loads everything in memory. Should not be a problem for diffraction files.
    out_file_content = h5todict(output_file, include_attributes=True)

    try:
        _reduce_integrated_stack(out_file_content, average_xy, output_file)
    except Exception as exc:
        print("Could not perform average_xy: %s" % str(exc))
        print(format_exc())

    dicttoh5(out_file_content, output_file_repacked)

    #
    # Cleanu-up
    #
    rename(output_file, output_file + ".virtual")

    def try_remove(fname):
        try:
            remove_file(fname)
        except Exception as exc:
            print("Cannot remove file %s: %s" % (fname, str(exc)))

    h5_files = glob(path.splitext(output_file)[0] + "*.h5")
    for fname in h5_files:
        if fname == output_file:
            continue
        try_remove(fname)

    rename(output_file_repacked, output_file)
    try_remove(output_file + ".virtual")

    print("done !", flush=True)
