import numpy
import copy
from collections import defaultdict

from chimerax.core.triggerset import TriggerSet
from chimerax.atomic import AtomicStructure, concatenate
from chimerax.geometry import Place, Places
from chimerax.geometry import find_close_points, find_close_points_sets
from chimerax.surface import zone
from chimerax.surface.shapes import sphere_geometry

from chimerax.core.models import Model, Drawing
from chimerax.std_commands import camera, cofr, cartoon
from chimerax.core.commands import atomspec
from chimerax.map_data import ArrayGridData
from chimerax.map import Volume, volumecommand

from .mousemodes import initialize_map_contour_mouse_modes
from .main import atom_list_from_sel
#from . import clipper
from .clipper_mtz import ReflectionDataContainer

DEFAULT_BOND_RADIUS = 0.2

def move_model(session, model, new_parent):
    '''
    Temporary method until something similar is added to the ChimeraX
    core. Picks up a model from the ChimeraX model tree and transplants
    it (with all its children intact) as the child of a different model.
    '''
    model._auto_style = False
    mlist = model.all_models()
    model_id = model.id
    if new_parent in mlist:
        raise RuntimeError('Target model cannot be one of the models being moved!')
    for m in mlist:
        m.removed_from_session(session)
        mid = m.id
        if mid is not None:
            del session.models._models[mid]
            m.id = None
    session.triggers.activate_trigger('remove models', mlist)
    if len(model_id) == 1:
        parent = session.models.drawing
    else:
        parent = session.models._models[model_id[:-1]]
    parent.remove_drawing(model, delete=False)
    parent._next_unused_id = None
    new_parent.add([model])

def symmetry_from_model_metadata(model):
    if 'CRYST1' in model.metadata.keys():
        return symmetry_from_model_metadata_pdb(model)
    elif 'cell' in model.metadata.keys():
        return symmetry_from_model_metadata_mmcif(model)
    raise TypeError('Model does not appear to have symmetry information!')

def symmetry_from_model_metadata_mmcif(model):
    metadata = model.metadata
    try:
        cell_headers = metadata['cell']
        cell_data = metadata['cell data']
        cell_dict = dict((key, data) for (key, data) in zip(cell_headers, cell_data))
        abc = [cell_dict['length_a'], cell_dict['length_b'], cell_dict['length_c']]
        angles = [cell_dict['angle_alpha'], cell_dict['angle_beta'], cell_dict['angle_gamma']]
    except:
        raise TypeError('No cell information available!')

    try:
        spgr_headers = metadata['symmetry']
        spgr_data = metadata['symmetry data']
    except:
        raise TypeError('No space group headers in metadata!')

    sprg_dict = dict((key, data) for (key, data) in zip(spgr_headers, spgr_data))
    spgr_str = spgr_dict['int_tables_number']
    if spgr_str == '?':
        spgr_str = spgr_dict['space_group_name_h-m']
    if spgr_str == '?':
        spgr_str = spgr_dict['space_group_name_hall']
    if spgr_str == '?':
        raise TypeError('No space group information available!')

    # TODO: ChimeraX does not currently keep refinement metadata (including resolution)
    res = 3.0
    from .clipper_python import Cell_descr, Cell, Spgr_descr, Spacegroup, Resolution, Grid_sampling
    cell_descr = Cell_descr(*abc, *angles)
    cell = Cell(cell_descr)
    spgr_descr = Spgr_descr(spgr_str)
    spacegroup = Spacegroup(spgr_descr)
    resolution = Resolution(res)
    grid_sampling = Grid_sampling(spacegroup, cell, resolution)
    return cell, spacegroup, grid_sampling

def symmetry_from_model_metadata_pdb(model):
    '''
    Generate Cell, Spacegroup and a default Grid_Sampling from the PDB
    CRYST1 card.
    '''
    cryst1 = model.metadata['CRYST1'][0].split()
    abc = [float(a) for a in cryst1[1:4]]
    angles = [float(c) for c in cryst1[4:7]]
    symstr = ' '.join(cryst1[7:])

    remarks = model.metadata['REMARK']
    i = 0

    '''
    Get the resolution. We need this to define a Grid_sampling
    for the unit cell (needed even in the absence of a map since
    atomic symmetry lookups are done with integerised symops for
    performance). We want to be as forgiving as possible at this
    stage - we'll use the official resolution if we find it, and
    set a default resolution if we don't. This will be overridden
    by the value from any mtz file that's loaded later.
    '''
    try:
        while 'REMARK   2' not in remarks[i]:
            i += 1
        # The first 'REMARK   2' line is blank by convention, and
        # resolution is on the second line
        i += 1
        line = remarks[i].split()
        res = line[3]
    except:
        res = 3.0

    # '''
    # The spacegroup identifier tends to be the most unreliable part
    # of the CRYST1 card, so it's considered safer to let Clipper
    # infer it from the list of symmetry operators at remark 290. This
    # typically looks something like the following:
    #
    # REMARK 290      SYMOP   SYMMETRY
    # REMARK 290     NNNMMM   OPERATOR
    # REMARK 290       1555   X,Y,Z
    # REMARK 290       2555   -X,-Y,Z+1/2
    # REMARK 290       3555   -Y+1/2,X+1/2,Z+1/4
    # REMARK 290       4555   Y+1/2,-X+1/2,Z+3/4
    # REMARK 290       5555   -X+1/2,Y+1/2,-Z+1/4
    # REMARK 290       6555   X+1/2,-Y+1/2,-Z+3/4
    # REMARK 290       7555   Y,X,-Z
    # REMARK 290       8555   -Y,-X,-Z+1/2
    #
    # Clipper is able to initialise a Spacegroup object from a
    # string containing a semicolon-delimited list of the symop
    # descriptors in the SYMMETRY OPERATOR column, so we need to
    # parse those out.
    # '''
    # # Find the start of the REMARK 290 section
    # while 'REMARK 290' not in remarks[i]:
    #     i += 1
    # while 'NNNMMM' not in remarks[i]:
    #     i += 1
    # i+=1
    # symstr = ''
    # thisline = remarks[i]
    # while 'X' in thisline and 'Y' in thisline and 'Z' in thisline:
    #     if len(symstr):
    #         symstr += ';'
    #     splitline = thisline.split()
    #     symstr += splitline[3]
    #     i+=1
    #     thisline = remarks[i]

    from .clipper_python import Cell_descr, Cell, Spgr_descr, Spacegroup, Resolution, Grid_sampling

    cell_descr = Cell_descr(*abc, *angles)
    cell = Cell(cell_descr)
    spgr_descr = Spgr_descr(symstr)
    spacegroup = Spacegroup(spgr_descr)
    resolution = Resolution(float(res))
    grid_sampling = Grid_sampling(spacegroup, cell, resolution)
    return cell, spacegroup, grid_sampling

