#! /usr/bin/env python

import os
import logging
import re
import math
import pathlib

from datetime import date
from datetime import datetime
from collections import namedtuple

import scipy.io as spio
import numpy as np
import warnings

import datajoint as dj

from pipeline import lab
from pipeline import experiment
from .. import get_schema_name

schema = dj.schema(get_schema_name('ingest_behavior'))

warnings.simplefilter(action='ignore', category=FutureWarning)

log = logging.getLogger(__name__)

# ================ PHOTOSTIM PROTOCOL ===============
photostim_duration = 0.5  # (s)
skull_ref = 'Bregma'
photostims = {4: {'photo_stim': 4, 'photostim_device': 'OBIS470', 'duration': photostim_duration,
                  'locations': [{'skull_reference': skull_ref, 'brain_area': 'ALM',
                                 'ap_location': 2500, 'ml_location': -1500, 'depth': 0,
                                 'theta': 15, 'phi': 15}]},
              5: {'photo_stim': 5, 'photostim_device': 'OBIS470', 'duration': photostim_duration,
                  'locations': [{'skull_reference': skull_ref, 'brain_area': 'ALM',
                                 'ap_location': 2500, 'ml_location': 1500, 'depth': 0,
                                 'theta': 15, 'phi': 15}]},
              6: {'photo_stim': 6, 'photostim_device': 'OBIS470', 'duration': photostim_duration,
                  'locations': [{'skull_reference': skull_ref, 'brain_area': 'ALM',
                                 'ap_location': 2500, 'ml_location': -1500, 'depth': 0,
                                 'theta': 15, 'phi': 15},
                                {'skull_reference': skull_ref, 'brain_area': 'ALM',
                                 'ap_location': 2500, 'ml_location': 1500, 'depth': 0,
                                 'theta': 15, 'phi': 15}
                                ]}}

def get_behavior_paths():
    '''
    retrieve behavior rig paths from dj.config
    config should be in dj.config of the format:

      dj.config = {
        ...,
        'custom': {
          'behavior_data_paths':
            [
                ["RRig", "/path/string", 0],
                ["RRig2", "/path2/string2", 1]
            ],
        }
        ...
      }

    where 'behavior_data_paths' is a list of multiple possible path for behavior data, each in format:
    [rig name, rig full path, search order]
    '''

    paths = dj.config.get('custom', {}).get('behavior_data_paths', None)
    if paths is None:
        raise ValueError("Missing 'behavior_data_paths' in dj.config['custom']")

    return sorted(paths, key=lambda x: x[-1])

