from __future__ import division

import numpy as np

import chainer
import chainer.functions as F
import chainer.links as L
from chainercv.links import Conv2DBNActiv
from chainercv.links.model.resnet.resblock import ResBlock
from chainercv.links import ResNet101
from chainercv import utils

from light_head_rcnn.functions import psroi_max_align_2d
from light_head_rcnn.links.model.global_context_module \
    import GlobalContextModule
from light_head_rcnn.links.model.light_head_rcnn_base import LightHeadRCNN
from light_head_rcnn.links.model.region_proposal_network \
    import RegionProposalNetwork

class LightHeadRCNNResNet101(LightHeadRCNN):

    """LightHead RCNN based on ResNet101.

    When you specify the path of a pre-trained chainer model serialized as
    a :obj:`.npz` file in the constructor, this chain model automatically
    initializes all the parameters with it.
    When a string in prespecified set is provided, a pretrained model is
    loaded from weights distributed on the Internet.
    The list of pretrained models supported are as follows:

    * :obj:`coco`: Loads weights trained with the trainval split of \
        COCO Detection Dataset.
    * :obj:`imagenet`: Loads weights trained with ImageNet Classfication \
        task for the feature extractor and the head modules. \
        Weights that do not have a corresponding layer in ResNet101 \
        will be randomly initialized.

    For descriptions on the interface of this model, please refer to
    :class:`~light_head_rcnn.links.model.light_head_rcnn_base.LightHeadRCNN`

    :class:`~light_head_rcnn.links.model.light_head_rcnn_base.LightHeadRCNN`
    supports finer control on random initializations of weights by arguments
    :obj:`vgg_initialW`, :obj:`rpn_initialW`, :obj:`loc_initialW` and
    :obj:`score_initialW`.
    It accepts a callable that takes an array and edits its values.
    If :obj:`None` is passed as an initializer, the default initializer is
    used.

    Args:
        n_fg_class (int): The number of classes excluding the background.
        pretrained_model (string): The destination of the pre-trained
            chainer model serialized as a :obj:`.npz` file.
            If this is one of the strings described
            above, it automatically loads weights stored under a directory
            :obj:`$CHAINER_DATASET_ROOT/pfnet/chainercv/models/`,
            where :obj:`$CHAINER_DATASET_ROOT` is set as
            :obj:`$HOME/.chainer/dataset` unless you specify another value
            by modifying the environment variable.
        min_size (int): A preprocessing paramter for :meth:`prepare`.
        max_size (int): A preprocessing paramter for :meth:`prepare`.
        ratios (list of floats): This is ratios of width to height of
            the anchors.
        anchor_scales (list of numbers): This is areas of anchors.
            Those areas will be the product of the square of an element in
            :obj:`anchor_scales` and the original area of the reference
            window.
        vgg_initialW (callable): Initializer for the layers corresponding to
            the VGG-16 layers.
        rpn_initialW (callable): Initializer for Region Proposal Network
            layers.
        loc_initialW (callable): Initializer for the localization head.
        score_initialW (callable): Initializer for the score head.
        proposal_creator_params (dict): Key valued paramters for
            :class:`~chainercv.links.model.faster_rcnn.ProposalCreator`.

    """

    _models = {
        'coco_converted': {
            'param': {'n_fg_class': 80},
            'url': 'https://github.com/knorth55/'
            'chainer-light-head-rcnn/releases/download/v0.0.0/'
            'light_head_rcnn_resnet101_converted_2018_07_12.npz',
            'cv2': True
        },
        'coco': {
            'param': {'n_fg_class': 80},
            'url': 'https://github.com/knorth55/'
            'chainer-light-head-rcnn/releases/download/v1.0.0/'
            'light_head_rcnn_resnet101_trained_2018_07_26.npz',
            'cv2': True
        },
    }
    feat_stride = 16
    proposal_creator_params = {
        'nms_thresh': 0.7,
        'n_train_pre_nms': 12000,
        'n_train_post_nms': 2000,
        'n_test_pre_nms': 6000,
        'n_test_post_nms': 1000,
        'force_cpu_nms': False,
        'min_size': 0,
    }

    def __init__(
            self,
            n_fg_class=None,
            pretrained_model=None,
            min_size=800, max_size=1333, roi_size=7,
            ratios=[0.5, 1, 2], anchor_scales=[2, 4, 8, 16, 32],
            loc_normalize_mean=(0., 0., 0., 0.),
            loc_normalize_std=(0.1, 0.1, 0.2, 0.2),
            resnet_initialW=None, rpn_initialW=None,
            global_module_initialW=None,
            loc_initialW=None, score_initialW=None,
            proposal_creator_params=None,
    ):

        param, path = utils.prepare_pretrained_model(
            {'n_fg_class': n_fg_class}, pretrained_model, self._models)

        if resnet_initialW is None and pretrained_model:
            resnet_initialW = chainer.initializers.HeNormal()
        if rpn_initialW is None:
            rpn_initialW = chainer.initializers.Normal(0.01)
        if global_module_initialW is None:
            global_module_initialW = chainer.initializers.Normal(0.01)
        if loc_initialW is None:
            loc_initialW = chainer.initializers.Normal(0.001)
        if score_initialW is None:
            score_initialW = chainer.initializers.Normal(0.01)
        if proposal_creator_params is not None:
            self.proposal_creator_params = proposal_creator_params

        extractor = ResNet101Extractor(
            initialW=resnet_initialW)
        rpn = RegionProposalNetwork(
            1024, 512,
            ratios=ratios,
            anchor_scales=anchor_scales,
            feat_stride=self.feat_stride,
            initialW=rpn_initialW,
            proposal_creator_params=self.proposal_creator_params,
        )

        head = LightHeadRCNNResNet101Head(
            param['n_fg_class'] + 1,
            roi_size=roi_size,
            spatial_scale=1. / self.feat_stride,
            global_module_initialW=global_module_initialW,
            loc_initialW=loc_initialW,
            score_initialW=score_initialW
        )
        mean = np.array([122.7717, 115.9465, 102.9801],
                        dtype=np.float32)[:, None, None]

        super(LightHeadRCNNResNet101, self).__init__(
            extractor, rpn, head, mean, min_size, max_size,
            loc_normalize_mean, loc_normalize_std)

        if path == 'imagenet':
            self._copy_imagenet_pretrained_resnet()
        elif path:
            chainer.serializers.load_npz(path, self)

    def _copy_imagenet_pretrained_resnet(self):
        def _copy_conv2dbn(src, dst):
            dst.conv.W.array = src.conv.W.array
            if src.conv.b is not None and dst.conv.b is not None:
                dst.conv.b.array = src.conv.b.array
            dst.bn.gamma.array = src.bn.gamma.array
            dst.bn.beta.array = src.bn.beta.array
            dst.bn.avg_var = src.bn.avg_var
            dst.bn.avg_mean = src.bn.avg_mean

        def _copy_bottleneck(src, dst):
            if hasattr(src, 'residual_conv'):
                _copy_conv2dbn(src.residual_conv, dst.residual_conv)
            _copy_conv2dbn(src.conv1, dst.conv1)
            _copy_conv2dbn(src.conv2, dst.conv2)
            _copy_conv2dbn(src.conv3, dst.conv3)

        def _copy_resblock(src, dst):
            for layer_name in src.layer_names:
                _copy_bottleneck(
                    getattr(src, layer_name), getattr(dst, layer_name))

        pretrained_model = ResNet101(arch='he', pretrained_model='imagenet')
        _copy_conv2dbn(pretrained_model.conv1, self.extractor.conv1)
        _copy_resblock(pretrained_model.res2, self.extractor.res2)
        _copy_resblock(pretrained_model.res3, self.extractor.res3)
        _copy_resblock(pretrained_model.res4, self.extractor.res4)
        _copy_resblock(pretrained_model.res5, self.extractor.res5)

class ResNet101Extractor(chainer.Chain):

    """ResNet101 Extractor for LightHeadRCNN ResNet101 implementation.

    This class is used as an extractor for LightHeadRCNNResNet101.
    This outputs feature maps.
    Dilated convolution is used in the C5 stage.

    Args:
        initialW: Initializer for ResNet101 extractor.
    """

    def __init__(self, initialW=None):
        super(ResNet101Extractor, self).__init__()

        if initialW is None:
            initialW = chainer.initializers.HeNormal()
        kwargs = {
            'initialW': initialW,
            'bn_kwargs': {'eps': 1e-5},
            'stride_first': True
        }

        with self.init_scope():
            # ResNet
            self.conv1 = Conv2DBNActiv(
                3, 64, 7, 2, 3, nobias=True, initialW=initialW)
            self.pool1 = lambda x: F.max_pooling_2d(x, ksize=3, stride=2)
            self.res2 = ResBlock(3, 64, 64, 256, 1, **kwargs)
            self.res3 = ResBlock(4, 256, 128, 512, 2, **kwargs)
            self.res4 = ResBlock(23, 512, 256, 1024, 2, **kwargs)
            self.res5 = ResBlock(3, 1024, 512, 2048, 1, 2, **kwargs)

    def __call__(self, x):
        """Forward the chain.

        Args:
            x (~chainer.Variable): 4D image variable.

       """

        with chainer.using_config('train', False):
            h = self.pool1(self.conv1(x))
            h = self.res2(h)
            h.unchain_backward()
            h = self.res3(h)
            res4 = self.res4(h)
            res5 = self.res5(res4)
        return res4, res5

class LightHeadRCNNResNet101Head(chainer.Chain):

    def __init__(
            self, n_class, roi_size, spatial_scale,
            global_module_initialW=None,
            loc_initialW=None, score_initialW=None
    ):

        super(LightHeadRCNNResNet101Head, self).__init__()
        self.n_class = n_class
        self.spatial_scale = spatial_scale
        self.roi_size = roi_size
        with self.init_scope():
            self.global_context_module = GlobalContextModule(
                2048, 256, self.roi_size * self.roi_size * 10, 15,
                initialW=global_module_initialW)
            self.fc1 = L.Linear(
                self.roi_size * self.roi_size * 10, 2048,
                initialW=score_initialW)
            self.score = L.Linear(2048, n_class, initialW=score_initialW)
            self.cls_loc = L.Linear(2048, 4 * n_class, initialW=loc_initialW)

    def __call__(self, x, rois, roi_indices):
        # global context module
        h = self.global_context_module(x)
        # psroi max align
        pool = psroi_max_align_2d(
            h, rois, roi_indices,
            10, self.roi_size, self.roi_size,
            self.spatial_scale, self.roi_size,
            sampling_ratio=2.)
        # fc
        fc1 = F.relu(self.fc1(pool))
        roi_cls_locs = self.cls_loc(fc1)
        roi_scores = self.score(fc1)
        return roi_cls_locs, roi_scores


