from cc3d.cpp import CompuCell
from deprecated import deprecated

# IMPORTANT: It is best to always provide hand-written iterators for STL
# containers even though swig generates them for you.
# with multiple swig modules  those autogenerated iterators will work on one platform and crash on another
# ones so best solution is to write iterators yourself

# this is used to iterate more easily over cells


class CellList:
    """
    Wraps current inventory of all :class:`~cc3d.cpp.CompuCell.CellG` instances in simulation
    """

    def __init__(self, inventory):
        self.inventory = inventory

    def __iter__(self):
        """

        :return: CellListIterator instance
        :rtype: CellListIterator
        """
        return CellListIterator(self)

    def __len__(self):
        """

        :return: number of cells
        :rtype: int
        """
        return int(self.inventory.getSize())


class CellListIterator:
    def __init__(self, cell_list):
        self.next = self.__next__
        self.inventory = cell_list.inventory
        self.invItr = CompuCell.STLPyIteratorCINV()
        self.invItr.initialize(self.inventory.getContainer())
        self.invItr.setToBegin()

    def __next__(self):
        """

        :return: a cell
        :rtype: cc3d.cpp.CompuCell.CellG
        """
        if not self.invItr.isEnd():
            self.cell = self.invItr.getCurrentRef()
            self.invItr.next()
            return self.cell
        else:
            raise StopIteration

    def __iter__(self):
        return self


#########################################################################
# iterating over inventory of cells of a given type
class CellListByType:
    """
    List of all cells in current inventory of a variable number of particular types
    """
    def __init__(self, inventory, *args):
        self.inventory = inventory

        self.types = CompuCell.vectorint()

        self.inventoryByType = CompuCell.mapLongCellGPtr()

        self.initTypeVec(args)
        self.inventory.initCellInventoryByMultiType(self.inventoryByType, self.types)

    def __iter__(self):
        """

        :return: CellListByTypeIterator instance
        :rtype: CellListByTypeIterator
        """
        return CellListByTypeIterator(self)

    def __len__(self):
        """

        :return: number of cells
        :rtype: int
        """
        return int(self.inventoryByType.size())

    def initTypeVec(self, _type_list):

        self.types.clear()
        if len(_type_list) <= 0:
            self.types.push_back(1)  # type 1
        else:
            for type in _type_list:
                self.types.push_back(type)

    def initializeWithType(self, _type):
        self.types.clear()
        self.types.push_back(_type)
        self.inventory.initCellInventoryByMultiType(self.inventoryByType, self.types)

    def refresh(self):
        self.inventory.initCellInventoryByMultiType(self.inventoryByType, self.types)


class CellListByTypeIterator:
    def __init__(self, _cellListByType):
        self.inventoryByType = _cellListByType.inventoryByType
        self.invItr = CompuCell.mapLongCellGPtrPyItr()
        self.invItr.initialize(self.inventoryByType)
        self.invItr.setToBegin()
        self.next = self.__next__

    def __next__(self):
        """

        :return: a cell
        :rtype: cc3d.cpp.CompuCell.CellG
        """
        if not self.invItr.isEnd():
            self.cell = self.invItr.getCurrentRef()
            # print 'self.idCellPair=',self.idCellPair
            # print 'dir(self.idCellPair)=',dir(self.idCellPair)
            self.invItr.next()
            return self.cell
        #
        else:
            raise StopIteration

    def __iter__(self):
        return self


#########################################################################
# this is used to iterate more easily over clusters
class ClusterList:
    def __init__(self, _inventory):
        self.inventory = _inventory.getClusterInventory().getContainer()

    def __iter__(self):
        """

        :return: ClusterListIterator instance
        :rtype: ClusterListIterator
        """
        return ClusterListIterator(self)

    def __len__(self):
        """

        :return: number of clusters
        :rtype: int
        """
        return int(self.inventory.size())


