# -*- coding: utf-8 -*-

# PyNX - Python tools for Nano-structures Crystallography
#   (c) 2023-present : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr

"""
Convert a series of cxi files to a single nxtomo file.
The images will be automatically cropped to the size of
the detector (= the size of the probe).h
"""

import os
import argparse
import warnings
import multiprocessing
import timeit
import psutil
import numpy as np
import nxtomo
# from tomoscan.esrf import NXtomoScan
# from tomoscan.validator import ReconstructionValidator
from ...utils import h5py

try:
    # Get the real number of processor cores available
    # os.sched_getaffinity is only available on some *nix platforms
    nproc = len(os.sched_getaffinity(0)) * psutil.cpu_count(logical=False) // psutil.cpu_count(logical=True)
except AttributeError:
    nproc = psutil.cpu_count(logical=False)


def make_parser():
    parser = argparse.ArgumentParser(prog="pynx-nfp2nxtomo",
                                     description=__doc__,
                                     formatter_class=argparse.RawTextHelpFormatter,
                                     epilog=None)
    parser.add_argument('--cxi', action='store', type=str, required=True,
                        help='input CXI pattern, e.g. "prefix%%04d_subtomo%%03d.cxi". '
                             'It is assumed that the first field is the subtomo, and '
                             'the second the projection index.')
    parser.add_argument('--nb', action='store', type=int, nargs=2, required=True,
                        help='Number of projections per subtomo and number of subtomo, '
                             'e.g. "--nb 400 4"')
    # parser.add_argument('--dxy', type=str, default=None,
    #                     help='File pattern for the subtomo files, e.g. "nfp_subtomo%04d.txt". '
    #                          '(UNUSED)')
    parser.add_argument('--extract', type=str, default=None,
                        choices=['phase_diff_complex', 'phase', 'phase_diff', 'amplitude'],
                        help='Value to extract from the CXI data, either the phase, amplitude, '
                             'or phase difference (using standard or complex differentiation. '
                             'Note that the phase gradient may already have been computed '
                             'by pynx using "cxi_output=object_phase_deriv", in which case '
                             'this can be ignored.')
    parser.add_argument('--output', action='store', type=str, required=False,
                        help='Output filename. If missing, will be deduced from the '
                             'input cxi files"')
    # parser.add_argument('--validate', type=bool, action='store_true',
    #                     help='Use this argument to check that the file has enough '
    #                          'information for processing with nabu & tomwer')
    return parser


def read1(i, isub, params):
    """Read the data for one projection and subtomo, and return the corresponding
    array, angle and positions

    :param i: the index of the projection
    :param isub: index of the subtomo
    :param params: the parameters from argparse
    :return: (d, angle, corner_positions) with the 2D object data, projection angle,
        and object corner coordinates.
    """
    with h5py.File(params.cxi % (isub, i)) as h:
        obj = h['entry_last/object/data'][()].squeeze()
        if params.extract is None:
            if np.iscomplexobj(obj):
                raise RuntimeError("CXI holds a complex object but no --extract value was given")
        elif params.extract == 'phase':
            if np.isrealobj(obj):
                warnings.warn('CXI object is real but "--extract phase" was given. Assuming the '
                              'phase was already extracted by pynx, but next time do not supply '
                              'a --extract option if already done in the CXI file.')
            else:
                obj = np.angle(obj)
        elif params.extract == 'phase_diff':
            if np.isrealobj(obj):
                raise RuntimeError('CXI holds a real object (phase ?) but '
                                   '"--extract phase_diff" was given')
            else:
                obj = np.angle(obj)
                obj -= np.roll(obj, -1, axis=-1)
        elif params.extract == 'phase_diff_complex':
            if np.isrealobj(obj):
                raise RuntimeError('CXI holds a real object (phase ?) but '
                                   '"--extract phase_diff_complex" was given')
            else:
                obj = np.angle(obj * np.roll(obj, -1, axis=-1).conj())
        elif params.extract == 'amplitude':
            if np.isrealobj(obj):
                raise RuntimeError('CXI holds a real object (phase ?) but '
                                   '"--extract amplitude" was given')
            else:
                obj = abs(obj)
        corner_positions = h['entry_last/object/corner_positions'][()]
        angle = float(h["/entry_last/process_1/configuration/tomo_metadata/angle"][()])
    return obj, angle, corner_positions


def read1_kw(kw):
    return read1(**kw)


