#! /opt/local/bin/python
# -*- coding: utf-8 -*-
import os.path
# PyNX - Python tools for Nano-structures Crystallography
#   (c) 2016-present : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr

import timeit
import numpy as np
from scipy.ndimage import fourier_shift
from ...ptycho import simulation, Ptycho, PtychoData, Calc2Obs, AP, DM, ScaleObjProbe, save_ptycho_data_cxi
from ...mpi import MPI

if MPI is not None:
    from ..mpi import PtychoSplit, ShowObjProbe
from .runner import PtychoRunner, PtychoRunnerScan, PtychoRunnerException, default_params as params0

helptext_epilog = " Examples:\n\n" \
                  " * ``pynx-ptycho-simulation``\n" \
                  " * ``pynx-ptycho-simulation --siemens``\n" \
                  " * ``pynx-ptycho-simulation --simul_object obj.npz --simul_probe probe.npz``\n" \
                  " * ``pynx-ptycho-simulation --near_field --direct_beam``\n" \
                  " * ``pynx-ptycho-simulation --frame_nb 128 --frame_size 256 " \
                  "--algorithm analysis,ML**100,DM**200,nbprobe=2,probe=1 " \
                  "--save_plot --live_plot``"

# NOTE: scripts to test absolute orientation:
# mpiexec -n 4 pynx-simulationpty.py logo asym liveplot saveplot mpi=split
#  frame_nb=400 frame_size=256 algorithm=analysis,AP**20,positions=1,AP**400,probe=1
# pynx-simulationpty.py logo asym liveplot saveplot frame_nb=200 frame_size=256
#   algorithm=analysis,AP**20,positions=1,AP**400,probe=1
#
# In both cases compare:
# - positions plot with lone illumination position
# - absolute orientation of object and probe in saved plot, silx view of object, illumination and probe


params_beamline = {
    'algorithm': 'ML**100,AP**200,DM**200,probe=1',
    'asym': False,
    'autocenter': False,
    'defocus': 100e-6,
    'detectordistance': 1,
    'direct_beam': False,
    'frame_nb': 128,
    'frame_size': 256,
    'instrument': 'simulation',
    'logo': False,
    'maxsize': None,
    'nrj': 8,
    'photons_per_frame': 1e8,
    'pixelsize': 55e-6,
    'probe': 'auto',
    'roi': 'full',
    'saveprefix': 'none',
    'siemens': False,
    'simul_background': 0,
}

default_params = params0.copy()
for k, v in params_beamline.items():
    default_params[k] = v