""" Classes for remixing audio files.
(c) 2017 - Dave Rensin - dave@rensin.com

This module contains classes for remixing audio files. It started
as an attempt to re-create the amazing Infinite Jukebox (http://www.infinitejuke.com)
created by Paul Lamere of Echo Nest.

The InfiniteJukebox class can do it's processing in a background thread and
reports progress via the progress_callback arg. To run in a thread, pass do_async=True
to the constructor. In that case, it exposes an Event named play_ready -- which will
be signaled when the processing is complete. The default mode is to run synchronously.

  Async example:

      def MyCallback(percentage_complete_as_float, string_message):
        print "I am now %f percent complete with message: %s" % (percentage_complete_as_float * 100, string_message)

      jukebox = InfiniteJukebox(filename='some_file.mp3', progress_callback=MyCallback, do_async=True)
      jukebox.play_ready.wait()

      <some work here...>

  Non-async example:

      def MyCallback(percentage_complete_as_float, string_message):
        print "I am now %f percent complete with message: %s" % (percentage_complete_as_float * 100, string_message)

      jukebox = InfiniteJukebox(filename='some_file.mp3', progress_callback=MyCallback, do_async=False)

      <blocks until completion... some work here...>

"""

import collections
import librosa
import math
import random
import scipy
import threading

import numpy as np
import sklearn.cluster
import sklearn.metrics

import multiprocessing
import functools
import time

import os
from pathlib import Path
import csv

from utils import CONFIG, CacheOptions

def smap(f):
    return f()