@schema
class BehaviorIngest(dj.Imported):
    definition = """
    -> experiment.Session
    """

    class BehaviorFile(dj.Part):
        ''' files in rig-specific storage '''
        definition = """
        -> BehaviorIngest
        behavior_file:              varchar(255)          # behavior file name
        """

    class CorrectedTrialEvents(dj.Part):
        ''' TrialEvents containing auto-corrected data '''
        definition = """
        -> BehaviorIngest
        -> experiment.TrialEvent
        """

    @staticmethod
    def get_session_user():
        '''
        Determine desired 'session user' for a session.

        - 1st, try dj.config['custom']['session.user']
        - 2nd, try dj.config['database.user']
        - else, use 'unknown'

        TODO: multi-user / bulk ingest support
        '''
        session_user = dj.config.get('custom', {}).get('session.user', None)

        session_user = (dj.config.get('database.user')
                        if not session_user else session_user)

        if len(lab.Person() & {'username': session_user}):
            return session_user
        else:
            return 'unknown'

    @property
    def key_source(self):

        # 2 letters, anything, _, anything, 8 digits, _, 6 digits, .mat
        # where:
        # (2 letters, anything): water restriction
        # (anything): task name
        # (8 digits): date YYYYMMDD
        # (6 digits): time HHMMSS

        rexp = '^[a-zA-Z]{2}.*_.*_[0-9]{8}_[0-9]{6}.mat$'

        # water_restriction_number -> subject
        h2os = {k: v for k, v in zip(*lab.WaterRestriction().fetch(
            'water_restriction_number', 'subject_id'))}

        def buildrec(rig, rigpath, root, f):

            if not re.match(rexp, f):
                log.debug("{f} skipped - didn't match rexp".format(f=f))
                return

            log.debug('found file {f}'.format(f=f))

            fullpath = pathlib.Path(root, f)
            subpath = fullpath.relative_to(rigpath)

            fsplit = subpath.stem.split('_')
            h2o = fsplit[0]
            ymd = fsplit[-2:-1][0]

            if h2o not in h2os:
                log.warning('{f} skipped - no animal for {h2o}'.format(
                    f=f, h2o=h2o))
                return

            animal = h2os[h2o]

            log.debug('animal is {animal}'.format(animal=animal))

            return {
                'subject_id': animal,
                'session_date': date(
                    int(ymd[0:4]), int(ymd[4:6]), int(ymd[6:8])),
                'rig': rig,
                'rig_data_path': rigpath.as_posix(),
                'subpath': subpath.as_posix()
            }

        recs = []
        found = set()
        known = set(BehaviorIngest.BehaviorFile().fetch('behavior_file'))
        rigs = get_behavior_paths()

        for (rig, rigpath, _) in rigs:
            rigpath = pathlib.Path(rigpath)

            log.info('RigDataFile.make(): traversing {}'.format(rigpath))
            for root, dirs, files in os.walk(rigpath):
                log.debug('RigDataFile.make(): entering {}'.format(root))
                for f in files:
                    log.debug('RigDataFile.make(): visiting {}'.format(f))
                    r = buildrec(rig, rigpath, root, f)
                    if not r:
                        continue
                    if f in set.union(known, found):
                        log.info('skipping already ingested file {}'.format(
                            r['subpath']))
                    else:
                        found.add(f)  # block duplicate path conf
                        recs.append(r)

        return recs

    def populate(self, *args, **kwargs):
        for k in self.key_source:
            with dj.conn().transaction:
                self.make(k)

    def make(self, key):
        log.info('BehaviorIngest.make(): key: {key}'.format(key=key))

        subject_id = key['subject_id']
        h2o = (lab.WaterRestriction() & {'subject_id': subject_id}).fetch1(
            'water_restriction_number')

        ymd = key['session_date']
        datestr = ymd.strftime('%Y%m%d')
        log.info('h2o: {h2o}, date: {d}'.format(h2o=h2o, d=datestr))

        # session record key
        skey = {}
        skey['subject_id'] = subject_id
        skey['session_date'] = ymd
        skey['username'] = self.get_session_user()
        skey['rig'] = key['rig']

        # File paths conform to the pattern:
        # dl7/TW_autoTrain/Session Data/dl7_TW_autoTrain_20180104_132813.mat
        # which is, more generally:
        # {h2o}/{training_protocol}/Session Data/{h2o}_{training protocol}_{YYYYMMDD}_{HHMMSS}.mat

        path = pathlib.Path(key['rig_data_path'], key['subpath'])

        if experiment.Session() & skey:
            log.info("note: session exists for {h2o} on {d}".format(
                h2o=h2o, d=ymd))

        trial = namedtuple(  # simple structure to track per-trial vars
            'trial', ('ttype', 'stim', 'free', 'settings', 'state_times',
                      'state_names', 'state_data', 'event_data',
                      'event_times', 'trial_start'))

        if os.stat(path).st_size/1024 < 1000:
            log.info('skipping file {} - too small'.format(path))
            return

        log.debug('loading file {}'.format(path))

        mat = spio.loadmat(path, squeeze_me=True)
        SessionData = mat['SessionData'].flatten()

        # parse session datetime
        session_datetime_str = str('').join((
            str(SessionData['Info'][0]['SessionDate']),
            ' ', str(SessionData['Info'][0]['SessionStartTime_UTC'])))

        session_datetime = datetime.strptime(
            session_datetime_str, '%d-%b-%Y %H:%M:%S')

        AllTrialTypes = SessionData['TrialTypes'][0]
        AllTrialSettings = SessionData['TrialSettings'][0]
        AllTrialStarts = SessionData['TrialStartTimestamp'][0]
        AllTrialStarts = AllTrialStarts - AllTrialStarts[0]  # real 1st trial

        RawData = SessionData['RawData'][0].flatten()
        AllStateNames = RawData['OriginalStateNamesByNumber'][0]
        AllStateData = RawData['OriginalStateData'][0]
        AllEventData = RawData['OriginalEventData'][0]
        AllStateTimestamps = RawData['OriginalStateTimestamps'][0]
        AllEventTimestamps = RawData['OriginalEventTimestamps'][0]

        # verify trial-related data arrays are all same length
        assert(all((x.shape[0] == AllStateTimestamps.shape[0] for x in
                    (AllTrialTypes, AllTrialSettings,
                     AllStateNames, AllStateData, AllEventData,
                     AllEventTimestamps, AllTrialStarts, AllTrialStarts))))

        # AllStimTrials optional special case
        if 'StimTrials' in SessionData.dtype.fields:
            log.debug('StimTrials detected in session - will include')
            AllStimTrials = SessionData['StimTrials'][0]
            assert(AllStimTrials.shape[0] == AllStateTimestamps.shape[0])
        else:
            log.debug('StimTrials not detected in session - will skip')
            AllStimTrials = np.array([
                None for _ in enumerate(range(AllStateTimestamps.shape[0]))])

        # AllFreeTrials optional special case
        if 'FreeTrials' in SessionData.dtype.fields:
            log.debug('FreeTrials detected in session - will include')
            AllFreeTrials = SessionData['FreeTrials'][0]
            assert(AllFreeTrials.shape[0] == AllStateTimestamps.shape[0])
        else:
            log.debug('FreeTrials not detected in session - synthesizing')
            AllFreeTrials = np.zeros(AllStateTimestamps.shape[0],
                                     dtype=np.uint8)

        trials = list(zip(AllTrialTypes, AllStimTrials, AllFreeTrials,
                          AllTrialSettings, AllStateTimestamps, AllStateNames,
                          AllStateData, AllEventData, AllEventTimestamps,
                          AllTrialStarts))

        if not trials:
            log.warning('skipping date {d}, no valid files'.format(d=date))
            return

        #
        # Trial data seems valid; synthesize session id & add session record
        # XXX: note - later breaks can result in Sessions without valid trials
        #

        assert skey['session_date'] == session_datetime.date()

        skey['session_date'] = session_datetime.date()
        skey['session_time'] = session_datetime.time()

        log.debug('synthesizing session ID')
        session = (dj.U().aggr(experiment.Session()
                               & {'subject_id': subject_id},
                               n='max(session)').fetch1('n') or 0) + 1

        log.info('generated session id: {session}'.format(session=session))
        skey['session'] = session
        key = dict(key, **skey)

        #
        # Actually load the per-trial data
        #
        log.info('BehaviorIngest.make(): trial parsing phase')

        # lists of various records for batch-insert
        rows = {k: list() for k in ('trial', 'behavior_trial', 'trial_note',
                                    'trial_event', 'corrected_trial_event',
                                    'action_event', 'photostim',
                                    'photostim_location', 'photostim_trial',
                                    'photostim_trial_event')}

        i = 0  # trial numbering starts at 1
        for t in trials:

            #
            # Misc
            #

            t = trial(*t)  # convert list of items to a 'trial' structure
            i += 1  # increment trial counter

            log.debug('BehaviorIngest.make(): parsing trial {i}'.format(i=i))

            # covert state data names into a lookup dictionary
            #
            # names (seem to be? are?):
            #
            # Trigtrialstart, PreSamplePeriod, SamplePeriod, DelayPeriod
            # EarlyLickDelay, EarlyLickSample, ResponseCue, GiveLeftDrop
            # GiveRightDrop, GiveLeftDropShort, GiveRightDropShort
            # AnswerPeriod, Reward, RewardConsumption, NoResponse
            # TimeOut, StopLicking, StopLickingReturn, TrialEnd
            #

            states = {k: (v+1) for v, k in enumerate(t.state_names)}
            required_states = ('PreSamplePeriod', 'SamplePeriod',
                               'DelayPeriod', 'ResponseCue', 'StopLicking',
                               'TrialEnd')

            missing = list(k for k in required_states if k not in states)

            if len(missing):
                log.warning('skipping trial {i}; missing {m}'
                            .format(i=i, m=missing))
                continue

            gui = t.settings['GUI'].flatten()

            # ProtocolType - only ingest protocol >= 3
            #
            # 1 Water-Valve-Calibration 2 Licking 3 Autoassist
            # 4 No autoassist 5 DelayEnforce 6 SampleEnforce 7 Fixed
            #

            if 'ProtocolType' not in gui.dtype.names:
                log.warning('skipping trial {i}; protocol undefined'
                            .format(i=i))
                continue

            protocol_type = gui['ProtocolType'][0]
            if gui['ProtocolType'][0] < 3:
                log.warning('skipping trial {i}; protocol {n} < 3'
                            .format(i=i, n=gui['ProtocolType'][0]))
                continue

            #
            # Top-level 'Trial' record
            #

            tkey = dict(skey)
            startindex = np.where(t.state_data == states['PreSamplePeriod'])[0]

            # should be only end of 1st StopLicking;
            # rest of data is irrelevant w/r/t separately ingested ephys
            endindex = np.where(t.state_data == states['StopLicking'])[0]

            log.debug('states\n' + str(states))
            log.debug('state_data\n' + str(t.state_data))
            log.debug('startindex\n' + str(startindex))
            log.debug('endindex\n' + str(endindex))

            if not(len(startindex) and len(endindex)):
                log.warning('skipping {}: start/end mismatch: {}/{}'.format(
                    i, str(startindex), str(endindex)))
                continue

            try:
                tkey['trial'] = i
                tkey['trial_uid'] = i
                tkey['start_time'] = t.trial_start
                tkey['stop_time'] = t.trial_start + t.state_times[endindex][0]
            except IndexError:
                log.warning('skipping {}: IndexError: {}/{} -> {}'.format(
                    i, str(startindex), str(endindex), str(t.state_times)))
                continue

            log.debug('tkey' + str(tkey))
            rows['trial'].append(tkey)

            #
            # Specific BehaviorTrial information for this trial
            #

            bkey = dict(tkey)
            bkey['task'] = 'audio delay'  # hard-coded here
            bkey['task_protocol'] = 1     # hard-coded here

            # determine trial instruction
            trial_instruction = 'left'    # hard-coded here

            if gui['Reversal'][0] == 1:
                if t.ttype == 1:
                    trial_instruction = 'left'
                elif t.ttype == 0:
                    trial_instruction = 'right'
            elif gui['Reversal'][0] == 2:
                if t.ttype == 1:
                    trial_instruction = 'right'
                elif t.ttype == 0:
                    trial_instruction = 'left'

            bkey['trial_instruction'] = trial_instruction

            # determine early lick
            early_lick = 'no early'

            if (protocol_type >= 5
                and 'EarlyLickDelay' in states
                    and np.any(t.state_data == states['EarlyLickDelay'])):
                    early_lick = 'early'
            if (protocol_type >= 5
                and ('EarlyLickSample' in states
                     and np.any(t.state_data == states['EarlyLickSample']))):
                    early_lick = 'early'

            bkey['early_lick'] = early_lick

            # determine outcome
            outcome = 'ignore'

            if ('Reward' in states
                    and np.any(t.state_data == states['Reward'])):
                outcome = 'hit'
            elif ('TimeOut' in states
                    and np.any(t.state_data == states['TimeOut'])):
                outcome = 'miss'
            elif ('NoResponse' in states
                    and np.any(t.state_data == states['NoResponse'])):
                outcome = 'ignore'

            bkey['outcome'] = outcome

            # Determine free/autowater (Autowater 1 == enabled, 2 == disabled)
            bkey['auto_water'] = True if gui['Autowater'][0] == 1 else False
            bkey['free_water'] = t.free

            rows['behavior_trial'].append(bkey)

            #
            # Add 'protocol' note
            #
            nkey = dict(tkey)
            nkey['trial_note_type'] = 'protocol #'
            nkey['trial_note'] = str(protocol_type)
            rows['trial_note'].append(nkey)

            #
            # Add 'autolearn' note
            #
            nkey = dict(tkey)
            nkey['trial_note_type'] = 'autolearn'
            nkey['trial_note'] = str(gui['Autolearn'][0])
            rows['trial_note'].append(nkey)

            #
            # Add 'bitcode' note
            #
            if 'randomID' in gui.dtype.names:
                nkey = dict(tkey)
                nkey['trial_note_type'] = 'bitcode'
                nkey['trial_note'] = str(gui['randomID'][0])
                rows['trial_note'].append(nkey)

            #
            # Add presample event
            #
            log.debug('BehaviorIngest.make(): presample')

            ekey = dict(tkey)
            sampleindex = np.where(t.state_data == states['SamplePeriod'])[0]

            ekey['trial_event_id'] = len(rows['trial_event'])
            ekey['trial_event_type'] = 'presample'
            ekey['trial_event_time'] = t.state_times[startindex][0]
            ekey['duration'] = (t.state_times[sampleindex[0]]
                                - t.state_times[startindex])[0]

            if math.isnan(ekey['duration']):
                log.debug('BehaviorIngest.make(): fixing presample duration')
                ekey['duration'] = 0.0  # FIXDUR: lookup from previous trial

            rows['trial_event'].append(ekey)

            #
            # Add other 'sample' events
            #

            log.debug('BehaviorIngest.make(): sample events')

            last_dur = None

            for s in sampleindex:  # in protocol > 6 ~-> n>1
                # todo: batch events
                ekey = dict(tkey)
                ekey['trial_event_id'] = len(rows['trial_event'])
                ekey['trial_event_type'] = 'sample'
                ekey['trial_event_time'] = t.state_times[s]
                ekey['duration'] = gui['SamplePeriod'][0]

                if math.isnan(ekey['duration']) and last_dur is None:
                    log.warning('... trial {} bad duration, no last_edur'
                                .format(i, last_dur))
                    ekey['duration'] = 0.0  # FIXDUR: cross-trial check
                    rows['corrected_trial_event'].append(ekey)

                elif math.isnan(ekey['duration']) and last_dur is not None:
                    log.warning('... trial {} duration using last_edur {}'
                                .format(i, last_dur))
                    ekey['duration'] = last_dur
                    rows['corrected_trial_event'].append(ekey)

                else:
                    last_dur = ekey['duration']  # only track 'good' values.

                rows['trial_event'].append(ekey)

            #
            # Add 'delay' events
            #

            log.debug('BehaviorIngest.make(): delay events')

            last_dur = None
            delayindex = np.where(t.state_data == states['DelayPeriod'])[0]

            for d in delayindex:  # protocol > 6 ~-> n>1
                ekey = dict(tkey)
                ekey['trial_event_id'] = len(rows['trial_event'])
                ekey['trial_event_type'] = 'delay'
                ekey['trial_event_time'] = t.state_times[d]
                ekey['duration'] = gui['DelayPeriod'][0]

                if math.isnan(ekey['duration']) and last_dur is None:
                    log.warning('... {} bad duration, no last_edur'
                                .format(i, last_dur))
                    ekey['duration'] = 0.0  # FIXDUR: cross-trial check
                    rows['corrected_trial_event'].append(ekey)

                elif math.isnan(ekey['duration']) and last_dur is not None:
                    log.warning('... {} duration using last_edur {}'
                                .format(i, last_dur))
                    ekey['duration'] = last_dur
                    rows['corrected_trial_event'].append(ekey)

                else:
                    last_dur = ekey['duration']  # only track 'good' values.

                log.debug('delay event duration: {}'.format(ekey['duration']))
                rows['trial_event'].append(ekey)

            #
            # Add 'go' event
            #
            log.debug('BehaviorIngest.make(): go')

            ekey = dict(tkey)
            responseindex = np.where(t.state_data == states['ResponseCue'])[0]

            ekey['trial_event_id'] = len(rows['trial_event'])
            ekey['trial_event_type'] = 'go'
            ekey['trial_event_time'] = t.state_times[responseindex][0]
            ekey['duration'] = gui['AnswerPeriod'][0]

            if math.isnan(ekey['duration']):
                log.debug('BehaviorIngest.make(): fixing go duration')
                ekey['duration'] = 0.0  # FIXDUR: lookup from previous trials
                rows['corrected_trial_event'].append(ekey)

            rows['trial_event'].append(ekey)

            #
            # Add 'trialEnd' events
            #

            log.debug('BehaviorIngest.make(): trialend events')

            last_dur = None
            trialendindex = np.where(t.state_data == states['TrialEnd'])[0]

            ekey = dict(tkey)
            ekey['trial_event_id'] = len(rows['trial_event'])
            ekey['trial_event_type'] = 'trialend'
            ekey['trial_event_time'] = t.state_times[trialendindex][0]
            ekey['duration'] = 0.0

            rows['trial_event'].append(ekey)

            #
            # Add lick events
            #

            lickleft = np.where(t.event_data == 69)[0]
            log.debug('... lickleft: {r}'.format(r=str(lickleft)))

            action_event_count = len(rows['action_event'])
            if len(lickleft):
                [rows['action_event'].append(
                    dict(tkey, action_event_id=action_event_count+idx,
                         action_event_type='left lick',
                         action_event_time=t.event_times[l]))
                 for idx, l in enumerate(lickleft)]

            lickright = np.where(t.event_data == 71)[0]
            log.debug('... lickright: {r}'.format(r=str(lickright)))

            action_event_count = len(rows['action_event'])
            if len(lickright):
                [rows['action_event'].append(
                    dict(tkey, action_event_id=action_event_count+idx,
                         action_event_type='right lick',
                         action_event_time=t.event_times[r]))
                    for idx, r in enumerate(lickright)]

            #
            # Photostim Events
            #

            if t.stim:
                log.debug('BehaviorIngest.make(): t.stim == {}'.format(t.stim))
                rows['photostim_trial'].append(tkey)
                delay_period_idx = np.where(
                    t.state_data == states['DelayPeriod'])[0][0]
                rows['photostim_trial_event'].append(
                    dict(tkey,
                         photo_stim=t.stim,
                         photostim_event_id=len(
                             rows['photostim_trial_event']),
                         photostim_event_time=t.state_times[delay_period_idx],
                         power=5.5))

            # end of trial loop.

        # Session Insertion

        log.info('BehaviorIngest.make(): adding session record')
        experiment.Session().insert1(skey)

        # Behavior Insertion

        log.info('BehaviorIngest.make(): bulk insert phase')

        log.info('BehaviorIngest.make(): saving ingest {d}'.format(d=key))
        self.insert1(key, ignore_extra_fields=True, allow_direct_insert=True)

        log.info('BehaviorIngest.make(): ... experiment.Session.Trial')
        experiment.SessionTrial().insert(
            rows['trial'], ignore_extra_fields=True, allow_direct_insert=True)

        log.info('BehaviorIngest.make(): ... experiment.BehaviorTrial')
        experiment.BehaviorTrial().insert(
            rows['behavior_trial'], ignore_extra_fields=True,
            allow_direct_insert=True)

        log.info('BehaviorIngest.make(): ... experiment.TrialNote')
        experiment.TrialNote().insert(
            rows['trial_note'], ignore_extra_fields=True,
            allow_direct_insert=True)

        log.info('BehaviorIngest.make(): ... experiment.TrialEvent')
        experiment.TrialEvent().insert(
            rows['trial_event'], ignore_extra_fields=True,
            allow_direct_insert=True, skip_duplicates=True)

        log.info('BehaviorIngest.make(): ... CorrectedTrialEvents')
        BehaviorIngest().CorrectedTrialEvents().insert(
            rows['corrected_trial_event'], ignore_extra_fields=True,
            allow_direct_insert=True)

        log.info('BehaviorIngest.make(): ... experiment.ActionEvent')
        experiment.ActionEvent().insert(
            rows['action_event'], ignore_extra_fields=True,
            allow_direct_insert=True)

        # Photostim Insertion

        photostim_ids = np.unique(
            [r['photo_stim'] for r in rows['photostim_trial_event']])

        unknown_photostims = np.setdiff1d(
            photostim_ids, list(photostims.keys()))

        if unknown_photostims:
            raise ValueError(
                'Unknown photostim protocol: {}'.format(unknown_photostims))

        if photostim_ids.size > 0:
            log.info('BehaviorIngest.make(): ... experiment.Photostim')
            for stim in photostim_ids:
                experiment.Photostim.insert1(
                    dict(skey, **photostims[stim]), ignore_extra_fields=True)

                experiment.Photostim.PhotostimLocation.insert(
                    (dict(skey, **loc,
                          photo_stim=photostims[stim]['photo_stim'])
                     for loc in photostims[stim]['locations']),
                    ignore_extra_fields=True)

        log.info('BehaviorIngest.make(): ... experiment.PhotostimTrial')
        experiment.PhotostimTrial.insert(rows['photostim_trial'],
                                         ignore_extra_fields=True,
                                         allow_direct_insert=True)

        log.info('BehaviorIngest.make(): ... experiment.PhotostimTrialEvent')
        experiment.PhotostimEvent.insert(rows['photostim_trial_event'],
                                         ignore_extra_fields=True,
                                         allow_direct_insert=True)

        # Behavior Ingest Insertion

        log.info('BehaviorIngest.make(): ... BehaviorIngest.BehaviorFile')
        BehaviorIngest.BehaviorFile().insert1(
            dict(key, behavior_file=os.path.basename(key['subpath'])),
            ignore_extra_fields=True, allow_direct_insert=True)