class CrystalStructure(Model):
    '''
    Master container class for a crystal structure, designed to act as the
    head node of a Model tree with the following general format:

    CrystalStructure
      |
      -- master model (AtomicStructure)
      |
      -- symmetry copies (SymModels)
      |    |
      |    -- AtomicStructure
      |    |
      |    -- AtomicStructure
      |    |
      |    -- ...
      |
      -- reciprocal space data (ReflectionDataContainer)
      |    |
      |    -- Free Flags (ReflectionDataFreeFlags)
      |    |
      |    -- Experimental (ReflectionDataNode)
      |    |    |
      |    |    -- F/SigF (ReflectionDataExp)
      |    |    |
      |    |    -- ...
      |    |
      |    -- Calculated (ReflectionDataCalc)
      |         |
      |         -- 2mFo-DFc (ReflectionDataCalc)
      |         |
      |         -- ...
      |
      |
      -- real-space maps (XMapSet)
           |
           -- 2mFo-Fc (Volume)
           |
           -- ...
    '''

    def __init__(self, session, model,
                 mtzfile = None, calculate_maps = True,
                 map_oversampling = 1.5,
                 live_map_scrolling = True, map_scrolling_radius = 12,
                 live_atomic_symmetry = True, atomic_symmetry_radius = 15,
                 show_nonpolar_H = False):
        '''
        Create a new crystal structure object from an atomic model and
        (optionally) a set of reciprocal-space data.
        Args:
            session:
                The ChimeraX session.
            model:
                A loaded AtomicStructure model. NOTE: this will be moved from
                its existing place in the session.models tree to become a
                child of this one.
            mtzfile (string):
                The name of an MTZ file containing experimental and/or
                calculated amplitudes and phases
            calculate_maps (bool, default True):
                If an MTZ file containing map structure factors is provided,
                generate an XmapSet containing the real-space maps.
            map_oversampling (float, default 1.5):
                The Shannon rate at which the map should be sampled.
                A rate of 1 means the grid spacing will be half the true
                resolution (the minimum required). The default value of
                1.5 gives a grid spacing of resolution/3.
            live_map_scrolling (bool, default True):
                If maps are generated, initialise live scrolling of a
                sphere of density around the centre of rotation.
            map_scrolling_radius (float, default 12):
                The radius (in Angstroms) of the live map sphere.
            live_atomic_symmetry (bool, default True):
                If true, switch to a mode in which atoms (including
                symmetry atoms) within the given radius are displayed as
                sticks, while the remainder of the structure is shown as
                a minimal cartoon. Cartoons of symmetry-related molecules
                will be automatically shown/hidden as they enter/leave
                the sphere.
            atomic_symmetry_radius (float, default 15):
                Radius (in Angstroms) of the live atomic symmetry sphere.
            show_nonpolar_H (bool, default False):
                Do you want non-polar hydrogens to be visible in all default
                visualisations?
        '''
        name = 'Crystal (' + model.name +')'
        Model.__init__(self, name, session)

        initialize_mouse_modes(session)
        volumecommand.volume(session, pickable=False)

        self._master_model = model
        self.mtzdata = None
        if mtzfile is not None:
            self.mtzdata = ReflectionDataContainer(self.session, mtzfile, shannon_rate = map_oversampling)
            self.add([self.mtzdata])
            self.cell = self.mtzdata.cell
            self.spacegroup = self.mtzdata.spacegroup
            self.grid = self.mtzdata.grid_sampling
            self.hklinfo = self.mtzdata.hklinfo
        else:
            self.cell, self.spacegroup, self.grid = symmetry_from_model_metadata(model)

        '''
        Named triggers to simplify handling of changes on key events (e.g. live
        updating of maps/symmetry).
        '''
        if not hasattr(self, 'triggers'):
            self.triggers = TriggerSet()
        trigger_names = (
            'map box changed',  # Changed shape of box for map viewing
            'map box moved',    # Changed location of box for map viewing
            'atom box changed', # Changed shape of box for showing symmetry atoms
            'atom box moved',   # Changed location of box for showing symmetry atoms
        )
        for t in trigger_names:
            self.triggers.add_trigger(t)

        self._voxel_size = self.cell.dim / self.grid.dim

        self.show_nonpolar_H = show_nonpolar_H

        ref = model.bounds().center()

        # Convert the atoms to a format recognised by Clipper
        self._clipper_atoms = atom_list_from_sel(model.atoms)
        # A C++ object holding the symmetry operations required to pack one
        # unit cell relative to the current atomic model, along with fast
        # functions returning the symops necessary to pack a given box in
        # xyz space.
        from .clipper_python import Unit_Cell
        self.unit_cell = Unit_Cell(ref,
                    self._clipper_atoms, self.cell, self.spacegroup, self.grid)

        # Container for managing all the symmetry copies
        self._sym_model_container = None

        # Container for drawing of special positions
        self._special_positions_model = None

        # Tool for "stepping" through the model by secondary structure
        self._stepper = None

        ###
        # LIVE ATOMIC SYMMETRY
        ###

        # Do we want to always have the reference model shown?
        self.sym_always_shows_reference_model = True
        # Trigger handler for live display of symmetry
        self._sym_handler = None
        # Centroid of the search space for symmetry equivalents
        self._sym_box_center = None
        # Half-width of the box in which we want to search
        self._sym_box_radius = atomic_symmetry_radius
        # Grid dimensions of the box
        self._sym_box_dimensions = None
        # Is the box already initialised?
        self._sym_box_initialized = False
        # Last grid coordinate for the cofr at which the symmetry was updated
        self._sym_last_cofr_grid = None
        # Factor defining the frequency of steps in the symmetry search
        # (larger number = more steps)
        self._sym_search_frequency = 2

        # Do we want to find and show atoms in the search radius at each iteration?
        self._live_atomic_symmetry = None
        self.live_atomic_symmetry = live_atomic_symmetry

        ###
        # REAL-SPACE MAPS / LIVE SCROLLING
        ###

        self.xmaps = None
        if calculate_maps:
            if self.mtzdata is not None:
                if self.mtzdata.calculated_data is not None:
                        # We want to add the capability to generate maps from
                        # the atomic model and experimental data here, but for
                        # now we'll only add maps if amplitudes and phases were
                        # provided in the MTZ file.
                        self.xmaps = XmapSet(self.session,
                            self.mtzdata.calculated_data, self,
                            live_scrolling = live_map_scrolling,
                            display_radius = map_scrolling_radius)
                        self.add([self.xmaps])
        self.add([model])
        #move_model(self.session, model, self)
        self.session.models.add([self])

    @property
    def master_model(self):
        return self._master_model

    def delete(self):
        if self._sym_handler is not None:
            self.session.triggers.remove_handler(self._sym_handler)
        super(CrystalStructure, self).delete()

    @property
    def stepper(self):
        '''
        Provides methods for "stepping" back and forth through the
        model according to secondary structure. For example, each call
        to stepper.step_forward() (with default arguments) will return
        an atom selection corresponding to the next pair of defined
        secondary structure elements plus their flanking loops.
        '''
        if self._stepper is None:
            from .structurestepper import StructureStepper
            self._stepper = StructureStepper(self.session, self.master_model)
        return self._stepper

    @property
    def live_atomic_symmetry_radius(self):
        '''
        Set the radius (in Angstroms) of the volume in which to show
        symmetry atoms.
        '''
        return self._sym_box_radius

    @live_atomic_symmetry_radius.setter
    def live_atomic_symmetry_radius(self, radius):
        self._change_sym_box_radius(radius)

    @property
    def live_atomic_symmetry(self):
        '''
        Turn live display of symmetry atoms on or off.
        '''
        return self._live_atomic_symmetry

    @live_atomic_symmetry.setter
    def live_atomic_symmetry(self, switch):
        if switch:
            if not self._sym_box_initialized:
                self.initialize_symmetry_display(self.live_atomic_symmetry_radius)
            set_to_default_cartoon(self.session, self.master_model)
            for key, m in self.items():
                m.bonds.radii = DEFAULT_BOND_RADIUS
            self._start_live_atomic_symmetry()
        else:
            self._stop_live_atomic_symmetry()
        self._live_atomic_symmetry = switch

    @property
    def sym_model_container(self):
        if self._sym_model_container is None:
            self._sym_model_container = SymModels(self.session, self)
        return self._sym_model_container

    @property
    def display(self):
        return super().display

    @display.setter
    def display(self, switch):
        if switch:
            if self.live_atomic_symmetry:
                self._start_live_atomic_symmetry()
            if self.xmaps.live_scrolling:
                self.xmaps._start_live_scrolling()
        else:
            self._stop_live_atomic_symmetry()
            self.xmaps._stop_live_scrolling()
        Model.display.fset(self, switch)

    def items(self):
        from .clipper import RTop_frac
        return ((RTop_frac.identity(), self.master_model), *self.sym_model_container.items())

    def sym_select_within(self, coords, radius):
        '''
        Given an array of (x,y,z) coordinates, return a list of Atoms lists
        (one per symmetry equivalent molecule) covering all atoms within
        radius of any input coordinate.
        Args:
          coords:
            An (n * [x,y,z)) array of coordinates
          radius:
            Search radius in Angstroms
        '''
        c = numpy.empty((len(coords), 3), numpy.float32)
        c[:] = coords
        master_atoms = self.master_model.atoms
        master_coords = master_atoms.coords.astype(numpy.float32)
        from .clipper_python import Util
        grid_minmax = Util.get_minmax_grid(coords, self.cell, self.grid)
        pad = calculate_grid_padding(radius, self.grid, self.cell)
        grid_minmax += numpy.array((-pad, pad))
        from .clipper_python import Coord_grid
        min_xyz = Coord_grid(grid_minmax[0]).coord_frac(self.grid).coord_orth(self.cell).xyz
        dim = grid_minmax[1] - grid_minmax[0]
        symops = self.unit_cell.all_symops_in_box(min_xyz, dim)
        symmats = symops.all_matrices_orth(self.cell, format = '3x4')
        target = [(c, Place().matrix.astype(numpy.float32))]
        search_list = []
        model_list = []
        for i, s in enumerate(symops):
            search_list.append((master_coords, symmats[i].astype(numpy.float32)))
            model_list.append(self.sym_model_container[s])
        i1, i2 = find_close_points_sets(search_list, target, radius)

        found = []
        for i, c in enumerate(i1):
            if len(c):
                found.append(model_list[i].atoms[c])
        return found

    def initialize_symmetry_display(self, radius = 14):
        '''
        Continually update the display to show all symmetry equivalents of
        the atomic model which enter a box of the given size, centred on the
        centre of rotation.

        Args:
          radius (float):
            The search volume is actually the smallest parallelepiped defined
            by the same angles as the unit cell, which will contain a sphere
            of the given radius.
        '''
        if self._sym_box_initialized:
            raise RuntimeError('''
              The symmetry box is already intialised for this structure. If you
              want to reset it, run sym_box_reset().
              ''')
        camera.camera(self.session, 'ortho')
        cofr.cofr(self.session, 'centerOfView', show_pivot = True)
        self.sym_always_shows_reference_model = True
        self._sym_box_radius = radius
        uc = self.unit_cell
        v = self.session.main_view
        c = self.cell
        g = self.grid
        self._sym_box_center = v.center_of_rotation
        from .clipper_python import Coord_orth
        self._sym_last_cofr_grid = Coord_orth(self._sym_box_center).coord_frac(c).coord_grid(g)
        box_corner_grid, box_corner_xyz, self._sym_box_dimensions = find_box_params(self._sym_box_center, c, g, radius)
        # self._sym_box_dimensions = (numpy.ceil(radius / self._voxel_size * 2)).astype(int)
        self._update_sym_box(force_update = True)
        self._sym_box_initialized = True

    def _change_sym_box_radius(self, radius):
        dim = (numpy.ceil(radius / self._voxel_size * 2)).astype(int)
        self._sym_box_dimensions = dim
        self._sym_box_radius = radius

    def _update_sym_box(self, *_, force_update = False):
        v = self.session.main_view
        uc = self.unit_cell
        cofr = v.center_of_rotation
        from .clipper_python import Coord_orth
        cofr_grid = Coord_orth(cofr).coord_frac(self.cell).coord_grid(self.grid)
        if not force_update:
            if self._sym_last_cofr_grid is not None:
                if cofr_grid == self._sym_last_cofr_grid:
                    return
        self._sym_last_cofr_grid = cofr_grid
        self._sym_box_center = cofr
        box_corner_grid, box_corner_xyz, dim = find_box_params(
                  self._sym_box_center, self.cell,
                  self.grid, self._sym_box_radius)
        symops = uc.all_symops_in_box(box_corner_xyz,
                  self._sym_box_dimensions.astype(numpy.int32), self.sym_always_shows_reference_model,
                  sample_frequency = self._sym_search_frequency)
        l1 = []
        search_entry = [(numpy.array([cofr], numpy.float32), Place().matrix.astype(numpy.float32))]
        coords = self.master_model.atoms.coords.astype(numpy.float32)
        display_mask = numpy.array([False]*len(coords))
        if not self.show_nonpolar_H:
            from .util import exclude_nonpolar_hydrogens
            h_mask = exclude_nonpolar_hydrogens(self.master_model.atoms)
            # h_mask = self.master_model.atoms.idatm_types != 'HC'
        for s in symops:
            this_model = self.sym_model_container[s]
            this_set = (coords, s.rtop_orth(self.cell).mat34.astype(numpy.float32))
            l1.append(this_set)
        i1, i2 = find_close_points_sets(l1, search_entry, self.live_atomic_symmetry_radius)
        for i, indices in enumerate(i1):
            this_model = self.sym_model_container[symops[i]]
            if len(indices):
                display_mask[:] = False
                a = this_model.atoms
                display_mask[a.indices(a[indices].residues.atoms)] = True
                if not self.show_nonpolar_H:
                    display_mask = numpy.logical_and(display_mask, h_mask)
                a.displays = display_mask
                #if this_model is not self.master_model:
                    #a.residues.ribbon_displays = display_mask
                this_model.display = True
            else:
                if this_model is not self.master_model:
                    this_model.display = False
        for s, m in self.sym_model_container.items():
            if s not in symops:
                m.display = False

    def show_large_scale_symmetry(self, box_width = 200):
        '''
        Show the model symmetry over a large volume by tiling the current
        representation of the master model. NOTE: this will automatically
        stop live atomic symmetry updating.
        '''
        self.live_atomic_symmetry = False
        box_center = self.master_model.bounds().center()
        uc = self.unit_cell
        # dim = (numpy.ceil(box_width / self._voxel_size)).astype(int)
        box_corner_grid, box_corner_xyz, dim = find_box_params(
                  box_center, self.cell,
                  self.grid, box_width/2)
        symops = uc.all_symops_in_box(box_corner_xyz, dim, True)
        num_symops = len(symops)
        sym_matrices = symops.all_matrices_orth(self.cell, format = '3x4')
        self.master_model.positions = Places(place_array=sym_matrices)

    def hide_large_scale_symmetry(self, restart_live_atomic_symmetry = True):
        self.master_model.position = Place()
        self.live_atomic_symmetry = restart_live_atomic_symmetry

    def _start_live_atomic_symmetry(self):
        self.hide_large_scale_symmetry(restart_live_atomic_symmetry = False)
        set_to_default_cartoon(self.session)
        if self._sym_handler is None:
            self._sym_handler = self.session.triggers.add_handler('new frame', self._update_sym_box)

    def _stop_live_atomic_symmetry(self):
        if self._sym_handler is not None:
            self.session.triggers.remove_handler(self._sym_handler)
            self._sym_handler = None
        for key, m in self.sym_model_container.items():
            m.display = False

    def isolate_and_cover_selection(self, atoms, include_surrounding_residues = 5,
                        show_context = 5, mask_radius = 3, extra_padding = 0, hide_surrounds = True, focus = True):
        '''
        Expand the map to cover a given atomic selection, then mask it to
        within a given distance of said atoms to reduce visual clutter. Adjust
        the atomic visualisation to show only the selected atoms, plus an
        optional surrounding buffer zone.
        Args:
          atoms (ChimeraX Atoms object):
            The main selection we're interested in. The existing selection will
            be expanded to include the whole residue for every selected atom.
          include_surrounding_residues (float):
            Any residue with an atom coming within this radius of the primary
            selection will be added to the selection covered by the map. To
            cover only the primary selection, set this value to zero.
          show_context (float):
            Any residue within an atom coming within this radius of the previous
            two selections will be displayed as a thinner stick representation,
            but will not be considered for the map masking calculation.
          mask_radius (float):
            Components of the map more than this distance from any atom will
            be hidden.
          pad (float):
            Optionally, further pad the volume by this distance. The extra
            volume will be hidden, but available for calculations.
          hide_surrounds (bool):
            If true, all residues outside the selection region will be hidden
          focus (bool):
            If true, the camera will be moved to focus on the selection
        '''
        # If we're in live mode, turn it off
        self.live_atomic_symmetry = False
        self.xmaps.live_scrolling = False
        orig_atoms = atoms
        atoms = atoms.residues.atoms
        coords = atoms.coords
        if include_surrounding_residues > 0:
            atoms = concatenate(
              self.sym_select_within(
                  coords, include_surrounding_residues)).residues.atoms
            coords = atoms.coords
        context_atoms = None
        if show_context > 0:
            context_atoms = concatenate(
              self.sym_select_within(
                  coords, show_context)).residues.atoms.subtract(atoms)
        pad = calculate_grid_padding(mask_radius, self.grid, self.cell)
        ep = calculate_grid_padding(extra_padding, self.grid, self.cell)
        from .clipper_python import Util
        box_bounds_grid = Util.get_minmax_grid(coords, self.cell, self.grid) \
                                + numpy.array((-pad, pad)) + numpy.array((-ep, ep))
        self.xmaps.set_box_limits(box_bounds_grid)

        self.xmaps._surface_zone.update(mask_radius, atoms, None)
        self.xmaps._reapply_zone()
        if context_atoms is None:
            found_models = atoms.unique_structures
        else:
            found_models = concatenate((context_atoms, atoms)).unique_structures
        self.master_model.bonds.radii = 0.05
        for key, m in self.items():
            if m not in found_models:
                m.display = False
            else:
                m.display = True
                m.atoms.displays = False
                m.residues.ribbon_displays = False
        self.master_model.atoms[numpy.in1d(
            self.master_model.atoms.names, numpy.array(
                ['N','C','CA']))].displays = True
        if not self.show_nonpolar_H:
            from .util import exclude_nonpolar_hydrogens
            atoms = atoms.filter(exclude_nonpolar_hydrogens(atoms))
        atoms.displays = True
        atoms.intra_bonds.radii = 0.2
        atoms.residues.ribbon_displays = True
        if context_atoms is not None:
            if not self.show_nonpolar_H:
                context_atoms = context_atoms.filter(context_atoms.idatm_types != 'HC')
            context_atoms.displays = True
            context_atoms.intra_bonds.radii = 0.1
        if focus:
            self.session.main_view.view_all(atoms.scene_bounds, 0.2)
        # return the original selection in case we want to re-run with modified settings
        return orig_atoms

    def draw_unit_cell_and_special_positions(self, offset = None):
        '''
        Quick-and-dirty drawing mapping out the special positions
        (positions which map back to themselves by at least one
        non-unity symop) within one unit cell. A sphere will be drawn
        at each grid-point with non-unit multiplicity, and colour-coded
        according to multiplicity:
            2-fold: white
            3-fold: cyan
            4-fold: yellow
            6-fold: magenta

        Ultimately it would be nice to replace this with something more
        elegant, that masks and scrolls continuously along with the model/
        map visualisation.

        Args:
            offset (1x3 numpy array, default = None):
                Optional (u,v,w) offset (in fractions of a unit cell axis)
        '''
        m = self._special_positions_model
        if not (m is None or m.deleted):
            # Just show the existing model
            m.display = True
            return
        model = self.master_model

        ref = model.bounds().center().astype(float)
        from .clipper_python import Coord_orth, Coord_frac
        frac_coords = Coord_orth(ref).coord_frac(self.cell).uvw
        if offset is None:
            offset = numpy.array([0,0,0],int)

        positions = []
        colors = []
        rgba_corner = numpy.array([255,0,255,128],numpy.int32)
        corners_frac = numpy.array([[0,0,0],[0,0,1],[0,1,0],[0,1,1],[1,0,0],[1,0,1],[1,1,0],[1,1,1]],numpy.double) + offset\
                        + self.unit_cell.min.coord_frac(self.grid).uvw

        corners = []
        for c in corners_frac:
            co = Coord_frac(c).coord_orth(self.cell)
            positions.append(Place(axes=numpy.identity(3)*4, origin=co.xyz))
            colors.append(rgba_corner)

        m = self.special_positions_model = Model('Special Positions',self.session)
        xmap = self.xmaps.child_models()[0].xmap
        uc = self.unit_cell
        spc = numpy.array(xmap.special_positions_unit_cell_xyz(uc, offset))
        d = Drawing('points')
        sphere = numpy.array(sphere_geometry(80))
        sphere[0]*=0.25
        d.set_geometry(*sphere)
        #d.vertices, d.normals, d.triangles = sphere

        if len(spc):

            coords = spc[:,0:3]
            multiplicity = spc[:,3].astype(int)
            scale_2fold = numpy.identity(3)
            scale_3fold = numpy.identity(3)* 1.5
            scale_4fold = numpy.identity(3)* 2
            scale_6fold = numpy.identity(3)* 3
            rgba_2fold = numpy.array([255,255,255,255],numpy.int32)
            rgba_3fold = numpy.array([0,255,255,255],numpy.int32)
            rgba_4fold = numpy.array([255,255,0,255],numpy.int32)
            rgba_6fold = numpy.array([255,0,0,255],numpy.int32)

            for coord, mult in zip(coords, multiplicity):
                if mult == 2:
                    positions.append(Place(axes=scale_2fold, origin=coord))
                    colors.append(rgba_2fold)
                elif mult == 3:
                    positions.append(Place(axes=scale_3fold, origin=coord))
                    colors.append(rgba_3fold)
                elif mult == 4:
                    positions.append(Place(axes=scale_4fold, origin=coord))
                    colors.append(rgba_4fold)
                elif mult == 6:
                    positions.append(Place(axes=scale_6fold, origin=coord))
                    colors.append(rgba_6fold)
            for c in corners:
                positions.append(Place(axes=scale_6fold, origin=c))
                colors.append(rgba_corner)

        d.positions = Places(positions)
        d.colors = numpy.array(colors)
        m.add_drawing(d)
        model.parent.add([m])
        m.display = True

