#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright © 2018 Michael J. Hayford
""" Manager class for a sequential optical model

.. codeauthor: Michael J. Hayford
"""

import itertools
import logging

from anytree import Node

from rayoptics.elem import surface
from . import gap
from . import medium as m
from rayoptics.raytr import raytrace as rt
from rayoptics.raytr import trace as trace
from rayoptics.raytr import analyses
from rayoptics.elem import transform as trns
from opticalglass import glassfactory as gfact
from opticalglass import glasserror as ge
import numpy as np
from math import copysign, sqrt
from rayoptics.util.misc_math import isanumber


class SequentialModel:
    """ Manager class for a sequential optical model

    A sequential optical model is a sequence of surfaces and gaps.

    The sequential model has this structure
    ::

        IfcObj  Ifc1  Ifc2  Ifc3 ... Ifci-1   IfcImg
             \  /  \  /  \  /             \   /
             GObj   G1    G2              Gi-1

    where

        - Ifc is a :class:`~rayoptics.seq.interface.Interface` instance
        - G   is a :class:`~rayoptics.seq.gap.Gap` instance

    There are N interfaces and N-1 gaps. The initial configuration has an
    object and image Surface and an object gap.

    The Interface API supports implementation of an optical action, such as
    refraction, reflection, scatter, diffraction, etc. The Interface may be
    realized as a physical profile separating the adjacent gaps or an idealized
    object, such as a thin lens or 2 point HOE.

    The Gap class maintains a simple separation (z translation) and the medium
    filling the gap. More complex coordinate transformations are handled
    through the Interface API.

    Attributes:
        opt_model: parent optical model
        ifcs: list of :class:`~rayoptics.seq.interface.Interface`
        gaps: list of :class:`~rayoptics.seq.gap.Gap`
        lcl_tfrms: forward transform, interface to interface
        rndx: a list with refractive indices for all **wvls**
        z_dir: -1 if gap follows an odd number of reflections, otherwise +1
        gbl_tfrms: global coordinates of each interface wrt the 1st interface
        stop_surface (int): index of stop interface
        cur_surface (int): insertion index for next interface
    """

    def __init__(self, opt_model, do_init=True, **kwargs):
        self.opt_model = opt_model
        self.ifcs = []
        self.gaps = []
        self.gbl_tfrms = []
        self.lcl_tfrms = []
        self.z_dir = []
        self.stop_surface = None
        self.cur_surface = None
        self.wvlns = []
        self.rndx = []
        self.do_apertures = True
        if do_init:
            self._initialize_arrays()

    def __json_encode__(self):
        attrs = dict(vars(self))
        del attrs['opt_model']
        del attrs['gbl_tfrms']
        del attrs['lcl_tfrms']
        del attrs['z_dir']
        del attrs['wvlns']
        del attrs['rndx']
        return attrs

    def _initialize_arrays(self):
        """ initialize object and image interfaces and intervening gap """
        # add object interface
        self.ifcs.append(surface.Surface('Obj', interact_mode='dummy'))

        tfrm = np.identity(3), np.array([0., 0., 0.])
        self.gbl_tfrms.append(tfrm)
        self.lcl_tfrms.append(tfrm)

        # add object gap
        self.gaps.append(gap.Gap())
        self.z_dir.append(1)
        self.rndx.append([1.0])

        # interfaces are inserted after cur_surface
        self.cur_surface = 0

        # add image interface
        self.ifcs.append(surface.Surface('Img', interact_mode='dummy'))
        self.gbl_tfrms.append(tfrm)
        self.lcl_tfrms.append(tfrm)

    def reset(self):
        self.__init__()

    def get_num_surfaces(self):
        return len(self.ifcs)

    def path(self, wl=None, start=None, stop=None, step=1):
        """ returns an iterable path tuple for a range in the sequential model

        Args:
            wl: wavelength in nm for path, defaults to central wavelength
            start: start of range
            stop: first value beyond the end of the range
            step: increment or stride of range

        Returns:
            (**ifcs, gaps, lcl_tfrms, rndx, z_dir**)
        """
        if wl is None:
            wl = self.central_wavelength()

        if step < 0:
            gap_start = start - 1 if start is not None else start
        else:
            gap_start = start

        wl_idx = self.index_for_wavelength(wl)
        try:
            rndx = [n[wl_idx] for n in self.rndx[start:stop:step]]
        except IndexError:
            self.wvlns = self.opt_model['osp']['wvls'].wavelengths
            self.rndx = self.calc_ref_indices_for_spectrum(self.wvlns)
            rndx = [n[wl_idx] for n in self.rndx[start:stop:step]]

        path = itertools.zip_longest(self.ifcs[start:stop:step],
                                     self.gaps[gap_start:stop:step],
                                     self.lcl_tfrms[start:stop:step],
                                     rndx,
                                     self.z_dir[start:stop:step])
        return path

    def reverse_path(self, wl=None, start=None, stop=None, step=-1):
        """ returns an iterable path tuple for a range in the sequential model
    
        Args:
            wl: wavelength in nm for path, defaults to central wavelength
            start: start of range
            stop: first value beyond the end of the range
            step: increment or stride of range
    
        Returns:
            (**ifcs, gaps, lcl_tfrms, rndx, z_dir**)
        """
        if wl is None:
            wl = self.central_wavelength()
    
        if step < 0:
            if start is not None:
                gap_start = start - 1
                rndx_start = start - 1
            else:
                gap_start = start
                rndx_start = -2
        else:
            gap_start = start
    
        tfrms = self.compute_local_transforms(step=-1)
        wl_idx = self.index_for_wavelength(wl)
        rndx = [n[wl_idx] for n in self.rndx[rndx_start:stop:step]]
        z_dir = [-z_dir for z_dir in self.z_dir[start:stop:step]]
        path = itertools.zip_longest(self.ifcs[start:stop:step],
                                     self.gaps[gap_start:stop:step],
                                     tfrms,
                                     rndx,
                                     z_dir)
        return path

    def calc_ref_indices_for_spectrum(self, wvls):
        """ returns a list with refractive indices for all **wvls**

        Args:
            wvls: list of wavelengths in nm
        """
        indices = []
        for g in self.gaps:
            ri = []
            mat = g.medium
            for w in wvls:
                rndx = mat.rindex(w)
                ri.append(rndx)
            indices.append(ri)

        return indices

    def central_wavelength(self):
        """ returns the central wavelength in nm of the model's ``WvlSpec`` """
        return self.opt_model.optical_spec.spectral_region.central_wvl

    def index_for_wavelength(self, wvl):
        """ returns index into rndx array for wavelength `wvl` in nm """
        self.wvlns = self.opt_model.optical_spec.spectral_region.wavelengths
        return self.wvlns.index(wvl)

    def central_rndx(self, i):
        """ returns the central refractive index of the model's ``WvlSpec`` """
        central_wvl = self.opt_model.optical_spec.spectral_region.reference_wvl
        return self.rndx[i][central_wvl]

    def get_surface_and_gap(self, srf=None):
        if srf is None:
            srf = self.cur_surface
        s = self.ifcs[srf]
        if srf == len(self.gaps):
            g = None
        else:
            g = self.gaps[srf]
        return s, g

    def set_cur_surface(self, s):
        self.cur_surface = s

    def set_stop(self):
        """ sets the stop surface to the current surface """
        self.stop_surface = self.cur_surface
        return self.stop_surface

    def __iadd__(self, node):
        if isinstance(node, gap.Gap):
            self.gaps.append(node)
        else:
            self.ifcs.insert(len(self.ifcs)-1, node)
        return self

    def insert(self, ifc, gap, prev=False):
        """ insert surf and gap at the cur_gap edge of the sequential model
            graph """
        if self.stop_surface is not None:
            num_ifcs = len(self.ifcs)
            if num_ifcs > 2:
                if self.stop_surface > self.cur_surface and \
                   self.stop_surface < num_ifcs - 2:
                    self.stop_surface += 1
        idx = self.cur_surface = (0 if self.cur_surface is None
                                   else self.cur_surface+1)
        self.ifcs.insert(idx, ifc)
        if gap is not None:
            idx_g = idx-1 if prev else idx
            self.gaps.insert(idx_g, gap)
        else:
            gap = self.gaps[idx]

        tfrm = np.identity(3), np.array([0., 0., 0.])
        self.gbl_tfrms.insert(idx, tfrm)
        self.lcl_tfrms.insert(idx, tfrm)

        new_z_dir = self.z_dir[idx-1] if idx > 1 else 1
        self.z_dir.insert(idx, new_z_dir)

        wvls = self.opt_model.optical_spec.spectral_region.wavelengths
        rindex = [gap.medium.rindex(w) for w in wvls]
        self.rndx.insert(idx, rindex)

        if ifc.interact_mode == 'reflect':
            self.update_reflections(start=idx)

    def remove(self, *args, prev=False):
        """Remove surf and gap at cur_surface or an input index argument.

        To avoid invalid sequence states, both an interface and a gap must be
        removed at the same time. The ``prev`` argument, if True, removes the
        gap preceding the interface. The default behavior is to remove the
        following gap.
        """
        if len(args) == 0:
            idx = self.cur_surface
        else:
            idx = args[0]

        num_ifcs = len(self.ifcs)
        # don't allow object or image interfaces to be removed
        if prev:
            if idx == 0 or idx == 1 or idx == num_ifcs:
                raise IndexError
        else:
            if idx == 0 or idx == -1 or idx == num_ifcs:
                raise IndexError

        if self.ifcs[idx].interact_mode == 'reflect':
            self.update_reflections(start=idx)

        # decrement stop surface as needed
        if self.stop_surface is not None:
            if num_ifcs > 2:
                if self.stop_surface > idx and self.stop_surface > 1:
                    self.stop_surface -= 1
        
        # remove the associated nodes
        pt = self.opt_model['part_tree']
        pt.trim_node(self.ifcs[idx])

        # interface related attribute lists
        del self.ifcs[idx]
        del self.gbl_tfrms[idx]
        del self.lcl_tfrms[idx]

        # gap node and related attribute lists
        idx = idx-1 if prev else idx
        pt.trim_node(self.gaps[idx])

        del self.gaps[idx]
        del self.z_dir[idx]
        del self.rndx[idx]

    def remove_node(self, e_node):
        part_tree = self.opt_model.part_tree
        ifcs = [n.id for n in part_tree.nodes_with_tag(tag='#ifc',
                                                       root=e_node)]
        for ifc in ifcs:
            idx = self.ifcs.index(ifc)
            if (self.stop_surface is not None and
                idx <= self.stop_surface and
                idx > 1):
                self.stop_surface -= 1
            del self.ifcs[idx]
            del self.gaps[idx]
            del self.z_dir[idx]
            del self.rndx[idx]
            del self.lcl_tfrms[idx]
            del self.gbl_tfrms[idx]

    def add_surface(self, surf_data, **kwargs):
        """ add a surface where `surf_data` is a list that contains:

        [curvature, thickness, refractive_index, v-number, semi-diameter]

        The `curvature` entry is interpreted as radius if `radius_mode` is **True**

        The `thickness` is the signed thickness

        The `refractive_index, v-number` entry can have several forms:

            - **refractive_index, v-number**
            - **refractive_index** only -> constant index model
            - **'REFL'** -> set interact_mode to 'reflect'
            - **glass_name, catalog_name** as 1 or 2 strings
            - blank -> defaults to air

        The `semi-diameter` entry is optional

        """
        radius_mode = self.opt_model.radius_mode
        mat = None
        if len(surf_data) > 2:
            if not isanumber(surf_data[2]):
                if surf_data[2].upper() == 'REFL':
                    mat = self.gaps[self.cur_surface].medium
        s, g, rn, tfrm = create_surface_and_gap(surf_data, prev_medium=mat,
                                                radius_mode=radius_mode,
                                                **kwargs)
        self.insert(s, g)

        root_node = self.opt_model['part_tree'].root_node
        idx = self.cur_surface
        Node(f'i{idx}', id=s, tag='#ifc', parent=root_node)
        if gap is not None:
            Node(f'g{idx}', id=(g, self.z_dir[idx]), tag='#gap', parent=root_node)

    def sync_to_restore(self, opt_model):
        self.opt_model = opt_model
        if hasattr(self, 'optical_spec'):
            opt_model.optical_spec = self.optical_spec
            delattr(self, 'optical_spec')
        init_z_dir = False
        if not hasattr(self, 'z_dir'):
            self.z_dir = []
            init_z_dir = True
            z_dir_work = 1
        for sg in itertools.zip_longest(self.ifcs, self.gaps):
            ifc, g = sg
            if hasattr(ifc, 'sync_to_restore'):
                ifc.sync_to_restore(self)
            if g:
                if hasattr(g, 'sync_to_restore'):
                    g.sync_to_restore(self)
                if init_z_dir:
                    if ifc.interact_mode == 'reflect':
                        z_dir_work = -z_dir_work
                    self.z_dir.append(z_dir_work)

        self.ifcs[0].interact_mode = 'dummy'
        self.ifcs[-1].interact_mode = 'dummy'
                
        if not hasattr(self, 'do_apertures'):
            self.do_apertures = True

    def update_model(self):
        # delta n across each surface interface must be set to some
        #  reasonable default value. use the index at the central wavelength
        osp = self.opt_model.optical_spec
        ref_wl = osp.spectral_region.reference_wvl

        self.wvlns = osp.spectral_region.wavelengths
        self.rndx = self.calc_ref_indices_for_spectrum(self.wvlns)
        n_before = self.rndx[0][ref_wl]

        self.z_dir = []
        z_dir_before = 1

        seq = itertools.zip_longest(self.ifcs, self.gaps)

        for i, sg in enumerate(seq):
            ifc, g = sg
            z_dir_after = copysign(1, z_dir_before)
            if ifc.interact_mode == 'reflect':
                z_dir_after = -z_dir_after

            # leave rndx data unsigned, track change of sign using z_dir
            if g is not None:
                n_after = self.rndx[i][ref_wl]
                if z_dir_after < 0:
                    n_after = -n_after
                ifc.delta_n = n_after - n_before
                n_before = n_after
    
                z_dir_before = z_dir_after
                self.z_dir.append(z_dir_after)

            # call update() on the surface interface
            ifc.update()

        self.gbl_tfrms = self.compute_global_coords()
        self.lcl_tfrms = self.compute_local_transforms()

        if self.do_apertures:
            if len(self.ifcs) > 2:
                osp.update_model()
    
                self.set_clear_apertures()

    def apply_scale_factor(self, scale_factor):
        for i, sg in enumerate(self.path()):
            ifc, g, lcl_tfrm, rndx, z_dir = sg
            ifc.apply_scale_factor(scale_factor)
            if g:
                g.apply_scale_factor(scale_factor)

        self.gbl_tfrms = self.compute_global_coords()
        self.lcl_tfrms = self.compute_local_transforms()

    def set_from_specsheet(self, specsheet):
        if self.opt_model.optical_spec.parax_data is None:
            return
        if len(specsheet.imager_inputs) == 2:
            fod = self.opt_model.optical_spec.parax_data.fod
            f_old = fod.efl
            f_new = specsheet.imager.f
            scale_factor = f_new/f_old
            if scale_factor != 1.0:
                self.apply_scale_factor(scale_factor)

            if specsheet.conjugate_type == 'finite':
                self.gaps[0].thi = scale_factor*fod.pp1 - specsheet.imager.s
                self.gaps[-1].thi = specsheet.imager.sp - scale_factor*fod.ppk
            elif specsheet.conjugate_type == 'infinite':
                self.gaps[-1].thi = specsheet.imager.sp - scale_factor*fod.ppk

    def insert_surface_and_gap(self):
        s = surface.Surface()
        g = gap.Gap()
        self.insert(s, g)
        return s, g

    def update_reflections(self, start):
        """ update interfaces and gaps following insertion of a mirror """

        for i, sg in enumerate(self.path(start=start), start=start):
            ifc, g, lcl_tfrm, rndx, z_dir = sg
            if i > start:
                ifc.apply_scale_factor(-1)
                if g:
                    g.apply_scale_factor(-1)
                    self.z_dir[i] = -z_dir

    def surface_label_list(self):
        """ list of surface labels or surface number, if no label """
        labels = []
        for i, s in enumerate(self.ifcs):
            if len(s.label) == 0:
                if i == self.stop_surface:
                    labels.append('Stop')
                else:
                    labels.append(str(i))
            else:
                labels.append(s.label)
        return labels

    def list_model(self, path=None):
        cvr = 'r' if self.opt_model.radius_mode else 'c'
        print("             {}            t        medium     mode   zdr"
              "      sd".format(cvr))
        labels = self.surface_label_list()
        path = self.path() if path is None else path
        for i, sg in enumerate(path):
            ifc, gap, _, _, _ = sg
            s = self.list_surface_and_gap(ifc, gp=gap)
            if gap is not None:
                s.append(self.z_dir[i])
            else:
                s.append(self.z_dir[-1])
            fmt = "{0:>4s}: {1:12.6f} {2:#12.6g} {3:>9s} {4:>10s} {6:2n}"
            if s[4] is not None:  # if the sd is not None...
                fmt += "  {5:#10.5g}"
            print(fmt.format(labels[i], *s))

    def list_model_old(self):
        cvr = 'r' if self.opt_model.radius_mode else 'c'
        print("           {}            t        medium     mode         sd"
              .format(cvr))
        for i, sg in enumerate(self.path()):
            ifc, g, _, _, _ = sg
            s = self.list_surface_and_gap(ifc, gp=g)
            fmt = "{0:2n}: {1:12.6f} {2:#12.6g} {3:>9s} {4.name:>10s}"
            if s[4] is not None:  # if the sd is not None...
                fmt += " {5:#10.5g}"
            print(fmt.format(i, *s))

    def list_gaps(self):
        for i, gp in enumerate(self.gaps):
            print(i, gp)

    def list_surfaces(self):
        for i, s in enumerate(self.ifcs):
            print(i, s)

    def list_surface_and_gap(self, ifc, gp=None):
        """Returns cvr, thi, med, imode, sd for input ifc and gap."""
        cvr = ifc.profile_cv
        if self.opt_model.radius_mode:
            if cvr != 0.0:
                cvr = 1.0/cvr
        sd = ifc.surface_od()
        imode = ifc.interact_mode if ifc.interact_mode == 'reflect' else ""

        if gp is not None:
            thi = gp.thi
            med = gp.medium.name()
        else:
            thi = 0.
            med = ''
        return [cvr, thi, med, imode, sd]

    def list_decenters(self, full=False):
        """List decenter data and gap separations.

        Arguments:
            full: lists all values if True, else only y offset and alpha tilt
        """
        fmt0a = ("             thi    medium/mode       type             x"
                 "          y       alpha      beta       gamma")
        fmt0b = ("             thi    medium/mode       type             y"
                 "       alpha")
        fmt1a = ("{:4n}:                {:>10s}  {:>14s} {:#10.5g} {:#10.5g}"
                 " {:#10.5g} {:#10.5g} {:#10.5g}")
        fmt1b = ("{:4n}:                {:>10s}  {:>14s} {:#10.5g}"
                 " {:#10.5g}")
        fmt1c = "{:4n}:                {:>10s}"
        fmt2 = "{:4n}: {:#12.6g}    {:>9s}"

        # print header
        if full:
            print(fmt0a)
        else:
            print(fmt0b)

        for i, sg in enumerate(self.path()):
            ifc, gap, lcl_tfrm, rndx, z_dir = sg
            imode = (ifc.interact_mode if ifc.interact_mode != 'transmit'
                     else "")

            if ifc.decenter is not None:
                d = ifc.decenter
                if full:
                    print(fmt1a.format(i, imode, d.dtype,
                                       d.dec[0], d.dec[1],
                                       d.euler[0], d.euler[1], d.euler[2]))
                else:
                    print(fmt1b.format(i, imode, d.dtype,
                                       d.dec[1], d.euler[0]))
            elif gap is None:  # final interface, just list interact_mode
                print(fmt1c.format(i, imode))

            if gap:
                print(fmt2.format(i, gap.thi, gap.medium.name()))

    def list_sg(self):
        """List decenter data and gap separations. """
        cvrd = 'r' if self.opt_model.radius_mode else 'c'
        fmt0a = ("              {}               mode           type"
                 "             y       alpha")
        fmt0b = ("                      t           medium")
        fmt1 = ("{:>4s}: {:#12.6g}       {:>10s}     {:>14s} {:#10.5g}"
                " {:#10.5g}")
        fmt2 = ("{:>4s}: {:#12.6g}       {:>10s}")
        fmt3 = ("               {:#12.6g}    {:>9s}")

        # print header
        print(fmt0a.format(cvrd))
        print(fmt0b)

        labels = self.surface_label_list()
        for i, sg in enumerate(self.path()):
            ifc, gap, lcl_tfrm, rndx, z_dir = sg
            s = self.list_surface_and_gap(ifc, gap)
            s.append(z_dir)
            cvr, thi, med, imode, sd, z_dir = s
            if ifc.decenter is not None:
                d = ifc.decenter
                print(fmt1.format(labels[i], cvr, imode, d.dtype,
                                  d.dec[1], d.euler[0]))
            else:
                print(fmt2.format(labels[i], cvr, imode))

            if gap:
                print(fmt3.format(gap.thi, gap.medium.name()))

    def list_elements(self):
        for i, gp in enumerate(self.gaps):
            if gp.medium.label.lower() != 'air':
                print(self.ifcs[i].profile,
                      self.ifcs[i+1].profile,
                      gp)

    def trace_fan(self, fct, fi, xy, num_rays=21, **kwargs):
        """ xy determines whether x (=0) or y (=1) fan """
        osp = self.opt_model.optical_spec
        fld = osp.field_of_view.fields[fi]
        wvl = self.central_wavelength()
        foc = osp.defocus.get_focus()

        rs_pkg, cr_pkg = trace.setup_pupil_coords(self.opt_model,
                                                  fld, wvl, foc)
        fld.chief_ray = cr_pkg
        fld.ref_sphere = rs_pkg

        # Use the chief ray intersection point for the central wavelength
        #  as the reference image point for wavefront error calculations
        ref_img_pt = cr_pkg[0].ray[-1][0]

        wvls = osp.spectral_region
        fans_x = []
        fans_y = []
        fan_start = np.array([0., 0.])
        fan_stop = np.array([0., 0.])
        fan_start[xy] = -1.0
        fan_stop[xy] = 1.0
        fan_def = [fan_start, fan_stop, num_rays]
        max_y_val = 0.0
        rc = []
        for wi, wvl in enumerate(wvls.wavelengths):
            rc.append(wvls.render_colors[wi])

            rs_pkg, cr_pkg = trace.setup_pupil_coords(self.opt_model,
                                                      fld, wvl, foc,
                                                      image_pt=ref_img_pt)
            fld.chief_ray = cr_pkg
            fld.ref_sphere = rs_pkg
            fan = trace.trace_fan(self.opt_model, fan_def, fld, wvl, foc,
                                  img_filter=lambda p, ray_pkg:
                                  fct(p, xy, ray_pkg, fld, wvl, foc), **kwargs)
            f_x = []
            f_y = []
            for p, y_val in fan:
                f_x.append(p[xy])
                f_y.append(y_val)
                if abs(y_val) > max_y_val:
                    max_y_val = abs(y_val)
            fans_x.append(f_x)
            fans_y.append(f_y)
        fans_x = np.array(fans_x)
        fans_y = np.array(fans_y)
        return fans_x, fans_y, max_y_val, rc

    def trace_grid(self, fct, fi, wl=None, num_rays=21, form='grid',
                   append_if_none=True, **kwargs):
        """ fct is applied to the raw grid and returned as a grid  """
        osp = self.opt_model.optical_spec
        wvls = osp.spectral_region
        wvl = self.central_wavelength()
        wv_list = wvls.wavelengths if wl is None else [wvl]
        fld = osp.field_of_view.fields[fi]
        foc = osp.defocus.get_focus()

        rs_pkg, cr_pkg = trace.setup_pupil_coords(self.opt_model,
                                                  fld, wvl, foc)
        fld.chief_ray = cr_pkg
        fld.ref_sphere = rs_pkg

        grids = []
        grid_start = np.array([-1., -1.])
        grid_stop = np.array([1., 1.])
        grid_def = [grid_start, grid_stop, num_rays]
        for wi, wvl in enumerate(wv_list):
            grid = trace.trace_grid(self.opt_model, grid_def, fld, wvl, foc,
                                    form=form, append_if_none=append_if_none,
                                    img_filter=lambda p, ray_pkg:
                                    fct(p, wi, ray_pkg, fld, wvl, foc),
                                    **kwargs)
            grids.append(grid)
        rc = wvls.render_colors
        return grids, rc

    def trace_wavefront(self, fld, wvl, foc, num_rays=32):

        def wave(p, ray_pkg, fld, wvl, foc):
            x = p[0]
            y = p[1]
            if ray_pkg is not None:
                fod = self.opt_model.optical_spec.parax_data.fod
                opd = analyses.wave_abr_full_calc(fod, fld, wvl, foc, ray_pkg,
                                                  fld.chief_ray,
                                                  fld.ref_sphere)
                opd = opd/self.opt_model.nm_to_sys_units(wvl)
            else:
                opd = 0.0
            return np.array([x, y, opd])

        rs_pkg, cr_pkg = trace.setup_pupil_coords(self.opt_model,
                                                  fld, wvl, foc)
        fld.chief_ray = cr_pkg
        fld.ref_sphere = rs_pkg

        grid_start = np.array([-1., -1.])
        grid_stop = np.array([1., 1.])
        grid_def = (grid_start, grid_stop, num_rays)

        grid = trace.trace_grid(self.opt_model, grid_def, fld, wvl, foc,
                                img_filter=lambda p, ray_pkg:
                                wave(p, ray_pkg, fld, wvl, foc), form='grid')
        return grid

    # def set_clear_apertures(self):
    #     def rd(v):
    #         """ take 2d length of input vector v """
    #         return np.sqrt(v[0]*v[0]+v[1]*v[1])

    #     if self.get_num_surfaces() > 2:
    #         fields_df = trace.trace_all_fields(self.opt_model)
    #         # a) Select the inc_pt data from the unstacked result and
    #         #    transpose so that intrfcs are the index
    #         inc_pts = fields_df.unstack()['inc_pt'].T
    #         # b) applymap() is used to apply the function rd() to each
    #         #    element in the dataframe
    #         inc_pts_rd = inc_pts.applymap(rd)
    #         # c) apply max() function to each row (i.e. across columns,
    #         #    axis=1)
    #         semi_ap = inc_pts_rd.max(axis=1)
    #         for s, max_ap in zip(self.ifcs[1:-1], semi_ap[1:-1]):
    #             s.set_max_aperture(max_ap)

    def set_clear_apertures_paraxial(self):
        ax_ray, pr_ray, _ = self.opt_model.optical_spec.parax_data
        for i, ifc in enumerate(self.ifcs):
            sd = abs(ax_ray[i][0]) + abs(pr_ray[i][0])
            ifc.set_max_aperture(sd)

    def set_clear_apertures(self):
        rayset = trace.trace_boundary_rays(self.opt_model,
                                           use_named_tuples=True)

        for i, s in enumerate(self.ifcs):
            max_ap = -1.0e+10
            update = True
            for f in rayset:
                for p in f:
                    ray = p.ray
                    if len(ray) > i:
                        ap = sqrt(ray[i].p[0]**2 + ray[i].p[1]**2)
                        if ap > max_ap:
                            max_ap = ap
                    else:  # ray failed before this interface, don't update
                        update = False
            if update:
                s.set_max_aperture(max_ap)

    def trace(self, pt0, dir0, wvl, **kwargs):
        return rt.trace(self, pt0, dir0, wvl, **kwargs)

    def compute_global_coords(self, glo=1):
        """ Return global surface coordinates (rot, t) wrt surface glo. """
        tfrms = []
        r, t = np.identity(3), np.array([0., 0., 0.])
        prev = r, t
        tfrms.append(prev)
        if glo > 0:
            # iterate in reverse over the segments before the
            #  global reference surface
            step = -1
            seq = itertools.zip_longest(self.ifcs[glo::step],
                                        self.gaps[glo-1::step])
            ifc, gap = after = next(seq)
            # loop of remaining surfaces in path
            while True:
                try:
                    b4_ifc, b4_gap = before = next(seq)
                    zdist = gap.thi
                    r, t = trns.reverse_transform(ifc, zdist, b4_ifc)
                    t = prev[0].dot(t) + prev[1]
                    r = prev[0].dot(r)
                    prev = r, t
                    tfrms.append(prev)
                    after, ifc, gap = before, b4_ifc, b4_gap
                except StopIteration:
                    break
            tfrms.reverse()
    
        seq = itertools.zip_longest(self.ifcs[glo:], self.gaps[glo:])
        b4_ifc, b4_gap = before = next(seq)
        prev = np.identity(3), np.array([0., 0., 0.])
        # loop forward over the remaining surfaces in path
        while True:
            try:
                ifc, gap = after = next(seq)
                zdist = b4_gap.thi
                r, t = trns.forward_transform(b4_ifc, zdist, ifc)
                t = prev[0].dot(t) + prev[1]
                r = prev[0].dot(r)
                prev = r, t
                tfrms.append(prev)
                before, b4_ifc, b4_gap = after, ifc, gap
            except StopIteration:
                break
    
        return tfrms

    def compute_local_transforms(self, seq=None, step=1):
        """ Return forward surface coordinates (r.T, t) for each interface. """
        tfrms = []
        if seq is None:
            seq = itertools.zip_longest(self.ifcs[::step],
                                        self.gaps[::step])
        b4_ifc, b4_gap = before = next(seq)
        while before is not None:
            try:
                ifc, gap = after = next(seq)
            except StopIteration:
                tfrms.append((np.identity(3), np.array([0., 0., 0.])))
                break
            else:
                zdist = step*b4_gap.thi
                r, t = trns.forward_transform(b4_ifc, zdist, ifc)
                rt = r.transpose()
                tfrms.append((rt, t))
                before, b4_ifc, b4_gap = after, ifc, gap

        return tfrms

    def find_matching_ifcs(self):
        rot_tols = dict(atol=1e-14, rtol=1e-8)
        tols = dict(atol=1e-14, rtol=1e-14)
        matches = []
        for i, gi in enumerate(self.gbl_tfrms):
            i1 = i+1
            for j, gj in enumerate(self.gbl_tfrms[i1:], start=i1):
                if (
                        np.allclose(gi[0], gj[0], **rot_tols) and
                        np.allclose(gi[1], gj[1], **tols)
                        ):
                    print(f'coincident surfs: {i} - {j}')
                    matches.append((i, j))
        return matches


