# ! /opt/local/bin/python
# -*- coding: utf-8 -*-
import argparse
import warnings

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

import numpy as np
from .runner import CDIRunner, CDIRunnerException, CDIRunnerScan, default_params as params0
from pynx.cdi import *

helptext_epilog = """
pynx-cdi-simulation examples:

  * pynx-cdi-simulation --liveplot: default simulation (3D, icosahedron) with live plot
"""

params_beamline = {
    'formula_phase': None,
    'instrument': 'Simulation',
    'nb_er': 200,
    'nb_hio': 0,
    'nb_ml': 0,
    'nb_photons': 1e8,
    'nb_raar': 600,
    'ndim': 3,
    'obj_shape': 'icosahedron',
    'obj_radius': 25,
    'save': None,
    'scan': 1,
    'size': 128,
}

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


class CDIRunnerScanSimul(CDIRunnerScan):
    def __init__(self, params, scan, timings=None):
        super(CDIRunnerScanSimul, self).__init__(params, scan, timings=timings)

    def load_data(self):
        sh = [self.params['size']] * self.params['ndim']
        pix = self.params['pixel_size_detector']
        w = self.params['wavelength']
        d = self.params['detector_distance']
        iobs = np.empty(sh)
        init_free_pixel_mask = False
        if self.cdi is None:
            init_free_pixel_mask = True
            self.cdi = CDI(iobs=iobs, pixel_size_detector=pix, wavelength=w, detector_distance=d)
        s = self.params['obj_shape']
        n = self.params['obj_radius']
        if n > sh[0] // 2:
            warnings.warn(f"CDIRunnerScanSimul: data size={sh[0]} and obj_radius={n} - there may "
                          f"be oversampling issues (2*radius > size/2)")
        fp = self.params['formula_phase']
        self.cdi = Calc2Obs() * InitObjShape(shape=s, radius=n, formula=None, formula_phase=fp) * self.cdi
        # need to shift so data is centred on the array as for any experimental data
        self.iobs = self.cdi.get_iobs(shift=True)
        # TODO: give an option for the number of photons, and do this on GPU
        self.iobs = np.random.poisson(self.params['nb_photons'] * self.iobs / self.iobs.sum())
        self.cdi.set_iobs(self.iobs, shift=True)
        if init_free_pixel_mask:
            # otherwise it is not initialised because we created self.cdi..
            self._init_free_pixel_mask(self.iobs)

    def prepare_cdi(self):
        """
        Prepare CDI object from input data.

        :return: nothing. Creates or updates self.cdi object.
        """
        super(CDIRunnerScanSimul, self).prepare_cdi()
        # Scale initial object (unnecessary if auto-correlation is used)
        if self.params['support'] != 'auto':
            self.cdi = ScaleObj(method='F', lazy=True) * self.cdi


class CDIRunnerSimul(CDIRunner):
    """
    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().__init__(argv, params)
        self.CDIRunnerScan = CDIRunnerScanSimul

    @classmethod
    def make_parser(cls, description=None, script_name="pynx-ptycho-simulation", epilog=None, default_par=None):
        if default_par is None:
            default_par = default_params
        if epilog is None:
            epilog = helptext_epilog
        if description is None:
            description = "Script to perform a CDI analysis on simulated data (mostly for tests)"

        parser = super().make_parser(script_name, description, epilog, default_par)
        grp = parser.add_argument_group("Simulation parameters")
        p = default_par

        grp.add_argument('--size', type=int, default=p['size'],
                         help="size of the dataset along each dimension")

        grp.add_argument('--ndim', type=int, default=p['ndim'],
                         choices=[2, 3],
                         help="number of dimensions for the data")

        grp.add_argument(
            '--obj_shape', '--objshape', type=str, default=p['obj_shape'],
            choices=['circle', 'sphere', 'square', 'cube', 'tetrahedron',
                     'octahedron', 'dodecahedron', 'icosahedron'],
            help="Simulated object shape, using smooth super-spheres "
                 "[Onaka, Phil. Mag. Lett. 86 (2006), 175] for "
                 "tetra/octa/dodeca/icosahedron")

        grp.add_argument('--obj_radius', type=int, default=p['obj_radius'],
                         help="object radius (half-size) in pixels")

        grp.add_argument('--nb_photons', '--nbphotons',
                         type=int, default=p['nb_photons'],
                         help="Total number of photons in the simulated dataset")

        grp.add_argument(
            '--formula_phase', '--formulaphase', type=str, default=p['formula_phase'],
            help="Optional formula (must be in quotes) to add a phase to the object.\n"
                 "In this formula, only sqrt, abs, sin, cos, tan, +, -, *, / "
                 "are allowed in addition to x,y,z, which"
                 "are coordinates relative to the object radius (i.e. varying normally from "
                 "-1 to +1). The computed phase is in radians.\n"
                 "Examples:\n\n"
                 "* ``--formula_phase '2 * (x*x-0.2*x*x*x + 0.5*y*y-0.1*y*y*y + 0.7*z*z)'``\n"
                 "* ``--formula_phase '1*cos(8*(x*x+0.5*y*y))'``\n\n"
                 "[default=None, the object is real and positive]")

        # Change default argument so --data is not required
        grp.add_argument('--data', type=str, default=None, required=False,
                         help=argparse.SUPPRESS)

        return parser


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