class SymModels(defaultdict):
    '''
    Handles creation, destruction and organisation of symmetry copies
    of an atomic model. Uses the Clipper RTop_frac object for the given
    symop as the dict key. If the key is not found, automatically creates
    a copy of the master model, sets colours, applies the Place transform
    for the symop, and adds the model to the session.
    '''
    def __init__(self, session, parent):
        '''
        Just create an empty dict.
        Args:
          session:
            The ChimeraX session.
          parent:
            The AtomicCrystalStructure describing the master model and symmetry
        '''
        self.session = session
        self.parent = parent
        self.master = parent.master_model

        # Add a sub-model to the master to act as a container for the
        # symmetry copies
        self._sym_container = None

    @property
    def sym_container(self):
        if self._sym_container is None or self._sym_container.deleted:
            self._sym_container = Model('symmetry equivalents', self.session)
            self.parent.add([self._sym_container])
            self.clear()
        return self._sym_container

    def __missing__(self, key):
        from .clipper_python import RTop_frac
        if type(key) is not RTop_frac:
            raise TypeError('Key must be a clipper.RTop_frac!')
        thisplace = Place(matrix=key.rtop_orth(self.parent.cell).mat34)
        if not thisplace.is_identity():
            thismodel = self.master.copy(name=key.format_as_symop)
            atoms = thismodel.atoms
            #thismodel.position = thisplace
            atoms.coords = thisplace.transform_points(atoms.coords)
            atom_colors = atoms.colors
            atom_colors[:,0:3] = (self.master.atoms.colors[:,0:3].astype(float)*0.6).astype(numpy.uint8)
            atoms.colors = atom_colors
            ribbon_colors = thismodel.residues.ribbon_colors
            ribbon_colors[:,0:3] = (self.master.residues.ribbon_colors[:,0:3].astype(float)*0.6).astype(numpy.uint8)
            thismodel.residues.ribbon_colors = ribbon_colors
            self.sym_container.add([thismodel])
            set_to_default_cartoon(self.session, thismodel)
            self[key] = thismodel
            thismodel.display = False
            return thismodel
        return self.master

    def __getitem__(self, key):
        s = self.sym_container
        m = super(SymModels, self).__getitem__(key)
        return m

def set_to_default_cartoon(session, model = None):
    '''
    Adjust the ribbon representation to provide information without
    getting in the way.
    '''
    try:
        if model is None:
            atoms = None
        else:
            arg = atomspec.AtomSpecArg('thearg')
            atoms = arg.parse('#' + model.id_string, session)[0]
        cartoon.cartoon(session, atoms = atoms, suppress_backbone_display=False)
        cartoon.cartoon_style(session, atoms = atoms, width=0.4, thickness=0.1, arrows_helix=True, arrow_scale = 2)
        cartoon.cartoon_tether(session, structures=atoms, opacity=0)
        from chimerax import nucleotides
        nucleotides.cmd.nucleotides(session, 'atoms')
    except:
        return

#TODO: update SurfaceZone class to handle symmetry atoms
class SurfaceZone:
    '''
    Add this as a property to a Volume object to provide it with the
    necessary information to update its triangle mask after re-contouring.
    '''
    def __init__(self, distance, atoms = None, coords = None):
        '''
        Args:
          distance (float in Angstroms):
            distance from points to which the map will be masked
          atoms:
            Atoms to mask to (coordinates will be updated upon re-masking)
          coords:
            (x,y,z) coordinates to mask to (will not be updated upon
            re-masking).

          Set both atoms and coords to None to disable automatic re-masking.
        '''
        self.update(distance, atoms, coords)

    def update(self, distance, atoms = None, coords = None):
        self.distance = distance
        self.atoms = atoms
        self.coords = coords

    @property
    def all_coords(self):
        if self.atoms is not None:
            if self.coords is not None:
                return numpy.concatenate(self.atoms.coords, self.coords)
            return self.atoms.coords
        return self.coords

# def surface_zones(models, points, distance):
#     '''
#     Essentially a copy of chimerax.surface.zone.surface_zone, but uses
#     find_close_points_sets to eke a little extra performance
#     '''
#     vlist = []
#     dlist = []
#     ident_matrix = Place().matrix.astype(numpy.float32)
#     search_entry = [(numpy.array(points, numpy.float32), Place().matrix.astype(numpy.float32))]
#     for m in models:
#         #for d in m.child_drawings():
#         for d in m.surfaces:
#             if not d.display:
#                 continue
#             if d.vertices is not None:
#                 dlist.append(d)
#                 vlist.append((d.vertices.astype(numpy.float32), ident_matrix))
#
#     i1, i2 = find_close_points_sets(vlist, search_entry, distance)
#
#     for vp, i, d in zip(vlist, i1, dlist):
#         v = vp[0]
#         nv = len(v)
#         mask = numpy.zeros((nv,), numpy.bool)
#         numpy.put(mask, i, 1)
#         t = d.triangles
#         if t is None:
#             return
#         tmask = numpy.logical_and(mask[t[:,0]], mask[t[:,1]])
#         numpy.logical_and(tmask, mask[t[:,2]], tmask)
#         d.triangle_mask = tmask