def gen_sequence(surf_data_list, **kwargs):
    """ create a sequence iterator from the surf_data_list

    Args:
        surf_data_list: a list of lists containing:
                        [curvature, thickness, refractive_index, v-number]
        **kwargs: keyword arguments

    Returns:
        (**ifcs**, **gaps**, **rndx**, **lcl_tfrms**, **z_dir**)
    """
    ifcs = []
    gaps = []
    rndx = []
    lcl_tfrms = []
    z_dir = []

    for surf_data in surf_data_list:
        s, g, rn, tfrm = create_surface_and_gap(surf_data, **kwargs)
        ifcs.append(s)
        gaps.append(g)
        rndx.append(rn)
        lcl_tfrms.append(tfrm)
        z_dir.append(1)
    ifcs[-1].interact_mode = 'dummy'

    n_before = 1.0
    z_dir_before = 1
    for i, s in enumerate(ifcs):
        z_dir_after = copysign(1, z_dir_before)
        n_after = np.copysign(rndx[i], n_before)
        if s.interact_mode == 'reflect':
            n_after = -n_after
            z_dir_after = -z_dir_after

        n_before = n_after
        rndx[i] = n_after
        z_dir_before = z_dir_after
        z_dir[i] = z_dir_after

    seq = itertools.zip_longest(ifcs, gaps[:-2], lcl_tfrms, rndx, z_dir)
    return seq