class ClusterListIterator:
    def __init__(self, _cellList):
        self.inventory = _cellList.inventory
        self.invItr = CompuCell.compartmentinventoryPtrPyItr()
        self.invItr.initialize(self.inventory)
        self.invItr.setToBegin()
        self.compartmentList = None
        self.next = self.__next__

    def __next__(self):
        """

        :return: a cell
        :rtype: cc3d.cpp.CompuCell.CellG
        """

        if not self.invItr.isEnd():
            self.compartmentList = self.invItr.getCurrentRef()
            self.invItr.next()
            return self.compartmentList
        #
        else:
            raise StopIteration


# this is used to iterate more easily over clusters and avoid strange looking Python syntax

class Clusters:
    def __init__(self, _inventory):
        self.inventory = _inventory.getClusterInventory().getContainer()

    def __iter__(self):
        """

        :return: ClustersIterator instance
        :rtype: ClustersIterator
        """
        return ClustersIterator(self)

    def __len__(self):
        """

        :return: number of clusters
        :rtype: int
        """
        return int(self.inventory.size())


class ClustersIterator:
    def __init__(self, _cellList):

        self.next = self.__next__
        self.inventory = _cellList.inventory

        self.invItr = CompuCell.compartmentinventoryPtrPyItr()
        self.invItr.initialize(self.inventory)
        self.invItr.setToBegin()

        self.compartmentList = None

    def __next__(self):
        """

        :return: CompartmentList instance
        :rtype: CompartmentList
        """

        if not self.invItr.isEnd():
            self.compartmentList = self.invItr.getCurrentRef()
            # print 'self.idCellPair=',self.idCellPair
            # print 'dir(self.idCellPair)=',dir(self.idCellPair)
            self.invItr.next()
            return CompartmentList(self.compartmentList)
        #
        else:
            raise StopIteration


# this is used to iterate more easily over list of compartments,
# notice regular map iteration will work too but this is more abstracted out and will work with other containers too

class CompartmentList:
    def __init__(self, _inventory):
        self.inventory = _inventory

    def __iter__(self):
        """

        :return: CompartmentListIterator instance
        :rtype: CompartmentListIterator
        """
        return CompartmentListIterator(self)

    def __len__(self):
        """

        :return: number of compartments
        :rtype: int
        """
        return int(self.inventory.size())

    def clusterId(self):
        return self.__iter__().next().clusterId


class CompartmentListIterator:
    def __init__(self, _cellList):
        self.next = self.__next__
        self.inventory = _cellList.inventory
        self.invItr = CompuCell.mapLongCellGPtrPyItr()
        self.invItr.initialize(self.inventory)
        self.invItr.setToBegin()

    def __next__(self):
        """

        :return: a cell
        :rtype: cc3d.cpp.CompuCell.CellG
        """
        if not self.invItr.isEnd():
            self.cell = self.invItr.getCurrentRef()
            # print 'self.idCellPair=',self.idCellPair
            # print 'dir(self.idCellPair)=',dir(self.idCellPair)
            self.invItr.next()
            return self.cell
        #
        else:
            raise StopIteration

    def __iter__(self):
        return self


# this is wrapper for std::vector<CellG*>
class ClusterCellList:
    def __init__(self, _inventory):
        self.inventory = _inventory

    def __iter__(self):
        """

        :return: ClusterCellListIterator instance
        :rtype: ClusterCellListIterator
        """
        return ClusterCellListIterator(self)

    def __getitem__(self, item):
        """

        :param int item:
        :return: a cell
        :rtype: cc3d.cpp.CompuCell.CellG
        """
        return self.inventory[item]

    def __len__(self):
        """

        :return: number of cells
        :rtype: int
        """
        return int(self.inventory.size())

    @deprecated(version='4.0.0', reason="You should use : len()")
    def size(self):
        return self.__len__()


class ClusterCellListIterator:
    def __init__(self, _cellList):
        self.next = self.__next__
        self.inventory = _cellList.inventory
        # print "dir(self.inventory)=",dir(self.inventory)
        self.currentIdx = 0
        self.cell = None
        # self.invItr.initialize(self.inventory.getContainer())
        # self.invItr.setToBegin()

    def __next__(self):
        """

        :return: a cell
        :rtype: cc3d.cpp.CompuCell.CellG
        """
        # if self.invItr !=  self.inventory.end():
        if self.currentIdx < self.inventory.size():
            # print "self.invItr=",dir(self.invItr)
            # print "self.invItr.next()=",self.invItr.next()
            # self.compartmentList = self.invItr.next()

            self.cell = self.inventory[self.currentIdx]
            self.currentIdx += 1

            return self.cell
        else:
            raise StopIteration

    def __iter__(self):
        return self

    # legacy code  - do not modify