import haiku as hk
import jax

from jdetr._typing import JaxArray

__all__ = ["ResNet50"]

class HeadlessResnet(hk.nets.ResNet):
    def __call__(
        self, inputs: JaxArray, is_training: bool, test_local_stats: bool = False
    ) -> JaxArray:
        out = inputs
        out = self.initial_conv(out)
        if not self.resnet_v2:
            out = self.initial_batchnorm(out, is_training, test_local_stats)
            out = jax.nn.relu(out)

        out = hk.max_pool(
            out, window_shape=(1, 3, 3, 1), strides=(1, 2, 2, 1), padding="SAME"
        )

        for block_group in self.block_groups:
            out = block_group(out, is_training, test_local_stats)

        return out

class ResNet50(HeadlessResnet, hk.nets.ResNet50):
    """
    >>> from jdetr.utils import Init
    >>> x = jax.numpy.zeros((1, 224, 224, 3))
    >>> y = Init(ResNet50, 0).__call__(x, is_training=True)
    >>> y.shape
    (1, 7, 7, 2048)
    """

    ...


"""Deep Q network with parameter sharing(PS) for multi-agent reinforcement learning"""
import numpy as np
import tensorflow as tf
#import tensorflow.compat.v1 as tf
#tf.disable_v2_behavior()

from base import TFBaseModel
#from marl_agent.utils.memory import ReplayBuffer
from utils.memory import ReplayBuffer
#from marl_agent.utils.policies import BoltzmanPolicy
from utils.policies import BoltzmanPolicy
import time
import os

class DeepQNetwork(TFBaseModel):
    def __init__(self, env, name, **kwargs):
                 #batch_size=64, learning_rate=1e-4, reward_decay=0.99,
                 #train_freq=1, target_update=2000, memory_size=2 ** 10, eval_obs=None,
                 #use_dueling=True, use_double=True, use_conv=False,
                 #custom_state_space=None, num_gpu=1, infer_batch_size=8192, network_type=0):
        """
        Init DQN
        :param env: Environment
            environment
        :param name: str
            name of this model
        :param batch_size: int
        :param learning_rate: float
        :param reward_decay: float
            reward_decay in TD
        :param train_freq: int
            mean training times of a sample
        :param target_update: int
            target will update every target_update batches
        :param memory_size: int
            weight of entropy loss in total loss
        :param eval_obs: numpy array
            evaluation set of observation
        :param use_dueling: bool
            whether use dueling q network
        :param use_double: bool
            whether use double q network
        :param use_conv: bool
            use convolution or fully connected layer as state encoder
        :param custom_state_space: tuple
        :param num_gpu: int
            number of gpu
        :param infer_batch_size: int
            batch size while inferring actions
        :param network_type:
        """
        TFBaseModel.__init__(self, env, name, "tfdqn")
        # ======================== set config  ========================
        self.env = env
        self.state_space = env.get_state_space()
        self.num_actions = env.get_action_space()

        self.batch_size = kwargs.setdefault("batch_size", 64)
        self.learning_rate = kwargs.setdefault("learning_rate", 1e4)
        self.training_freq = kwargs.setdefault("training_freq", 1)  # train time of every sample (s,a,r,s')
        self.target_update = kwargs.setdefault("target_update", 1000)   # target network update frequency
        self.eval_obs = kwargs.setdefault("eval_obs", None)
       # self.infer_batch_size = kwargs.setdefault("infer_batch_size", 8192)  # maximum batch size when infer actions,
       # change this to fit your GPU memory if you meet a OOM

        self.network_param = kwargs.setdefault("network", False)
        self.use_dueling = self.network_param["use_dueling"]
        self.use_double = self.network_param["use_double"]
        self.num_gpu = self.network_param["num_gpu"]
        self.use_conv = self.network_param["use_conv"]

        self.nn_layers = self.network_param["layers"]
        self.activation = self.network_param["activation"]
        self.train_ct = 0

        # ======================= build network =======================
        tf.reset_default_graph()
        # input place holder
        self.target = tf.placeholder(tf.float32, [None], name="target")
        self.weight = tf.placeholder(tf.float32, [None], name="weight")

        self.input_state = tf.placeholder(tf.float32, [None, self.state_space], name="input_state")
        self.action = tf.placeholder(tf.int32, [None], name="action")
        self.mask = tf.placeholder(tf.float32, [None], name="mask")
        self.eps = tf.placeholder(tf.float32, name="eps")  # e-greedy

        # build a graph
        with tf.variable_scope(self.name):
            with tf.variable_scope("eval_net_scope"):
                self.eval_scope_name = tf.get_variable_scope().name
                self.qvalues = self._create_network(self.input_state, self.use_conv)

            if self.num_gpu > 1: # build inference graph for multiple gpus
                self._build_multi_gpu_infer(self.num_gpu)

            with tf.variable_scope("target_net_scope"):
                self.target_scope_name = tf.get_variable_scope().name
                self.target_qvalues = self._create_network(self.input_state, self.use_conv)

        # loss
        self.gamma = kwargs.setdefault("reward_decay", 0.99)
        self.actions_onehot = tf.one_hot(self.action, self.num_actions)
        td_error = tf.square(self.target - tf.reduce_sum(tf.multiply(self.actions_onehot, self.qvalues), axis=1))
        self.loss = tf.reduce_sum(td_error * self.mask) / tf.reduce_sum(self.mask)
        self.loss_summary = tf.summary.scalar(name='Loss_summary', tensor=self.loss)

        # train op(clip gradient)
        optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate, name="ADAM")
        gradients, variables = zip(*optimizer.compute_gradients(self.loss))
        gradients, _ = tf.clip_by_global_norm(gradients, 5.0)
        self.train_op = optimizer.apply_gradients(zip(gradients, variables), name="train_op")
#        self.train_summary = tf.summary.scalar(name='Train_summary', tensor=self.train_op)

        # output action
        def out_action(qvalues):
            best_action = tf.argmax(qvalues, axis=1)
            best_action = tf.to_int32(best_action)
            random_action = tf.random_uniform(tf.shape(best_action), 0, self.num_actions, tf.int32, name="random_action")
            should_explore = tf.random_uniform(tf.shape(best_action), 0, 1) < self.eps
            return tf.where(should_explore, random_action, best_action)

        self.output_action = out_action(self.qvalues)
        if self.num_gpu > 1:
            self.infer_out_action = [out_action(qvalue) for qvalue in self.infer_qvalues]

        # target network update op
        self.update_target_op = []
        t_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, self.target_scope_name)
        e_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, self.eval_scope_name)
        for i in range(len(t_params)):
            self.update_target_op.append(tf.assign(t_params[i], e_params[i]))

        # Initialize the tensor board
        if not os.path.exists('summaries'):
            os.mkdir('summaries')
        if not os.path.exists(os.path.join('summaries', 'first')):
            os.mkdir(os.path.join('summaries', 'first'))

        # init tensorflow session
        config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)
        config.gpu_options.allow_growth = True
        self.sess = tf.Session(config=config)
        self.summ_writer = tf.summary.FileWriter(os.path.join('summaries', 'first'), self.sess.graph)
        self.sess.run(tf.global_variables_initializer())

        # init replay buffers
        self.replay_buffer_len = 0
        self.memory_size = int(kwargs.setdefault("memory_size", 2**10))
        print("Memory size ", self.memory_size)
        self.replay_buf_state = ReplayBuffer(shape=(self.memory_size, self.state_space))
        self.replay_buf_action = ReplayBuffer(shape=(self.memory_size,), dtype=np.int32)
        self.replay_buf_reward = ReplayBuffer(shape=(self.memory_size,))
        self.replay_buf_terminal = ReplayBuffer(shape=(self.memory_size,), dtype=np.bool)
        self.replay_buf_mask = ReplayBuffer(shape=(self.memory_size,))
        # if mask[i] == 0, then the item is used for padding, not for training
        self.policy = BoltzmanPolicy(action_space=self.num_actions)

    def _create_network(self, input_state, use_conv=False, reuse=None):
        """
        Define computation graph of network
        :param input_state: tf.tensor
        :param use_conv: bool
        :param reuse: bool
        :return:
        """
        kernel_num = [32, 32]
        hidden_size = self.nn_layers.values()
        if len(hidden_size) is 1:
            if self.activation == "Linear":
                print("1 NN with Linear activation, " + str(hidden_size[0]) + " neurons")
                h_state = tf.layers.dense(input_state,  units=hidden_size[0], activation=None,
                                    name="h_state", reuse=reuse)
            else:
                print("1 NN with RELU activation, " + str(hidden_size[0]) + " neurons")
                h_state = tf.layers.dense(input_state,  units=hidden_size[0], activation=tf.nn.relu,
                                    name="h_state", reuse=reuse)
        elif len(hidden_size) is 2:
            print("2 Layers NN with " + str(hidden_size[0]) + " and " + str(hidden_size[1]) + " neurons")
            activation = None
            if self.activation != "Linear":
                activation = tf.nn.relu

            h_state_0 = tf.layers.dense(input_state,  units=hidden_size[0], activation=activation,
                                    name="h_state_0", reuse=reuse)

            h_state = tf.layers.dense(h_state_0,  units=hidden_size[1], activation=activation,
                                    name="h_state", reuse=reuse)

        if self.use_dueling:
            value = tf.layers.dense(h_state, units=1, name="value", reuse=reuse)
            advantage = tf.layers.dense(h_state, units=self.num_actions, use_bias=False,
                                        name="advantage", reuse=reuse)

            qvalues = value + advantage - tf.reduce_mean(advantage, axis=1, keep_dims=True)
        else:
            qvalues = tf.layers.dense(h_state, units=self.num_actions, name="value", reuse=reuse)

        return qvalues

    def infer_action(self, user_id, sn, obs, step, policy="e_greedy", eps=0):
        """
        infer action for the given agent.
        :param raw_obs:
        :param user_id: int
            id of the user
        :param policy:
            can be eps-greedy or greedy
        :param eps: float
            used when policy is eps-greedy
        :return:
        """

        if policy == 'e_greedy':
            eps = eps
        elif policy == 'greedy':
            eps = 0