class InfiniteJukebox(object):

    """ Class to "infinitely" remix a song.

    This class will take an audio file (wav, mp3, ogg, etc) and
    (a) decompose it into individual beats, (b) find the tempo
    of the track, and (c) create a play path that you can use
    to play the song approx infinitely.

    The idea is that it will find and cluster beats that are
    musically similar and return them to you so you can automatically
    'remix' the song.

    Attributes:

     play_ready: an Event that triggers when the processing/clustering is complete and
                 playback can begin. This is only defined if you pass do_async=True in the
                 constructor.

    start_index: the start index of the original track before trimming (i.e. leading silence is before this start index)

       duration: the duration (in seconds) of the track after the leading and trailing silences
                 have been removed.

      raw_audio: an array of numpy.Int16 that is suitable for using for playback via pygame
                 or similar modules. If the audio is mono then the shape of the array will
                 be (bytes,). If it's stereo, then the shape will be (2,bytes).

    sample_rate: the sample rate from the audio file. Usually 44100 or 48000

       clusters: the number of clusters used to group the beats. If you pass in a value, then
                 this will be reflected here. If you let the algorithm decide, then auto-generated
                 value will be reflected here.

          beats: a dictionary containing the individual beats of the song in normal order. Each
                 beat will have the following keys:

                         id: the ordinal position of the beat in the song
                      start: the time (in seconds) in the song where this beat occurs
                   duration: the duration (in seconds) of the beat
                     buffer: an array of audio bytes for this beat. it is just raw_audio[start:start+duration]
                    cluster: the cluster that this beat most closely belongs. Beats in the same cluster
                             have similar harmonic (timbre) and chromatic (pitch) characteristics. They
                             will "sound similar"
                    segment: the segment to which this beat belongs. A 'segment' is a contiguous block of
                             beats that belong to the same cluster.
                  amplitude: the loudness of the beat
                       next: the next beat to play after this one, if playing sequentially
            jump_candidates: a list of the other beats in the song to which it is reasonable to jump. Those beats
                             (a) are in the same cluster as the NEXT oridnal beat, (b) are of the same segment position
                             as the next ordinal beat, (c) are in the same place in the measure as the NEXT beat,
                             (d) but AREN'T the next beat.

                 An example of playing the first 32 beats of a song:

                    from Remixatron import InfiniteJukebox
                    from pygame import mixer
                    import time

                    jukebox = InfiniteJukebox('some_file.mp3')

                    pygame.mixer.init(frequency=jukebox.sample_rate)
                    channel = pygame.mixer.Channel(0)

                    for beat in jukebox.beats[0:32]:
                        snd = pygame.Sound(buffer=beat['buffer'])
                        channel.queue(snd)
                        time.sleep(beat['duration'])

    play_vector: a beat play list of 1024^2 items. This represents a pre-computed
                 remix of this song that will last beat['duration'] * 1024 * 1024
                 seconds long. A song that is 120bpm will have a beat duration of .5 sec,
                 so this playlist will last .5 * 1024 * 1024 seconds -- or 145.67 hours.

                 Each item contains:

                    beat: an index into the beats array of the beat to play
                 seq_len: the length of the musical sequence being played
                          in this part of play_vector.
                 seq_pos: this beat's position in seq_len. When
                          seq_len - seq_pos == 0 the song will "jump"

    """

    def __init__(self, filepath, start_beat=1, use_cache = False, clusters=0, max_clusters = 48, progress_callback=None,
                 do_async=False, use_v1=False):

        """ The constructor for the class. Also starts the processing thread.

            Args:

                filepath: the path to the audio file to process
              start_beat: the first beat to play in the file. Should almost always be 1,
                          but you can override it to skip into a specific part of the song.
               use_cache: use cache if it exists
                clusters: the number of similarity clusters to compute. The DEFAULT value
                          of 0 means that the code will try to automatically find an optimal
                          cluster. If you specify your own value, it MUST be non-negative. Lower
                          values will create more promiscuous jumps. Larger values will create higher quality
                          matches, but run the risk of jumps->0 -- which will just loop the
                          audio sequentially ~forever.
            max_clusters: number of clusters to try
       progress_callback: a callback function that will get periodic satatus updates as
                          the audio file is processed. MUST be a function that takes 2 args:

                             percent_complete: FLOAT between 0.0 and 1.0
                                      message: STRING with the progress message
                  use_v1: set to True if you want to use the original auto clustering algorithm.
                          Otherwise, it will use the newer silhouette-based scheme.
        """
        self.__progress_callback = progress_callback
        self.filepath = filepath
        self.__start_beat = start_beat
        self.clusters = clusters
        self.__max_clusters = max_clusters
        self._extra_diag = ""
        self._use_v1 = use_v1
        self.cache_option = CacheOptions.DISCARD

        if use_cache and os.path.isfile(os.path.join(CONFIG['cacheDir'], Path(filepath).stem + '.csv')):
            self.evecs = np.array([])
            self.__load_cache()
        else:
            if do_async == True:
                self.play_ready = threading.Event()
                self.__thread = threading.Thread(target=self.__process_audio)
                self.__thread.start()
            else:
                self.play_ready = None
                self.__process_audio()

    def save_cache(self, cache_evecs = False):
        os.makedirs(CONFIG['cacheDir'], exist_ok=True)
        with open(os.path.join(CONFIG['cacheDir'], Path(self.filepath).stem + '.csv'), 'w', newline='') as csvfile:
            fieldnames = ['start_index', 'cluster']  # , 'stop_index', 'start', 'duration' ]
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerow({'start_index': "",  # jukebox.filename,
                             'cluster': self.avg_amplitude})
            writer.writerow({'start_index': self.__start_beat,
                             'cluster': self.clusters})
            # ,
            # 'stop_index': filepath,
            # 'start': 0,
            # 'duration': jukebox.duration})
            for beat in self.beats:
                writer.writerow({'start_index': beat['start_index'],
                                 'cluster': beat['cluster']})  # ,
                # 'stop_index': beat['stop_index'],
                # 'start': beat['start'],
                # 'duration': beat['duration']})

        if cache_evecs:
            np.save(os.path.join(CONFIG['cacheDir'], Path(self.filepath).stem + '.npy'), self.evecs)

    def remove_cache(self):
        songname = Path(self.filepath).stem
        if self.cache_option == self.cache_option.DISCARD:
            if os.path.isfile(os.path.join(CONFIG['cacheDir'], songname + '.csv')):
                os.remove(os.path.join(CONFIG['cacheDir'], songname + '.csv'))
            if os.path.isfile(os.path.join(CONFIG['cacheDir'], songname + '.npy')):
                os.remove(os.path.join(CONFIG['cacheDir'], songname + '.npy'))
        elif self.cache_option == self.cache_option.KEEP_CACHE:
            if os.path.isfile(os.path.join(CONFIG['cacheDir'], songname + '.npy')):
                os.remove(os.path.join(CONFIG['cacheDir'], songname + '.npy'))

    def __load_cache(self):
        self.beats = []
        with open(os.path.join(CONFIG['cacheDir'], Path(self.filepath).stem + '.csv'), newline='') as csvfile:
            reader = csv.DictReader(csvfile)
            for i, beat in enumerate(reader):
                if i == 0:
                    #if beat['start_index'] != self.filepath:
                    #    break
                    self.__report_progress(.8, "loading from cache...")
                    self.avg_amplitude = float(beat['cluster'])

                elif i == 1:
                    start_index_diff = max(0, int(beat['start_index']) - self.__start_beat)
                    self.__start_beat += start_index_diff
                    clusters = int(beat['cluster'])
                elif i >= start_index_diff + 2:
                    self.beats.append({'start_index': int(beat['start_index']),
                                  'cluster': int(beat['cluster'])})#,
                                  #'id': i})  # ,
                    # 'stop_index': beat['stop_index'],
                    # 'start': beat['start'],
                    # 'duration': beat['duration']})

        y, sr = librosa.core.load(self.filepath, mono=False, sr=None)
        y, index = librosa.effects.trim(y)

        self.start_index = index[0]

        self.duration = librosa.core.get_duration(y, sr)
        self.raw_audio = (y * np.iinfo(np.int16).max).astype(np.int16).T.copy(order='C')
        self.sample_rate = sr

        if os.path.isfile(os.path.join(CONFIG['cacheDir'], Path(self.filepath).stem + '.npy')):
            self.evecs = np.load(os.path.join(CONFIG['cacheDir'], Path(self.filepath).stem + '.npy'))

        if self.clusters == 0: # if 0 in config, use what was saved
            self.clusters = clusters
        self.recompute_beat_array(clusters)

        self.time_elapsed = -1

    def recompute_beat_array(self, clusters):

        if (self.clusters != clusters): # if num of clusters different from what was saved

            seg_ids, self.clusters = self.__compute_cluster(self.evecs, clusters, self._use_v1)

            for i, beat in enumerate(self.beats):
                beat['cluster'] = seg_ids[i]
        else:
            self.clusters = clusters

        total_indices = self.raw_audio.shape[0]

        for i, beat in enumerate(self.beats):
            beat['id'] = i
            beat['start'] = (beat['start_index']/total_indices) * self.duration
            beat['quartile'] = beat['id'] // (len(self.beats) / 4.0)

            if i == (len(self.beats) - 1):
                beat['stop_index'] = total_indices
                beat['next'] = self.beats[0]['id']
                beat['duration'] = self.duration - beat['start']
            else:
                beat['stop_index']  = self.beats[i + 1]['start_index']
                beat['next'] = i + 1
                beat['duration'] = ((self.beats[i + 1]['start_index'] - beat['start_index']) / total_indices) * self.duration

            if i == 0:
                beat['segment'] = 0
                beat['is'] = 0
            else:
                if beat['cluster'] != self.beats[i - 1]['cluster']:
                    beat['segment'] = self.beats[i - 1]['segment'] + 1
                    beat['is'] = 0
                else:
                    beat['segment'] = self.beats[i - 1]['segment']
                    beat['is'] = self.beats[i - 1]['is'] + 1

            beat['buffer'] = self.raw_audio[beat['start_index']: beat['stop_index']]

        for beat in self.beats[:-1]:
            jump_candidates = [bx['id'] for bx in self.beats[:beat['id']] if  # only consider beats that are earlier
                               (bx['cluster'] == self.beats[beat['next']]['cluster']) and
                               (bx['is'] == self.beats[beat['next']]['is']) and
                               # (bx['id'] % 4 == beats[beat['next']]['id'] % 4) and # removed as was limiting loop points
                               (bx['segment'] != beat['segment']) and
                               (bx['id'] != beat['next'])]

            if jump_candidates:
                beat['jump_candidates'] = jump_candidates
            else:
                beat['jump_candidates'] = []

        self.beats[-1]['jump_candidates'] = []

    def __process_audio(self):

        """ The main audio processing routine for the thread.

        This routine uses Laplacian Segmentation to find and
        group similar beats in the song.

        This code has been adapted from the sample created by Brian McFee at
        https://librosa.github.io/librosa_gallery/auto_examples/plot_segmentation.html#sphx-glr-auto-examples-plot-segmentation-py
        and is based on his 2014 paper published at http://bmcfee.github.io/papers/ismir2014_spectral.pdf

        I have made some performance improvements, but the basic parts remain (mostly) unchanged
        """

        start = time.time()

        self.__report_progress( .1, "loading file and extracting raw audio")

        #
        # load the file as stereo with a high sample rate and
        # trim the silences from each end
        #

        y, sr = librosa.core.load(self.filepath, mono=False, sr=None)
        y, index = librosa.effects.trim(y)

        self.start_index = index[0]

        self.duration = librosa.core.get_duration(y,sr)
        self.raw_audio = (y * np.iinfo(np.int16).max).astype(np.int16).T.copy(order='C')
        self.sample_rate = sr

        # after the raw audio bytes are saved, convert the samples to mono
        # because the beat detection algorithm in librosa requires it.

        y = librosa.core.to_mono(y)

        self.__report_progress( .2, "computing pitch data and finding beats...")

        # Compute the constant-q chromagram for the samples.

        BINS_PER_OCTAVE = 12 * 3
        N_OCTAVES = 7

        # If multiple cores exist, process cqt and beat track at same time to cut computation time
        if multiprocessing.cpu_count() > 1:
            f_cqt = functools.partial(librosa.cqt, y=y, sr=sr, bins_per_octave=BINS_PER_OCTAVE, n_bins=N_OCTAVES * BINS_PER_OCTAVE)
            f_beat_track = functools.partial(librosa.beat.beat_track, y=y, sr=sr, trim=False)

            with multiprocessing.Pool(processes=2) as pool:
                res = pool.map(smap, [f_cqt, f_beat_track])
                cqt = res[0]
                tempo = res[1][0]
                btz = res[1][1]
        else:
            cqt = librosa.cqt(y=y, sr=sr, bins_per_octave=BINS_PER_OCTAVE, n_bins=N_OCTAVES * BINS_PER_OCTAVE) ######
            tempo, btz = librosa.beat.beat_track(y=y, sr=sr, trim=False)  ##########

        # Cynthia
        # single core: 21s, 23s, 28s, 26s
        # multi core: 19s

        # 08 HOHR
        # single core: 45s, 46s
        # multi core: 37s, 37s

        # Hornet
        # single core: 14s, 12s, 10s
        # multi core: 11s, 11s

        C = librosa.amplitude_to_db( np.abs(cqt), ref=np.max)

        #self.__report_progress( .3, "Finding beats..." )

        ##########################################################
        # To reduce dimensionality, we'll beat-synchronous the CQT
        #tempo, btz = librosa.beat.beat_track(y=y, sr=sr, trim=False) ##########
        # tempo, btz = librosa.beat.beat_track(y=y, sr=sr)
        Csync = librosa.util.sync(C, btz, aggregate=np.median)

        self.tempo = tempo

        # For alignment purposes, we'll need the timing of the beats
        # we fix_frames to include non-beat frames 0 and C.shape[1] (final frame)
        beat_times = librosa.frames_to_time(librosa.util.fix_frames(btz,
                                                                    x_min=0,
                                                                    x_max=C.shape[1]),
                                            sr=sr)

        self.__report_progress( .4, "building recurrence matrix...")
        #####################################################################
        # Let's build a weighted recurrence matrix using beat-synchronous CQT
        # (Equation 1)
        # width=3 prevents links within the same bar
        # mode='affinity' here implements S_rep (after Eq. 8)
        R = librosa.segment.recurrence_matrix(Csync, width=3, mode='affinity',
                                              sym=True)

        # Enhance diagonals with a median filter (Equation 2)
        df = librosa.segment.timelag_filter(scipy.ndimage.median_filter)
        Rf = df(R, size=(1, 7))

        ###################################################################
        # Now let's build the sequence matrix (S_loc) using mfcc-similarity
        #
        #   :math:`R_\text{path}[i, i\pm 1] = \exp(-\|C_i - C_{i\pm 1}\|^2 / \sigma^2)`
        #
        # Here, we take :math:`\sigma` to be the median distance between successive beats.
        #
        mfcc = librosa.feature.mfcc(y=y, sr=sr)
        Msync = librosa.util.sync(mfcc, btz)

        path_distance = np.sum(np.diff(Msync, axis=1)**2, axis=0)
        sigma = np.median(path_distance)
        path_sim = np.exp(-path_distance / sigma)

        R_path = np.diag(path_sim, k=1) + np.diag(path_sim, k=-1)

        ##########################################################
        # And compute the balanced combination (Equations 6, 7, 9)

        deg_path = np.sum(R_path, axis=1)
        deg_rec = np.sum(Rf, axis=1)

        mu = deg_path.dot(deg_path + deg_rec) / np.sum((deg_path + deg_rec)**2)

        A = mu * Rf + (1 - mu) * R_path

        #####################################################
        # Now let's compute the normalized Laplacian (Eq. 10)
        L = scipy.sparse.csgraph.laplacian(A, normed=True)

        # and its spectral decomposition
        _, evecs = scipy.linalg.eigh(L)

        # We can clean this up further with a median filter.
        # This can help smooth over small discontinuities
        evecs = scipy.ndimage.median_filter(evecs, size=(9, 1))

        self.evecs = evecs #save intermediate step for caching

        # Cluster beats
        seg_ids, self.clusters = self.__compute_cluster(self.evecs, self.clusters, self._use_v1)

        # Get the amplitudes and beat-align them
        self.__report_progress( .6, "getting amplitudes")

        # newer versions of librosa have renamed the rmse function

        if hasattr(librosa.feature,'rms'):
            amplitudes = librosa.feature.rms(y=y)
        else:
            amplitudes = librosa.feature.rmse(y=y)

        ampSync = librosa.util.sync(amplitudes, btz)

        # create a list of tuples that include the ordinal position, the start time of the beat,
        # the cluster to which the beat belongs and the mean amplitude of the beat

        zbeat_tuples = zip(range(0,len(btz)), beat_times, seg_ids, ampSync[0].tolist())
        beat_tuples =tuple(zbeat_tuples)

        info = []

        bytes_per_second = int(round(len(self.raw_audio) / self.duration))

        last_cluster = -1
        current_segment = -1
        segment_beat = 0

        for i in range(0, len(beat_tuples)):
            final_beat = {}
            final_beat['start'] = float(beat_tuples[i][1])
            final_beat['cluster'] = int(beat_tuples[i][2])
            final_beat['amplitude'] = float(beat_tuples[i][3])

            if final_beat['cluster'] != last_cluster:
                current_segment += 1
                segment_beat = 0
            else:
                segment_beat += 1

            final_beat['segment'] = current_segment
            final_beat['is'] = segment_beat

            last_cluster = final_beat['cluster']

            if i == len(beat_tuples) - 1:
                final_beat['duration'] = self.duration - final_beat['start']
            else:
                final_beat['duration'] = beat_tuples[i+1][1] - beat_tuples[i][1]

            if ( (final_beat['start'] * bytes_per_second) % 2 > 1.5 ):
                final_beat['start_index'] = int(math.ceil(final_beat['start'] * bytes_per_second))
            else:
                final_beat['start_index'] = int(final_beat['start'] * bytes_per_second)

            final_beat['stop_index'] = int(math.ceil((final_beat['start'] + final_beat['duration']) * bytes_per_second))

            # save pointers to the raw bytes for each beat with each beat.
            final_beat['buffer'] = self.raw_audio[ final_beat['start_index'] : final_beat['stop_index'] ]

            info.append(final_beat)

        #self.__report_progress( .7, "truncating to fade point...")

        # get the max amplitude of the beats
        # max_amplitude = max([float(b['amplitude']) for b in info])
        avg_amplitude = sum([float(b['amplitude']) for b in info]) / len(info)

        # assume that the fade point of the song is the last beat of the song that is >= 75% of
        # the max amplitude.

        self.avg_amplitude = avg_amplitude

        fade = len(info) - 1

        #for b in reversed(info):
        #    if b['amplitude'] >= (.75 * avg_amplitude):
        #        fade = info.index(b)
        #        break

        # truncate the beats to [start:fade + 1]
        beats = info[self.__start_beat:fade + 1]

        loop_bounds_begin = self.__start_beat

        self.__report_progress( .8, "computing final beat array and finding loops...")

        # assign final beat ids
        for beat in beats:
            beat['id'] = beats.index(beat)
            beat['quartile'] = beat['id'] // (len(beats) / 4.0)

        # compute a coherent 'next' beat to play. This is always just the next ordinal beat
        # unless we're at the end of the song. Then it gets a little trickier.

        for beat in beats[:-1]:
            #if beat == beats[-1]:

                # if we're at the last beat, then we want to find a reasonable 'next' beat to play. It should (a) share the
                # same cluster, (b) be in a logical place in its measure, (c) be after the computed loop_bounds_begin, and
                # is in the first half of the song. If we can't find such an animal, then just return the beat
                # at loop_bounds_begin

                # beat['next'] = next( (b['id'] for b in beats if b['cluster'] == beat['cluster'] and
                #                       b['id'] % 4 == (beat['id'] + 1) % 4 and
                #                       b['id'] <= (.5 * len(beats)) and
                #                       b['id'] >= loop_bounds_begin), loop_bounds_begin )
                #beat['next'] = beats[0]['id']
            #else:
            beat['next'] = beat['id'] + 1

            # find all the beats that (a) are in the same cluster as the NEXT oridnal beat, (b) are of the same
            # cluster position as the next ordinal beat, (c) are in the same place in the measure as the NEXT beat,
            # (d) but AREN'T the next beat, and (e) AREN'T in the same cluster as the current beat.
            #
            # THAT collection of beats contains our jump candidates

            jump_candidates = [bx['id'] for bx in beats[:beat['id']] if # only consider beats that are earlier
                               (bx['cluster'] == beats[beat['next']]['cluster']) and
                               (bx['is'] == beats[beat['next']]['is']) and
                               #(bx['id'] % 4 == beats[beat['next']]['id'] % 4) and # removed as was limiting loop points
                               (bx['segment'] != beat['segment']) and
                               (bx['id'] != beat['next'])]

            if jump_candidates:
                beat['jump_candidates'] = jump_candidates
            else:
                beat['jump_candidates'] = []

        beats[-1]['jump_candidates'] = []
        beats[-1]['next'] = beats[0]['id']

        # save off the segment count

        self.segments = max([b['segment'] for b in beats]) + 1

        # we don't want to ever play past the point where it's impossible to loop,
        # so let's find the latest point in the song where there are still jump
        # candidates and make sure that we can't play past it.

        # last_chance = len(beats) - 1
        #
        # for b in reversed(beats):
        #     if len(b['jump_candidates']) > 0:
        #         last_chance = beats.index(b)
        #         break

        # if we play our way to the last beat that has jump candidates, then just skip
        # to the earliest jump candidate rather than enter a section from which no
        # jumping is possible.

        #beats[last_chance]['next'] = min(beats[last_chance]['jump_candidates'])

        # store the beats that start after the last jumpable point. That's
        # the outro to the song. We can use these
        # beasts to create a sane ending for a fixed-length remix

        # outro_start = last_chance + 1 + self.__start_beat
        #
        # if outro_start >= len(info):
        #     self.outro = []
        # else:
        #     self.outro = info[outro_start:]

        #
        # Computing play_vector is removed in this as it loops live
        #

        random.seed()

        play_vector = []

        # save off the beats array and play_vector. Signal
        # the play_ready event (if it's been set)

        self.beats = beats
        self.play_vector = play_vector

        self.__report_progress(1.0, "finished processing")

        if self.play_ready:
            self.play_ready.set()

        self.time_elapsed = time.time() - start

    def __report_progress(self, pct_done, message):

        """ If a reporting callback was passed, call it in order
            to mark progress.
        """
        if self.__progress_callback:
            self.__progress_callback(pct_done, message, self.filepath)

    def __compute_cluster(self, evecs, clusters, use_v1 = False):

        # cumulative normalization is needed for symmetric normalize laplacian eigenvectors
        Cnorm = np.cumsum(evecs ** 2, axis=1) ** 0.5

        # If we want k clusters, use the first k normalized eigenvectors.
        # Fun exercise: see how the segmentation changes as you vary k

        self.__report_progress(.5, "clustering...")

        # if a value for clusters wasn't passed in, then we need to auto-cluster
        if clusters <= 0:

            # if we've been asked to use the original auto clustering algorithm, otherwise
            # use the new and improved one that accounts for silhouette scores.

            if use_v1:
                clusters, seg_ids = self.__compute_best_cluster(evecs, Cnorm)
            else:
                clusters, seg_ids = self.__compute_best_cluster_with_sil(evecs, Cnorm)
        else:
            # otherwise, just use the cluster value passed in

            self.__report_progress(.51, "using %d clusters" % clusters)

            X = evecs[:, :clusters] / Cnorm[:, clusters - 1:clusters]

            while (np.any(np.isnan(X))) and (not np.all(np.isfinite(X))): # if input is invalid, increment until a valid input is calculated
                clusters += 1
                X = evecs[:, :clusters] / Cnorm[:, clusters - 1:clusters]

            seg_ids = sklearn.cluster.KMeans(n_clusters=clusters, max_iter=1000,
                                             random_state=0, n_init=1000, n_jobs=-1).fit_predict(X)

        return seg_ids, clusters

    def __compute_best_cluster_with_sil(self, evecs, Cnorm):

        ''' Attempts to compute optimum clustering

            Uses the the silhouette score to pick the best number of clusters.
            See: https://en.wikipedia.org/wiki/Silhouette_(clustering)

            PARAMETERS:
                evecs: Eigen-vectors computed from the segmentation algorithm
                Cnorm: Cumulative normalization of evecs. Easier to pass it in than
                       compute it from scratch here.

            KEY DEFINITIONS:

                  Clusters: buckets of musical similarity
                  Segments: contiguous blocks of beats belonging to the same cluster
                Silhouette: A score given to a cluster that measures how well the cluster
                            members fit together. The value is from -1 to +1. Higher values
                            indicate higher quality.
                   Orphans: Segments with only one beat. The presence of orphans is a potential
                            sign of overfitting.

            SUMMARY:

                There are lots of things that might indicate one cluster count is better than another.
                High silhouette scores for the candidate clusters mean that the jumps will be higher
                quality.

                On the other hand, we could easily choose so many clusters that everyone has a great
                silhouette score but none of the beats have other segments into which they can jump.
                That will be a pretty boring result!

                So, the cluster/segment ratio matters, too The higher the number, the more places (on average)
                a beat can jump. However, if the beats aren't very similar (low silhouette scores) then
                the jumps won't make any musical sense.

                So, we can't just choose the cluster count with the highest average silhouette score or the
                highest cluster/segment ratio.

                Instead, we comput a simple fitness score of:
                        cluster_count * ratio * average_silhouette

                Finally, segments with only one beat are a potential (but not definite) sign of overfitting.
                We call these one-beat segments 'orphans'. We want to keep an eye out for those and slightly
                penalize any candidate cluster count that contains orphans.

                If we find an orphan, we scale the fitness score by .8 (ie. penalize it 20%). That's
                enough to push any candidate cluster count down the stack rank if orphans aren't
                otherwise very common across most of the other cluster count choices.

        '''

        self._clusters_list = []

        best_cluster_size = 0
        best_labels = None
        best_cluster_score = 0

        # we need at least 3 clusters for any song and shouldn't need to calculate more than
        # 48 clusters for even a really complicated peice of music.

        for n_clusters in range(self.__max_clusters, 2, -1):

            self.__report_progress(.51, "Testing a cluster value of %d..." % n_clusters)

            # compute a matrix of the Eigen-vectors / their normalized values
            X = evecs[:, :n_clusters] / Cnorm[:, n_clusters-1:n_clusters]

            if (not np.any(np.isnan(X))) and (np.all(np.isfinite(X))): # ensure input is valid

                # create the candidate clusters and fit them
                clusterer = sklearn.cluster.KMeans(n_clusters=n_clusters, max_iter=300,
                                                   random_state=0, n_init=20, n_jobs=-1) # multiprocess

                cluster_labels = clusterer.fit_predict(X)

                # get some key statistics, including how well each beat in the cluster resemble
                # each other (the silhouette average), the ratio of segments to clusters, and the
                # length of the smallest segment in this cluster configuration

                silhouette_avg = sklearn.metrics.silhouette_score(X, cluster_labels)

                ratio, min_segment_len = self.__segment_stats_from_labels(cluster_labels.tolist())

                # We need to grade each cluster according to how likely it is to produce a good
                # result. There are a few factors to look at.
                #
                # First, we can look at how similar the beats in each cluster (on average) are for
                # this candidate cluster size. This is known as the silhouette score. It ranges
                # from -1 (very bad) to 1 (very good).
                #
                # Another thing we can look at is the ratio of clusters to segments. Higher ratios
                # are preferred because they afford each beat in a cluster the opportunity to jump
                # around to meaningful places in the song.
                #
                # All other things being equal, we prefer a higher cluster count to a lower one
                # because it will tend to make the jumps more selective -- and therefore higher
                # quality.
                #
                # Lastly, if we see that we have segments equal to just one beat, that might be
                # a sign of overfitting. We call these one beat segments 'orphans'. Some songs,
                # however, will have orphans no matter what cluster count you use. So, we don't
                # want to throw out a cluster count just because it has orphans. Instead, we
                # just de-rate its fitness score. If most of the cluster candidates have orphans
                # then this won't matter in the overall scheme because everyone will be de-rated
                # by the same scaler.
                #
                # Putting this all together, we muliply the cluster count * the average
                # silhouette score for the clusters in this candidate * the ratio of clusters to
                # segments. Then we scale (or de-rate) the fitness score by whether or not is has
                # orphans in it.

                orphan_scaler = .8 if min_segment_len == 1 else 1

                cluster_score = n_clusters * silhouette_avg * ratio * orphan_scaler
                #cluster_score = ((n_clusters/48.0) * silhouette_avg * (ratio/10.0)) * orphan_scaler

                # if this cluster count has a score that's better than the best score so far, store
                # it for later.

                if cluster_score >= best_cluster_score:
                    best_cluster_score = cluster_score
                    best_cluster_size = n_clusters
                    best_labels = cluster_labels

        # return the best results
        return (best_cluster_size, best_labels)

    @staticmethod
    def __segment_count_from_labels(labels):

        ''' Computes the number of unique segments from a set of ordered labels. Segements are
            contiguous beats that belong to the same cluster. '''

        segment_count = 0
        previous_label = -1

        for label in labels:
            if label != previous_label:
                previous_label = label
                segment_count += 1

        return segment_count

    def __segment_stats_from_labels(self, labels):
        ''' Computes the segment/cluster ratio and min segment size value given an array
            of labels. '''

        segment_count = 0.0
        segment_length = 0
        clusters = max(labels) + 1

        previous_label = -1

        segment_lengths = []

        for label in labels:
            if label != previous_label:
                previous_label = label
                segment_count += 1.0

                if segment_length > 0:
                    segment_lengths.append(segment_length)

                segment_length = 1
            else:
                segment_length +=1

        # self.__report_progress( .52, "clusters: %d,  ratio: %f,  min_seg: %d" % (clusters, segment_count/len(labels), segment_length) )

        return float(segment_count) / float(clusters), min(segment_lengths)

    def __compute_best_cluster(self, evecs, Cnorm):

        ''' Attempts to compute optimum clustering from a set of simplified
            hueristics. This method has been deprecated in favor of code above that takes into
            account the average silhouette score of each cluster. You can force the code to use
            this method by passing in use_v1=True in the constructor.

            PARAMETERS:
                evecs: Eigen-vectors computed from the segmentation algorithm
                Cnorm: Cumulative normalization of evecs. Easier to pass it in than
                       compute it from scratch here.

            KEY DEFINITIONS:

                Clusters: buckets of musical similarity
                Segments: contiguous blocks of beats belonging to the same cluster
                 Orphans: clusters that only belong to one segment
                    Stub: a cluster with less than N beats. Stubs are a sign of
                          overfitting

            SUMMARY:

                Group the beats in [8..64] clusters. They key metric is the segment:cluster ratio.
                This value gives the avg number of different segments to which a cluster
                might belong. The higher the value, the more diverse the playback because
                the track can jump more freely. There is a balance, however, between this
                ratio and the number of clusters. In general, we want to find the highest
                numeric cluster that has a ratio of segments:clusters nearest 4.
                That ratio produces the most musically pleasing results.

                Basically, we're looking for the highest possible cluster # that doesn't
                obviously overfit.

                Someday I'll implement a proper RMSE algorithm...
        '''

        self._clusters_list = []

        # We compute the clusters between 4 and 64. Owing to the inherent
        # symmetry of Western popular music (including Jazz and Classical), the most
        # pleasing musical results will often, though not always, come from even cluster values.

        for ki in range(4,64, 2):

            # compute a matrix of the Eigen-vectors / their normalized values
            X = evecs[:, :ki] / Cnorm[:, ki-1:ki]

            if (not np.any(np.isnan(X))) and (np.all(np.isfinite(X))): # ensure input is valid
                # cluster with candidate ki
                labels = sklearn.cluster.KMeans(n_clusters=ki, max_iter=1000,
                                                random_state=0, n_init=20, n_jobs=-1).fit_predict(X)

                entry = {'clusters':ki, 'labels':labels}

                # create an array of dictionary entries containing (a) the cluster label,
                # (b) the number of total beats that belong to that cluster, and
                # (c) the number of segments in which that cluster appears.

                lst = []

                for i in range(0,ki):
                    lst.append( {'label':i, 'beats':0, 'segs':0} )

                last_label = -1

                for l in labels:

                    if l != last_label:
                        lst[l]['segs'] += 1
                        last_label = l

                    lst[l]['beats'] += 1

                entry['cluster_map'] = lst

                # get the average number of segments to which a cluster belongs
                entry['seg_ratio'] = np.mean([l['segs'] for l in entry['cluster_map']])

                self._clusters_list.append(entry)

        # get the max cluster with the segments/cluster ratio nearest to 4. That
        # will produce the most musically pleasing effect

        max_seg_ratio = max( [cl['seg_ratio'] for cl in self._clusters_list] )
        max_seg_ratio = min( max_seg_ratio, 4 )

        final_cluster_size = max(cl['clusters'] for cl in self._clusters_list if cl['seg_ratio'] >= max_seg_ratio)

        # compute a very high fidelity set of clusters using our selected cluster size.
        X = evecs[:, :final_cluster_size] / Cnorm[:, final_cluster_size-1:final_cluster_size]
        labels = sklearn.cluster.KMeans(n_clusters=final_cluster_size, max_iter=1000,
                                        random_state=0, n_init=1000, n_jobs=-1).fit_predict(X)

        # labels = next(c['labels'] for c in self._clusters_list if c['clusters'] == final_cluster_size)

        # return a tuple of (winning cluster size, [array of cluster labels for the beats])
        return (final_cluster_size, labels)

    def __add_log(self, line):
        """Convenience method to add debug logging info for later"""

        self._extra_diag += line + "\n"