def surface_zones(models, points, distance):
    from chimerax.surface import zone
    for m in models:
        for s in m.surfaces:
            spoints = s.position.inverse(is_orthonormal=True) * points
            zone.surface_zone(s, spoints, distance, auto_update=True)

class XmapSet(Model):
    '''
    Each crystal will typically have multiple maps associated with it -
    the standard 2mFo-DFc and mFo-DFc maps, for a start, but also
    potentially sharpened/smoothed versions of these, anomalous difference
    maps, omit maps, etc. etc. etc.

    XmapSet is designed as a class to contain and organise these, control
    their display and recalculation, etc.
    '''

    STANDARD_LOW_CONTOUR = numpy.array([1.5])
    STANDARD_HIGH_CONTOUR = numpy.array([2.0])
    STANDARD_DIFFERENCE_MAP_CONTOURS = numpy.array([-3.0, 3.0])

    DEFAULT_MESH_MAP_COLOR = [0,1.0,1.0,1.0] # Solid cyan
    DEFAULT_SOLID_MAP_COLOR = [0,1.0,1.0,0.4] # Transparent cyan
    DEFAULT_DIFF_MAP_COLORS = [[1.0,0,0,1.0],[0,1.0,0,1.0]] #Solid red and green

    SESSION_SAVE=False

    def __init__(self, session, crystal, datasets = None,
                 live_scrolling = True, display_radius = 12,
                 atom_selection = None, padding = 3):
        '''
        Args:
            session:
                The ChimeraX session
            datasets:
                An iterable of ReflectionDataCalc objects
            crystal:
                The CrystalStructure object this object will be subordinate to.
            live_scrolling:
                If True, the maps will be initialised in live scrolling
                mode, displaying a sphere of density centred on the
                centre of rotation. This option takes precedence over
                atom_selection.
            display_radius:
                The radius (in Angstroms) of the display sphere used in
                live scrolling mode.
            atom_selection:
                If live_scrolling is False, this argument must provide a
                selection of atoms around which the maps will be masked.
            padding:
                The radius (in Angstroms) of the mask surrounding each
                atom in atom_selection.
        '''
        if not live_scrolling and not atom_selection:
            raise TypeError('''
            If live_scrolling is False, you must provide a set of atoms
            to mask the maps to!
            ''')
        Model.__init__(self, 'Real-space maps', session)
        self.crystal = crystal
        if not hasattr(self, 'triggers'):
            from chimerax.core.triggerset import TriggerSet
            self.triggers = TriggerSet()
        trig = self.triggers
        trigger_names = (
            'map box changed',  # Changed shape of box for map viewing
            'map box moved',    # Just changed the centre of the box
        )
        for t in trigger_names:
            trig.add_trigger(t)
        #############
        # Variables involved in handling live redrawing of maps in a box
        # centred on the cofr
        #############

        # Handler for live box update
        self._box_update_handler = None
        # Is the box already initialised?
        self._box_initialized = False
        # Object storing the parameters required for masking (used after
        # adjusting contours)
        self._surface_zone = SurfaceZone(display_radius, None, None)
        # Is the map box moving with the centre of rotation?
        self._live_scrolling = False
        # Radius of the sphere in which the map will be displayed when
        # in live-scrolling mode
        self.display_radius = display_radius
        # Actual box dimensions in (u,v,w) grid coordinates
        self._box_dimensions = None
        # Centre of the box (used when tracking the centre of rotation)
        self._box_center = None
        # Last grid coordinate of the box centre. We only need to update
        # the map if the new centre maps to a different integer grid point
        self._box_center_grid = None
        # Minimum corner of the box in (x,y,z) coordinates. This must
        # correspond to the grid coordinate in self._box_corner_grid
        self._box_corner_xyz = None
        # Minimum corner of the box in grid coordinates
        self._box_corner_grid = None

        self.live_scrolling = live_scrolling

        if not live_scrolling:
            # Get the initial box parameters based on atom_selection and padding
            self._box_corner_grid, self._box_corner_xyz, self._box_dimensions = \
                _get_bounding_box(atom_selection.coords, padding, self.grid, self.cell)
            self._surface_zone.update(padding, atoms = atom_selection)

        self._box_initialized = True

        if datasets is not None:
            for dataset in datasets:
                print('Working on dataset: {}'.format(dataset.name))
                self.add_xmap_handler(dataset)

        self.display=False
        # Apply the surface mask
        self.session.triggers.add_handler('frame drawn', self._rezone_once_cb)
        # self._reapply_zone()

    def _rezone_once_cb(self, *_):
        self.display = True
        from chimerax.core.triggerset import DEREGISTER
        return DEREGISTER

    @property
    def hklinfo(self):
        return self.crystal.hklinfo

    @property
    def spacegroup(self):
        return self.crystal.spacegroup

    @property
    def cell(self):
        return self.crystal.cell

    @property
    def res(self):
        return self.hklinfo.resolution

    @property
    def grid(self):
        return self.crystal.grid

    @property
    def voxel_size(self):
        return self.cell.dim / self.grid.dim

    @property
    def voxel_size_frac(self):
        return 1/ self.grid.dim

    @property
    def unit_cell(self):
        return self.crystal.unit_cell

    @property
    def display_radius(self):
        return self._display_radius

    @display_radius.setter
    def display_radius(self, radius):
        '''Set the radius (in Angstroms) of the live map display sphere.'''
        self._display_radius = radius
        v = self.session.main_view
        cofr = self._box_center = v.center_of_rotation
        from .clipper_python import Coord_orth
        self._box_center_grid = Coord_orth(cofr).coord_frac(self.cell).coord_grid(self.grid)
        # dim = self._box_dimensions = \
        #     2 * calculate_grid_padding(radius, self.grid, self.cell)
        self._box_corner_grid, self._box_corner_xyz, self._box_dimensions = find_box_params(
            cofr, self.cell, self.grid, radius)
        self.triggers.activate_trigger('map box changed',
            (self._box_corner_xyz, self._box_corner_grid, self._box_dimensions))
        self._surface_zone.update(radius, coords = numpy.array([cofr]))
        self._reapply_zone()

    @property
    def live_scrolling(self):
        '''Turn live map scrolling on and off.'''
        return self._live_scrolling

    @live_scrolling.setter
    def live_scrolling(self, switch):
        if switch:
            self.position = Place()
            if not self._live_scrolling:
                '''
                Set the box dimensions to match the stored radius.
                '''
                self.display_radius = self._display_radius
            self._start_live_scrolling()
        else:
            self._stop_live_scrolling()

    @property
    def display(self):
        return super().display

    @display.setter
    def display(self, switch):
        if switch:
            if self.live_scrolling:
                self._start_live_scrolling()
        Model.display.fset(self, switch)

    def _start_live_scrolling(self):
        if self._box_update_handler is None:
            self._box_update_handler = self.crystal.triggers.add_handler(
                'box center moved', self.update_box)
        if self._box_center is not None:
            self.update_box(None, (self._box_center, self._box_center_grid), force=True)
        self.positions = Places()
        self._live_scrolling = True

    def _stop_live_scrolling(self):
        if self._box_update_handler is not None:
            self.crystal.triggers.remove_handler(self._box_update_handler)
            self._box_update_handler = None
        self._live_scrolling = False

    def __getitem__(self, name_or_index):
        '''Get one of the child maps by name or index.'''
        if type(name_or_index) == str:
            for m in self.child_models():
                if m.name == name_or_index:
                    return m
            raise KeyError('No map with that name!')
        else:
            return self.child_models()[name_or_index]

    def set_box_limits(self, minmax):
        '''
        Set the map box to fill a volume encompassed by the provided minimum
        and maximum grid coordinates. Automatically turns off live scrolling.
        '''
        self.live_scrolling = False
        from .clipper_python import Coord_grid
        cmin = Coord_grid(minmax[0])
        cmin_xyz = cmin.coord_frac(self.grid).coord_orth(self.cell).xyz
        dim = (minmax[1]-minmax[0])
        self.triggers.activate_trigger('map box changed',
            (cmin_xyz, cmin, dim))

    def cover_unit_cells(self, nuvw = [1,1,1], offset = [0,0,0]):
        '''
        Expand the map(s) to cover multiple unit cells. In order to
        maintain reasonable performance, this method cheats a little by
        filling just one unit cell and then tiling it using the graphics
        engine. This leaves some minor artefacts at the cell edges, but
        is a worthwhile tradeoff.
        Automatically turns off live scrolling.
        Args:
            nuvw (array of 3 positive integers):
                Number of unit cells to show in each direction.
            offset (array of 3 integers):
                Shifts the starting corner of the displayed volume by
                this number of unit cells in each direction.
        '''
        self.live_scrolling = False
        uc = self.unit_cell
        box_min_grid = uc.min.uvw
        # Add a little padding to the max to create a slight overlap between copies
        from .clipper_python import Coord_grid
        box_max_grid = (uc.max+Coord_grid([2,2,2])).uvw
        minmax = [box_min_grid, box_max_grid]
        self.set_box_limits(minmax)
        self._surface_zone.update(None, None)
        # Tile by the desired number of cells
        places = []
        grid_dim = self.grid.dim
        nu, nv, nw = nuvw
        ou, ov, ow = offset
        for i in range(ou, nu+ou):
            for j in range(ov, nv+ov):
                for k in range(ow, nw+ow):
                    thisgrid = Coord_grid(numpy.array([i,j,k])*grid_dim)
                    thisorigin = thisgrid.coord_frac(self.grid).coord_orth(self.cell).xyz
                    places.append(Place(origin = thisorigin))
        self.positions = Places(places)

    def add_nxmap_handler(self, volume):
        from .real_space_map import NXmapHandler
        m = NXmapHandler(self.session, self, volume)
        self.add([m])

    def add_xmap_handler(self, dataset, is_difference_map = None,
                color = None, style = None, contour = None):
        '''
        Add a new XmapHandler based on the given reflections and phases.
        Args:
            dataset:
                a ReflectionDataCalc object.
            is_difference_map:
                Decides whether this map is to be treated as a difference
                map (with positive and negative contours) or a standard
                map with positive contours only. Leave as None to allow
                this to be determined automatically from the dataset,
                or set it to True or False to force an explicit choice.
            color:
                an array of [r, g, b, alpha] as integers in the range 0-255
                for a positive density map, or
                [[r,g,b,alpha],[r,g,b,alpha]] representing the colours of
                the positive and negative contours respectively if the
                map is a difference map
            style:
                one of 'mesh' or 'surface'
            contour:
                The value(s) (in sigma units) at which to contour the map
                on initial loading. For a standard map this should be a
                single value; for a difference map it should be
                [negative contour, positive contour]
        '''
        data = dataset.data
        new_xmap = Xmap(self.spacegroup, self.cell, self.grid, name = dataset.name, hkldata = data)
        if is_difference_map is None:
            is_difference_map = dataset.is_difference_map
        new_xmap.is_difference_map = is_difference_map
        if is_difference_map and color is not None and len(color) != 2:
            err_string = '''
            ERROR: For a difference map you need to define colours for
            both positive and negative contours, as:
            [[r,g,b,a],[r,g,b,a]] in order [positive, negative].
            '''
            raise TypeError(err_string)
        new_handler = XmapHandler(self.session, self, dataset.name, new_xmap,
            self._box_corner_xyz, self._box_corner_grid, self._box_dimensions)
        if style is None:
            style = 'mesh'
        if color is None:
            if is_difference_map:
                color = self.DEFAULT_DIFF_MAP_COLORS
            elif style == 'mesh':
                color = [self.DEFAULT_MESH_MAP_COLOR]
            else:
                color = [self.DEFAULT_SOLID_MAP_COLOR]
        if contour is None:
            if is_difference_map:
                contour = self.STANDARD_DIFFERENCE_MAP_CONTOURS
            else:
                contour = self.STANDARD_LOW_CONTOUR
        elif not hasattr(contour, '__len__'):
                contour = numpy.array([contour])
        else:
            contour = numpy.array(contour)
        contour = contour * new_xmap.sigma
        self.add([new_handler])
        new_handler.set_representation(style)
        new_handler.set_parameters(**{'cap_faces': False,
                                  'surface_levels': contour,
                                  'show_outline_box': False,
                                  'surface_colors': color,
                                  'square_mesh': True})
        # new_handler.update_surface()
        new_handler.show()

    def update_box(self, trigger_name, new_center, force=True):
        '''Update the map box to surround the current centre of rotation.'''
        center, center_grid = new_center
        self._box_center = center
        self._box_center_grid = center_grid
        if not self.visible and not force:
            # Just store the box parameters for when we're re-displayed
            return
        if self.live_scrolling:
            box_corner_grid, box_corner_xyz, dim = find_box_params(center, self.cell, self.grid, self.display_radius)
            self.triggers.activate_trigger('map box moved', (box_corner_xyz, box_corner_grid, self._box_dimensions))
            self._surface_zone.update(self.display_radius, coords = numpy.array([center]))
            self._reapply_zone()

    def _reapply_zone(self):
        '''
        Reapply any surface zone applied to the volume after changing box
        position.
        '''
        coords = self._surface_zone.all_coords
        radius = self._surface_zone.distance
        if coords is not None:
            surface_zones(self, coords, radius)

    def delete(self):
        self.live_scrolling = False
        super(XmapSet, self).delete()