#        if self.num_gpu > 1 and n > batch_size:  # infer by multi gpu in parallel
#            ret = self._infer_multi_gpu(view, feature, ids, eps)
        qvalues = self.sess.run(self.qvalues, feed_dict={self.input_state: obs})
        best_actions = np.argmax(qvalues, axis=1)

        n = 1  # Since we take an action only for 1 user.
        random = np.random.randint(self.num_actions, size=(n,))
        cond = np.random.uniform(0, 1, size=(n,)) < eps
        ret = np.where(cond, random, best_actions)
        action = ret.astype(np.int32)
        # TODO: enable this later.
        #action = self.policy.take_action(qvalues, step)
        #actions.append(action)
        # action = self.sess.run(self.output_action, feed_dict={
        #     self.input_state: obs,
        #     self.eps: eps
        # })
        return action

    def _calc_target(self, next_state, reward, terminal):
        """
        Calculate target value
        :param next_state: next_state of the user.
        :param reward: reward of the previous action
        :param terminal:
        :return:
        """
        n = len(reward)
        if self.use_double:
            t_qvalues, qvalues = self.sess.run([self.target_qvalues, self.qvalues],
                                               feed_dict={self.input_state: next_state})
            next_value = t_qvalues[np.arange(n), np.argmax(qvalues, axis=1)]
        else:
            t_qvalues = self.sess.run(self.target_qvalues, feed_dict={self.input_state: next_state})
            next_value = np.max(t_qvalues, axis=1)

        target = np.where(terminal, reward, reward + self.gamma * next_value)

        return target

    def _add_to_replay_buffer(self, sample_buffer):
        """
        Add stored episode buffers to replay buffer.
        :param sample_buffer:
        :return:
        """
        n = 0
        for episode in sample_buffer.episodes(): # Each user has its own episode.
            s, a, r = [], [], []
            for step in range(len(episode.states)):  # Step represent the sequence number of the transmitted packet.
                if (episode.states[step] is not -1)and (episode.actions[step] is not -1)and (episode.rewards[step]is not -1):
                    # This part is required to make sure we synchronize the s,a and reward.
                    # in order words, to alleviate the effect of delayed rewards.
                    s.append(episode.states[step])
                    a.append(episode.actions[step])
                    r.append(episode.rewards[step])

            m = len(r)
            if m is 0:
                continue
            mask = np.ones((m,))
            terminal = np.zeros((m,), dtype=np.bool)
            if episode.terminal:
                terminal[-1] = True
            else:
                mask[-1] = 0

            self.replay_buf_state.put(s)
            self.replay_buf_action.put(a)
            self.replay_buf_reward.put(r)
            self.replay_buf_terminal.put(terminal)
            self.replay_buf_mask.put(mask)

            n += m

        self.replay_buffer_len = min(self.memory_size, self.replay_buffer_len + n)
        return n

    def train(self, sample_buffer, print_every=1000, **kwargs):
        """
        Add new samples in sample buffer to replay buffer and train
        Parameters
        ----------
        :param sample_buffer: memory.EpisodesBuffer
        :param print_every: int
            print log every print_every batches
        :param kwargs:
        :return:
        loss: float
            bellman residual loss
        value: float
            estimated state value
        """
        add_num = self._add_to_replay_buffer(sample_buffer)
        batch_size = self.batch_size
        total_loss = 0

        n_batches = int(self.training_freq * add_num / batch_size)
        if n_batches == 0:
            return 0, 0

        print("batch number: %d  add: %d  batch_size: %d training_freq: %d  replay_len: %d/%d" %
              (n_batches, add_num, batch_size, self.training_freq, self.replay_buffer_len, self.memory_size))

        start_time = time.time()
        ct = 0
        for i in range(n_batches):
            # fetch a batch
            index = np.random.choice(self.replay_buffer_len - 1, batch_size)

            batch_state = self.replay_buf_state.get(index)
            batch_action = self.replay_buf_action.get(index)
            batch_reward = self.replay_buf_reward.get(index)
            batch_terminal = self.replay_buf_terminal.get(index)
            batch_mask = self.replay_buf_mask.get(index)

            batch_next_state = self.replay_buf_state.get(index+1)

            batch_target = self._calc_target(batch_next_state, batch_reward, batch_terminal)

            ret = self.sess.run([self.train_op, self.loss], feed_dict = {
                self.input_state: batch_state,
                self.action: batch_action,
                self.target: batch_target,
                self.mask: batch_mask
            })
            loss = ret[1]
            total_loss += loss

            if ct % self.target_update == 0:
                print("Target Q update ct " + str(ct))
                self.sess.run(self.update_target_op)

            if ct % print_every == 0:
                print("batch %5d,  loss %.6f, eval %.6f" % (ct, loss, self._eval(batch_target)))

            ct += 1
            self.train_ct += 1

        total_time = time.time() - start_time
        step_average = total_time / max(1.0, (ct / 1000.0))
        print("batches: %d,  total time: %.2f,  1k average: %.2f" % (ct, total_time, step_average))

        return total_loss / ct if ct != 0 else 0, self._eval(batch_target)

    def _eval(self, target):
        """ Evaluate estimated q value"""
        if self.eval_obs is None:
            return np.mean(target)
        else:
            return np.mean(self.sess.run([self.qvalues], feed_dict = {
                self.input_state: self.eval_obs[0]
            }))

    def clean_buffer(self):
        """ Clean replay buffer """
        self.replay_buf_len = 0
        self.replay_buf_view.clear()
        self.replay_buf_feature.clear()
        self.replay_buf_action.clear()
        self.replay_buf_reward.clear()
        self.replay_buf_terminal.clear()
        self.replay_buf_mask.clear()


from collections.abc import Mapping
from collections import defaultdict

import numba
from numba import generated_jit, types
import numpy as np

from africanus.util.patterns import Multiton
from africanus.experimental.rime.fused.arguments import ArgumentDependencies
from africanus.experimental.rime.fused.intrinsics import IntrinsicFactory
from africanus.experimental.rime.fused.specification import RimeSpecification

DATASET_TYPES = []

try:
    from daskms.dataset import Dataset as dmsds
except ImportError:
    pass
else:
    DATASET_TYPES.append(dmsds)

try:
    from xarray import Dataset as xrds
except ImportError:
    pass
else:
    DATASET_TYPES.append(xrds)

def rime_impl_factory(terms, transformers, ncorr):
    @generated_jit(nopython=True, nogil=True, cache=True)
    def rime(names, *inargs):
        if len(inargs) != 1 or not isinstance(inargs[0], types.BaseTuple):
            raise TypeError(f"{inargs[0]} must be be a Tuple")

        if not isinstance(names, types.BaseTuple):
            raise TypeError(f"{names} must be a Tuple of strings")

        if len(names) != len(inargs[0]):
            raise ValueError(f"len(names): {len(names)} "
                             f"!= {len(inargs[0])}")

        if not all(isinstance(n, types.Literal) for n in names):
            raise TypeError(f"{names} must be a Tuple of strings")

        if not all(n.literal_type is types.unicode_type for n in names):
            raise TypeError(f"{names} must be a Tuple of strings")

        # Get literal argument names
        names = tuple(n.literal_value for n in names)

        # Generate intrinsics
        argdeps = ArgumentDependencies(names, terms, transformers)
        factory = IntrinsicFactory(argdeps)
        out_names, pack_opts_indices = factory.pack_optionals_and_indices_fn()
        out_names, pack_transformed = factory.pack_transformed_fn(out_names)
        term_state = factory.term_state_fn(out_names)
        term_sampler = factory.term_sampler_fn()

        try:
            lm_i = out_names.index("lm")
            uvw_i = out_names.index("uvw")
            chan_freq_i = out_names.index("chan_freq")
        except ValueError as e:
            raise ValueError(f"{str(e)} is required")

        def impl(names, *inargs):
            args_opt_idx = pack_opts_indices(inargs)
            args = pack_transformed(args_opt_idx)
            state = term_state(args)

            nsrc, _ = args[lm_i].shape
            nrow, _ = args[uvw_i].shape
            nchan, = args[chan_freq_i].shape

            vis = np.zeros((nrow, nchan, ncorr), np.complex128)
            # Kahan summation compensation
            compensation = np.zeros_like(vis)

            for s in range(nsrc):
                for r in range(nrow):
                    t = state.time_index[r]
                    a1 = state.antenna1[r]
                    a2 = state.antenna2[r]
                    f1 = state.feed1[r]
                    f2 = state.feed2[r]

                    for ch in range(nchan):
                        X = term_sampler(state, s, r, t, f1, f2, a1, a2, ch)

                        for c, value in enumerate(numba.literal_unroll(X)):
                            # Kahan summation
                            y = value - compensation[r, ch, c]
                            current = vis[r, ch, c]
                            x = current + y
                            compensation[r, ch, c] = (x - current) - y
                            vis[r, ch, c] = x

            return vis

        return impl

    return rime

class RimeFactory(metaclass=Multiton):
    REQUIRED_ARGS = ArgumentDependencies.REQUIRED_ARGS
    REQUIRED_ARGS_LITERAL = tuple(types.literal(n) for n in REQUIRED_ARGS)
    DEFAULT_SPEC = "(Kpq, Bpq): [I, Q, U, V] -> [XX, XY, YX, YY]"

    def __reduce__(self):
        return (RimeFactory, (self.rime_spec,))

    def __hash__(self):
        return hash(self.rime_spec)

    def __eq__(self, rhs):
        return (isinstance(rhs, RimeFactory) and
                self.rime_spec == rhs.rime_spec)

    def __init__(self, rime_spec=DEFAULT_SPEC):
        if isinstance(rime_spec, RimeSpecification):
            pass
        elif isinstance(rime_spec, (list, tuple)):
            rime_spec = RimeSpecification(*rime_spec)
        elif isinstance(rime_spec, str):
            rime_spec = RimeSpecification(rime_spec)

        self.rime_spec = rime_spec
        self.impl = rime_impl_factory(
            rime_spec.terms,
            rime_spec.transformers,
            len(rime_spec.corrs))

    def dask_blockwise_args(self, **kwargs):
        """ Get the dask schema """
        argdeps = ArgumentDependencies(
            tuple(kwargs.keys()),
            self.rime_spec.terms,
            self.rime_spec.transformers)
        # Holds kwargs + any dummy outputs from transformations
        dummy_kw = kwargs.copy()

        dask_schema = defaultdict(list)
        for a in argdeps.REQUIRED_ARGS:
            dask_schema[a].append(("internal", ("row",)))

        POISON = object()

        for transformer in argdeps.can_create.values():
            kw = {}

            for a in transformer.ARGS:
                v = dummy_kw.get(a, None if a in argdeps.KEY_ARGS else POISON)
                kw[a] = v

            for a, d in transformer.KWARGS.items():
                kw[a] = dummy_kw.get(a, d)

            inputs, outputs = transformer.dask_schema(**kw)

            for k, schema in inputs.items():
                dask_schema[k].append((transformer, schema))

            dummy_kw.update(outputs)

        for term in self.rime_spec.terms:
            kw = {a: dummy_kw[a] for a in term.ALL_ARGS if a in dummy_kw}

            for k, v in term.dask_schema(**kw).items():
                dask_schema[k].append((term, v))

        merged_schema = {}

        for a, candidates in dask_schema.items():
            dims = set(pair[1] for pair in candidates)
            if len(dims) != 1:
                raise ValueError(
                    f"Multiple candidates provided conflicting "
                    f"dimension definitions for {a}: {candidates}.")

            merged_schema[a] = dims.pop()

        names = list(sorted(argdeps.valid_inputs & set(kwargs.keys())))
        blockwise_args = [e for n in names
                          for e in (kwargs[n], merged_schema.get(n, None))]

        assert 2 * len(names) == len(blockwise_args)
        return names, blockwise_args

    def __call__(self, time, antenna1, antenna2, feed1, feed2, **kwargs):
        keys = (self.REQUIRED_ARGS_LITERAL +
                tuple(map(types.literal, kwargs.keys())))

        return self.impl(keys, time,
                         antenna1, antenna2,
                         feed1, feed2,
                         *kwargs.values())