class CellNeighborListFlex:
    def __init__(self, _neighborPlugin, _cell):
        self.neighborPlugin = _neighborPlugin
        self.neighborTrackerAccessor = self.neighborPlugin.getNeighborTrackerAccessorPtr()
        self.cell = _cell

        # legacy API
        self.commonSurfaceAreaWithCellTypes = self.common_surface_area_with_cell_types
        self.commonSurfaceAreaByType = self.common_surface_area_by_type
        self.neighborCountByType = self.neighbor_count_by_type

    def __len__(self):
        """

        :return: number of cells
        :rtype: int
        """
        neighborTracker = self.neighborTrackerAccessor.get(self.cell.extraAttribPtr)
        return int(neighborTracker.cellNeighbors.size())

    def __getitem__(self, idx):
        """

        :param int idx:
        :raises IndexError: Out of bounds index
        :return: neighbor, common surface area
        :rtype: (cc3d.cpp.CompuCell.CellG, int)
        """
        if idx > self.__len__() - 1: raise IndexError(
            "Out of bounds index: CellNeighborListAuto index = %s is out of bounds" % str(idx))
        for counter, data in enumerate(self.__iter__()):
            if idx == counter: return data

    def common_surface_area_with_cell_types(self, cell_type_list):
        """
        Returns common surface area with a variable number of cell types

        :param cell_type_list: list or tuple of cell types
        :type cell_type_list: list of int or tuple of int
        :return: total common surface area with all given types
        :rtype: int
        """
        area = 0
        for neighbor, commonSurfaceArea in self.__iter__():
            cell_type = 0 if not neighbor else neighbor.type
            if cell_type in cell_type_list:
                area += commonSurfaceArea
        return area

    def common_surface_area_by_type(self):
        """
        Returns common surface area with all cell types

        :return: total common surface area by cell type
        :rtype: collections.defaultdict of {int: int}
        """
        from collections import defaultdict
        area_dict = defaultdict(int)
        for neighbor, commonSurfaceArea in self.__iter__():
            cell_type = 0 if not neighbor else neighbor.type
            area_dict[cell_type] += commonSurfaceArea
        return area_dict

    def neighbor_count_by_type(self):
        """
        Returns number of neighbors of all cell types

        :return: number of neighbors by type
        :rtype: collections.defaultdict of {int: int}
        """
        from collections import defaultdict
        neighbor_counter_dict = defaultdict(int)

        for neighbor, commonSurfaceArea in self.__iter__():
            cell_type = 0 if not neighbor else neighbor.type
            neighbor_counter_dict[cell_type] += 1
        return neighbor_counter_dict

    def __iter__(self):
        """

        :return: CellNeighborIteratorFlex instance
        :rtype: CellNeighborIteratorFlex
        """
        return CellNeighborIteratorFlex(self)


class CellNeighborIteratorFlex:
    def __init__(self, _cellNeighborList):
        self.neighborTrackerAccessor = _cellNeighborList.neighborTrackerAccessor
        self.cell = _cellNeighborList.cell
        self.nsdItr = CompuCell.nsdSetPyItr()
        self.nTracker = self.neighborTrackerAccessor.get(self.cell.extraAttribPtr)
        self.nsdItr.initialize(self.nTracker.cellNeighbors)
        self.nsdItr.setToBegin()

    def __next__(self):
        """

        :return: neighbor, common surface area
        :rtype: (cc3d.cpp.CompuCell.CellG, int)
        """
        if not self.nsdItr.isEnd():
            self.neighborCell = self.nsdItr.getCurrentRef().neighborAddress
            self.currentNsdItr = self.nsdItr.current
            self.currentNeighborSurfaceData = self.nsdItr.getCurrentRef()
            self.nsdItr.next()
            # return self.currentNeighborSurfaceData.neighborAddress,self.currentNeighborSurfaceData.commonSurfaceArea
            return self.currentNeighborSurfaceData.neighborAddress, self.currentNeighborSurfaceData.commonSurfaceArea
        else:
            raise StopIteration

    def __iter__(self):
        return self