def calculate_grid_padding2(radius, grid, cell):
    angles = cell.angles
    d = 2*radius
    # a is always parallel to x
    da = d
    # gamma is angle between a and b
    db = d/sin(angles[2])
    # beta is the angle between a and c
    dc = d/sin(angles[1])
    dabc = numpy.array([da, db, dc])
    dabc_grid = numpy.ceil(dabc/(cell.dim/grid.dim)).astype(numpy.int32)
    return dabc_grid

def calculate_grid_padding(radius, grid, cell):
    '''
    Calculate the number of grid steps needed on each crystallographic axis
    in order to capture at least radius angstroms in x, y and z.
    '''
    corner_mask = numpy.array([[0,0,0],[0,0,1],[0,1,0],[0,1,1],[1,0,0],[1,0,1],[1,1,0],[1,1,1]])
    corners = (corner_mask * radius).astype(float)
    grid_upper = numpy.zeros([8,3], int)
    grid_lower = numpy.zeros([8,3], int)
    from .clipper_python import Coord_orth
    for i, c in enumerate(corners):
        co = Coord_orth(c)
        cm = co.coord_frac(cell).coord_map(grid).uvw
        grid_upper[i,:] = numpy.ceil(cm).astype(int)
        grid_lower[i,:] = numpy.floor(cm).astype(int)
    return grid_upper.max(axis=0) - grid_lower.min(axis=0)

def find_box_params(center, cell, grid, radius = 20):
    '''
    Return the origin (in grid and cartesian coordinates) and dimensions of a
    rhombohedral box big enough to hold a sphere of the desired radius.
    '''
    from .clipper_python import Coord_orth, Coord_frac, Coord_grid
    radii_frac = Coord_frac(radius/cell.dim)
    center_frac = Coord_orth(center).coord_frac(cell)
    half_dim = calculate_grid_padding(radius, grid, cell)
    bottom_corner_grid = center_frac.coord_grid(grid) \
                - Coord_grid(half_dim)
    bottom_corner_orth = bottom_corner_grid.coord_frac(grid).coord_orth(cell)
    return bottom_corner_grid, bottom_corner_orth.xyz, half_dim*2

def _get_bounding_box(coords, padding, grid, cell):
    '''
    Find the minimum and maximum grid coordinates of a box which will
    encompass the given (x,y,z) coordinates plus padding (in Angstroms).
    '''
    from .clipper_python import Util, Coord_grid
    grid_pad = calculate_grid_padding(padding, grid, cell)
    box_bounds_grid = Util.get_minmax_grid(coords, cell, grid)\
                        + numpy.array((-grid_pad, grid_pad))
    box_origin_grid = box_bounds_grid[0]
    box_origin_xyz = Coord_grid(box_origin_grid).coord_frac(grid).coord_orth(cell)
    dim = box_bounds_grid[1] - box_bounds_grid[0]
    return [box_origin_grid, box_origin_xyz, dim]

# from .clipper_python import Xmap_float
# class Xmap(Xmap_float):
#     def __init__(self, spacegroup, cell, grid_sampling,
#                  name = None, hkldata = None, is_difference_map = None):
#         super().__init__(spacegroup, cell, grid_sampling)
#         self.name = name
#         self.is_difference_map = is_difference_map
#         if hkldata is not None:
#             self.fft_from(hkldata)
#         from .clipper_python import Map_stats
#         self._stats = Map_stats(self)
#
#     @property
#     def stats(self):
#         if self._stats is None:
#             from .clipper_python import Map_stats
#             self._stats = Map_stats(self)
#         return self._stats
#
#     @property
#     def mean(self):
#         return self.stats.mean
#
#     @property
#     def std_dev(self):
#         return self.stats.std_dev
#
#     @property
#     def sigma(self):
#         return self.stats.std_dev
#
#     @property
#     def min(self):
#         return self.stats.min
#
#     @property
#     def max(self):
#         return self.stats.max
#
#     @property
#     def range(self):
#         return self.stats.range
#
#
#
#
# class XmapHandler(Volume):
#     '''
#     An XmapHandler is in effect a resizable window into a periodic
#     crystallographic map. The actual map data (a clipper Xmap object) is
#     held within, and filled into the XmapWindow.data array as needed.
#     Methods are included for both live updating (e.g. tracking and filling
#     a box centred on the centre of rotation) and static display of a
#     given region.
#     '''
#     def __init__(self, session, manager, name, xmap, origin, grid_origin, dim):
#         '''
#         Args:
#             sesssion:
#                 The ChimeraX session
#             crystal:
#                 The CrystalStructure object this belongs to
#             name:
#                 A descriptive name for this map
#             xmap:
#                 A clipper.Xmap
#             origin:
#                 The (x,y,z) coordinates of the bottom left corner of the
#                 volume.
#             grid_origin:
#                 The (u,v,w) integer grid coordinates corresponding to
#                 origin.
#             dim:
#                 The shape of the box in (u,v,w) grid coordinates.
#         '''
#         self.box_params = (origin, grid_origin, dim)
#         self.xmap = xmap
#         self.manager = manager
#         darray = self._generate_and_fill_data_array(origin, grid_origin, dim)
#         Volume.__init__(self, darray, session)
#
#         self.is_difference_map = xmap.is_difference_map
#         self.name = name
#         self.initialize_thresholds()
#
#         # If the box shape changes while the volume is hidden, the change
#         # will not be applied until it's shown again.
#         self._needs_update = True
#         self.show()
#         self._box_shape_changed_cb_handler = self.manager.triggers.add_handler(
#             'map box changed', self._box_changed_cb)
#         self._box_moved_cb_handler = self.manager.triggers.add_handler(
#             'map box moved', self._box_moved_cb)
#
#
#
#     def show(self, *args, **kwargs):
#         if self._needs_update:
#             self._swap_volume_data(self.box_params, force_update = True)
#             self._needs_update = False
#         else:
#             # Just set the origin and fill the box with the data for
#             # the current location
#             origin, grid_origin, ignore = self.box_params
#             self._fill_volume_data(self._data_fill_target, grid_origin)
#         super(XmapHandler, self).show(*args, **kwargs)
#
#     @property
#     def hklinfo(self):
#         return self.manager.hklinfo
#
#     @property
#     def spacegroup(self):
#         return self.manager.spacegroup
#
#     @property
#     def cell(self):
#         return self.manager.cell
#
#     @property
#     def res(self):
#         return self.hklinfo.resolution
#
#     @property
#     def grid(self):
#         return self.manager.grid
#
#     @property
#     def voxel_size(self):
#         return self.cell.dim / self.grid.dim
#
#     @property
#     def voxel_size_frac(self):
#         return 1/ self.grid.dim
#
#     @property
#     def unit_cell(self):
#         return self.manager.unit_cell
#
#     @property
#     def _surface_zone(self):
#         return self.manager._surface_zone
#
#     def mean_sd_rms(self):
#         '''
#         Overrides the standard Volume method to give the overall values
#         from the Clipper object.
#         '''
#         x = self.xmap
#         # RMS is not currently calculated by Clipper, so we'll just return
#         # the sigma twice.
#         return (x.mean, x.sigma, x.sigma)
#
#
#     def _box_changed_cb(self, name, params):
#         self.box_params = params
#         self._needs_update = True
#         if not self.display:
#             # No sense in wasting cycles on this if the volume is hidden.
#             # We'll just store the params and apply them when we show the
#             # volume.
#             # NOTE: this means we need to over-ride show() to ensure
#             # it's updated before re-displaying.
#             return
#         self._swap_volume_data(params)
#         self.data.values_changed()
#         self.show()
#
#     def _box_moved_cb(self, name, params):
#         self.box_params = params
#         if not self.display:
#             return
#         self.data.set_origin(params[0])
#         self._fill_volume_data(self._data_fill_target, params[1])
#         self.data.values_changed()
#
#     def delete(self):
#         bh = self._box_shape_changed_cb_handler
#         if bh is not None:
#             self.manager.triggers.remove_handler(bh)
#             self._box_shape_changed_cb_handler = None
#         bm = self._box_moved_cb_handler
#         if bm is not None:
#             self.manager.triggers.remove_handler(bm)
#             self._box_moved_cb_handler = None
#         super(XmapHandler, self).delete()
#
#
#
#
#     def _swap_volume_data(self, params, force_update = False):
#         '''
#         Replace this Volume's data array with one of a new shape/size
#         Args:
#             params:
#                 A tuple of (new_origin, new_grid_origin, new_dim)
#         '''
#         if not self._needs_update and not force_update:
#             # Just store the parameters
#             self.box_params = params
#             return
#         new_origin, new_grid_origin, new_dim = params
#         darray = self._generate_and_fill_data_array(new_origin, new_grid_origin, new_dim)
#         self._box_dimensions = new_dim
#         self.replace_data(darray)
#         self.new_region(ijk_min=(0,0,0), ijk_max=darray.size, ijk_step=(1,1,1), adjust_step=False)
#
#     def _generate_and_fill_data_array(self, origin, grid_origin, dim):
#         data = self._data_fill_target = numpy.empty(dim, numpy.float32)
#         self._fill_volume_data(data, grid_origin)
#         order = numpy.array([2,1,0], int)
#         darray = ArrayGridData(data.transpose(), origin = origin,
#             step = self.voxel_size, cell_angles = self.cell.angles_deg)
#         return darray
#
#
#     def _fill_volume_data(self, target, start_grid_coor):
#         #shape = (numpy.array(target.shape)[::-1] - 1)
#         #end_grid_coor = start_grid_coor + clipper.Coord_grid(shape)
#         #self.data.set_origin(origin_xyz)
#         from .clipper_python import Coord_grid
#         xmap = self.xmap
#         xmap.export_section_numpy(Coord_grid(start_grid_coor), target)
#
#     # def update_drawings(self):
#     #     super().update_drawings()
#     #     if hasattr(self, '_surface_zone'):
#     #         sz = self._surface_zone
#     #         coords = sz.all_coords
#     #         distance = sz.distance
#     #         if coords is not None:
#     #             from chimerax.surface.zone import surface_zone
#     #             surface_zone(self, coords, distance)