import numpy as np
import cv2
import json
import os
from time import time
from scipy import ndimage
from functools import wraps

from utils.config import Config
from model import ICNet, ICNet_BN

model_config = {'train': ICNet, 'trainval': ICNet, 'train_bn': ICNet_BN, 'trainval_bn': ICNet_BN, 'others': ICNet_BN}

# Choose dataset here, but remember to use `script/downlaod_weight.py` first
dataset = 'cityscapes'
filter_scale = 1

SIZE_WEBOT = (64, 128, 3)
SIZE_HD = (1080,1920, 3)

def timer(func):
    @wraps(func)
    def wrapper(instance, *args, **kwargs):
        start_time = time()
        result = func(instance, *args, **kwargs)
        instance.cycle_time = time() - start_time
        if instance.average_cycle_time is None:
            instance.average_cycle_time = instance.cycle_time
        instance.average_cycle_time = instance.average_cycle_time * 0.95 + instance.cycle_time * 0.05
        return result
    return wrapper

class Brain:
    class InferenceConfig(Config):
        def __init__(self, dataset, is_training, filter_scale):
            Config.__init__(self, dataset, is_training, filter_scale)
        model_type = 'trainval'
        model_weight = './model/cityscapes/icnet_cityscapes_train_30k.npy'

    class HotReloadConfigs:
        def __init__(self, config_file_path='config.json'):
            self.last_modified = 0
            self.config = None
            self.path = config_file_path
            self.reload()

        def reload(self):
            stat = os.stat(self.path)
            if stat.st_mtime != self.last_modified:
                print ('reloading')
                try:
                    with open(self.path) as f:
                        self.config = json.load(f)
                    self.last_modified = stat.st_mtime
                    return True
                except (IOError, ValueError, KeyError):
                    print('can not reload config')
            return False

        def __getattr__(self, name):
            return self.config[name]

    def __init__(self, inner_size=SIZE_HD, exponential_filter_alpha=0.9):
        self.last_theta = 0
        self.alpha = exponential_filter_alpha
        self.cycle_time = None
        self.average_cycle_time = None
        self.config = Brain.HotReloadConfigs()

        print('Configuring Brain')
        self.cfg = Brain.InferenceConfig(dataset, is_training=False, filter_scale=filter_scale)
        self.cfg.INFER_SIZE = inner_size
        self.cfg.display()

        print('Creating Model Graph')
        self.model = model_config[self.cfg.model_type]
        self.net = self.model(cfg=self.cfg, mode='inference')

        print('Loading Network')
        self.net.create_session()
        self.net.restore(self.cfg.model_weight)

    def exp_filter(self, theta):
        self.last_theta = self.last_theta * (1.0 - self.alpha) + theta * self.alpha
        return self.last_theta

    @timer
    def process(self, image):
        if image.shape != self.cfg.INFER_SIZE:
            print('[WARNING]: Image not in the correct size. Resizing...')
            image = cv2.resize(image, (self.cfg.INFER_SIZE[1], self.cfg.INFER_SIZE[0]))

        self.config.reload()
        c = self.run_network(image)

        # calculate theta
        theta = np.arctan2((int(c[0]) - image.shape[0]), (int(c[1]) - int(image.shape[1] / 2)))
        if self.config.output_type == "one_to_one":
            theta = (np.arctan2((int(c[0]) - image.shape[0]), (int(c[1]) - int(image.shape[1] / 2))) / np.pi)
        elif self.config.output_type == "degree":
            theta = (theta / np.pi) * 180.0
            theta = (theta + 90.0) / 90.0
        elif self.config.output_type == "radian":
            pass

        # apply gain and offset
        theta = theta * self.config.gain + self.config.offset

        if self.config.smooth_result:
            return self.exp_filter(theta)
        return theta

    def run_network(self, image):
        results1 = self.net.predict(image)
        a = results1[0].copy() / 255
        a[a[:, :, 0] > 0.502] = 0
        a[a[:, :, 0] < 0.501] = 0
        a[a[:, :, 1] > 0.26] = 0
        a[a[:, :, 1] < 0.24] = 0
        a[a[:, :, 2] > 0.502] = 0
        a[a[:, :, 2] < 0.501] = 0
        c = ndimage.measurements.center_of_mass(a)
        return c