class _FocalPointPlasticityDataListBase:
    def __init__(self, _fpp_plugin, _cell):
        # Maintaining legacy feature: cell is attached to data list
        self.cell = _cell

        self._data = self._get_inventory(_fpp_plugin).getFPPTrackerDataSet(_cell)

    def __len__(self):
        """

        :return: number of links
        :rtype: int
        """
        return len(self._data)

    def __getitem__(self, idx):
        """

        :param idx: {int}
        :return: {:class:`cc3d.cpp.CompuCell.FocalPointPlasticityTrackerData`}
        """
        return self._data[idx]

    def __iter__(self):
        """

        :return: FocalPointPlasticityDataIterator instance
        :rtype: FocalPointPlasticityDataIterator
        """
        return self._data.__iter__()

    @staticmethod
    def _get_inventory(_fpp_plugin):
        """
        Get link inventory container

        :param _fpp_plugin: focal point plasticity plugin
        :return: link inventory container
        """
        raise NotImplementedError


class _FocalPointPlasticityDataListIteratorBase:
    def __init__(self, _fpp_data_list: _FocalPointPlasticityDataListBase):
        # Maintaining legacy feature: cell is attached to data list iterator
        self.cell = _fpp_data_list.cell

        self._val = iter(_fpp_data_list._data)

    def __next__(self):
        """

        :return: FocalPointPlasticityTrackerData instance
        :rtype: cc3d.cpp.CompuCell.FocalPointPlasticityTrackerData
        """
        return self._val.__next__()

    def __iter__(self):
        return self._val


class FocalPointPlasticityDataList(_FocalPointPlasticityDataListBase):
    def __init__(self, _fpp_plugin, _cell):
        super().__init__(_fpp_plugin, _cell)

    @staticmethod
    def _get_inventory(_fpp_plugin):
        return _fpp_plugin.getLinkInventory()


class FocalPointPlasticityDataIterator(_FocalPointPlasticityDataListIteratorBase):
    def __init__(self, _fpp_data_list: FocalPointPlasticityDataList):
        super().__init__(_fpp_data_list)


class InternalFocalPointPlasticityDataList(_FocalPointPlasticityDataListBase):
    def __init__(self, _fpp_plugin, _cell):
        super().__init__(_fpp_plugin, _cell)

    @staticmethod
    def _get_inventory(_fpp_plugin):
        return _fpp_plugin.getInternalLinkInventory()


class InternalFocalPointPlasticityDataIterator(_FocalPointPlasticityDataListIteratorBase):
    def __init__(self, _fpp_data_list: InternalFocalPointPlasticityDataList):
        super().__init__(_fpp_data_list)


class AnchorFocalPointPlasticityDataList(_FocalPointPlasticityDataListBase):
    def __init__(self, _fpp_plugin, _cell):
        super().__init__(_fpp_plugin, _cell)

    @staticmethod
    def _get_inventory(_fpp_plugin):
        return _fpp_plugin.getAnchorInventory()


class AnchorFocalPointPlasticityDataIterator(_FocalPointPlasticityDataListIteratorBase):
    def __init__(self, _fpp_data_list: AnchorFocalPointPlasticityDataList):
        super().__init__(_fpp_data_list)


class _FocalPointPlasticityLinkListBase:

    # Python iterator type
    inv_itr_t = None

    def __init__(self, _fpp_plugin, cell: CompuCell.CellG = None):
        self.__cell = cell
        self._inv = self._get_inventory(_fpp_plugin)

    @property
    def cell(self):
        """

        :return: cell of this list, if any
        """
        return self.__cell

    def __len__(self):
        """

        :return: number of links
        :rtype: int
        """
        return int(self._inv.getLinkInventorySize())

    def __iter__(self):
        return self.inv_itr_t(self)

    def _get_inventory(self, _fpp_plugin):
        """
        Get link inventory container

        :param _fpp_plugin: focal point plasticity plugin
        :return: link inventory container
        """
        raise NotImplementedError