def consolidate_args(args, kw):
    mapping = {}
    oargs = []

    for element in args:
        if isinstance(element, tuple(DATASET_TYPES)):
            mapping.update((k.lower(), v.data) for k, v in element.items())
        elif isinstance(element, Mapping):
            mapping.update(element)
        else:
            oargs.append(element)

    mapping.update(zip(oargs, RimeFactory.REQUIRED_ARGS))
    mapping.update(kw)

    return mapping

def rime(rime_spec, *args, **kw):
    """
    Evaluates the Radio Interferometer Measurement Equation (RIME), given
    the Specification of the RIME :code:`rime_spec`, as well as the
    inputs to the RIME given in :code:`*args` and :code:`**kwargs`.
    """
    mapping = consolidate_args(args, kw)
    factory = RimeFactory(rime_spec=rime_spec)
    return factory(**mapping)


# -*- coding: utf-8 -*-
"""
===============================================================================
Network.tools.topology: Assorted topological manipulation methods
===============================================================================

"""
import scipy as _sp
import numpy as _np
import scipy.ndimage as _spim
from OpenPNM.Base import logging as _logging
from OpenPNM.Base import Workspace as _workspace
logger = _logging.getLogger(__name__)
_mgr = _workspace()

def extend(network, pore_coords=[], throat_conns=[], labels=[]):
    r'''
    Add individual pores and/or throats to the network from a list of coords
    or conns.  This is an in-place operation, meaning the received Network
    object will be altered directly.

    Parameters
    ----------
    network : OpenPNM Network Object
        The Network to which pores or throats should be added
    pore_coords : array_like
        The coordinates of the pores to add
    throat_conns : array_like
        The throat connections to add
    labels : string, or list of strings, optional
        A list of labels to apply to the new pores and throats

    Notes
    -----
    This needs to be enhanced so that it increases the size of all pore
    and throat props and labels on ALL associated Phase objects.  At the
    moment it throws an error is there are any associated Phases.

    '''
    if (network._phases != []):
        raise Exception('Network has active Phases, cannot proceed')

    logger.info('Extending network')
    Np_old = network.num_pores()
    Nt_old = network.num_throats()
    Np = Np_old + int(_sp.size(pore_coords)/3)
    Nt = Nt_old + int(_sp.size(throat_conns)/2)
    # Adjust 'all' labels
    del network['pore.all'], network['throat.all']
    network['pore.all'] = _sp.ones((Np,), dtype=bool)
    network['throat.all'] = _sp.ones((Nt,), dtype=bool)
    # Add coords and conns
    if _sp.size(pore_coords) > 0:
        coords = _sp.vstack((network['pore.coords'], pore_coords))
        network['pore.coords'] = coords
    if _sp.size(throat_conns) > 0:
        conns = _sp.vstack((network['throat.conns'], throat_conns))
        network['throat.conns'] = conns
    # Increase size of any prop or label arrays on Network
    for item in list(network.keys()):
        if item.split('.')[1] not in ['coords', 'conns', 'all']:
            if item.split('.')[0] == 'pore':
                N = Np
            else:
                N = Nt
            if network[item].dtype == bool:
                temp = _sp.where(network[item])[0]
                network[item] = _sp.zeros((N,), dtype=bool)
                network[item][temp] = True
            elif network[item].dtype == object:
                temp = network[item]
                network[item] = _sp.ndarray((N,), dtype=object)
                network[item][_sp.arange(0, _sp.shape(temp)[0])] = temp
            else:
                temp = network[item]
                try:
                    network[item] = _sp.ones((N, _sp.shape(temp)[1]),
                                             dtype=float)*_sp.nan
                except:
                    network[item] = _sp.ones((N,), dtype=float)*_sp.nan
                network[item][_sp.arange(0, _sp.shape(temp)[0])] = temp
    # Apply labels, if supplied
    if labels != []:
        # Convert labels to list if necessary
        if type(labels) is str:
            labels = [labels]
        for label in labels:
            # Remove pore or throat from label, if present
            label = label.split('.')[-1]
            if _sp.size(pore_coords) > 0:
                Ps = _sp.r_[Np_old:Np]
                if 'pore.'+label not in network.labels():
                    network['pore.'+label] = False
                network['pore.'+label][Ps] = True
            if _sp.size(throat_conns) > 0:
                Ts = _sp.r_[Nt_old:Nt]
                if 'throat.'+label not in network.labels():
                    network['throat.'+label] = False
                network['throat.'+label][Ts] = True
    # Regnerate the adjacency matrices
    network._update_network()

def trim(network, pores=[], throats=[]):
    '''
    Remove pores or throats from the network.  This is an in-place operation,
    meaning the received Network object will be altered directly.

    Parameters
    ----------
    network : OpenPNM Network Object
        The Network from which pores or throats should be removed
    pores (or throats) : array_like
        A boolean mask of length Np (or Nt) or a list of indices of the
        pores (or throats) to be removed.

    Notes
    -----
    Trimming only adjusts Phase, Geometry, and Physics objects. Trimming a
    Network that has already been used to run simulations will break those
    simulation objects.

    Examples
    --------
    >>> import OpenPNM
    >>> pn = OpenPNM.Network.TestNet()
    >>> pn.Np
    125
    >>> pn.Nt
    300
    >>> pn.trim(pores=[1])
    >>> pn.Np
    124
    >>> pn.Nt
    296

    '''
    mgr = network.workspace
    for net in mgr.networks():
        if net._parent is network:
            raise Exception('This Network has been cloned, cannot trim')
    if (_sp.size(pores) > 0) and (_sp.size(throats) > 0):
        raise Exception('Cannot delete pores and throats simultaneously')
    elif _sp.size(pores) > 0:
        pores = _sp.array(pores, ndmin=1)
        Pkeep = _sp.ones((network.num_pores(),), dtype=bool)
        Pkeep[pores] = False
        Tkeep = _sp.ones((network.num_throats(),), dtype=bool)
        Ts = network.find_neighbor_throats(pores)
        if len(Ts) > 0:
            Tkeep[Ts] = False
    elif _sp.size(throats) > 0:
        throats = _sp.array(throats, ndmin=1)
        Tkeep = _sp.ones((network.num_throats(),), dtype=bool)
        Tkeep[throats] = False
        Pkeep = network['pore.all'].copy()
    else:
        logger.warning('No pores or throats recieved')
        return

    # Trim all associated objects
    for item in network._geometries+network._physics+network._phases:
        Pnet = network['pore.'+item.name]*Pkeep
        Tnet = network['throat.'+item.name]*Tkeep
        temp = network.map_pores(pores=_sp.where(Pnet)[0],
                                 target=item,
                                 return_mapping=True)
        Ps = temp['target']
        temp = network.map_throats(throats=_sp.where(Tnet)[0],
                                   target=item,
                                   return_mapping=True)
        Ts = temp['target']
        # Then resize 'all
        item.update({'pore.all': _sp.ones((_sp.sum(Pnet),), dtype=bool)})
        item.update({'throat.all': _sp.ones((_sp.sum(Tnet),), dtype=bool)})
        # Overwrite remaining data and info
        for key in list(item.keys()):
            if key.split('.')[1] not in ['all']:
                temp = item.pop(key)
                if key.split('.')[0] == 'throat':
                    logger.debug('Trimming {a} from {b}'.format(a=key,
                                                                b=item.name))
                    item[key] = temp[Ts]
                if key.split('.')[0] == 'pore':
                    logger.debug('Trimming {a} from {b}'.format(a=key,
                                                                b=item.name))
                    item[key] = temp[Ps]

    # Remap throat connections
    Pmap = _sp.ones((network.Np,), dtype=int)*-1
    Pmap[Pkeep] = _sp.arange(0, _sp.sum(Pkeep))
    tpore1 = network['throat.conns'][:, 0]
    tpore2 = network['throat.conns'][:, 1]
    Tnew1 = Pmap[tpore1[Tkeep]]
    Tnew2 = Pmap[tpore2[Tkeep]]
    # Write 'all' label specifically
    network.update({'throat.all': _sp.ones((_sp.sum(Tkeep),), dtype=bool)})
    network.update({'pore.all': _sp.ones((_sp.sum(Pkeep),), dtype=bool)})
    # Write throat connections specifically
    network.update({'throat.conns': _sp.vstack((Tnew1, Tnew2)).T})
    # Overwrite remaining data and info
    for item in list(network.keys()):
        if item.split('.')[-1] not in ['conns', 'all']:
            temp = network.pop(item)
            if item.split('.')[0] == 'throat':
                logger.debug('Trimming {a} from {b}'.format(a=item,
                                                            b=network.name))
                network[item] = temp[Tkeep]
            if item.split('.')[0] == 'pore':
                logger.debug('Trimming {a} from {b}'.format(a=item,
                                                            b=network.name))
                network[item] = temp[Pkeep]

    # Reset network graphs
    network._update_network(mode='regenerate')

    # Check Network health
    health = network.check_network_health()
    if health['trim_pores'] != []:
        logger.warning('Isolated pores exist!  Run check_network_health to ID \
                        which pores to remove.')
        pass

def clone_pores(network, pores, apply_label=['clone'], mode='parents'):
    r'''
    Clones the specified pores and adds them to the network

    Parameters
    ----------
    network : OpenPNM Network Object
        The Network object to which the new pores are to be added
    pores : array_like
        List of pores to clone
    apply_labels : string, or list of strings
        The labels to apply to the clones, default is 'clone'
    mode : string
        Controls the connections between parents and clones.  Options are:

        - 'parents': (Default) Each clone is connected only to its parent
        - 'siblings': Clones are only connected to each other in the same
                      manner as parents were connected
        - 'isolated': No connections between parents or siblings
    '''
    if (network._geometries != []):
        logger.warning('Network has active Geometries, new pores must be \
                        assigned a Geometry')
    if (network._phases != []):
        raise Exception('Network has active Phases, cannot proceed')

    logger.debug('Cloning pores')
    apply_label = list(apply_label)
    # Clone pores
    Np = network.num_pores()
    Nt = network.num_throats()
    parents = _sp.array(pores, ndmin=1)
    pcurrent = network['pore.coords']
    pclone = pcurrent[pores, :]
    pnew = _sp.concatenate((pcurrent, pclone), axis=0)
    Npnew = _sp.shape(pnew)[0]
    clones = _sp.arange(Np, Npnew)
    # Add clone labels to network
    for item in apply_label:
        if 'pore.' + item not in network.keys():
            network['pore.'+item] = False
        if 'throat.' + item not in network.keys():
            network['throat.'+item] = False
    # Add connections between parents and clones
    if mode == 'parents':
        tclone = _sp.vstack((parents, clones)).T
        extend(network=network, pore_coords=pclone, throat_conns=tclone)
    if mode == 'siblings':
        ts = network.find_neighbor_throats(pores=pores, mode='intersection')
        tclone = network['throat.conns'][ts] + network.num_pores()
        extend(network=network, pore_coords=pclone, throat_conns=tclone)
    if mode == 'isolated':
        extend(network=network, pore_coords=pclone)
    # Apply provided labels to cloned pores
    for item in apply_label:
        network['pore.'+item][network.pores('all') >= Np] = True
        network['throat.'+item][network.throats('all') >= Nt] = True

    # Any existing adjacency and incidence matrices will be invalid
    network._update_network()