import re
import sys
import string
import os
from sklearn.cluster import KMeans
from sklearn.ensemble import RandomForestClassifier
#from numpy import array
import numpy as np
from argparse import Namespace

def find_training_genome(trainingFlag,INSTALLATION_DIR):
    try:
        f = open(INSTALLATION_DIR+"data/trainingGenome_list.txt","r")
    except:
        print('cannot open '+INSTALLATION_DIR+'data/trainingGenome_list.txt')
        return ''

    for line in f:
        temp = re.split('\t',line.strip())
        if int(temp[0]) == trainingFlag:
            f.close()
            return temp[1].strip()
    return ''

def call_randomforest(**kwargs):
    output_dir = kwargs['output_dir']
    trainingFile = kwargs['training_set']
    bin_path = os.path.join(os.path.dirname(os.path.dirname(os.path.relpath(__file__))),'bin')
    infile = os.path.join(output_dir, "testSet.txt")
    outfile = os.path.join(output_dir, "classify.tsv")
    #trainingFile = find_training_genome(trainingFlag,INSTALLATION_DIR)
    #train_data = np.genfromtxt(fname=trainingFile, delimiter="\t", skip_header=1, filling_values=1)
    #test_data = np.genfromtxt(fname=infile, delimiter="\t", skip_header=1, filling_values=1)
    #clf = RandomForestClassifier()
    #clf.fit(train_data[:, :-1], train_data[:, -1].astype('int'))
    #print(clf.predict(test_data))
    #exit()
    cmd = "Rscript " + bin_path + "/randomForest.r " + trainingFile + " " + infile + " " + outfile
    os.system(cmd)