"""
.. _ref_any:

Any
====================
"""
import traceback
import warnings

import ansys.dpf.core.server_types
from ansys.dpf.core import server as server_module
from ansys.dpf.core import errors
from ansys.dpf.core.common import type_to_internal_object_keyword
from ansys.dpf.gate import any_abstract_api

class Any:
    """Common wrapper representing any supported DPF Data Types.

    Parameters
    ----------
    any_dpf : ctypes.c_void_p, ansys.grpc.dpf.any_pb2.Any message, optional
    server : DPFServer, optional
        Server with channel connected to the remote or local instance.
        The default is ``None``, in which case an attempt is made to use the
        global server.

    Notes
    -----
    Class available with server's version starting at 7.0 (Ansys 2024 R1 pre0).
    """

    def __init__(self, any_dpf=None, server=None):
        # step 1: get server
        self._server = server_module.get_or_create_server(server)

        if any_dpf is None and not self._server.meet_version("7.0"):
            raise errors.DpfVersionNotSupported("7.0")

        self._api_instance = None

        # step 2: if object exists, take the instance, else create it
        if any_dpf is not None:
            self._internal_obj = any_dpf

        self._api.init_any_environment(self)  # creates stub when gRPC

        self._internal_type = None
        self._get_as_method = None

    @staticmethod
    def _type_to_new_from_get_as_method(any_dpf):
        from ansys.dpf.core import (
            field,
            property_field,
            generic_data_container,
            string_field,
            scoping,
            data_tree,
        )

        return [
            (
                int,
                any_dpf._api.any_new_from_int,
                any_dpf._api.any_get_as_int,
                any_dpf._api.any_new_from_int_on_client,
            ),
            (
                str,
                any_dpf._api.any_new_from_string,
                any_dpf._api.any_get_as_string,
                any_dpf._api.any_new_from_string_on_client,
            ),
            (
                float,
                any_dpf._api.any_new_from_double,
                any_dpf._api.any_get_as_double,
                any_dpf._api.any_new_from_double_on_client,
            ),
            (field.Field, any_dpf._api.any_new_from_field, any_dpf._api.any_get_as_field),
            (
                property_field.PropertyField,
                any_dpf._api.any_new_from_property_field,
                any_dpf._api.any_get_as_property_field,
            ),
            (
                string_field.StringField,
                any_dpf._api.any_new_from_string_field,
                any_dpf._api.any_get_as_string_field,
            ),
            (
                generic_data_container.GenericDataContainer,
                any_dpf._api.any_new_from_generic_data_container,
                any_dpf._api.any_get_as_generic_data_container,
            ),
            (
                scoping.Scoping,
                any_dpf._api.any_new_from_scoping,
                any_dpf._api.any_get_as_scoping,
            ),
            (
                data_tree.DataTree,
                any_dpf._api.any_new_from_data_tree,
                any_dpf._api.any_get_as_data_tree,
            ),
        ]

    @staticmethod
    def new_from(obj, server=None):
        """Return an Any instance from the given object.

        Parameters
        ----------
        obj : Object wrap as an Any

        Returns
        -------
        any : Any
            Wrapped any type.
        """

        innerServer = server if server is not None else obj._server

        if not innerServer.meet_version("7.0"):
            raise errors.DpfVersionNotSupported("7.0")

        any_dpf = Any(server=innerServer)

        for type_tuple in Any._type_to_new_from_get_as_method(any_dpf):
            if isinstance(obj, type_tuple[0]):
                # call respective new_from function
                if isinstance(server, ansys.dpf.core.server_types.InProcessServer) or not (
                    isinstance(obj, int) or isinstance(obj, str) or isinstance(obj, float)
                ):
                    any_dpf._internal_obj = type_tuple[1](obj)
                else:
                    any_dpf._internal_obj = type_tuple[3](innerServer.client, obj)
                # store get_as & type for casting back to original type
                any_dpf._internal_type = type_tuple[0]
                any_dpf._get_as_method = type_tuple[2]

                return any_dpf

        raise TypeError(f"{obj.__class__} is not currently supported by the Any class.")

    @property
    def _api(self) -> any_abstract_api.AnyAbstractAPI:
        from ansys.dpf.gate import any_capi, any_grpcapi

        if not self._api_instance:
            self._api_instance = self._server.get_api_for_type(
                capi=any_capi.AnyCAPI, grpcapi=any_grpcapi.AnyGRPCAPI
            )

        return self._api_instance

    def __str__(self):
        """Describe the entity.

        Returns
        -------
        description : str
            Description of the entity.
        """
        from ansys.dpf.core.core import _description

        return _description(self._internal_obj, self._server)

    def cast(self, output_type=None):
        """Cast the Any back to its original type.

        Parameters
        ----------
        output_type: output_type, optional
            Used when the Any instance was retrieved from the server.
            Not necessary when the instance was created using the
            :func:`ansys.dpf.core.Any.new_from`

        Returns
        -------
        type
            Original object instance
        """

        self._internal_type = output_type if output_type is not None else self._internal_type

        for type_tuple in Any._type_to_new_from_get_as_method(self):
            if self._internal_type == type_tuple[0]:
                # call the get_as function for the appropriate type
                internal_obj = type_tuple[2](self)
                if (
                    self._internal_type is int
                    or self._internal_type is str
                    or self._internal_type is float
                ):
                    obj = internal_obj
                else:
                    # get current type's constructors' variable keyword for passing the internal_obj
                    internal_obj_keyword = type_to_internal_object_keyword()[type_tuple[0]]

                    # wrap parameters in a dictionary for parameters expansion when calling
                    # constructor
                    keyword_args = {internal_obj_keyword: internal_obj, "server": self._server}
                    # call constructor
                    obj = type_tuple[0](**keyword_args)

                return obj

        raise TypeError(f"{output_type} is not currently supported by the Any class.")

    def __del__(self):
        try:
            if hasattr(self, "_deleter_func"):
                obj = self._deleter_func[1](self)
                if obj is not None:
                    self._deleter_func[0](obj)
        except Exception:
            warnings.warn(traceback.format_exc())


import re
from ...feats.filter_banks import FilterBankFactory as FBF
from .audio_feats import *

FFT = 'fft'
SPEC = 'spec'
LOG_SPEC = 'log_spec'
LOG_FB = 'logfb'
MFCC = 'mfcc'

FEAT_TYPES = [FFT, SPEC, LOG_SPEC, LOG_FB, MFCC]

class AudioFeatsFactory(object):

    @staticmethod
    def create(audio_feat, fs=16000, frame_length=25, frame_shift=10,
               fft_length=512,
               remove_dc_offset=True, preemph_coeff=0.97,
               window_type='povey', use_fft_mag=False, dither=1,
               fb_type='mel_kaldi',
               low_freq=20, high_freq=0, num_filters=23, norm_filters=False,
               num_ceps=13, snip_edges=True, cepstral_lifter=22,
               energy_floor=0, raw_energy=True, use_energy=True):

        if audio_feat == FFT:
            return Wav2FFT(
                fs, frame_length, frame_shift, fft_length,
                remove_dc_offset=remove_dc_offset, preemph_coeff=preemph_coeff,
                window_type=window_type, dither=dither,
                snip_edges=snip_edges,
                energy_floor=energy_floor, raw_energy=raw_energy,
                use_energy=use_energy)

        if audio_feat == SPEC:
            return Wav2Spec(
                fs, frame_length, frame_shift, fft_length,
                remove_dc_offset=remove_dc_offset, preemph_coeff=preemph_coeff,
                window_type=window_type, use_fft_mag=use_fft_mag, dither=dither,
                snip_edges=snip_edges,
                energy_floor=energy_floor, raw_energy=raw_energy,
                use_energy=use_energy)

        if audio_feat == LOG_SPEC:
            return Wav2LogSpec(
                fs, frame_length, frame_shift, fft_length,
                remove_dc_offset=remove_dc_offset, preemph_coeff=preemph_coeff,
                window_type=window_type, use_fft_mag=use_fft_mag, dither=dither,
                snip_edges=snip_edges,
                energy_floor=energy_floor, raw_energy=raw_energy,
                use_energy=use_energy)

        if audio_feat == LOG_FB:
            return Wav2LogFilterBank(
                fs, frame_length, frame_shift, fft_length,
                remove_dc_offset=remove_dc_offset, preemph_coeff=preemph_coeff,
                window_type=window_type, use_fft_mag=use_fft_mag, dither=dither,
                fb_type=fb_type,
                low_freq=low_freq, high_freq=high_freq,
                num_filters=num_filters, norm_filters=norm_filters,
                snip_edges=snip_edges,
                energy_floor=energy_floor, raw_energy=raw_energy,
                use_energy=use_energy)

        if audio_feat == MFCC:
            return Wav2MFCC(
                fs, frame_length, frame_shift, fft_length,
                remove_dc_offset=remove_dc_offset, preemph_coeff=preemph_coeff,
                window_type=window_type, use_fft_mag=use_fft_mag, dither=dither,
                fb_type=fb_type,
                low_freq=low_freq, high_freq=high_freq,
                num_filters=num_filters, norm_filters=norm_filters,
                num_ceps=num_ceps, snip_edges=snip_edges,
                cepstral_lifter=cepstral_lifter,
                energy_floor=energy_floor, raw_energy=raw_energy,
                use_energy=use_energy)

    @staticmethod
    def filter_args(prefix=None, **kwargs):
        """Filters MFCC args from arguments dictionary.

           Args:
             prefix: Options prefix.
             kwargs: Arguments dictionary.

           Returns:
             Dictionary with MFCC options.
        """
        if prefix is None:
            p = ''
        else:
            p = re.sub('-','_',prefix) + '_'

        valid_args = ('fs', 'frame_length', 'frame_shift',
                      'fft_length', 'remove_dc_offset', 'preemph_coeff',
                      'window_type', 'blackman_coeff', 'use_fft_mag', 'dither',
                      'fb_type', 'low_freq', 'high_freq', 'num_filters', 'norm_filters',
                      'num_ceps', 'snip_edges',
                      'energy_floor', 'raw_energy', 'use_energy', 'cepstral_lifter',
                      'audio_feat')

        d = dict((k, kwargs[p+k])
                 for k in valid_args if p+k in kwargs)
        return d

    @staticmethod
    def add_argparse_args(parser, prefix=None):
        """Adds MFCC options to parser.

           Args:
             parser: Arguments parser
             prefix: Options prefix.
        """

        if prefix is None:
            p1 = '--'
            p2 = ''
        else:
            p1 = '--' + prefix + '-'
            p2 = re.sub('-','_', prefix) + '_'

        parser.add_argument(p1+'sample-frequency', dest=(p2+'fs'),
                            default=16000, type=int,
                            help='Waveform data sample frequency (must match the waveform file, if specified there)')

        parser.add_argument(p1+'frame-length', dest=(p2+'frame_length'), type=int,
                            default=25,
                            help='Frame length in milliseconds')
        parser.add_argument(p1+'frame-shift', dest=(p2+'frame_shift'), type=int,
                            default=10,
                            help='Frame shift in milliseconds')
        parser.add_argument(p1+'fft-length', dest=(p2+'fft_length'), type=int,
                            default=512,
                            help='Length of FFT')

        parser.add_argument(p1+'remove-dc-offset', dest=(p2+'remove_dc_offset'),
                            default=True, type=str2bool,
                            help='Subtract mean from waveform on each frame')

        parser.add_argument(
            p1+'preemphasis-coeff', dest=(p2+'preemph_coeff'), type=float,
            default=0.97,
            help='Coefficient for use in signal preemphasis')

        parser.add_argument(
            p1+'window-type', dest=(p2+'window_type'),
            default='povey',
            choices=['hamming', 'hanning', 'povey', 'rectangular', 'blackman'],
            help=('Type of window ("hamming"|"hanning"|"povey"|"rectangular"|"blackmann")'))

        parser.add_argument(p1+'use-fft-mag', dest=(p2+'use_fft_mag'),
                            default=False, action='store_true',
                            help='If true, it uses |X(f)|, if false, it uses |X(f)|^2')

        parser.add_argument(
            p1+'dither', dest=(p2+'dither'), type=float,
            default=1,
            help='Dithering constant (0.0 means no dither)')

        FBF.add_argparse_args(parser, prefix)

        parser.add_argument(p1+'num-ceps', dest=(p2+'num_ceps'), type=int,
                            default=13,
                            help='Number of cepstra in MFCC computation (including C0)')

        parser.add_argument(p1+'snip-edges', dest=(p2+'snip_edges'),
                            default=True, type=str2bool,
                            help='If true, end effects will be handled by outputting only frames that completely fit in the file, and the number of frames depends on the frame-length.  If false, the number of frames depends only on the frame-shift, and we reflect the data at the ends.')

        parser.add_argument(
            p1+'energy-floor', dest=(p2+'energy_floor'), type=float,
            default=0,
            help='Floor on energy (absolute, not relative) in MFCC computation')

        parser.add_argument(p1+'raw-energy', dest=(p2+'raw_energy'),
                            default=True, type=str2bool,
                            help='If true, compute energy before preemphasis and windowing')
        parser.add_argument(p1+'use-energy', dest=(p2+'use_energy'),
                            default=True, type=str2bool,
                            help='Use energy (not C0) in MFCC computation')

        parser.add_argument(
            p1+'cepstral-lifter', dest=(p2+'cepstral_lifter'), type=float,
            default=22,
            help='Constant that controls scaling of MFCCs')

        parser.add_argument(
            p1+'audio-feat', dest=(p2+'audio_feat'),
            default='cepstrum',
            choices=['fft', 'spec', 'log_spec', 'logfb', 'mfcc' ],
            help='It can return intermediate result: fft, spec, log_spec, logfb, mfcc')