class PtychoRunnerScanSimul(PtychoRunnerScan):

    def load_scan(self):
        nb = self.params['frame_nb']
        n = self.params['frame_size']
        pixel_size_detector = self.params['pixelsize']
        wavelength = 12.3984e-10 / self.params['nrj']
        detector_distance = self.params['detectordistance']
        # 50 scan positions correspond to 4 turns, 78 to 5 turns, 113 to 6 turns
        scan_info = {'type': 'raster' if self.params['raster_scan'] else 'spiral',
                     'scan_step_pix': self.params['step_size_window'] * n, 'n_scans': nb}

        s = simulation.Simulation(obj_info=None, probe_info=None, scan_info=scan_info, data_info=None,
                                  verbose=self.mpi_master)
        s.make_scan()

        if self.params['near_field']:
            pixel_size_object = pixel_size_detector
        else:
            pixel_size_object = wavelength * detector_distance / pixel_size_detector / n
        self.x, self.y = s.scan.values[0] * pixel_size_object, s.scan.values[1] * pixel_size_object
        self.imgn = np.arange(nb, dtype=np.int32)

        if self.params['moduloframe'] is not None:
            n1, n2 = self.params['moduloframe']
            idx = np.where(self.imgn % n1 == n2)[0]
            self.imgn = self.imgn.take(idx)
            self.x = self.x.take(idx)
            self.y = self.y.take(idx)

        if self.params['maxframe'] is not None:
            N = self.params['maxframe']
            if len(self.imgn) > N:
                print("MAXFRAME: only using first %d frames" % (N))
                self.imgn = self.imgn[:N]
                self.x = self.x[:N]
                self.y = self.y[:N]

        if self.params['asym']:
            idx = np.where(np.logical_or(self.x <= self.x.mean(), self.y >= self.y.mean()))[0]
            # Keep one illumination farthest in the corner
            idx = np.append(idx, [np.argmax(self.x - self.y)])
            self.x, self.y = np.take(self.x, idx), np.take(self.y, idx)
            # Append existing positions to keep total number of frames
            dn = nb - len(self.x)
            idx = np.random.randint(0, len(self.x), dn)
            self.x = np.append(self.x, self.x[idx])
            self.y = np.append(self.y, self.y[idx])

            # Also off-center positions to check absolute position
            self.x += 2 * np.abs(self.x).max()
            self.y += 3 * np.abs(self.y).max()
        if self.params['direct_beam']:
            # The last frame will be with the direct beam.
            self.x[-1] = 1e12
            self.y[-1] = 1e12

    def load_data(self):
        t0 = timeit.default_timer()
        n = self.params['frame_size']
        pixel_size_detector = self.params['pixelsize']
        wavelength = 12.3984e-10 / self.params['nrj']
        detector_distance = self.params['detectordistance']
        obj0 = None  # Object loaded from file ?
        if self.params['siemens']:
            obj_info = {'type': 'siemens', 'alpha_win': .2}
        elif self.params['logo']:
            obj_info = {'type': 'logo', 'phase_stretch': 1, 'ampl_range': (0.9, 0.1)}
        elif self.params['simul_object'] is not None:
            if 'split' in self.params['mpi']:
                raise PtychoRunnerException('--simul_object is not compatible with --mpi=split')
            if not os.path.isfile(self.params['simul_object']):
                raise PtychoRunnerException(f"--simul_object {self.params['simul_object']}: "
                                            f"file does not exist")
            if os.path.splitext(self.params['simul_object'])[-1] == ".npy":
                obj0 = np.load(self.params['simul_object']).squeeze()
            elif os.path.splitext(self.params['simul_object'])[-1] == ".npz":
                # Take the first array with a size>1000
                for k, v in np.load(self.params['simul_object']).items():
                    if v.size > 1000:
                        obj0 = v.squeeze()
                        break
            else:
                raise PtychoRunnerException(f"--simul_object {self.params['simul_object']}: "
                                            f"need an .npy or .npz file")
            obj_info = None
        else:
            obj_info = {'type': 'phase_ampl', 'phase_stretch': np.pi / 2, 'alpha_win': .2}

        near_field = self.params['near_field']
        if near_field:
            pixel_size_object = pixel_size_detector
        else:
            # The probe is calculated so that it will defocus to about 40% of the object frame size,
            # and the step size is adapted accordingly
            pixel_size_object = wavelength * detector_distance / (n * pixel_size_detector)
        aperture = 400e-6
        focal_length = 0.1
        if self.params['probe_size_window'] is None:
            defocus = 0.4 * pixel_size_object * n / aperture * focal_length
        else:
            defocus = self.params['probe_size_window'] * pixel_size_object * n / aperture * focal_length
        probe0 = None
        if self.params['simul_probe'] is not None:
            if not os.path.isfile(self.params['simul_probe']):
                raise PtychoRunnerException(f"--simul_probe {self.params['simul_probe']}: file does not exist")
            if os.path.splitext(self.params['simul_probe'])[-1] == ".npy":
                probe0 = np.load(self.params['simul_probe'])
            elif os.path.splitext(self.params['simul_probe'])[-1] == ".npz":
                # Take the first array with a size>1000
                for k, v in np.load(self.params['simul_probe']).items():
                    if v.size > 1000:
                        probe0 = v
                        break
            probe0 = probe0.squeeze()  # Just in case it's 3D
            if probe0.shape != (n, n):
                raise PtychoRunnerException(f"--simul_probe: shape {probe0.shape} != ({n}, {n})")
            probe_info = None
        elif near_field:
            ap = pixel_size_detector * n * 0.7
            probe_info = {'type': 'near_field', 'aperture': (ap, ap), 'defocus': 0.3, 'shape': (n, n)}
        else:
            probe_info = {'type': 'focus', 'aperture': (aperture, aperture), 'focal_length': focal_length,
                          'defocus': defocus, 'shape': (n, n)}

        if self.params['probe'] == 'auto' and not near_field:
            # Put some randomness in the starting probe
            r = np.random.uniform(0.9, 1.1, 4)
            self.params['probe'] = 'focus,%ex%e,%e' % (aperture * r[0], aperture * r[1], focal_length * r[2])
            self.params['defocus'] = defocus * r[3]

        # 50 scan positions correspond to 4 turns, 78 to 5 turns, 113 to 6 turns
        scan_info = {'type': 'custom', 'x': self.x / pixel_size_object, 'y': self.y / pixel_size_object}
        data_info = {'num_phot_max': 1e6, 'bg': 0, 'wavelength': wavelength, 'detector_distance': detector_distance,
                     'detector_pixel_size': pixel_size_detector, 'noise': 'poisson', 'near_field': near_field}

        s = simulation.Simulation(obj_info=obj_info, probe_info=probe_info, scan_info=scan_info,
                                  data_info=data_info, verbose=self.mpi_master)
        if probe0 is None:
            s.make_probe()
            probe0 = s.probe.values

        t1 = timeit.default_timer()
        if self.timings is not None:
            dt = t1 - t0
            if "load_data_simul_probe" in self.timings:
                self.timings["load_data_simul_probe"] += dt
            else:
                self.timings["load_data_simul_probe"] = dt

        if self.params['multiscan_reuse_ptycho']:
            # Randomly displace the scan positions on the object by 10%
            # See corresponding shift of the object array below
            if self.mpi_master:
                a = np.random.uniform(0, 2 * np.pi)
                dx, dy = n / 10 * np.cos(a), n / 10 * np.sin(a)
            if MPI is not None:
                if self.mpi_master:
                    dx, dy = self.mpic.bcast((dx, dy), root=0)
                else:
                    dx, dy = self.mpic.bcast(None, root=0)
            self.x += dx * pixel_size_object
            self.y += dy * pixel_size_object
            self.print(f"multiscan_reuse_ptycho: <x>={self.x.mean() / pixel_size_object} "
                       f"<y>={self.y.mean() / pixel_size_object} pixels")

        data = PtychoData(iobs=np.ones((len(self.x), n, n), dtype=np.float32), positions=(self.x, self.y),
                          detector_distance=detector_distance, mask=None, pixel_size_detector=pixel_size_detector,
                          wavelength=wavelength, near_field=near_field)

        if 'split' not in self.params['mpi']:
            if obj0 is None:
                s.make_obj()
                s.make_obj_true(data.get_required_obj_shape(margin=2))
                o = s.obj.values
            else:
                o = obj0
                nyo, nxo = data.get_required_obj_shape(margin=2)
                if obj0.shape != (nyo, nxo):
                    print(f"Adapting object shape {o.shape} to the desired one ({nyo}, {nxo}) "
                          f"by cropping or zero-padding)")
                if o.shape[0] < nyo:
                    dn = nyo - o.shape[0]
                    o = np.pad(o, (dn // 2, dn - dn // 2, 0, 0))
                elif o.shape[0] > nyo:
                    dn = o.shape[0] - nyo
                    o = o[dn // 2:-(dn - dn // 2)]
                if o.shape[1] < nxo:
                    dn = nxo - o.shape[1]
                    o = np.pad(o, (0, 0, dn // 2, dn - dn // 2))
                elif o.shape[1] > nxo:
                    dn = o.shape[1] - nxo
                    o = o[:, dn // 2:-(dn - dn // 2)]

            if self.params['multiscan_reuse_ptycho']:
                # Shift object array
                o = np.pad(s.obj.values, n // 10, constant_values=0)
                o = np.fft.ifftn(fourier_shift(np.fft.fftn(o), (-dy, -dx)))
            print(probe0.shape, o.shape)
            p = Ptycho(probe=probe0, obj=o, data=data, background=None)
        else:
            p = PtychoSplit(probe=probe0, obj=None, data=data, background=None,
                            mpi_neighbour_xy=self.mpi_neighbour_xy)
            if self.mpi_master:
                s.make_obj()
                # The object shape actually needed is computed in PtychoSplit.init_mpi_obj()
                s.make_obj_true(p.mpi_obj.shape[-2:])
            p.set_mpi_obj(s.obj.values)

        if 'asym' in self.params:
            if self.params['asym']:
                # Use asymmetry to check the absolute orientation of the reconstruction
                x, y = p.get_probe_coord()
                y, x = np.meshgrid(y, x, indexing='ij')
                px, py = data.pixel_size_object()
                tmp = np.logical_or(x <= 0, y >= 0).astype(np.int8)
                probe0 *= tmp + (1 - tmp) * np.exp(np.maximum(y, -x) / (px * len(x) / 32))
                p.set_probe(probe0)

        t2 = timeit.default_timer()
        if self.timings is not None:
            dt = t2 - t1
            if "load_data_simul_obj" in self.timings:
                self.timings["load_data_simul_obj"] += dt
            else:
                self.timings["load_data_simul_obj"] = dt

        p = Calc2Obs(nb_photons_per_frame=1e8) * p
        iobs = np.fft.fftshift(p.data.iobs, axes=(1, 2))

        dark = None
        if self.params['simul_background'] > 0:
            x, y = p.get_probe_coord()
            dx, dy = np.meshgrid((x - x.mean()) / (x.max() - x.min()), (y - y.mean()) / (y.max() - y.min()))
            tmp = np.exp(-5 * (dx ** 2 + dy ** 2))
            dark = iobs.sum() / (tmp.sum() * len(iobs)) * self.params['simul_background'] * tmp
            iobs += dark

        t3 = timeit.default_timer()
        if self.timings is not None:
            dt = t3 - t2
            if "load_data_simul_iobs" in self.timings:
                self.timings["load_data_simul_iobs"] += dt
            else:
                self.timings["load_data_simul_iobs"] = dt

        if 'split' in self.params['mpi']:
            vnb = self.mpic.gather(len(iobs), root=0)
            viobs_sum = self.mpic.gather(iobs.sum(), root=0)
            if self.mpi_master:
                print(vnb, viobs_sum)
                scale = self.mpic.bcast(np.array(viobs_sum).sum() / np.array(vnb).sum(), root=0)
            else:
                scale = self.mpic.bcast(None, root=0)
        else:
            scale = iobs.sum() / len(iobs)

        # 1e8 photons on average
        self.raw_data = np.random.poisson(iobs / scale * 1e8)
        if dark is not None and not self.params['unknown_background']:
            dark = dark * (1e8 / scale)
            print("Including dark with <dark>=%6.0f ph/pixel, <Iobs>=%6.0f" % (dark.mean(), self.raw_data.mean()))
            self.dark = dark

        if self.params['random_position_errors'] > 0:
            dx = self.params['frame_size'] * self.params['random_position_errors'] * 0.5
            self.x += (np.random.uniform(-dx, dx, self.params['frame_nb']) * pixel_size_object).astype(np.float32)
            self.y += (np.random.uniform(-dx, dx, self.params['frame_nb']) * pixel_size_object).astype(np.float32)

        t4 = timeit.default_timer()
        if self.timings is not None:
            dt = t4 - t3
            if "load_data_simul_noise" in self.timings:
                self.timings["load_data_simul_noise"] += dt
            else:
                self.timings["load_data_simul_noise"] = dt

        if self.params['direct_beam']:
            # Extract the last simulated frame as reference data
            self.x = self.x[:-1]
            self.y = self.y[:-1]
            self.data_ref = self.raw_data[-1]
            self.raw_data = self.raw_data[:-1]

        # if self.mpi_master:
        #     print("RunnerScanSimul: elapsed time for load_data()= %5.2fs" % (time.time()-t0))
        self.load_data_post_process()


class PtychoRunnerSimul(PtychoRunner):
    """
    Class to process a series of scans with a series of algorithms, given from the command-line
    """

    def __init__(self, argv, params, *args, **kwargs):
        super(PtychoRunnerSimul, self).__init__(argv, default_params if params is None else params)
        self.PtychoRunnerScan = PtychoRunnerScanSimul

    @classmethod
    def make_parser(cls, default_par, description=None, script_name="pynx-ptycho-simulation", epilog=None):
        if epilog is None:
            epilog = helptext_epilog
        if description is None:
            description = "Script to perform a ptychography analysis on a simulated dataset"
        p = default_par
        parser = super().make_parser(default_par, script_name, description, epilog)
        grp = parser.add_argument_group("Simulation parameters")
        grp.add_argument(
            '--frame_nb', '--framenb', '--nb_frame', '--nbframe', type=int, default=p['frame_nb'],
            help='number of frames for the simulation')
        grp.add_argument('--frame_size', '--framesize', type=int, default=p['frame_size'],
                         help='size of the frame for the simulation')
        grp.add_argument('--siemens', '--siemens_star', action='store_true', default=p['siemens'],
                         help='use a siemens star for the simulation, instead of two different '
                              'images for the amplitude and phase')
        grp.add_argument('--logo', action='store_true', default=p['logo'],
                         help='use the PyNX logo for the simulation')
        grp.add_argument('--asym', action='store_true', default=p['asym'],
                         help='if given on the command-line, all but 1 scan positions'
                              'in the lower right quadrant (x>x.mean(), y<y.mean()) will '
                              'be removed. Positions are added in the other parts to '
                              'keep the number of frames as expected. The probe used '
                              'for simulation will also have the lower right quadrant '
                              'set to zero.')
        grp.add_argument('--simul_background', action='store', type=float,
                         default=p['simul_background'],
                         help='add incoherent (gaussian-shaped) background, with an '
                              'integrated intensity equal to the given factor multiplied '
                              'by the simulated intensity. Reasonable values up to 0.2. '
                              'By default the background will be transferred to the '
                              'optimisation algorithm, unless ``unknown_background`` '
                              'is used')
        grp.add_argument('--unknown_background', action='store_true',
                         default=False,
                         help='Use this not to give the initial background to '
                              'the optimisation algorithm, when ``--simul_background`` '
                              'is used. You should then activate the background optmisation '
                              'in the algorithm chain to recover it.')
        grp.add_argument('--direct_beam', '--directbeam', action='store_true',
                         default=p['direct_beam'],
                         help='if this option is used, one frame will be with '
                              'the direct beam (no sample), allowing '
                              'an absolute scale of the object and probe')
        grp.add_argument('--probe_size_window', type=float,
                         help="Probe size, relative to the Fourier window. By default it "
                              "is 0.4 for far-field, and this will be ignored for near-field "
                              "where the probe fills the whole field of view. This can be used "
                              "in combination to ``--step_size_window`` to modify the overlap "
                              "and the extent of the scanned area.")
        grp.add_argument('--step_size_window', type=float, default=0.1,
                         help="Step size, relative to the frame. The default is "
                              "0.1. This can be used  in combination to ``--probe_size_window`` "
                              "to modify the overlap and the extent of the scanned area.")
        grp.add_argument('--random_position_errors', type=float, default=0,
                         help="Random positional errors, given as a fraction of the "
                              "frame size in pixels.")
        grp.add_argument('--photons_per_frame', type=float, default=p['photons_per_frame'],
                         help="Average number of photons per frame")
        grp.add_argument('--raster_scan', action='store_true',
                         default=False,
                         help='Use a raster scan, instead of a spiral. Raster scan are '
                              'strongly discouraged since they can lead to correlation '
                              'between the object & probe structure and the scan positions')
        grp.add_argument('--simul_object', type=str,
                         help="Load an object from a file to simulate the observed data. "
                              "Accepted formats are numpy files (npy or npz), which should "
                              "contain only one 2D array. The object will be zero-padded "
                              "if it is smaller than the required shape based on the scan "
                              "and the probe size.")
        grp.add_argument('--simul_probe', type=str,
                         help="Load a probe from a file to simulate the observed data. "
                              "Accepted formats are numpy files (npy or npz), which should "
                              "contain only one 2D array, which should have the same size "
                              "as the give frame_size.")
        return parser

    def check_params_beamline(self):
        """
        Check if self.params includes a minimal set of valid parameters, specific to a beamiline
        Returns: Nothing. Will raise an exception if necessary
        """
        if self.params['pixelsize'] == 55e-6 and self.params['near_field']:
            self.params['pixelsize'] = 100e-9
        if self.params['direct_beam']:
            self.params['use_direct_beam'] = True
        if self.params['probe_size_window'] is not None:
            if self.params['probe_size_window'] >= 1 or self.params['probe_size_window'] < 0.001:
                raise PtychoRunnerException(f"--probe_size_window={self.params['probe_size_window']} "
                                            f"must be within [0.001, 1[")
        if self.params['step_size_window'] >= 0.5 or self.params['step_size_window'] < 0.01:
            raise PtychoRunnerException(f"--probe_size_window={self.params['probe_size_window']} "
                                        f"must be within [0.01, 0.5[")
        if self.params['random_position_errors'] > 2 * self.params['step_size_window']:
            raise PtychoRunnerException(f"--random_position_errors={self.params['random_position_errors']} "
                                        f"must be <= 2 * step_size_window={self.params['step_size_window']}")


def make_parser_sphinx():
    """Returns the argparse for sphinx documentation"""
    return PtychoRunnerSimul.make_parser(default_params)