def create_surface_and_gap(surf_data, radius_mode=False, prev_medium=None,
                           wvl=550.0, **kwargs):
    """ create a surface and gap where `surf_data` is a list that contains:

    [curvature, thickness, refractive_index, v-number, semi-diameter]
    
    The `curvature` entry is interpreted as radius if `radius_mode` is **True**

    The `thickness` is the signed thickness

    The `refractive_index, v-number` entry can have several forms:
        
        - **refractive_index, v-number**
        - **refractive_index** only -> constant index model
        - **'REFL'** -> set interact_mode to 'reflect'
        - **glass_name, catalog_name** as 1 or 2 strings
        - blank -> defaults to air

    The `semi-diameter` entry is optional
    """
    s = surface.Surface()

    if radius_mode:
        if surf_data[0] != 0.0:
            s.profile.cv = 1.0/surf_data[0]
        else:
            s.profile.cv = 0.0
    else:
        s.profile.cv = surf_data[0]

    if len(surf_data) > 2:
        if isanumber(surf_data[2]):  # assume all args are numeric
            if len(surf_data) < 3:
                if surf_data[2] == 1.0:
                    mat = m.Air()
                else:
                    mat = m.Medium(surf_data[2])
            else:
                if surf_data[2] == 1.0:
                    mat = m.Air()
                else:
                    mat = m.Glass(surf_data[2], surf_data[3], '')

        else:  # string args
            if surf_data[2].upper() == 'REFL':
                s.interact_mode = 'reflect'
                mat = prev_medium
            else:
                num_str_args = 0
                for tkn in surf_data[2:]:
                    if isinstance(tkn, str) and len(tkn) > 0:
                        num_str_args += 1
                if num_str_args == 2:
                    name, cat = surf_data[2], surf_data[3]
                elif num_str_args == 1:
                    name, cat = surf_data[2].split(',')
                elif num_str_args == 0:
                    mat = m.Air()

                if num_str_args > 0:
                    try:
                        mat = gfact.create_glass(name, cat)
                    except ge.GlassNotFoundError as gerr:
                        logging.info('%s glass data type %s not found',
                                     gerr.catalog,
                                     gerr.name)
                        logging.info('Replacing material with air.')
                        mat = m.Air()

        if len(surf_data) >= 5:
            s.set_max_aperture(surf_data[4])

    else:  # only curvature and thickness entered, set material to air
        mat = m.Air()

    thi = surf_data[1]
    g = gap.Gap(thi, mat)
    rndx = mat.rindex(wvl)
    tfrm = np.identity(3), np.array([0., 0., thi])

    return s, g, rndx, tfrm