def my_sort(orf_list):
     n = len(orf_list)
     i = 1
     while( i <= n ):
          j = i + 1
          while( j < n ):
               flag = 0
               #direction for both
               if( orf_list[i]['start'] < orf_list[i]['stop'] ):
                    dir_i = 1
               else:
                    dir_i = -1
               if( orf_list[j]['start'] < orf_list[j]['stop'] ):
                    dir_j = 1
               else:
                    dir_j = -1

               #check whether swap need or not
               if dir_i == dir_j:
                    if orf_list[i]['start']>orf_list[j]['start']:
                         flag = 1
               else:
                    if dir_i == 1:
                         if orf_list[i]['start']>orf_list[j]['stop']:
                              flag = 1
                    else:
                         if orf_list[i]['stop']>orf_list[j]['start']:
                              flag = 1
               #swap
               if flag == 1:
                    temp = orf_list[i]
                    orf_list[i] = orf_list[j]
                    orf_list[j] = temp
               j = j+1
          i = i+1
     return orf_list

def find_mean(all_len):
     sum = 0.0
     for i in all_len:
          sum = sum + i
     return float(sum)/len(all_len)

def calc_pp(func):
    x = 0
    func = func.replace('-',' ')
    func = func.replace(',',' ')
    a = re.split(' ',func)
    if (
       'phage'         in a or
       'lysin'         in a or
       'endolysin'     in a or
       'holin'         in a or
       'capsid'        in a or
       'tail'          in a or
       'bacteriophage' in a or
       'prophage'      in a or
       'portal'        in a or
       'terminase'     in a or
       'tapemeasure'   in a or
       'baseplate'     in a or
       'virion'        in a or
       'antirepressor' in a or
       'excisionase'   in a or
       'mobile element protein' == func or
       re.search(r"\b%s\b" % "tape measure", func) or
       re.search(r"\b%s\b" % "Cro-like repressor", func) or
       re.search(r"\b%s\b" % "CI-like repressor", func) or
       re.search(r"\b%s\b" % "rIIA lysis", func) or
       re.search(r"\b%s\b" % "rI lysis", func) or
       re.search(r"\b%s\b" % "rIIB lysis", func) or
       re.search(r"\b%s\b" % "base plate", func)
       ):
           x = 1
    elif ('unknown' in func) or ('hypothetical' in func):
           x = 0.5
    else:
           x = 0
    if 'recombinase' in func or 'integrase' in func:
           x = 1.5
    if ('phage' in func) and ('shock' in func):
           x = 0
    if "dna binding domain" in func:
           x = 0
    return x

def calc_function_3files(organism):
    my_func = {}
    x = 0 #no need it for computation.. just a flag in "except"
    try:
        f_fun = open(organism+'/proposed_non_ff_functions','r')
        for line in f_fun:
            temp = re.split('\t',line.strip())
            if len(temp)>=2:
                my_func[temp[0]] = temp[1]
        f_fun.close()
    except:
        x = x + 1
    try:
        f_fun = open(organism+'/proposed_functions','r')
        for line in f_fun:
            temp = re.split('\t',line.strip())
            if len(temp)>=2:
                my_func[temp[0]] = temp[1]
        f_fun.close()
    except:
        x = x + 1
    try:
        f_fun = open(organism+'/assigned_functions','r')
        for line in f_fun:
            temp = re.split('\t',line.strip())
            if len(temp)>=2:
                my_func[temp[0]] = temp[1]
        f_fun.close()
    except:
        x = x + 1

    return my_func

def input_bactpp(**kwargs):
    print(kwargs)
    exit()

    #bact_file = organism+'/Features/peg/tbl'
    #try:
    #    fh = open(bact_file,'r')
    #except:
    #    print('cant open file- assigned functions/tbl file:',organism)
    #    return {}
    #my_func = calc_function_3files(organism)
    all_orf_list = {}
    for i in fh:
        temp = re.split('\t',i.strip())
        temp1 = re.split('_',temp[1])
        if ',' in temp[1]:
            ttemp = re.split(',',temp[1])
            temp[1] = ttemp[len(ttemp)-1]
        temp1 = re.split('_',temp[1])
        contig = temp[1][:temp[1][:temp[1].rfind('_')].rfind('_')]
        start = int(temp1[len(temp1)-2])
        stop = int(temp1[len(temp1)-1])
        #save info for sorting orf
        if contig in all_orf_list:
            x = len(all_orf_list[contig]) + 1
        else:
            x = 1
            all_orf_list[contig]={}
        all_orf_list[contig][x]={}
        all_orf_list[contig][x]['fig'] = temp[0]
        all_orf_list[contig][x]['contig'] = str(contig)
        all_orf_list[contig][x]['start'] = start
        all_orf_list[contig][x]['stop'] = stop
        if temp[0] in my_func:
            all_orf_list[contig][x]['function'] = my_func[temp[0]]
            all_orf_list[contig][x]['pp'] = calc_pp(my_func[temp[0]].lower(),INSTALLATION_DIR)
        else:
            all_orf_list[contig][x]['function'] = "-"
            all_orf_list[contig][x]['pp'] = 0.5
    fh.close()
    all = {}
    index = 1
    for mycontig in all_orf_list:
        orf_list = my_sort(all_orf_list[mycontig])
        i = 1
        while i <= len(orf_list):
            all[index] = {}
            all[index]['fig'] = orf_list[i]['fig']
            all[index]['function'] = orf_list[i]['function']
            all[index]['contig'] = orf_list[i]['contig']
            all[index]['start'] = orf_list[i]['start']
            all[index]['stop'] = orf_list[i]['stop']
            all[index]['rank'] = 0.0
            all[index]['status'] = 0
            all[index]['pp'] = orf_list[i]['pp']
            i = i+1
            index = index+1
    return all

def make_initial_tbl(**kwargs): #organismPath, output_dir, window, INSTALLATION_DIR):
    self = Namespace(**kwargs)
    x = []
    for entry in self.record:
        for feature in entry.get_features('CDS'):
            all = {}
            all['fig'] = feature.id
            all['function'] = feature.function
            all['contig'] = entry.id
            all['start'] = feature.start
            all['stop'] = feature.stop
            all['rank'] = 0.0
            all['status'] = 0
            all['pp'] = calc_pp(feature.function)
            x.append(all)
    try:
        infile = open(os.path.join(self.output_dir, 'classify.tsv'), 'r')
        outfile = open(os.path.join(self.output_dir, 'initial_tbl.tsv'), 'w')
    except:
        sys.exit('ERROR: Cannot open classify.tsv')
    #x = input_bactpp(**kwargs)
    j = 1
    ranks = [[] for n in range(len(x))]
    for line in infile:
        val = float(line.strip())
        for k in range(j-int(self.window_size/2), j+int(self.window_size/2)):
            if k <= 0 or k >= len(x) or j >= len(x) or x[k]['contig'] != x[j]['contig']:
                continue
            ranks[k].append(val)
        j += 1
    infile.close()
    #calculate threshold
    y = []
    j = 1
    while j < len(x):
        x[j]['rank'] = sum(ranks[j]) / len(ranks[j])
        x[j]['extra'] =  ranks[j]
        y.append(x[j]['rank'])
        #y.append([x[j]['rank']])
        j = j+1
    #threshold = max(y)/2
    y2 = np.array(y).reshape(-1, 1)
    km = KMeans(n_clusters = 2)
    km.fit(y2)
    centers = km.cluster_centers_
    threshold = max(centers[0][0], centers[1][0])
    """
    Note added by Rob:
    At this point we have the classifications for each ORF and we want to take a sliding window and decide where the phage should
    start. We have two calculations for a threshold for the rank: either the kmeans centers and finding things above the larger center
    or just a plain threshold.

    """
    j = 1
    outfile.write('fig_no\tfunction\tcontig\tstart\tstop\tposition\trank\tmy_status\tpp\tFinal_status\tstart of attL\tend of attL\tstart of attR\tend of attR\tsequence of attL\tsequence of attR\tReason for att site\n')
    while j < len(x):
        if x[j]['rank'] > threshold:
            x[j]['status'] = 1
        outfile.write(str(x[j]['fig']))
        outfile.write('\t')
        outfile.write(str(x[j]['function']))
        outfile.write('\t')
        outfile.write(str(x[j]['contig']))
        outfile.write('\t')
        outfile.write(str(x[j]['start']))
        outfile.write('\t')
        outfile.write(str(x[j]['stop']))
        outfile.write('\t')
        outfile.write(str(j))
        outfile.write('\t')
        outfile.write(str(x[j]['rank']))
        outfile.write('\t')
        outfile.write(str(x[j]['status']))
        outfile.write('\t')
        outfile.write(str(x[j]['pp']))
        outfile.write('\n')
        j = j+1
    outfile.close()


import skimage.io as io
import skimage.transform as skt
import numpy as np
from PIL import Image, ImageOps
from src.models.class_patcher import patcher
from src.utils.imgproc import *

class patcher(patcher):
    def __init__(self, body='./body/body_salty.png', **options):
        super().__init__('ソルティ', body=body, pantie_position=[1044, 736], **options)

    def convert(self, image):
        # image = Image.open('./dream/0101.png')
        pantie = np.array(image)
        patch = np.copy(pantie[-140:-5, 546:, :])
        pantie[-140:, 546:, :] = 0
        patch = skt.resize(patch[::-1, ::-1, :], (160, 70), anti_aliasing=True, mode='reflect')
        [pr, pc, d] = patch.shape
        pantie[115:115 + pr, :pc, :] = np.uint8(patch * 255)

        arrx = np.zeros(25)
        arry = np.sin(np.linspace(0, np.pi, 25)) * 120 - 20
        arry[3:11] -= np.sin(np.linspace(0, np.pi, 8)) * 100
        pantie = affine_transform_by_arr(pantie, arrx, arry, smoothy=True, mvy=2)
        pantie = np.uint8(pantie[:280, 18:-60] * 255)
        return Image.fromarray(pantie)


"""The module contains everything to handle cross section interfaces."""

from abc import ABCMeta, abstractmethod
from os.path import isfile, join

import numpy as np

import prince_cr.decays as decs
from prince_cr.data import spec_data
from prince_cr.util import (
    get_2Dinterp_object, get_interp_object, info, bin_widths, get_AZN)
import prince_cr.config as config