import numpy as np
import tensorflow as tf
from tensorflow import keras
from keras import layers as lr
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.utils import np_utils
from keras.datasets import mnist
from matplotlib import pyplot as plt
from tensorflow.keras.utils import plot_model

def create_model(save_plot=False):
    INPUT_IMAGE_SHAPE = 50, 60, 3
    INPUT_AUDIO_SHAPE = 43, 20  # (time, rows)
    OUTUT_SHAPE = 24, 50, 60, 3
    ACT = 'elu'

    im_inputs = lr.Input(INPUT_IMAGE_SHAPE, name='image')
    au_inputs = lr.Input(INPUT_AUDIO_SHAPE, name='audio')

    # x = (im_inputs - keras.backend.mean(im_inputs, axis=(0, 1, 2), keepdims=2)) / keras.backend.std(im_inputs, axis=(0, 1, 2), keepdims=2)

    x = lr.Conv2D(16, (3, 3), padding='same')(im_inputs)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)

    x = lr.Conv2D(16, (3, 3), padding='same')(x)
    x = lr.BatchNormalization()(x)
    inner_link = lr.Activation(ACT)(x)

    x = lr.Conv2D(32, (3, 3), padding='same', strides=2)(inner_link)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)

    x = lr.Conv2D(64, (3, 3), padding='same')(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)

    x = lr.Conv2D(128, (3, 3), padding='same')(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)

    z = lr.Reshape(au_inputs.shape[1:] + [1, 1])(au_inputs)
    x = lr.Reshape([1] + x.shape[1:])(x)

    x = lr.UpSampling3D((z.shape[1] + 5, 1, 1))(x)
    z = lr.UpSampling3D((1, 1, x.shape[3] - 5))(z)

    z = lr.Conv3DTranspose(1, 5, padding='valid')(z)
    z = lr.BatchNormalization()(z)
    z = lr.Conv3DTranspose(1, 2, padding='valid')(z)
    z = lr.BatchNormalization()(z)
    z = lr.TimeDistributed(lr.Conv2D(32, 3, padding='same'))(z)
    z = lr.BatchNormalization()(z)
    z = lr.Conv3D(32, 3, padding='same')(z)
    z = lr.BatchNormalization()(z)
    z = lr.ConvLSTM2D(64, 3, padding='same', return_sequences=True)(z)
    z = lr.BatchNormalization()(z)

    x = lr.Concatenate(axis=-1)([x, z])

    x = lr.Conv2D(128, (3, 3), padding='same')(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)

    # x = lr.ConvLSTM2D(128, 3, padding='same', return_sequences=True)(x)

    x = lr.TimeDistributed(lr.Conv2DTranspose(128, 3, strides=(2, 2), padding='same'))(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)
    x = lr.SpatialDropout3D(0.2)(x)

    # x = lr.ConvLSTM2D(128, 3, padding='same', return_sequences=True)(x)
    x = lr.AveragePooling3D((2, 1, 1))(x)

    # inner_link = lr.Reshape([1] + inner_link.shape[1:])(inner_link)
    # inner_link = lr.UpSampling3D((24, 1, 1))(inner_link)

    x = lr.TimeDistributed(lr.Conv2D(128, 3, padding='same'))(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)
    x = lr.TimeDistributed(lr.Conv2D(64, 3, padding='same'))(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)

    x = lr.Conv3D(64, 3, padding='same')(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)

    # x = lr.Concatenate(axis=-1)([x, inner_link])

    x = lr.ConvLSTM2D(64, 3, padding='same', return_sequences=True)(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)
    x = lr.ConvLSTM2D(32, 3, padding='same', return_sequences=True)(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)
    x = lr.ConvLSTM2D(3, 3, padding='same', return_sequences=True)(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)

    x = lr.TimeDistributed(lr.Conv2D(3, 2, padding='same'))(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)
    x = lr.TimeDistributed(lr.Conv2D(3, 2, padding='same'))(x)
    x = lr.Activation(ACT)(x)
    x = lr.BatchNormalization()(x)
    outputs = lr.Activation('sigmoid')(x)

    model = tf.keras.Model(inputs={'image': im_inputs,
                                   'audio': au_inputs},
                           outputs=outputs)

    if save_plot:
        plot_model(model, dpi=70, show_shapes=True)

    return model


import rhinoscriptsyntax as rs
import joinCrvsByLayer as jc

def make2dByLayer(layer):
    objs = rs.ObjectsByLayer(layer)
    if objs is None:
        return
    #Make2d

    rs.SelectObjects(objs)
    rs.Command("-_make2d _D _U _Enter")
    projLines = rs.GetObjects("Press escape", preselect = True)
    finalCrvs = rs.JoinCurves(projLines)
    rs.DeleteObjects(projLines)
    rs.MatchObjectAttributes(finalCrvs, objs[0])
    return finalCrvs

def main():
    layers = rs.GetLayers()
    for layer in layers:
        make2dByLayer(layer)

if __name__ ==  "__main__":
    main()


import pulp

def abs_load_diff(ilp_vars, x, y, vnfs, vims, dec_vars, types=['cpu', 'ram']):
    """
    Calculating the actual load difference between two vims x and y.
    """

    load_x = 0.0
    load_y = 0.0

    for res_type in types:
        if res_type == 'cpu':
            type_used = 'core_used'
            type_tot = 'core_total'
        if res_type == 'ram':
            type_used = 'memory_used'
            type_tot = 'memory_total'

        used_x = sum(vnfs[k[0]][res_type] * ilp_vars[k] for k in dec_vars if k[1] == x)
        load_x = load_x + (vims[x][type_used] + used_x) / vims[x][type_tot]

        used_y = sum(vnfs[k[0]][res_type] * ilp_vars[k] for k in dec_vars if k[1] == y)
        load_y = load_y + (vims[y][type_used] + used_y) / vims[y][type_tot]

    return load_x / len(types) - load_y / len(types)

def number_of_vnfs_mapped_to_pop(ilp_vars, x, dec_vars):

    number_of_vnf = 0
    for combo in dec_vars:
        if combo[1] == x:
            number_of_vnf = number_of_vnf + ilp_vars[combo]
    return number_of_vnf

def calc_op_obj(ilp_vars, vnfs, vims, dec_vars, op_pol, LOG):

    LOG.info("Active operator Policy: " + str(op_pol))
    op_obj = 0
    add_constr = []

    # Setting up the operator_factor. When the load is balanced, the difference
    # in load between vims needs to be minimized, hence factor 1. When vims
    # need to be used in a fill first capacity, the difference in load between
    # vims should be maximised, or the inverse minimised, hence factor -1
    op_factor = 0
    if op_pol['policy'] == 'load balanced':
        op_factor = 1
    if op_pol['policy'] == 'fill first':
        op_factor = -1

    if op_factor != 0:
        # Creating expressions for fill first or load balanced objectives, we
        # have to deal with absolute values, which require the introduction of
        # new variables to be described as an ilp constraint, as well as new
        # constraints
        constr_to_add = {}
        new_floats1 = []
        new_floats2 = []
        new_bins = []
        for x in range(len(vims)):
            for y in range(len(vims)):
                if x < y:
                    # expression of which absolute value should be minimized
                    expr = abs_load_diff(ilp_vars, y, x, vnfs, vims, dec_vars)
                    # new dec vars to problem to bypass absolute value issue
                    nameFloat1 = "extraFloat1_" + str(x) + "_" + str(y)
                    nameFloat2 = "extraFloat2_" + str(x) + "_" + str(y)
                    nameBin = "extraBin_" + str(x) + "_" + str(y)
                    new_floats1.append(nameFloat1)
                    new_floats2.append(nameFloat2)
                    new_bins.append(nameBin)
                    # we need to store the new constraints to be added after
                    # object function is defined
                    constr_to_add[nameFloat1] = expr

        flts1 = pulp.LpVariable.dicts('extra_Float1',
                                      new_floats1,
                                      lowBound=0,
                                      upBound=1)
        flts2 = pulp.LpVariable.dicts('extra_Float2',
                                      new_floats2,
                                      lowBound=0,
                                      upBound=1)
        bins = pulp.LpVariable.dicts('extra_bin',
                                     new_bins,
                                     lowBound=0,
                                     upBound=1,
                                     cat=pulp.LpInteger)

        # Combine all objectives in object function
        rn = range(len(new_floats1))
        op_obj = sum(flts1[new_floats1[x]] + flts2[new_floats2[x]] for x in rn)
        op_obj = op_factor * op_obj

        # Add the constraints related to the new variables
        for index in range(len(new_floats1)):
            bn = bins[new_bins[index]]
            ft1 = flts1[new_floats1[index]]
            ft2 = flts2[new_floats2[index]]
            add_constr.append(ft1 - ft2 == constr_to_add[new_floats1[index]])
            add_constr.append(ft1 <= bn)
            add_constr.append(ft2 <= 1 - bn)

    elif op_pol['policy'] == 'priority':
        number_of_vnfs = []
        priority_log = []
        priorityList = op_policy['policy_list']
        LOG.info("Priority list: " + str(priorityList))

        for index in range(len(vims)):
            vim_name = vims[index]['id']

            if vim_name in priorityList:
                priority = len(vims) - priorityList.index(vim_name)
            else:
                LOG.info("PoP not in priority list: " + str(vim_name))
                priority = 1

            x = number_of_vnfs_mapped_to_pop(ilp_vars, index, dec_vars)
            number_of_vnfs.append(x)
            priority_log.append(priority)

        rn = range(len(priority_log))
        op_obj = (-1) * sum(number_of_vnfs[x] * priority_log[x] for x in rn)

    else:
        LOG.info("No supported operator policy set.")

    return op_obj, add_constr