class _FocalPointPlasticityDataIteratorBase:

    # C++ iterator type
    py_itr_t = None

    def __init__(self, _link_list: _FocalPointPlasticityLinkListBase):
        self._link_list = _link_list
        self._inv = self._link_list._inv
        self._itr = self.py_itr_t()
        self._itr.initialize(self._inv.getContainer())
        self._itr.setToBegin()

    def __next__(self):
        """

        :return: FocalPointPlasticityTrackerData instance
        :rtype: cc3d.cpp.CompuCell.FocalPointPlasticityTrackerData
        """
        if self._itr.isEnd():
            raise StopIteration
        else:
            _itr_val = self._itr.getCurrentRef()
            self._itr.next()
            return _itr_val

    def __iter__(self):
        return self


class FocalPointPlasticityLinkListItr(_FocalPointPlasticityDataIteratorBase):

    py_itr_t = CompuCell.mapFPPLinkIDFPPLinkPyItr


class FocalPointPlasticityLinkList(_FocalPointPlasticityLinkListBase):

    inv_itr_t = FocalPointPlasticityLinkListItr

    def _get_inventory(self, _fpp_plugin):
        inv = _fpp_plugin.getLinkInventory()
        if self.cell is None:
            return inv
        return inv.getCellLinkInventory(self.cell)


class FocalPointPlasticityInternalLinkListItr(_FocalPointPlasticityDataIteratorBase):

    py_itr_t = CompuCell.mapFPPLinkIDFPPInternalLinkPyItr


class FocalPointPlasticityInternalLinkList(_FocalPointPlasticityLinkListBase):

    inv_itr_t = FocalPointPlasticityInternalLinkListItr

    def _get_inventory(self, _fpp_plugin):
        inv = _fpp_plugin.getInternalLinkInventory()
        if self.cell is None:
            return inv
        return inv.getCellLinkInventory(self.cell)


class FocalPointPlasticityAnchorListItr(_FocalPointPlasticityDataIteratorBase):

    py_itr_t = CompuCell.mapFPPLinkIDFPPAnchorPyItr


class FocalPointPlasticityAnchorList(_FocalPointPlasticityLinkListBase):

    inv_itr_t = FocalPointPlasticityAnchorListItr

    def _get_inventory(self, _fpp_plugin):
        inv = _fpp_plugin.getAnchorInventory()
        if self.cell is None:
            return inv
        return inv.getCellLinkInventory(self.cell)


class CellPixelList:
    def __init__(self, _pixelTrackerPlugin, _cell):
        self.pixelTrackerPlugin = _pixelTrackerPlugin
        self.pixelTrackerAccessor = self.pixelTrackerPlugin.getPixelTrackerAccessorPtr()
        self.cell = _cell

    def __iter__(self):
        """

        :return: CellPixelIterator instance
        :rtype: CellPixelIterator
        """
        return CellPixelIterator(self)

    @deprecated(version='4.0.0', reason="You should use : number_of_pixels")
    def numberOfPixels(self):
        return self.number_of_pixels()

    def number_of_pixels(self):
        """

        :return: number of pixels
        :rtype: int
        """
        return self.pixelTrackerAccessor.get(self.cell.extraAttribPtr).pixelSet.size()


class CellPixelIterator:
    def __init__(self, _cellPixelList):

        self.pixelTrackerAccessor = _cellPixelList.pixelTrackerAccessor
        self.pixelTrackerPlugin = _cellPixelList.pixelTrackerPlugin
        self.cell = _cellPixelList.cell
        self.pixelItr = CompuCell.pixelSetPyItr()
        self.pixelTracker = self.pixelTrackerAccessor.get(self.cell.extraAttribPtr)
        self.pixelItr.initialize(self.pixelTracker.pixelSet)
        self.pixelItr.setToBegin()

    def __next__(self):
        """

        :return: PixelTrackerData instance
        :rtype: cc3d.cpp.CompuCell.PixelTrackerData
        """
        if not self.pixelItr.isEnd():
            #             self.neighborCell = self.nsdItr.getCurrentRef().neighborAddress
            #             self.currentNsdItr = self.nsdItr.current
            self.currentPixelTrackerData = self.pixelItr.getCurrentRef()
            self.pixelItr.next()

            return self.pixelTrackerPlugin.getPixelTrackerData(self.currentPixelTrackerData)
            # return self.currentPixelTrackerData
        else:
            raise StopIteration

    def __iter__(self):
        return self