class CrossSectionBase(object, metaclass=ABCMeta):
    """Base class for cross section interfaces to tabulated models.

    The class is abstract and it is not inteded to be instantiated.

    Child Classes either define the tables:
        self._egrid_tab
        self._nonel_tab
        self._incl_tab
        self._incl_diff

    Or directly reimplememnt the functions
        nonel(self, mother, daughter)
        incl(self, mother, daughter)
        incl_diff(self, mother, daughter)

    The flag self.supports_redistributions = True/False should be set
    To tell the class to include/ignore incl_diff
    """
    def __init__(self):
        # Tuple, defining min and max energy cuts on the grid
        self._range = None
        # Energy grid, as defined in files
        self._egrid_tab = None
        # Dictionary of nonel. cross sections on egrid, indexed by (mother)
        self._nonel_tab = {}
        # Dictionary of incl. cross sections on egrid, indexed by (mother, daughter)
        self._incl_tab = {}
        # Dictionary of incl. diff. cross sections on egrid, indexed by (mother, daughter)
        self._incl_diff_tab = {}
        # List of available mothers for nonel cross sections
        self.nonel_idcs = []
        # List of available (mothers,daughter) reactions in incl. cross sections
        self.incl_idcs = []
        # List of available (mothers,daughter) reactions in incl. diff. cross sections
        self.incl_diff_idcs = []
        # Common grid in x (the redistribution variable)
        self.xbins = None

        # Flag, which tells if the model supports secondary redistributions
        if not hasattr(self, 'supports_redistributions'):
            self.supports_redistributions = None  # JH: to differ from explicitly set False
        # List of all known particles (after optimization)
        self.known_species = []
        # List of all boost conserving inclusive channels (after optimization)
        self.known_bc_channels = []
        # List of all differential inclusive channels (after optimization)
        self.known_diff_channels = []
        # Dictionary of (mother, daughter) reactions for each mother
        self.reactions = {}

        # Class name of the model
        self.mname = self.__class__.__name__

    def set_range(self, e_min=None, e_max=None):
        """Set energy range within which to return tabulated data.

        Args:
            e_min (float): minimal energy in GeV
            e_max (float): maximal energy in GeV
        """
        if e_min is None:
            e_min = np.min(self._egrid_tab)
        if e_max is None:
            e_max = np.max(self._egrid_tab)

        info(5, "Setting range to {0:3.2e} - {1:3.2e}".format(e_min, e_max))
        self._range = np.where((self._egrid_tab >= e_min)
                               & (self._egrid_tab <= e_max))[0]
        info(
            2, "Range set to {0:3.2e} - {1:3.2e}".format(
                np.min(self._egrid_tab[self._range]),
                np.max(self._egrid_tab[self._range])))

    @property
    def egrid(self):
        """Returns energy grid of the tabulated data in selected range.

        Returns:
            (numpy.array): Energy grid in GeV
        """

        return self._egrid_tab[self._range]

    @property
    def xcenters(self):
        """Returns centers of the grid in x.

        Returns:
            (numpy.array): x grid
        """

        return 0.5 * (self.xbins[1:] + self.xbins[:-1])

    @property
    def xwidths(self):
        """Returns bin widths of the grid in x.

        Returns:
            (numpy.array): x widths
        """

        return self.xbins[1:] - self.xbins[:-1]

    @property
    def resp(self):
        """Return ResponseFunction corresponding to this cross section
        Will only create the Response function once.
        """
        if not hasattr(self, '_resp'):
            info(2, 'First Call, creating instance of ResponseFunction now')
            from .response import ResponseFunction
            self._resp = ResponseFunction(self)
        return self._resp

    def is_differential(self, mother, daughter):
        """Returns true if the model supports redistributions and requested
        mother/daughter combination should return non-zero redistribution matrices.

        Args:
            mother (bool): Neucosma ID of mother particle
            daughter (bool): Neucosma ID of daughter particle

        Returns:
            (bool): ``True`` if the model has this particular redistribution function
        """
        # info(10, mother, daughter, " asking for redist")
        # if not self.supports_redistributions:
        #     info(10, mother, daughter, " model doesn't support redist")
        #     return False
        if (daughter <= config.redist_threshold_ID
                or (mother, daughter) in self.incl_diff_idcs):
            info(60, 'Daughter requires redistribution.', mother, daughter)
            return True
        info(60, 'Daughter conserves boost.', mother, daughter)
        return False

    def _update_indices(self):
        """Updates the list of indices according to entries in the
        _tab variables"""

        self.nonel_idcs = sorted(self._nonel_tab.keys())
        self.incl_idcs = sorted(self._incl_tab.keys())
        self.incl_diff_idcs = sorted(self._incl_diff_tab.keys())

    def generate_incl_channels(self, mo_indices):
        """Generates indices for all allowed channels given mo_indices
            Note: By default this returns an empty list,
                  meant to be overwritten in cases where
                  the child class needs to dynamically generate indices

        Args:
            mo_indices (list of ints): list of indices for mother nuclei

        Returns:
           Returns:
            list of tuples: list of allowed channels given as (mo_idx, da_idx)
        """
        incl_channels = []

        return incl_channels

    def _optimize_and_generate_index(self):
        """Construct a list of mothers and (mother, daughter) indices.

        Args:
            just_reactions (bool): If True then fill just the reactions index.
        """

        # Integrate out short lived processes and leave only stable particles
        # in the databases
        self._reduce_channels()

        # Go through all three cross section categories
        # index contents in the ..known..variable
        self.reactions = {}

        self._update_indices()

        for mo, da in self.incl_idcs:
            if da >= 100 and get_AZN(da)[0] > get_AZN(mo)[0]:
                raise Exception(
                    'Daughter {0} heavier than mother {1}. Physics??'.format(
                        da, mo))

            if mo not in self.reactions:
                self.reactions[mo] = []
                self.known_species.append(mo)

            if (mo, da) not in self.reactions[mo]:
                # Make sure it's a unique list
                self.reactions[mo].append((mo, da))
            if self.is_differential(mo, da):
                # Move the distributions which are expected to be differential
                # to _incl_diff_tab
                self._incl_diff_tab[(mo, da)] = self._arange_on_xgrid(
                    self._incl_tab.pop((mo, da)))
                info(10, "Channel {0} -> {1} forced to be differential.")
            else:
                self.known_bc_channels.append((mo, da))
                self.known_species.append(da)

        for mo, da in list(self._incl_diff_tab.keys()):
            if da >= 100 and get_AZN(da)[0] > get_AZN(mo)[0]:
                raise Exception(
                    'Daughter {0} heavier than mother {1}. Physics??'.format(
                        da, mo))

            if mo not in self.reactions:
                self.reactions[mo] = []
                self.known_species.append(mo)

            if (mo, da) not in self.reactions[mo]:
                # Make sure it's a unique list to avoid unnecessary loops
                self.reactions[mo].append((mo, da))
                self.known_diff_channels.append((mo, da))
                self.known_species.append(da)

        # Remove duplicates
        self.known_species = sorted(list(set(self.known_species)))
        self.known_bc_channels = sorted(list(set(self.known_bc_channels)))
        self.known_diff_channels = sorted(list(set(self.known_diff_channels)))

        for sp in self.known_species:
            if sp >= 100 and (sp, sp) not in self.known_diff_channels:
                self.known_bc_channels.append((mo, mo))
            if (mo, mo) not in self.reactions[mo]:
                self.reactions[mo].append((mo, mo))

        # Make sure the indices are up to date
        self._update_indices()

        # Count numbers of channels for statistics
        # Count number of incl channels for activated nuclear species
        # n_incl = np.sum([
        #     len(self.reactions[mother])
        #     for mother in self.spec_man.known_species if mother >= 100
        # ])

    def _reduce_channels(self):
        """Follows decay chains until all inclusive reactions point to
        stable final state particles.

        The "tau_dec_threshold" parameter in the config controls the
        definition of stable. Unstable nuclei for which no decay channels
        are known, will be forced to beta-decay until they reach a stable
        element.
        """
        from prince_cr.util import AdditiveDictionary
        # TODO: check routine, how to avoid empty channels and
        # mothers with zero nonel cross sections

        # The new dictionary that will replace _incl_tab
        new_incl_tab = AdditiveDictionary()
        new_dec_diff_tab = AdditiveDictionary()

        threshold = config.tau_dec_threshold

        # How to indent debug printout for recursion
        dbg_indent = lambda lev: 4 * lev * "-" + ">" if lev else ""

        info(2, "Integrating out species with lifetime smaller than",
             threshold)
        info(3, (
            "Before optimization, the number of known primaries is {0} with " +
            "in total {1} inclusive channels").format(len(self._nonel_tab),
                                                      len(self._incl_tab)))

        if self.xbins is None:
            info(
                4,
                'Model does not provide a native xbins. Assuming JH special sophia',
                'binning.')
            from .photo_meson import SophiaSuperposition
            self.xbins = SophiaSuperposition().xbins

        bc = self.xcenters
        bw = bin_widths(self.xbins)
        # The x_mu/x_pi grid
        # dec_grid = np.fromfunction(
        #     lambda j, i: 10**(np.log10(bc[1] / bc[0]) * (j - i)), (len(bc),
        #                                                            len(bc)))

        # dec_grid = np.outer(bc, 1 / bc)

        dec_bins = np.outer(self.xbins, 1 / bc)
        dec_bins_lower = dec_bins[:-1]
        dec_bins_upper = dec_bins[1:]

        # dec_grid[dec_grid > 1.] *= 0.
        # The differential element dx_mu/x_pi
        int_scale = np.tile(bw / bc, (len(bc), 1))

        from functools import lru_cache
        @lru_cache(maxsize=512, typed=False)
        def decay_cached(mother,daughter):
            dec_dist = int_scale * decs.get_decay_matrix_bin_average(
                mother, daughter, dec_bins_lower, dec_bins_upper)

            return dec_dist

        def convolve_with_decay_distribution(diff_dist, mother, daughter,
                                             branching_ratio):
            r"""Computes the prompt decay xdist by convolving the x distribution
            of the unstable particle with the decay product distribution.

            :math:`\frac{{\rm d}N^{A\gamma \to \mu}}{{\rm d}x_j} =
            \sum_{i=0}^{N_x}~\Delta x_i
            \frac{{\rm d}N^{A\gamma \to \pi}}{{\rm d} x_i}~
            \frac{{\rm d}N^{\pi \to \mu}}{{\rm d} x_j}`
            """
            # dec_dist = int_scale * decs.get_decay_matrix(
            #     mother, daughter, dec_grid)
            dec_dist = decay_cached(mother,daughter)

            info(20, 'convolving with decay dist', mother, daughter)
            # Handle the case where table entry is (energy_grid, matrix)
            if not isinstance(diff_dist, tuple):
                return branching_ratio * dec_dist.dot(diff_dist)
            else:
                return diff_dist[0], branching_ratio * dec_dist.dot(
                    diff_dist[1])

        def follow_chain(first_mo, da, csection, reclev):
            """Recursive function to follow decay chains until all
            final state particles are stable.

            The result is saved in two dictionaries; one for the boost
            conserving inclusive channels and the other one collects
            channels with meson or lepton decay products, which will
            need special care due to energy redistributions of these
            secondaries.
            """

            info(10, dbg_indent(reclev), 'Entering with', first_mo, da)

            if da not in spec_data:
                info(
                    3, dbg_indent(reclev),
                    'daughter {0} unknown, forcing beta decay. Not Implemented yet!!'
                    .format(da))
                return

            # Daughter is stable. Add it to the new dictionary and terminate
            # recursion
            if spec_data[da]["lifetime"] >= threshold:
                if self.is_differential(None, da):
                    # If the daughter is a meson or lepton, use the dictionary for
                    # differential channels
                    info(
                        20, dbg_indent(reclev),
                        'daughter {0} stable and differential. Adding to ({1}, {2})'
                        .format(da, first_mo, da))
                    new_dec_diff_tab[(first_mo, da)] = csection
                else:
                    info(
                        20, dbg_indent(reclev),
                        'daughter {0} stable. Adding to ({1}, {2})'.format(
                            da, first_mo, da))
                    new_incl_tab[(first_mo, da)] = csection
                return

            # ..otherwise follow decay products of this daughter, tracking the
            # original mother particle (first_mo). The cross section (csection) is
            # reduced by the branching ratio (br) of this particular channel
            for br, daughters in spec_data[da]["branchings"]:
                info(10, dbg_indent(reclev),
                     ("{3} -> {0:4d} -> {2:4.2f}: {1}").format(
                         da, ", ".join(map(str, daughters)), br, first_mo))

                for chained_daughter in daughters:
                    # Follow each secondary and increment the recursion level by one
                    if self.is_differential(None, chained_daughter):
                        info(10, 'daughter', chained_daughter, 'of', da,
                             'is differential')
                        follow_chain(
                            first_mo, chained_daughter,
                            convolve_with_decay_distribution(
                                self._arange_on_xgrid(csection), da,
                                chained_daughter, br), reclev + 1)
                    else:
                        follow_chain(first_mo, chained_daughter, br * csection,
                                     reclev + 1)

        # Remove all unstable particles from the dictionaries
        for mother in sorted(self._nonel_tab.keys()):
            if mother not in spec_data or spec_data[mother][
                    "lifetime"] < threshold:
                info(
                    20,
                    "Primary species {0} does not fulfill stability criteria.".
                    format(mother))
                _ = self._nonel_tab.pop(mother)
        # Only stable (interacting) mother particles are left
        self._update_indices()

        for (mother, daughter) in self.incl_idcs:

            if mother not in self.nonel_idcs:
                info(
                    30, "Removing {0}/{1} from incl, since mother not stable ".
                    format(mother, daughter))
                _ = self._incl_tab.pop((mother, daughter))

            elif self.is_differential(mother, daughter):
                # Move the distributions which are expected to be differential
                # to _incl_diff_tab
                self._incl_diff_tab[(mother,
                                     daughter)] = self._arange_on_xgrid(
                                         self._incl_tab.pop(
                                             (mother, daughter)))

        self._update_indices()

        for (mother, daughter) in self.incl_diff_idcs:

            if mother not in self.nonel_idcs:
                info(
                    30,
                    "Removing {0}/{1} from diff incl, since mother not stable "
                    .format(mother, daughter))
                _ = self._incl_diff_tab.pop((mother, daughter))

        self._update_indices()

        # Launch the reduction for each inclusive channel
        for (mo, da), value in list(self._incl_tab.items()):
            #print mo, da, value
            #print '---'*30
            follow_chain(mo, da, value, 0)

        for (mo, da), value in list(self._incl_diff_tab.items()):
            #print mo, da, value
            #print '---'*30
            follow_chain(mo, da, value, 0)

        # Overwrite the old incl dictionary
        self._incl_tab = dict(new_incl_tab)
        # Overwrite the old incl_diff dictionary
        self._incl_diff_tab = dict(new_dec_diff_tab)
        # Reduce also the incl_diff_tab by removing the unknown mothers. At this stage
        # of the code, the particles with redistributions are
        info(
            3,
            ("After optimization, the number of known primaries is {0} with " +
             "in total {1} inclusive channels").format(
                 len(self._nonel_tab),
                 len(self._incl_tab) + len(self._incl_diff_tab)))
        info(2, f'Cache used for decays, {decay_cached.cache_info()}') # pylint:disable=no-value-for-parameter

    def nonel_scale(self, mother, scale='A'):
        """Returns the nonel cross section scaled by `scale`.

        Convenience funtion for plotting, where it is important to
        compare the cross section per nucleon.

        Args:
            mother (int): Mother nucleus(on)
            scale (float): If `A` then nonel/A is returned, otherwise
                           scale can be any float.

        Returns:
            (numpy.array, numpy.array): Tuple of Energy grid in GeV,
                                        scale * inclusive cross section
                                        in :math:`cm^{-2}`
        """

        egr, csection = self.nonel(mother)

        if scale == 'A':
            scale = 1. / get_AZN(mother)[0]

        return egr, scale * csection

    def incl_scale(self, mother, daughter, scale='A'):
        """Same as :func:`~cross_sections.CrossSectionBase.nonel_scale`,
        just for inclusive cross sections.
        """

        egr, csection = self.incl(mother, daughter)

        if scale == 'A':
            scale = 1. / get_AZN(mother)[0]

        return egr, scale * csection

    def nonel(self, mother):
        """Returns non-elastic cross section.

        Absorption cross section of `mother`, which is
        the total minus elastic, or in other words, the inelastic
        cross section.

        Args:
            mother (int): Mother nucleus(on)

        Returns:
            (numpy.array, numpy.array): Tuple of Energy grid in GeV, inclusive cross
                                        section in :math:`cm^{-2}`
        """

        if mother not in self._nonel_tab:
            raise Exception('Mother {0} unknown.'.format(mother))

        if isinstance(self._nonel_tab[mother], tuple):
            return self._nonel_tab[mother]
        else:
            return self.egrid, self._nonel_tab[mother][self._range]

    def incl(self, mother, daughter):
        """Returns inclusive cross section.

        Inclusive cross section for daughter in photo-nuclear
        interactions of `mother`.

        Args:
            mother (int): Mother nucleus(on)
            daughter (int): Daughter nucleus(on)

        Returns:
            (numpy.array): Inclusive cross section in :math:`cm^{-2}`
                           on self._egrid_tab
        """

        from scipy.integrate import trapz

        if (mother, daughter) in self._incl_diff_tab:
            # Return the integral of the differential for the inclusive
            egr_incl, cs_diff = self.incl_diff(mother, daughter)
            # diff_mat = diff_mat.transpose()
            cs_incl = trapz(cs_diff,
                            x=self.xcenters,
                            dx=bin_widths(self.xbins),
                            axis=0)

            if isinstance(self._incl_diff_tab[(mother, daughter)], tuple):
                return egr_incl, cs_incl

            return self.egrid, cs_incl[self._range]

        elif (mother, daughter) not in self._incl_tab:
            raise Exception(
                self.__class__.__name__ + '::'
                '({0},{1}) combination not in inclusive cross sections'.format(
                    mother, daughter))

        # If _nonel_tab contains tuples of (egrid, cs) return tuple
        # otherwise return (egrid, cs) in range defined by self.range

        if isinstance(self._incl_tab[(mother, daughter)], tuple):
            return self._incl_tab[(mother, daughter)]
        return self.egrid, self._incl_tab[(mother, daughter)][self._range]

    def incl_diff(self, mother, daughter):
        """Returns inclusive cross section.

        Inclusive differential cross section for daughter in photo-nuclear
        interactions of `mother`. Only defined, if the daughter is distributed
        in :math:`x = E_{da} / E_{mo}`

        Args:
            mother (int): Mother nucleus(on)
            daughter (int): Daughter nucleus(on)

        Returns:
            (numpy.array): Inclusive cross section in :math:`cm^{-2}`
                           on self._egrid_tab
        """

        if (mother, daughter) not in self._incl_diff_tab:
            raise Exception(
                self.__class__.__name__ +
                '({0},{1}) combination not in inclusive differential cross sections'
                .format(mother, daughter))

        # If _nonel_tab contains tuples of (egrid, cs) return tuple
        # otherwise return (egrid, cs) in range defined by self.range

        if isinstance(self._incl_diff_tab[(mother, daughter)], tuple):
            return self._incl_diff_tab[(mother, daughter)]
        return self.egrid, self._incl_diff_tab[(mother,
                                                daughter)][:, self._range]

    def _arange_on_xgrid(self, incl_cs):
        """Returns the inclusive cross section on an xgrid at x=1."""

        egr, cs = None, None

        if isinstance(incl_cs, tuple):
            egr, cs = incl_cs
        else:
            cs = incl_cs

        nxbins = len(self.xbins) - 1
        if len(cs.shape) > 1 and cs.shape[0] != nxbins:
            raise Exception(
                'One dimensional cross section expected, instead got',
                cs.shape, '\n', cs)
        elif len(cs.shape) == 2 and cs.shape[0] == nxbins:
            info(20, 'Supplied 2D distribution seems to be distributed in x.')
            if isinstance(incl_cs, tuple):
                return egr, cs
            return cs

        csec = np.zeros((nxbins, cs.shape[0]))
        # NOTE: The factor 2 in the following line is a workarround to account for the latter linear interpolation
        #       This is needed because linear spline integral will result in a trapz,
        #       which has half the area of the actual first bin
        corr_factor = 2 * self.xwidths[-1] / (self.xcenters[-1] -
                                              self.xcenters[-2])
        csec[-1, :] = cs / self.xwidths[-1] * corr_factor
        info(
            4,
            'Warning! Workaround to account for linear interpolation in x, factor 2 added!'
        )
        if isinstance(incl_cs, tuple):
            return egr, csec
        return csec

    def multiplicities(self, mother, daughter):
        '''Return the multiplicities from either the inclusive channels, or the
        differential ones integrated by x, as a function of Energy.
        '''
        egrid_incl, cs_incl = self.incl(mother, daughter)
        egrid_nonel, cs_nonel = self.nonel(mother)

        if egrid_incl.shape != egrid_nonel.shape:
            raise Exception('Problem with different grid shapes')

        multiplicities = cs_incl / np.where(cs_nonel == 0, np.inf, cs_nonel)

        return egrid_nonel, multiplicities

if __name__ == "__main__":
    pass