def stitch(network, donor, P_network, P_donor, method='nearest',
           len_max=_sp.inf, label_suffix=''):
    r'''
    Stitches a second a network to the current network.

    Parameters
    ----------
    networK : OpenPNM Network Object
        The Network that will to which to donor Network will be attached

    donor : OpenPNM Network Object
        The Network to stitch on to the current Network

    P_network : array_like
        The pores on the current Network

    P_donor : array_like
        The pores on the donor Network

    label_suffix : string or None
        Some text to append to each label in the donor Network before
        inserting them into the recipient.  The default is to append no
        text, but a common option would be to append the donor Network's
        name. To insert none of the donor labels, use None.

    len_max : float
        Set a length limit on length of new throats

    method : string (default = 'delaunay')
        The method to use when making pore to pore connections. Options are:

        - 'delaunay' : Use a Delaunay tessellation
        - 'nearest' : Connects each pore on the receptor network to its nearest
                      pore on the donor network

    Notes
    -----
    Before stitching it is necessary to translate the pore coordinates of
    one of the Networks so that it is positioned correctly relative to the
    other.

    Examples
    --------
    >>> import OpenPNM
    >>> pn = OpenPNM.Network.TestNet()
    >>> pn2 = OpenPNM.Network.TestNet()
    >>> [pn.Np, pn.Nt]
    [125, 300]
    >>> [pn2.Np, pn2.Nt]
    [125, 300]
    >>> pn2['pore.coords'][:, 2] += 5.0
    >>> pn.stitch(donor=pn2, P_network=pn.pores('top'),
    ...           P_donor=pn2.pores('bottom'), method='nearest', len_max=1.0)
    >>> [pn.Np, pn.Nt]
    [250, 625]

    '''
    # Ensure Networks have no associated objects yet
    if (len(network._simulation()) > 1) or (len(donor._simulation()) > 1):
        raise Exception('Cannot stitch a Network with active sibling objects')
    network['throat.stitched'] = False
    # Get the initial number of pores and throats
    N_init = {}
    N_init['pore'] = network.Np
    N_init['throat'] = network.Nt
    if method == 'nearest':
        P1 = P_network
        P2 = P_donor + N_init['pore']  # Increment pores on donor
        C1 = network['pore.coords'][P_network]
        C2 = donor['pore.coords'][P_donor]
        D = _sp.spatial.distance.cdist(C1, C2)
        [P1_ind, P2_ind] = _sp.where(D <= len_max)
        conns = _sp.vstack((P1[P1_ind], P2[P2_ind])).T
    else:
        raise RuntimeError('<{}> method not supported'.format(method))

    # Enter donor's pores into the Network
    extend(network=network, pore_coords=donor['pore.coords'])

    # Enter donor's throats into the Network
    extend(network=network, throat_conns=donor['throat.conns'] +
           N_init['pore'])

    # Trim throats that are longer then given len_max
    C1 = network['pore.coords'][conns[:, 0]]
    C2 = network['pore.coords'][conns[:, 1]]
    L = _sp.sum((C1 - C2)**2, axis=1)**0.5
    conns = conns[L <= len_max]

    # Add donor labels to recipient network
    if label_suffix is not None:
        if label_suffix != '':
            label_suffix = '_'+label_suffix
        for label in donor.labels():
            element = label.split('.')[0]
            locations = _sp.where(network._get_indices(element) >=
                                  N_init[element])[0]
            try:
                network[label + label_suffix]
            except:
                network[label + label_suffix] = False
            network[label+label_suffix][locations] = donor[label]

    # Add the new stitch throats to the Network
    extend(network=network, throat_conns=conns, labels='stitched')

    # Remove donor from Workspace, if present
    # This check allows for the reuse of a donor Network multiple times
    if donor in _mgr.values():
        _mgr.purge_object(donor)

def connect_pores(network, pores1, pores2, labels=[], add_conns=True):
    r'''
    Returns the possible connections between two group of pores, and optionally
    makes the connections.

    Parameters
    ----------
    network : OpenPNM Network Object

    pores1 : array_like
        The first group of pores on the network

    pores2 : array_like
        The second group of pores on the network

    labels : list of strings
        The labels to apply to the new throats.  This argument is only needed
        if ``add_conns`` is True.

    add_conns : bool
        Indicates whether the connections should be added to the supplied
        network (default is True).  Otherwise, the connections are returned
        as an Nt x 2 array that can be passed directly to ``extend``.

    Notes
    -----
    It creates the connections in a format which is acceptable by
    the default OpenPNM connection ('throat.conns') and either adds them to
    the network or returns them.

    Examples
    --------
    >>> import OpenPNM
    >>> pn = OpenPNM.Network.TestNet()
    >>> pn.Nt
    300
    >>> pn.connect_pores(pores1=[22, 32], pores2=[16, 80, 68])
    >>> pn.Nt
    306
    >>> pn['throat.conns'][300:306]
    array([[16, 22],
           [22, 80],
           [22, 68],
           [16, 32],
           [32, 80],
           [32, 68]])

    '''
    size1 = _sp.size(pores1)
    size2 = _sp.size(pores2)
    array1 = _sp.repeat(pores1, size2)
    array2 = _sp.tile(pores2, size1)
    conns = _sp.vstack([array1, array2]).T
    if add_conns:
        extend(network=network, throat_conns=conns, labels=labels)
    else:
        return conns

def find_centroid(coords=None, mode='geometric'):
    r'''
    Finds the coordinates of the centroid of the specified pores.

    Parameters
    ----------
    coords : array_like
        An Np x 3 list of [x, y, z] coordinates for the pores of which the
        centroid is sought

    mode : string
        Controls how the centroid is calculated.  Options are:

        **'geometric'** : (default) Simply calcuates the average [X, Y, Z] of
        the given coordinates using scipy.mean(coords, axis=0)

        **'center_of_mass'** : not implemented yet

        **'inscribed_sphere'** : not implemented yet
    '''
    if mode == 'geometric':
        c = _sp.mean(coords, axis=0)
    return c

def find_pores_distance(network, pores1=None, pores2=None):
    r'''
    It finds the distance between two group of pores.
    '''
    from scipy.spatial.distance import cdist
    p1 = _sp.array(pores1, ndmin=1)
    p2 = _sp.array(pores2, ndmin=1)
    coords = network['pore.coords']
    return cdist(coords[p1], coords[p2])

def subdivide(network, pores, shape, labels=[]):
    r'''
    It trim the pores and replace them by cubic networks with the sent shape.

    Parameters
    ----------
    network : OpenPNM Network Object

    pores : array_like
        The first group of pores to be replaced

    shape : array_like
        The shape of cubic networks in the target locations

    Notes
    -----
    - It works only for cubic networks.

    Examples
    --------
    >>> import OpenPNM
    >>> pn = OpenPNM.Network.Cubic(shape=[5,6,5], spacing=0.001)
    >>> pn.Np
    150
    >>> nano_pores = [2,13,14,15]
    >>> pn.subdivide(pores=nano_pores, shape=[4,7,3], labels='nano')
    >>> pn.Np
    482
    >>> assert pn.Np == (150+4*(4*7*3)-4)

    '''
    mro = [item.__name__ for item in network.__class__.__mro__]
    if 'Cubic' not in mro:
        raise Exception('Subdivide is only supported for Cubic Networks')
    from OpenPNM.Network import Cubic
    pores = _sp.array(pores, ndmin=1)

    # Checks to find boundary pores in the selected pores
    if 'pore.boundary' in network.labels():
        if (_sp.in1d(pores, network.pores('boundary'))).any():
            raise Exception('boundary pores cannot be subdivided!')
    if not hasattr(network, '_subdivide_flag'):
        network._subdivide_flag = True
    else:
        raise Exception('The network has subdivided pores, so the method ' +
                        'does not support another subdivision.')
    # Assigning right shape and division
    if _sp.size(shape) != 2 and _sp.size(shape) != 3:
        raise Exception('Subdivide not implemented for Networks other than 2D \
                         and 3D')
    elif _sp.size(shape) == 3 and 1 not in shape:
        div = _sp.array(shape, ndmin=1)
        single_dim = None
    else:
        single_dim = _sp.where(_sp.array(network._shape) == 1)[0]
        if _sp.size(single_dim) == 0:
            single_dim = None
        if _sp.size(shape) == 3:
            div = _sp.array(shape, ndmin=1)
        else:
            div = _sp.zeros(3, dtype=_sp.int32)
            if single_dim is None:
                dim = 2
            else:
                dim = single_dim
            div[dim] = 1
            div[-_sp.array(div, ndmin=1, dtype=bool)] = _sp.array(shape,
                                                                  ndmin=1)

    # Creating small network and handling labels
    network_spacing = network._spacing
    new_net_spacing = network_spacing/div
    new_net = Cubic(shape=div, spacing=new_net_spacing)
    main_labels = ['left', 'right', 'front', 'back', 'top', 'bottom']
    if single_dim is not None:
        label_groups = _sp.array([['front', 'back'],
                                  ['left', 'right'],
                                  ['top', 'bottom']])
        non_single_labels = label_groups[_sp.array([0, 1, 2]) != single_dim]
    for l in main_labels:
        new_net['pore.surface_' + l] = False
        network['pore.surface_' + l] = False
        if single_dim is None:
            new_net['pore.surface_' + l][new_net.pores(labels=l)] = True
        else:
            for ind in [0, 1]:
                loc = (non_single_labels[ind] == l)
                temp_pores = new_net.pores(non_single_labels[ind][loc])
                new_net['pore.surface_' + l][temp_pores] = True

    old_coords = _sp.copy(new_net['pore.coords'])
    if labels == []:
        labels = ['pore.subdivided_' + new_net.name]
    for P in pores:
        # Shifting the new network to the right location and attaching it to
        # the main network
        shift = network['pore.coords'][P] - network_spacing/2
        new_net['pore.coords'] += shift
        Pn = network.find_neighbor_pores(pores=P)
        try:
            Pn_new_net = network.pores(labels)
        except:
            Pn_new_net = []
        Pn_old_net = Pn[~_sp.in1d(Pn, Pn_new_net)]
        Np1 = network.Np
        extend(pore_coords=new_net['pore.coords'],
               throat_conns=new_net['throat.conns'] + Np1,
               labels=labels, network=network)

        # Moving the temporary labels to the big network
        for l in main_labels:
            network['pore.surface_'+l][Np1:] = new_net['pore.surface_'+l]

        # Stitching the old pores of the main network to the new extended pores
        surf_pores = network.pores('surface_*')
        surf_coord = network['pore.coords'][surf_pores]
        for neighbor in Pn:
            neighbor_coord = network['pore.coords'][neighbor]
            dist = [round(_sp.inner(neighbor_coord-x, neighbor_coord-x),
                          20) for x in surf_coord]
            nearest_neighbor = surf_pores[dist == _sp.amin(dist)]
            if neighbor in Pn_old_net:
                coplanar_labels = network.labels(pores=nearest_neighbor)
                new_neighbors = network.pores(coplanar_labels,
                                              mode='intersection')
                # This might happen to the edge of the small network
                if _sp.size(new_neighbors) == 0:
                    labels = network.labels(pores=nearest_neighbor,
                                            mode='intersection')
                    common_label = [l for l in labels if 'surface_' in l]
                    new_neighbors = network.pores(common_label)
            elif neighbor in Pn_new_net:
                new_neighbors = nearest_neighbor
            connect_pores(network=network, pores1=neighbor,
                          pores2=new_neighbors, labels=labels)

        # Removing temporary labels
        for l in main_labels:
            network['pore.surface_' + l] = False
        new_net['pore.coords'] = _sp.copy(old_coords)

    network._label_surfaces()
    for l in main_labels:
        del network['pore.surface_'+l]
    trim(network=network, pores=pores)
    _mgr.purge_object(obj=new_net, mode='complete')