class CellBoundaryPixelList:
    def __init__(self, boundary_pixel_tracker_plugin, cell, neighbor_order=-1):
        self.neighbor_order = neighbor_order
        self.boundary_pixel_tracker_plugin = boundary_pixel_tracker_plugin
        self.boundary_pixel_tracker_accessor = self.boundary_pixel_tracker_plugin.getBoundaryPixelTrackerAccessorPtr()
        self.cell = cell

    def __iter__(self):
        """

        :return: CellBoundaryPixelIterator instance
        :rtype: CellBoundaryPixelIterator
        """
        return CellBoundaryPixelIterator(self, self.neighbor_order)

    @deprecated(version='4.0.0', reason="You should use : number_of_pixels")
    def numberOfPixels(self):
        return self.number_of_pixels()

    def number_of_pixels(self):
        """

        :return: number of pixels
        :rtype: int
        """
        return self.boundary_pixel_tracker_accessor.get(self.cell.extraAttribPtr).pixelSet.size()


class CellBoundaryPixelIterator:
    def __init__(self, cell_pixel_list, neighbor_order=-1):
        self.boundary_pixel_tracker_accessor = cell_pixel_list.boundary_pixel_tracker_accessor
        self.boundary_pixel_tracker_plugin = cell_pixel_list.boundary_pixel_tracker_plugin
        self.cell = cell_pixel_list.cell
        self.boundary_pixel_itr = CompuCell.boundaryPixelSetPyItr()
        self.boundary_pixel_tracker = self.boundary_pixel_tracker_accessor.get(self.cell.extraAttribPtr)
        if neighbor_order <= 0:
            self.pixelSet = self.boundary_pixel_tracker.pixelSet
        else:
            self.pixelSet = self.boundary_pixel_tracker_plugin.getPixelSetForNeighborOrderPtr(self.cell, neighbor_order)
            if not self.pixelSet:
                raise LookupError('LookupError: CellBoundaryPixelIterator could not locate pixel set '
                                  'for neighbor order = %s. Make sure your BoundaryPixelTracker plugin definition '
                                  'requests tracking of neighbor order =%s boundary' % (neighbor_order, neighbor_order))

        # self.boundaryPixelItr.initialize(self.boundaryPixelTracker.pixelSet)
        self.boundary_pixel_itr.initialize(self.pixelSet)
        self.boundary_pixel_itr.setToBegin()

    def __next__(self):
        """

        :return: BoundaryPixelTrackerData instance
        :rtype: cc3d.cpp.CompuCell.BoundaryPixelTrackerData
        """
        if not self.boundary_pixel_itr.isEnd():
            self.current_boundary_pixel_tracker_data = self.boundary_pixel_itr.getCurrentRef()
            self.boundary_pixel_itr.next()
            return self.boundary_pixel_tracker_plugin.getBoundaryPixelTrackerData(
                self.current_boundary_pixel_tracker_data)
        else:
            raise StopIteration

    def __iter__(self):
        return self


class ElasticityDataList:
    def __init__(self, elasticity_tracker_plugin, _cell):
        self.elasticity_tracker_plugin = elasticity_tracker_plugin
        self.elasticity_tracker_accessor = self.elasticity_tracker_plugin.getElasticityTrackerAccessorPtr()
        self.cell = _cell

    def __iter__(self):
        """

        :return: ElasticityDataIterator instance
        :rtype: ElasticityDataIterator
        """
        return ElasticityDataIterator(self)