def calc_dev_obj(constr, dus, src, dest, vims, ilp_vars, dec_vars, LOG):

    # If no soft constraints have been defined, there is no objective
    if not constr:
        return 0, {}, {}, {}

    # TODO: code needs revision
    return 0, {}, {}, {}

    # dev_obj = 0
    # add_constr = []

    # # Account for the distance factor
    # dist_constr = []
    # for constraint in constr:
    #     if constraint['objective'] in ['proximity', 'affinity']:
    #         dist_constr.append(constraint['involved_actors'])
    # LOG.info("Distance constr: " + str(dist_constr))

    # # Check if IPs are provided for source and destination if they are used
    # # in the soft constraints
    # act_constr = []
    # for constraint in dist_constr:
    #     constr_complete = True
    #     if 'source' in constraint:
    #         if not src:
    #             constraint_complete = False
    #             LOG.info("constraint useless due missing src IP.")
    #     if 'destination' in constraint:
    #         if not dest:
    #             constraint_complete = False
    #             LOG.info("constraint useless due missing dest IP.")
    #     if constraint_complete:
    #         act_constr.append(constraint)

    # dist_matrix = calc_dist(src, dest, vims, LOG)
    # LOG.info("Distance matrix: " + str(dist_matrix))

    # for constraint in act_constr:
    #     indexes_first_vnf = find_images_vnf_indexes(nsd, vnfs, images_to_map, constr[0])
    #     indexes_second_vnf = find_images_vnf_indexes(nsd, vnfs, images_to_map, constr[1])

    #     if indexes_first_vnf != [] and indexes_second_vnf != []:
    #         new_bins = []
    #         matrix_equivalent = []
    #         old_var_1 = []
    #         old_var_2 = []
    #         for index_1 in indexes_first_vnf:
    #             for index_2 in indexes_second_vnf:
    #                 for pop_1 in range(len(top)):
    #                     for pop_2 in range(len(top)):
    #                         new_name = "extraDeveloperBin_" + str(index_1) + "_" + str(index_2) + "_" + str(pop_1) + "_" + str(pop_2)
    #                         new_bins.append(new_name)
    #                         matrix_equivalent.append(distance_matrix[pop_1][pop_2])
    #                         old_var_1.append(variables[(index_1, pop_1)])
    #                         old_var_2.append(variables[(index_2, pop_2)])

    #         new_bins_lp = pulp.LpVariable.dicts('extra_bins', new_bins, lowBound=0, upBound=1, cat=pulp.LpInteger)

    #         # Combine all objectives in object function
    #         LOG.info(str(matrix_equivalent))
    #         developer_objective += sum((-1) * new_bins_lp[new_bins[x]] * matrix_equivalent[x] / 4 for x in range(len(new_bins)))

    #         for new_constraint in range(len(new_bins)):
    #             additional_constraints.append(new_bins_lp[new_bins[new_constraint]] <= (old_var_1[new_constraint] + old_var_2[new_constraint])/2.0)

    #     else:
    #         if constr[0] in ['source', 'destination']:
    #             indexes_vnf = indexes_second_vnf
    #             if constr[0] == 'source':
    #                 index_endpoint = len(top)
    #             if constr[0] == 'destination':
    #                 index_endpoint = len(top) + 1
    #         else:
    #             indexes_vnf = indexes_first_vnf
    #             if constr[1] == 'source':
    #                 index_endpoint = len(top)
    #             if constr[1] == 'destination':
    #                 index_endpoint = len(top) + 1

    #         for index in indexes_vnf:
    #             for pop in range(len(top)):
    #                 developer_objective += (-1) * variables[(index, pop)] * distance_matrix[pop][index_endpoint]
    #         LOG.info(developer_objective)

    # return developer_objective, additional_constraints

def find_images_vnf_indexes(nsd, vnfs, images_to_map, vnf_id):

    if vnf_id in ['source', 'destination']:
        return []

    for vnf_desc in nsd['network_functions']:
        if vnf_desc['vnf_id'] == vnf_id:
            name = vnf_desc["vnf_name"]
            version = vnf_desc["vnf_version"]
            vendor = vnf_desc["vnf_vendor"]

    for vnf in vnfs:
        vnfd = vnf['vnfd']
        if vnfd['vendor'] == vendor:
            if vnfd['name'] == name:
                if vnfd['version'] == version:
                    function_id = vnf['id']

    indexes = []
    for index in range(len(images_to_map)):
        if images_to_map[index]['function_id'] == function_id:
            indexes.append(index)

    return indexes

def calculate_distances(source, destination, pop, LOG, weightfactor=1):

    matrix = []

    for index_x in range(len(pop) + 2):
        matrix.append([])
        for index_y in range(len(pop) + 2):
            if index_x == len(pop):
                ip_x = source
            elif index_x == len(pop) + 1:
                ip_x = destination
            else:
                ip_x = pop[index_x]['vim_endpoint']

            if index_y == len(pop):
                ip_y = source
            elif index_y == len(pop) + 1:
                ip_y = destination
            else:
                ip_y = pop[index_y]['vim_endpoint']
            matrix[index_x].append((find_number_of_common_segments(ip_x, ip_y, LOG)) ** weightfactor)

    return matrix

def find_number_of_common_segments(first_ip, second_ip, LOG):

    if not first_ip:
        return 0

    if not second_ip:
        return 0

    first_segments = first_ip.split('.')
    second_segments = second_ip.split('.')

    if first_segments == second_segments:
        return 4

    elif first_segments[:3] == second_segments[:3]:
        return 3

    elif first_segments[:2] == second_segments[:2]:
        return 2

    elif first_segments[0] == second_segments[0]:
        return 1
    else:
        return 0

def convert_name_to_index(name, vnfs):

    for index in range(len(vnfs)):
        if vnfs[index]['name'] == name:
            return index

def map_ref_on_id(ref, nsd, vnfds, eps):
    """
    This method is used to map a reference in a virtual link descriptor on
    a deployment unit.
    """
    nodes_list = []
    ref_list = []
    in_out = None
    if ':' in ref:
        # reference is towards a vnf
        for vnf in nsd['network_functions']:
            if vnf['vnf_id'] == ref.split(':')[0]:
                vnf_name = vnf['vnf_name']
                vnf_version = vnf['vnf_version']
                vnf_vendor = vnf['vnf_vendor']
                break
        for vnfd in vnfds:
            if vnfd['name'] == vnf_name and vnfd['version'] == vnf_version and \
               vnfd['vendor'] == vnf_vendor:

                for vl in vnfd['virtual_links']:
                    cpr = vl['connection_points_reference']
                    if ref.split(':')[1] in cpr:
                        ref_i = cpr.index(ref.split(':')[1])
                        du_int_ref = cpr[len(cpr) - 1 - ref_i]
                        du_ref = du_int_ref.split(':')[0]
                        break

                if 'cloudnative_deployment_units' in vnfd.keys():
                    for cdu in vnfd['cloudnative_deployment_units']:
                        if cdu['id'].startswith(du_ref):
                            nodes_list.append(cdu['id'])
                            ref_list.append(ref)

                if 'virtual_deployment_units' in vnfd.keys():
                    for vdu in vnfd['virtual_deployment_units']:
                        if vdu['id'].startswith(du_ref):
                            nodes_list.append(vdu['id'])
                            ref_list.append(ref)

    else:

        index_e = 0
        for ep in eps:
            if ep['type'] == 'egress':
                index_e = eps.index(ep)
                break

        # reference is towards an endpoint
        # TODO: find a solid way to differentiate between ingress and egress,
        # currently, we assume that ingress comes before egress in the nsd cp
        # list. If there is only one, we assume that it is ingress
        cps = nsd['connection_points']
        for cp in cps:
            if cp['id'] == ref:
                cp_i = cps.index(cp)
        if len(cps) == 1:
            for ep in eps:
                if ep['type'] == 'ingress':
                    nodes_list.append(ep['id'].split('_')[0])
                    ref_list.append(ref + '+' + str(eps.index(ep)))
                    in_out = 'ingress'
        if len(cps) == 3 and cp_i == 1:
            for ep in eps:
                if ep['type'] == 'ingress':
                    nodes_list.append(ep['id'].split('_')[0])
                    ref_list.append(ref + '+' + str(eps.index(ep)))
                    in_out = 'ingress'
        if len(cps) == 3 and cp_i == 2:
            for ep in eps:
                if ep['type'] == 'egress':
                    nodes_list.append(ep['id'].split('_')[0])
                    ref_list.append(ref + '+' + str(eps.index(ep) - index_e))
                    in_out = 'egress'
        if len(cps) == 2 and cp_i == 1:
            for ep in eps:
                if ep['type'] == 'ingress':
                    nodes_list.append(ep['id'].split('_')[0])
                    ref_list.append(ref + '+' + str(eps.index(ep)))
                    in_out = 'ingress'
            return nodes_list, ref_list, in_out
        if len(cps) == 2 and cp_i == 1:
            for ep in eps:
                if ep['type'] == 'egress':
                    nodes_list.append(ep['id'].split('_')[0])
                    ref_list.append(ref + '+' + str(eps.index(ep) - index_e))
                    in_out = 'egress'

    return nodes_list, ref_list, in_out


import random
start = input("請決定隨機數字範圍的開始值： ")
end = input("請決定隨機數字範圍的結束值： ")
start = int(start)
end = int(end)
r = random.randint(start, end)
count = 0
while True:
	count += 1 # 跟 count = count + 1
	num = input("請猜數字： ")
	num = int(num)
	if num == r:
		print("你猜對了！")
		print("你猜的次數是", count, "次")
		break
	elif num > r:
		print("比答案大")
	elif num < r:
		print("比答案小")
	print("你猜的次數是", count, "次")