def trim_occluded_throats(network, mask='all'):
    r"""
    Remove throats with zero area from the network and also remove
    pores that are isolated (as a result or otherwise)

    Parameters
    ----------
    network : OpenPNM Network Object

    mask : string
        Applies routine only to pores and throats with this label
    """
    occluded_ts = network['throat.area'] == 0
    if _sp.sum(occluded_ts) > 0:
        # Apply mask
        occluded_ts *= network["throat."+mask]
        trim(network=network, throats=occluded_ts)
    # Also get rid of isolated pores
    isolated_ps = network.check_network_health()['isolated_pores']
    if _sp.size(isolated_ps) > 0:
        # Convert to Bool array and apply mask
        temp_array = _sp.zeros(network.num_pores()).astype(bool)
        temp_array[isolated_ps] = True
        isolated_ps = temp_array * network["pore."+mask]
        trim(network=network, pores=isolated_ps)

def merge_pores(network, pores, labels=['merged']):
    r"""
    Combines a selection of pores into a new single pore located at the
    centroid of the selected pores and connected to all of their neighbors.

    Parameters
    ----------
    network : OpenPNM Network Object

    pores : array_like
        The list of pores which are to be combined into a new single pore

    labels : string or list of strings
        The labels to apply to the new pore and new throat connections

    Notes
    -----
    The selection of pores should be chosen carefully, preferrable so that
    they all form a continuous cluster.  For instance, it is recommended
    to use the ``find_nearby_pores`` method to find all pores within a
    certain distance of a given pore, and these can then be merged without
    causing any abnormal connections.

    Examples
    --------
    >>> import OpenPNM as op
    >>> pn = op.Network.Cubic(shape=[20,20,1])
    >>> topo = op.Utilities.topology()
    >>> P = pn.find_nearby_pores(pores=111, distance=5, flatten=True)
    >>> topo.merge_pores(network=pn, pores=P, labels=['merged'])
    >>> print(pn.Np)
    321
    >>> pn.pores('merged')
    array([320])
    >>> pn.num_throats('merged')
    32

    """
    Pn = network.find_neighbor_pores(pores=pores,
                                     mode='union',
                                     flatten=True,
                                     excl_self=True)
    xyz = _sp.mean(network['pore.coords'][pores], axis=0)
    extend(network, pore_coords=xyz, labels=labels)
    Pnew = network.Ps[-1]
    connect_pores(network, pores1=Pnew, pores2=Pn, labels=labels)
    trim(network=network, pores=pores)

def _template_sphere_disc(dim, outer_radius, inner_radius):
    r"""
    This private method generates an image array of a sphere/shell-disc/ring.
    It is useful for passing to Cubic networks as a ``template`` to make
    networks with desired shapes.

    Parameters
    ----------
    dim : int
        Network dimension

    outer_radius : int
        Number of the nodes in the outer radius of the network

    inner_radius : int
        Number of the nodes in the inner radius of the network

    Returns
    -------
    A Numpy array containing 1's to demarcate the desired shape, and 0's
    elsewhere.

    """
    rmax = _np.array(outer_radius, ndmin=1)
    rmin = _np.array(inner_radius, ndmin=1)
    ind = 2 * rmax - 1
    coord = _np.indices((ind * _np.ones(dim, dtype=int)))
    coord = coord - (ind - 1)/2
    x = coord[0, :]
    y = coord[1, :]
    if dim == 2:
        img = (x ** 2 + y ** 2) < rmax ** 2
    elif dim == 3:
        z = coord[2, :]
        img = (x ** 2 + y ** 2 + z ** 2) < rmax ** 2
    if rmin[0] != 0:
        if dim == 2:
            img_min = (x ** 2 + y ** 2) > rmin ** 2
        elif dim == 3:
            img_min = (x ** 2 + y ** 2 + z ** 2) > rmin ** 2
        img = img * img_min
    return (img)

def template_sphere_shell(outer_radius, inner_radius=0):
    r"""
    This method generates an image array of a sphere-shell. It is useful for
    passing to Cubic networks as a ``template`` to make spherical shaped
    networks.

    Parameters
    ----------
    outer_radius : int
        Number of the nodes in the outer radius of the shell

    inner_radius : int
        Number of the nodes in the inner radius of the shell

    Returns
    -------
    A Numpy array containing 1's to demarcate the sphere-shell, and 0's
    elsewhere.

    """
    img = _template_sphere_disc(dim=3, outer_radius=outer_radius,
                                inner_radius=inner_radius)
    return(img)

def template_disc_ring(outer_radius, inner_radius=0):
    r"""
    This method generates an image array of a disc-ring.  It is useful for
    passing to Cubic networks as a ``template`` to make circular-shaped 2D
    networks.

    Parameters
    ----------
    outer_radius : int
        Number of the nodes in the outer radius of the disc

    inner_radius : int
        Number of the nodes in the inner radius of the disc

    Returns
    -------
    A Numpy array containing 1's to demarcate the disc-ring, and 0's
    elsewhere.

    """

    img = _template_sphere_disc(dim=2, outer_radius=outer_radius,
                                inner_radius=inner_radius)
    return(img)

def find_surface_pores(network, markers=None, label='surface'):
    r"""
    Find the pores on the surface of the domain by performing a Delaunay
    triangulation between the network pores and some external ``markers``. All
    pores connected to these external marker points are considered surface
    pores.

    Parameters
    ----------
    network: OpenPNM Network Object
        The network for which the surface pores are to be found

    markers: array_like
        3 x N array of the marker coordiantes to use in the triangulation.  The
        labeling is performed in one step, so all points are added, and then
        any pores connected to at least one marker is given the provided label.
        By default, this function will automatically generate 6 points outside
        each axis of the network domain.

        Users may wish to specify a single external marker point and provide an
        appropriate label in order to identify specific faces.  For instance,
        the marker may be *above* the domain, and the label might be
        'top_surface'.

    label : string
        The label to apply to the pores.  The default is 'surface'.

    Notes
    -----
    This function does not check whether the given markers actually lie outside
    the domain, allowing the labeling of *internal* sufaces.

    If this method fails to mark some surface pores, consider sending more
    markers on each face.

    Examples
    --------
    >>> import OpenPNM as op
    >>> net = op.Network.Cubic(shape=[5, 5, 5])
    >>> op.Network.tools.find_surface_pores(network=net)
    >>> net.num_pores('surface')
    98

    When cubic networks are created, the surfaces are already labeled:

    >>> net.num_pores(['top','bottom', 'left', 'right', 'front','back'])
    98

    This function is mostly useful for unique networks such as spheres, random
    topology, or networks that have been subdivied.

    """
    import scipy.spatial as sptl
    if markers is None:
        (xmax, ymax, zmax) = _sp.amax(network['pore.coords'], axis=0)
        (xmin, ymin, zmin) = _sp.amin(network['pore.coords'], axis=0)
        xave = (xmin+xmax)/2
        yave = (ymin+ymax)/2
        zave = (zmin+zmax)/2
        markers = [[xmax + xave, yave, zave],
                   [xmin - xave, yave, zave],
                   [xave, ymax + yave, zave],
                   [xave, ymin - yave, zave],
                   [xave, yave, zmax + zave],
                   [xave, yave, zmin - zave]]
    markers = _sp.atleast_2d(markers)
    tri = sptl.Delaunay(network['pore.coords'], incremental=True)
    tri.add_points(markers)
    (indices, indptr) = tri.vertex_neighbor_vertices
    for k in range(network.Np, tri.npoints):
        neighbors = indptr[indices[k]:indices[k+1]]
        inds = _sp.where(neighbors < network.Np)
        neighbors = neighbors[inds]
        if 'pore.'+label not in network.keys():
            network['pore.'+label] = False
        network['pore.'+label][neighbors] = True

def label_faces(network):
    r"""
    Finds pores on the surface of the network and labels them according to
    whether they are on the *top*, *bottom*, etc.  This function assumes the
    network is cubic in shape (i.e. with six flat sides)

    Parameters
    ----------
    network : OpenPNM Network object
        The network to apply the labels

    """
    find_surface_pores(network)
    Ps = network['pore.surface']
    crds = network['pore.coords']
    network['pore.top'] = (crds[:, 2] > 0.9*_sp.amax(crds[:, 2])) * Ps
    network['pore.bottom'] = (crds[:, 2] < 0.1*_sp.amax(crds[:, 2])) * Ps
    network['pore.right'] = (crds[:, 1] > 0.9*_sp.amax(crds[:, 1])) * Ps
    network['pore.left'] = (crds[:, 1] < 0.1*_sp.amax(crds[:, 1])) * Ps
    network['pore.back'] = (crds[:, 0] > 0.9*_sp.amax(crds[:, 0])) * Ps
    network['pore.front'] = (crds[:, 0] < 0.1*_sp.amax(crds[:, 0])) * Ps

def plot_connections(network, throats=None, fig=None, **kwargs):
    r"""
    Produces a 3D plot of the network topology showing how throats connect
    for quick visualization without having to export data to veiw in Paraview.

    Parameters
    ----------
    network : OpenPNM Network Object
        The network whose topological connections to plot

    throats : array_like (optional)
        The list of throats to plot if only a sub-sample is desired.  This is
        useful for inspecting a small region of the network.  If no throats are
        specified then all throats are shown.

    fig and **kwargs: Matplotlib figure handle and line property arguments
        If a ``fig`` is supplied, then the topology will be overlaid.  By also
        passing in different line properties such as ``color`` and limiting
        which ``throats`` are plots, this makes it possible to plot different
        types of throats on the same plot.

        For information on available line style options, visit the Matplotlib
        documentation at:

        http://matplotlib.org/api/lines_api.html#matplotlib.lines.Line2D

    Notes
    -----
    The figure handle returned by this method can be passed into
    ``plot_coordinates`` to create a plot that combines pore coordinates and
    throat connections, and vice versa.

    Examples
    --------
    >>> import OpenPNM as op
    >>> pn = op.Network.Cubic(shape=[10, 10, 3])
    >>> pn.add_boundaries()
    >>> Ts = pn.throats('*boundary', mode='not')
    >>> # Create figure showing boundary throats
    >>> fig = op.Network.tools.plot_connections(network=pn, throats=Ts)
    >>> Ts = pn.throats('*boundary')
    >>> # Pass existing fig back into function to plot additional throats
    >>> fig = op.Network.tools.plot_connections(network=pn, throats=Ts,
    ...                                         fig=fig, color='r')

    """
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D

    if throats is None:
        Ts = network.Ts
    else:
        Ts = network._parse_locations(locations=throats)

    if fig is None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
    else:
        ax = fig.get_axes()[0]

    # Create dummy indexing to sp.inf
    i = -1*_sp.ones((_sp.size(Ts)*3, ), dtype=int)
    i[0::3] = network['throat.conns'][Ts, 0]
    i[1::3] = network['throat.conns'][Ts, 1]

    # Collect coordinates and scale axes to fit
    Ps = _sp.unique(network['throat.conns'][Ts])
    X = network['pore.coords'][Ps, 0]
    Y = network['pore.coords'][Ps, 1]
    Z = network['pore.coords'][Ps, 2]
    _scale_3d_axes(ax=ax, X=X, Y=Y, Z=Z)

    # Add sp.inf to the last element of pore.coords (i.e. -1)
    inf = _sp.array((_sp.inf,))
    X = _sp.hstack([network['pore.coords'][:, 0], inf])
    Y = _sp.hstack([network['pore.coords'][:, 1], inf])
    Z = _sp.hstack([network['pore.coords'][:, 2], inf])
    ax.plot(xs=X[i], ys=Y[i], zs=Z[i], **kwargs)

    return fig