class ElasticityDataIterator:
    def __init__(self, elasticity_data_list):
        self.elasticity_tracker_accessor = elasticity_data_list.elasticity_tracker_accessor
        self.cell = elasticity_data_list.cell
        self.elasticity_tracker_plugin = elasticity_data_list.elasticity_tracker_plugin
        self.elasticity_tracker = self.elasticity_tracker_accessor.get(self.cell.extraAttribPtr)
        self.elasticity_data_set_itr = CompuCell.elasticitySetPyItr()
        self.elasticity_data_set_itr.initialize(self.elasticity_tracker.elasticityNeighbors)
        self.elasticity_data_set_itr.setToBegin()

    def __next__(self):
        """

        :return: ElasticityTrackerData instance
        :rtype: cc3d.cpp.CompuCell.ElasticityTrackerData
        """
        if not self.elasticity_data_set_itr.isEnd():
            self.current_elasticity_data_set_itr = self.elasticity_data_set_itr.current
            self.elasticity_data = self.elasticity_data_set_itr.getCurrentRef()
            self.elasticity_data_set_itr.next()
            return self.elasticity_tracker_plugin.getElasticityTrackerData(self.elasticity_data)
        else:
            raise StopIteration

    def __iter__(self):
        return self


class EnergyDataList:
    def __init__(self, _potts):
        self.function_names = _potts.getEnergyFunctionNames()
        self.energy_changes = _potts.getCurrentEnergyChanges()
        self.flip_results = _potts.getCurrentFlipResults()

    def __iter__(self):
        """

        :return: EnergyDataListIterator instance
        :rtype: EnergyDataListIterator
        """
        return EnergyDataListIterator(self)


class EnergyDataListIterator:
    def __init__(self, _energy_data_list: EnergyDataList):
        self.__function_names = _energy_data_list.function_names
        self.__energy_changes = _energy_data_list.energy_changes
        self.__flip_results = _energy_data_list.flip_results

        self.__flip_idx = 0
        self.__idx_functions = range(self.__function_names.__len__())
        self.__num_flips = self.__flip_results.__len__()

    def __next__(self):
        """

        :return: flip result, energy change by energy function
        :rtype: (bool, dict of {str: float})
        """
        if self.__flip_idx < self.__num_flips:
            flip_result = self.__flip_results[self.__flip_idx]
            data_dict = {self.__function_names[idx]: self.__energy_changes[self.__flip_idx][idx]
                         for idx in self.__idx_functions}
            self.__flip_idx += 1
            return flip_result, data_dict
        else:
            raise StopIteration

    def __iter__(self):
        return self


class PlasticityDataList:
    def __init__(self, plasticity_tracker_plugin, cell):
        self.plasticity_tracker_plugin = plasticity_tracker_plugin
        self.plasticityTrackerAccessor = self.plasticity_tracker_plugin.getPlasticityTrackerAccessorPtr()
        self.cell = cell

    def __iter__(self):
        """

        :return: PlasticityDataIterator instance
        :rtype: PlasticityDataIterator
        """
        return PlasticityDataIterator(self)


class PlasticityDataIterator:
    def __init__(self, plasticity_data_list):
        self.plasticityTrackerAccessor = plasticity_data_list.plasticityTrackerAccessor
        self.cell = plasticity_data_list.cell
        self.plasticity_tracker_plugin = plasticity_data_list.plasticity_tracker_plugin
        self.plasticityTracker = self.plasticityTrackerAccessor.get(self.cell.extraAttribPtr)
        self.plasticityDataSetItr = CompuCell.plasticitySetPyItr()
        self.plasticityDataSetItr.initialize(self.plasticityTracker.plasticityNeighbors)
        self.plasticityDataSetItr.setToBegin()

    def next(self):
        """

        :return: PlasticityTrackerData instance
        :rtype: cc3d.cpp.CompuCell.PlasticityTrackerData
        """
        if not self.plasticityDataSetItr.isEnd():
            self.currentPlasticityDataSetItr = self.plasticityDataSetItr.current
            self.plasticityData = self.plasticityDataSetItr.getCurrentRef()
            self.plasticityDataSetItr.next()
            return self.plasticity_tracker_plugin.getPlasticityTrackerData(self.plasticityData)
        #             return self.plasticityData
        else:
            raise StopIteration

    def __iter__(self):
        return self

