"""Functions involved in the propagation of a wavefront."""

import numpy as np

from . import srwlib


[docs]
def groupOpticalSystem(opticalSystem, groupSeq):
    """Inspect the elements of an optical system and group a given sequence of 
    elements. Useful if the user is not interested in the intensity on some 
    intermediate planes (for instance just after a lens). The computation of a 
    sequence is more efficicent because it is performed at the SRW level and 
    requires also less memory cause the wavefront is overridden.

    Args:
        opticalSystem: the optical system dictionary.
        groupSeq: list of element types to group.

    Returns:
        list: list of optical containers where elements are grouped according 
            to the specified sequence.
    """
    
    seqLen = len(groupSeq)
    optElTypes = [optEl["type"] for optEl in opticalSystem.values()]
    optEls = [{key: opticalSystem[key]} for key in opticalSystem.keys()]

    if len(optElTypes) < len(groupSeq):
        # nothing to group
        optCnts = [optEls]
    else:
        optCnts = []
        i = 0
        while i < len(optEls):
            if seqLen != 0 and optElTypes[i:i + seqLen] == groupSeq:
                # sequence detected, group elements in one container
                optCnts.append([optEl for optEl in optEls[i:i + seqLen]])
                i += seqLen
            else:
                # spare element, add to list of optical containers
                optCnts.append([optEls[i]])
                i += 1

    return optCnts



[docs]
def cropOpticalSystem(opticalSystem, startPropAt=None, stopPropAt=None):
    """Crop an optical system to propagate a wavefront only within a 
    sub-sequence of the initial system.

    Args:
        opticalSystem: the dictionary of the whole optical system.
        startPropAt (optional): the name of the optical element to start from.
            If None, the first element is picked as start.
            Defaults to None.
        stopPropAt (optional): the name of the optical element to end with.
            If None, the last element is picked as stop.
            Defaults to None.
    Returns:
        dict: the dictionary of the cropped optical system.
    """

    optElKeys = list(opticalSystem.keys())

    startPropAt = optElKeys[0] if startPropAt is None else startPropAt
    stopPropAt = optElKeys[-1] if stopPropAt is None else stopPropAt
    startIdx = optElKeys.index(startPropAt)
    stopIdx = optElKeys.index(stopPropAt) + 1

    opticalSystemCropped = {
        key: opticalSystem[key] for key in optElKeys[startIdx:stopIdx]
    }

    return opticalSystemCropped



# =============================== PARSE TO SRW ================================


[docs]
def parsePropPar(autoResizeBefore=False, autoResizeAfter=False,
                 autoResizePrecision=1.0,
                 quadPhaseSemiAnal=False,
                 resizeInFourier=False,
                 hRangeChange=1.0, hResChange=1.0,
                 vRangeChange=1.0, vResChange=1.0):
    """Get the list of parameters for the propagation through an optical 
    element.

    Args:
        autoResizeBefore (optional): if True, auto-resize the wfr before the 
            propagation through the element. 
            Defaults to False.
        autoResizeBefore (optional): if True, auto-resize the wfr after the 
            propagation through the element. 
            Defaults to False.
        autoResizePrecision (optional): relative precision for propagation with 
            auto-resizing.
            Defaults to 1.0.
        quadPhaseSemiAnal (optional): use semi-analytical treatment of the
             quadratic (leading) phase. 
             Defaults to False.
        resizeInFourier (optional): resize in frequency domain using FTT.
            Defaults to False.
        hRangeChange (optional): resize horizontal extension.
            Defaults to 1.0.
        hResChange (optional): resize horizontal resolution.
            Defaults to 1.0.
        vRangeChange (optional): resize vertical extension.
            Defaults to 1.0.
        vResChange (optional): resize vertical resolution.
            Defaults to 1.0.

    Returns:
        list: [description]
    """

    SRWpropPar = [int(autoResizeBefore), int(autoResizeAfter),
                  float(autoResizePrecision),
                  int(quadPhaseSemiAnal),
                  int(resizeInFourier),
                  hRangeChange, hResChange,
                  vRangeChange, vResChange]

    return SRWpropPar



[docs]
def parseToSRWelement(optEl):
    """Get the SRW element and propagation parameter from the PySRW definition.

    Args:
        optEl: dictionary with the definition of a valid PySRW optical element.

    Returns:
        list: a list with the SRW element and the propagation parameters.
    """

    optElType = optEl["type"]
    
    if optElType == "drift":
        SRWoptEl, SRWpropPar = _parseDrift(optEl)
    elif optElType in ["rectAp", "rectOb", "circAp", "circOb"]:
        SRWoptEl, SRWpropPar = _parseApOrObs(optEl)
    elif optElType == "lens":
        SRWoptEl, SRWpropPar = _parseLens(optEl)
    elif optElType == "transmission":
        SRWoptEl, SRWpropPar = _parseTransmission(optEl)
    else:
        raise Exception(f"optical element type '{optElType}' not recognized")

    return SRWoptEl, SRWpropPar


def _parseDrift(optEl):
    # convert PySRW drift element to SRWLOptD
    SRWoptEl = srwlib.SRWLOptD(optEl["length"])

    if "propParam" in optEl.keys():
        SRWpropPar = parsePropPar(**optEl["propParam"])
    else:
        SRWpropPar = parsePropPar() # leave default values, mesh not adapted

    return SRWoptEl, SRWpropPar

def _parseApOrObs(optEl):
    # convert PySRW aperture/obstacle element to SRWLOptA
    elType = optEl["type"].lower()
    shape = "c" if "circ" in elType else "r"
    apOrObs = "a" if "ap" in elType else "o" 

    SRWoptEl = srwlib.SRWLOptA(shape, apOrObs, 
                               _Dx=optEl["extension"][0], 
                               _Dy=optEl["extension"][1],
                               _x=optEl["centre"][0], 
                               _y=optEl["centre"][1])

    if "propParam" in optEl.keys():
        SRWpropPar = parsePropPar(**optEl["propParam"])
    else:
        SRWpropPar = parsePropPar()

    return SRWoptEl, SRWpropPar

def _parseLens(optEl):
    # convert PySRW lens element to SRWLOptL
    if isinstance(optEl["f"], (list, tuple)):
        # cylindrical lens
        focal = optEl["f"]
    else:
        # axial lens
        focal = [optEl["f"], optEl["f"]]

    SRWoptEl = srwlib.SRWLOptL(_Fx=focal[0], _Fy=focal[1], 
                               _x=optEl["centre"][0], _y=optEl["centre"][1])

    if "propParam" in optEl.keys():
        SRWpropPar = parsePropPar(**optEl["propParam"])
    else:
        SRWpropPar = parsePropPar()

    return SRWoptEl, SRWpropPar

def _parseTransmission(optEl):
    # convert PySRW transmission element to SRWLOptT
    opd = optEl["opd"] * 1e-9
    ampl = optEl["amp"]
    extension = optEl["extension"]
    centre = optEl["centre"]

    nx, ny = np.shape(opd)
    Tr = np.zeros(2*nx*ny)
    Tr[::2] = np.reshape(ampl, (nx*ny))
    Tr[1::2] = np.reshape(opd, (nx*ny))
    arTr = srwlib.array("d", Tr)

    SRWoptEl = srwlib.SRWLOptT(_nx= nx, _ny=ny, 
                               _rx= extension[0], _ry=extension[1],
                               _arTr=arTr, _x=centre[0], _y=centre[1])
                        
    if "propParam" in optEl.keys():
        SRWpropPar = parsePropPar(**optEl["propParam"])
    else:
        SRWpropPar = parsePropPar()

    return SRWoptEl, SRWpropPar