def main():
    params = make_parser().parse_args()
    if params.output is None:
        params.output = os.path.basename(params.cxi).split('%')[0] + '.nx'

    angles = []
    vd = []
    corner_coordinates = []
    t0 = timeit.default_timer()
    print(f"Reading {params.nb[1]} x {params.nb[0]} projections...")
    if False:
        for isub in range(1, params.nb[1] + 1):
            for i in range(params.nb[0]):
                print("Reading: " + params.cxi % (isub, i))
                d, a, c = read1(i, isub, params)
                angles.append(a)
                vd.append(d)
                corner_coordinates.append(c[:2])
    else:
        vkw = [{'i': i, 'isub': isub, 'params': params}
               for isub in range(1, params.nb[1] + 1) for i in range(params.nb[0])]
        with multiprocessing.Pool(nproc) as pool:
            results = pool.imap(read1_kw, vkw)  # , chunksize=1
            for i in range(len(vkw)):
                d, a, c = results.next(timeout=20)
                angles.append(a)
                vd.append(d)
                corner_coordinates.append(c[:2])
    dt = timeit.default_timer() - t0
    print(f"Finished loading projections ["
          f"{vd[0].nbytes * params.nb[0] * params.nb[1] / dt / 1024 ** 3:.3f} GB/s]")

    with h5py.File(params.cxi % (1, 0)) as h:
        pixel_size = h['/entry_last/process_1/configuration/pixelsize'][()]
        # Original image size, for cropping
        roi = h['/entry_last/process_1/configuration/roi_actual'][()]
        ny, nx = roi[1] - roi[0], roi[3] - roi[2]
        distance = h['/entry_last/instrument_1/detector_1/distance'][()]
        nrj_j = h['/entry_last/instrument_1/source_1/energy'][()]
        sample_name = None
        if '/entry_last/sample/name' in h:
            try:
                sample_name = h['/entry_last/sample/name'][()].decode()
            except:
                # KeyError: 'Unable to synchronously open object (invalid identifier type to function)' ??
                pass

    # dxy = []
    # if params.dxy is not None:
    #     # Random motion from beamline
    #     for isub in range(1, params.nb[1] + 1):
    #         dxy.append(np.loadtxt(params.dxy % (isub + 1)))
    #     dxy = np.concatenate(dxy).astype(np.float32)

    # Actual corner positions as pynx output [dx, dy, dz], in pixel units
    dxy_pynx = np.array(corner_coordinates, dtype=np.float32) / pixel_size
    dxy_pynx = dxy_pynx * np.array([-1, 1], dtype=np.float32) + np.array((-ny / 2, -nx / 2))
    # Crop to center arrays as best as possible
    for i in range(len(vd)):
        # crop array
        dx, dy = np.round(dxy_pynx[i]).astype(int)
        vd[i] = vd[i][dy:dy + ny, dx:dx + nx]
        # keep residual dx dy
        dxy_pynx[i] -= np.array([dx, dy])

    # Sort by angular values
    idx = np.argsort(angles)
    vd = [vd[i] for i in idx]
    angles = [angles[i] for i in idx]
    dxy_pynx = np.take(dxy_pynx, idx, axis=0)

    vd = np.array(vd, dtype=np.float32)
    angles = np.array(angles, dtype=np.float32)

    # Make NXtomo file
    nxt = nxtomo.NXtomo()

    nxt.instrument.detector.image_key_control = np.concatenate([
        [nxtomo.nxobject.nxdetector.ImageKey.PROJECTION] * np.product(params.nb)])

    nxt.instrument.detector.data = vd
    nxt.sample.rotation_angle = angles
    if sample_name is not None:
        nxt.sample.name = sample_name

    # Energy and distance metadata - not useful for volume reconstruction
    nxt.energy = nrj_j / 1.602e-16  # Convert to keV
    nxt.instrument.detector.distance = distance

    nxt.sample.x_translation = dxy_pynx[:, 0] * pixel_size
    nxt.sample.y_translation = dxy_pynx[:, 1] * pixel_size

    nxt.instrument.detector.x_pixel_size = \
        nxt.instrument.detector.y_pixel_size = pixel_size

    print(f"Saving to: {params.output}")
    nxt.save(file_path=params.output, data_path="entry", overwrite=True)
    del nxt
    # Fix .nx file (nxtomo 1.0.4, https://gitlab.esrf.fr/tomotools/nxtomo/-/issues/1)
    with h5py.File(params.output, 'a') as h:
        for e in ['beam', 'data', 'instrument', 'sample']:
            if e in h and f'entry/{e}' not in h:
                print(f"Fix nxtomo format: moving {e} to entry/{e}")
                h.move(e, f"entry/{e}")

    dxy_out = os.path.splitext(params.output)[0] + "_dxy.txt"
    print(f"Saving residual displacements (x,y): {dxy_out}")
    np.savetxt(dxy_out, dxy_pynx)

    # if params.validate:
    #     scan = NXtomoScan(params.output, entry="entry")
    #     validator = ReconstructionValidator(scan, check_phase_retrieval=False, check_values=True)
    #     assert validator.is_valid()


if __name__ == '__main__':
    main()