def plot_coordinates(network, pores=None, fig=None, **kwargs):
    r"""
    Produces a 3D plot showing specified pore coordinates as markers

    Parameters
    ----------
    network : OpenPNM Network Object
        The network whose topological connections to plot

    pores : array_like (optional)
        The list of pores to plot if only a sub-sample is desired.  This is
        useful for inspecting a small region of the network.  If no pores are
        specified then all are shown.

    fig and **kwargs: Matplotlib figure handle and line property arguments
        If a ``fig`` is supplied, then the topology will be overlaid.  By also
        passing in different marker properties such as size (``s``) and
        limiting which ``pores`` are plotted, this makes it possible to plot
        different types of pores on the same plot.

        For information on available marker style options, visit the Matplotlib
        documentation at:

        http://matplotlib.org/api/lines_api.html#matplotlib.lines.Line2D

    Notes
    -----
    The figure handle returned by this method can be passed into
    ``plot_topology`` to create a plot that combines pore coordinates and
    throat connections, and vice versa.

    Examples
    --------
    >>> import OpenPNM as op
    >>> pn = op.Network.Cubic(shape=[10, 10, 3])
    >>> pn.add_boundaries()
    >>> Ps = pn.pores('internal')
    >>> # Create figure showing internal pores
    >>> fig = op.Network.tools.plot_coordinates(network=pn, pores=Ps, c='b')
    >>> Ps = pn.pores('*boundary')
    >>> # Pass existing fig back into function to plot boundary pores
    >>> fig = op.Network.tools.plot_coordinates(network=pn, pores=Ps, fig=fig,
    ...                                         c='r')

    """
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D

    if pores is None:
        Ps = network.Ps
    else:
        Ps = network._parse_locations(locations=pores)

    if fig is None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
    else:
        ax = fig.get_axes()[0]

    # Collect specified coordinates
    X = network['pore.coords'][Ps, 0]
    Y = network['pore.coords'][Ps, 1]
    Z = network['pore.coords'][Ps, 2]
    _scale_3d_axes(ax=ax, X=X, Y=Y, Z=Z)

    ax.scatter(xs=X, ys=Y, zs=Z, **kwargs)
    return fig

def _scale_3d_axes(ax, X, Y, Z):
    if hasattr(ax, '_scaled'):
        logger.warning('Axes is already scaled to previously plotted data')
    else:
        ax._scaled = True
        max_range = _sp.array([X.max()-X.min(), Y.max()-Y.min(),
                               Z.max()-Z.min()]).max() / 2.0
        mid_x = (X.max()+X.min()) * 0.5
        mid_y = (Y.max()+Y.min()) * 0.5
        mid_z = (Z.max()+Z.min()) * 0.5
        ax.set_xlim(mid_x - max_range, mid_x + max_range)
        ax.set_ylim(mid_y - max_range, mid_y + max_range)
        ax.set_zlim(mid_z - max_range, mid_z + max_range)

def generate_base_points(num_points, domain_size, prob=None):
    r"""
    Generates a set of base points for passing into the DelaunayVoronoiDual
    class.  The points can be distributed in spherical, cylindrical, or
    rectilinear patterns.

    Parameters
    ----------
    num_points : scalar
        The number of base points that lie within the domain.  Note that the
        actual number of points returned will be larger, with the extra points
        lying outside the domain.

    domain_size : list or array
        Controls the size and shape of the domain, as follows:

        **sphere** : If a single value is received, its treated as the radius
        [r] of a sphere centered on [0, 0, 0].

        **cylinder** : If a two-element list is received it's treated as the
        radius and height of a cylinder [r, z] positioned at [0, 0, 0] and
        extending in the positive z-direction.

        **rectangle** : If a three element list is received, it's treated
        as the outer corner of rectangle [x, y, z] whose opposite corner lies
        at [0, 0, 0].

    prob : 3D array, optional
        A 3D array that contains fractional (0-1) values indicating the
        liklihood that a point in that region should be kept.  If not specified
        an array containing 1's in the shape of a sphere, cylinder, or cube is
        generated, depnending on the give ``domain_size`` with zeros outside.
        When specifying a custom probabiliy map is it recommended to also set
        values outside the given domain to zero.  If not, then the correct
        shape will still be returned, but with too few points in it.

    Notes
    -----
    This method places the given number of points within the specified domain,
    then reflects these points across each domain boundary.  This results in
    smooth flat faces at the boundaries once these excess pores are trimmed.

    The reflection approach tends to create larger pores near the surfaces, so
    it might be necessary to use the ``prob`` argument to specify a slightly
    higher density of points near the surfaces.

    For rough faces, it is necessary to define a larger than desired domain
    then trim to the desired size.  This will discard the reflected points
    plus some of the original points.

    Examples
    --------
    The following generates a spherical array with higher values near the core.
    It uses a distance transform to create a sphere of radius 10, then a
    second distance transform to create larger values in the center away from
    the sphere surface.  These distance values could be further skewed by
    applying a power, with values higher than 1 resulting in higher values in
    the core, and fractional values smoothinging them out a bit.

    >>> import OpenPNM as op
    >>> import scipy as sp
    >>> import scipy.ndimage as spim
    >>> im = sp.ones([21, 21, 21], dtype=int)
    >>> im[10, 10, 10] = 0
    >>> im = spim.distance_transform_edt(im) <= 20  # Create sphere of 1's
    >>> prob = spim.distance_transform_edt(im)
    >>> prob = prob / sp.amax(prob)  # Normalize between 0 and 1
    >>> pts = op.Network.tools.generate_base_points(num_points=50,
    ...                                             domain_size=[2],
    ...                                             prob=prob)
    >>> net = op.Network.DelaunayVoronoiDual(points=pts, domain_size=[2])
    """
    def _try_points(num_points, prob):
        prob = _sp.array(prob)/_sp.amax(prob)  # Ensure prob is normalized
        base_pts = []
        N = 0
        while N < num_points:
            pt = _sp.random.rand(3)  # Generate a point
            # Test whether to keep it or not
            [indx, indy, indz] = _sp.floor(pt*_sp.shape(prob)).astype(int)
            if _sp.random.rand(1) <= prob[indx][indy][indz]:
                base_pts.append(pt)
                N += 1
        base_pts = _sp.array(base_pts)
        return base_pts
    if len(domain_size) == 1:  # Spherical
        domain_size = _sp.array(domain_size)
        if prob is None:
            prob = _sp.ones([41, 41, 41])
            prob[20, 20, 20] = 0
            prob = _spim.distance_transform_bf(prob) <= 20
        base_pts = _try_points(num_points, prob)
        # Convert to spherical coordinates
        [X, Y, Z] = _sp.array(base_pts - [0.5, 0.5, 0.5]).T  # Center at origin
        r = 2*_sp.sqrt(X**2 + Y**2 + Z**2)*domain_size[0]
        theta = 2*_sp.arctan(Y/X)
        phi = 2*_sp.arctan(_sp.sqrt(X**2 + Y**2)/Z)
        # Trim points outside the domain (from improper prob images)
        inds = r <= domain_size[0]
        [r, theta, phi] = [r[inds], theta[inds], phi[inds]]
        # Reflect base points across perimeter
        new_r = 2*domain_size - r
        r = _sp.hstack([r, new_r])
        theta = _sp.hstack([theta, theta])
        phi = _sp.hstack([phi, phi])
        # Convert to Cartesean coordinates
        X = r*_sp.cos(theta)*_sp.sin(phi)
        Y = r*_sp.sin(theta)*_sp.sin(phi)
        Z = r*_sp.cos(phi)
        base_pts = _sp.vstack([X, Y, Z]).T
    elif len(domain_size) == 2:  # Cylindrical
        domain_size = _sp.array(domain_size)
        if prob is None:
            prob = _sp.ones([41, 41, 41])
            prob[20, 20, :] = 0
            prob = _spim.distance_transform_bf(prob) <= 20
        base_pts = _try_points(num_points, prob)
        # Convert to cylindrical coordinates
        [X, Y, Z] = _sp.array(base_pts - [0.5, 0.5, 0]).T  # Center on z-axis
        r = 2*_sp.sqrt(X**2 + Y**2)*domain_size[0]
        theta = 2*_sp.arctan(Y/X)
        z = Z*domain_size[1]
        # Trim points outside the domain (from improper prob images)
        inds = r <= domain_size[0]
        [r, theta, z] = [r[inds], theta[inds], z[inds]]
        inds = ~((z > domain_size[1]) + (z < 0))
        [r, theta, z] = [r[inds], theta[inds], z[inds]]
        # Reflect base points about faces and perimeter
        new_r = 2*domain_size[0] - r
        r = _sp.hstack([r, new_r])
        theta = _sp.hstack([theta, theta])
        z = _sp.hstack([z, z])
        r = _sp.hstack([r, r, r])
        theta = _sp.hstack([theta, theta, theta])
        z = _sp.hstack([z, -z, 2-z])
        # Convert to Cartesean coordinates
        X = r*_sp.cos(theta)
        Y = r*_sp.sin(theta)
        Z = z
        base_pts = _sp.vstack([X, Y, Z]).T
    elif len(domain_size) == 3:  # Rectilinear
        domain_size = _sp.array(domain_size)
        Nx, Ny, Nz = domain_size
        if prob is None:
            prob = _sp.ones([10, 10, 10], dtype=float)
        base_pts = _try_points(num_points, prob)
        base_pts = base_pts*domain_size
        # Reflect base points about all 6 faces
        orig_pts = base_pts
        base_pts = _sp.vstack((base_pts, [-1, 1, 1]*orig_pts +
                                         [2.0*Nx, 0, 0]))
        base_pts = _sp.vstack((base_pts, [1, -1, 1]*orig_pts +
                                         [0, 2.0*Ny, 0]))
        base_pts = _sp.vstack((base_pts, [1, 1, -1]*orig_pts +
                                         [0, 0, 2.0*Nz]))
        base_pts = _sp.vstack((base_pts, [-1, 1, 1]*orig_pts))
        base_pts = _sp.vstack((base_pts, [1, -1, 1]*orig_pts))
        base_pts = _sp.vstack((base_pts, [1, 1, -1]*orig_pts))
    return base_pts


def predict1(model_final,img):
    ya={0: 'Tomato___Late_blight', 1:'Tomato___Leaf_Mold' ,2 : 'Tomato___Spider_mites Two-spotted_spider_mite', 3:'Tomato___Tomato_mosaic_virus' , 4:'Tomato___healthy' }
    import cv2
    import numpy as np
    img=cv2.resize(img,(224,224))
    img=np.expand_dims(img,axis=0)
    y1=model_final.predict(img)
    #print(y1)
    y1[y1>0.5]=1
    #print(y1)
    for i in range(5):
    #print(y5[i])
        if y1[0,i]==1:
            return ya[i]