"""
Create tidy df from raw dfs of cell features, drug features, and drug responses.
Yitan provided the raw dfs and the data splits.
"""
from __future__ import print_function, division

import warnings
warnings.filterwarnings('ignore')

import os
import sys
import platform
from pathlib import Path
import psutil
import argparse
from datetime import datetime
from time import time
from pprint import pprint, pformat
from collections import OrderedDict
from glob import glob

import sklearn
import numpy as np
import pandas as pd

from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

SEED = 42

# File path
file_path = Path(__file__).resolve().parent

# Utils
#utils_path = file_path / '../../utils'
#sys.path.append(str(utils_path))
#import utils
#from classlogger import Logger
import utils

utils_path = file_path / '../../utils'
sys.path.append(str(utils_path))
#import utils
from classlogger import Logger

# Path
OUTDIR = file_path/'../../data/yitan/Data/tidy'

def parse_args(args):
    parser = argparse.ArgumentParser(description="Create tidy df from raw (Yitan's) data.")
    parser.add_argument('--src', default='GDSC', type=str, help='Data source (default: GDSC).')
    # parser.add_argument('--fold', default=0, type=int, help='Fold (default: 0).')

    parser.add_argument('-cf', '--ccl_fea', nargs='+', default=['geneGE'], choices=['geneGE'],
                        help='Cell line features (default: `geneGE`).')
    parser.add_argument('-df', '--drg_fea', nargs='+', default=['DD'], choices=['DD'],
                        help='Drug features (default: `DD`).')

    parser.add_argument('-dk', '--drg_pca_k', default=None, type=int, help='Number of drug PCA components (default: None).')
    parser.add_argument('-ck', '--ccl_pca_k', default=None, type=int, help='Number of cell PCA components (default: None).')

    parser.add_argument('--n_jobs', default=4, type=int, help='Number of cpu workers (default: 4).')

    args = parser.parse_args(args)
    return args

def create_outdir(outdir, args):
    if any([True for i in [args['ccl_pca_k'], args['drg_pca_k']] if i is not None]):
        l = [args['src']] + args['ccl_fea'] + args['drg_fea'] + [f'drg_pca{drg_pca_k}'] + [f'ccl_pca{ccl_pca_k}']
    else:
        l = [args['src']] + args['ccl_fea'] + args['drg_fea']

    name_sffx = '.'.join( l )
    # outdir = Path(outdir) / (name_sffx + '_' + t)
    outdir = Path(outdir)/name_sffx
    os.makedirs(outdir)
    return outdir

def cnt_fea(df, fea_sep='_', verbose=True, logger=None):
    """ Count the number of features per feature type. """
    dct = {}
    unq_prfx = df.columns.map(lambda x: x.split(fea_sep)[0]).unique() # unique feature prefixes
    for prfx in unq_prfx:
        fea_type_cols = [c for c in df.columns if (c.split(fea_sep)[0]) in prfx] # all fea names of specific type
        dct[prfx] = len(fea_type_cols)

    if verbose and logger is not None:
        logger.info(pformat(dct))
    elif verbose:
        pprint(dct)

    return dct

def extract_subset_fea(df, fea_list, fea_sep='_'):
    """ Extract features based feature prefix name. """
    fea = [c for c in df.columns if (c.split(fea_sep)[0]) in fea_list]
    df = df[fea]
    return df

def run(args):
    src = args['src']
    ccl_fea_list = args['ccl_fea']
    drg_fea_list = args['drg_fea']
    n_jobs = args['n_jobs']
    drg_pca_k = args['drg_pca_k']
    ccl_pca_k = args['ccl_pca_k']
    fea_sep = '_'

    # =====================================================
    #       Logger
    # =====================================================
    run_outdir = create_outdir(OUTDIR, args)
    logfilename = run_outdir/'logfile.log'
    lg = Logger(logfilename)
    lg.logger.info(datetime.now())
    lg.logger.info(f'\nFile path: {file_path}')
    lg.logger.info(f'Machine: {platform.node()} ({platform.system()}, {psutil.cpu_count()} CPUs)')
    lg.logger.info(f'\n{pformat(args)}')

    # Dump args to file
    utils.dump_dict(args, run_outdir/'args.txt')

    # =====================================================
    #       Load data
    # =====================================================
    datadir = Path(file_path/'../../data/yitan/Data')
    ccl_folds_dir = Path(file_path/'../../data/yitan/CCL_10Fold_Partition')
    pdm_folds_dir = Path(file_path/'../../data/yitan/PDM_10Fold_Partition')
    fea_data_name = 'CCL_PDM_TransferLearningData_rmFactor_0.0_ddNorm_std.pkl'
    #fea_data_name = 'CCL_PDM_TransferLearningData_rmFactor_0.1_ddNorm_quantile.pkl'

    # Un-pickle files
    import _pickle as cp
    pkl_file = open(datadir/fea_data_name, 'rb')
    res = cp.load(pkl_file)
    ccl = cp.load(pkl_file)
    drg = cp.load(pkl_file)
    pkl_file.close()

    lg.logger.info('\n{}'.format('=' * 50))
    lg.logger.info('res: {}'.format(res.shape))
    lg.logger.info('ccl: {}'.format(ccl.shape))
    lg.logger.info('drg: {}'.format(drg.shape))

    # Update resp
    res = res.reset_index()
    res = res.rename(columns={'index': 'idx', 'ccl_name': 'cclname', 'SOURCE': 'src', 'area_under_curve': 'auc'})

    tmp = res.groupby('src').agg({'cclname': 'nunique', 'ctrpDrugID': 'nunique'}).reset_index()
    lg.logger.info(tmp)

    # =====================================================
    #       Prepare dfs for merging
    # =====================================================
    # Retain specific source
    lg.logger.info('\nFull dataset: {}'.format(res.shape))
    res = res[ res['src'].isin([src]) ]
    lg.logger.info('Only {}: {}'.format(src, res.shape))

    # Extract specific types of features
    lg.logger.info('\nExtract fea types ...')
    cnt_fea(ccl, logger=lg.logger);
    ccl = extract_subset_fea(df=ccl, fea_list=ccl_fea_list, fea_sep=fea_sep)
    cnt_fea(ccl, logger=lg.logger);

    cnt_fea(drg, logger=lg.logger);
    drg = extract_subset_fea(df=drg, fea_list=drg_fea_list, fea_sep=fea_sep)
    cnt_fea(drg, logger=lg.logger);

    # ---------------------------------------------
    #     PCA
    # ---------------------------------------------
    def pca_exp_var(pca, k):
        return np.sum(pca.explained_variance_ratio_[:k])

    def dump_pca_data(datadir, pca_x, k, fea_name):
        df = pd.DataFrame(pca_x.iloc[:, :k])
        df.to_csv(datadir/f'{fea_name}_pca{k}.csv', index=False)

    def pca_on_fea(df, fea_name, k, datadir):
        lg.logger.info(f'Calc PCA for {fea_name} features.')
        index = df.index
        n = 1400
        pca = PCA(n_components=n, random_state=SEED, svd_solver='auto')
        pca_x = pca.fit_transform(df)
        pca_x = StandardScaler().fit_transform(pca_x)
        pca_x = pd.DataFrame(pca_x, index=index,
                             columns=[f'{fea_name}_PC'+str(i+1) for i in range(n)])
        exp_var = pca_exp_var(pca, k)
        # dump_pca_data(datadir, pca_x=pca_x, k=k, fea_name=fea_name);
        lg.logger.info(f'{fea_name} PCA variance explained {exp_var:.4f} (k={k}).')
        return pca_x.iloc[:, :k]

    if drg_pca_k is not None:
        ccl = pca_on_fea(df=ccl, fea_name='ccl', k=ccl_pca_k, datadir=datadir)

    if ccl_pca_k is not None:
        drg = pca_on_fea(df=drg, fea_name='drg', k=drg_pca_k, datadir=datadir)

    # Bring the labels in, in order to merge on
    ccl = ccl.reset_index().rename(columns={'index': 'cclname'})
    drg = drg.reset_index().rename(columns={'index': 'ctrpDrugID'})

    # =====================================================
    #       Merge
    # =====================================================
    def merge_dfs(res_df, ccl_df, drg_df):
        """ Merge the following dfs: response, ccl fea, drug fea """
        mrg_df = pd.merge(res_df, ccl_df, on='cclname', how='inner')
        mrg_df = pd.merge(mrg_df, drg_df, on='ctrpDrugID', how='inner')
        return mrg_df

    lg.logger.info('\nMerge ...')
    mrg_df = merge_dfs(res, ccl, drg)
    lg.logger.info('mrg_df.shape: {}'.format(mrg_df.shape))

    # =====================================================
    #       Extract xdata, ydata, meta and  dump to file
    # =====================================================
    if any([True for i in [ccl_pca_k, drg_pca_k] if i is not None]):
        prfx = '.'.join( [src.lower()] + ccl_fea_list + drg_fea_list + [f'drg_pca{drg_pca_k}'] + [f'ccl_pca{ccl_pca_k}'] )
    else:
        prfx = '.'.join( [src.lower()] + ccl_fea_list + drg_fea_list )

    mrg_df.to_parquet(run_outdir/(prfx+'_data.parquet'), index=False)
    lg.logger.info('mrg_df usage: {:.2f} GB\n'.format(sys.getsizeof(mrg_df)/1e9))

    lg.kill_logger()

def main(args):
    args = parse_args(args)
    args = vars(args)
    ret = run(args)

if __name__ == '__main__':
    """ __name__ == '__main__' explained: www.youtube.com/watch?v=sugvnHA7ElY """
    main(sys.argv[1:])


import cv2
import numpy as np

#img = cv2.imread("./output/Image_Filtered_HP.png",0)

img = cv2.imread("./imgs/img5.jpeg",0)

kernel =np.ones((5,5), np.uint8)

#kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5,5))

ero = cv2.erode(img,kernel,iterations=1)

cv2.imshow('dilata', ero)
cv2.imshow('sem dilata', img)

cv2.waitKey(0)
cv2.destroyAllWindows()


import datetime
import functools

from mypy_boto3_sns.client import SNSClient
from mypy_boto3_dynamodb.service_resource import Table
from threatexchange.signal_type.pdq import PdqSignal
from hmalib.common.models.pipeline import MatchRecord
import typing as t
import time

from threatexchange.signal_type.index import IndexMatch, SignalTypeIndex
from threatexchange.signal_type.signal_base import SignalType

from hmalib import metrics
from hmalib.common.logging import get_logger
from hmalib.common.mappings import INDEX_MAPPING
from hmalib.common.messages.match import BankedSignal, MatchMessage
from hmalib.common.configs.fetcher import (
    ThreatExchangeConfig,
    AdditionalMatchSettingsConfig,
)
from hmalib.common.models.bank import BankMember, BanksTable
from hmalib.common.models.signal import ThreatExchangeSignalMetadata
from hmalib.indexers.metadata import (
    BANKS_SOURCE_SHORT_CODE,
    THREAT_EXCHANGE_SOURCE_SHORT_CODE,
    BaseIndexMetadata,
    BankedSignalIndexMetadata,
)

logger = get_logger(__name__)

PG_CONFIG_CACHE_TIME_SECONDS = 300

@functools.lru_cache(maxsize=128)
def _get_privacy_group_matcher_active(privacy_group_id: str, cache_buster) -> bool:
    config = ThreatExchangeConfig.get(privacy_group_id)
    if not config:
        logger.warning("Privacy group %s is not found!", privacy_group_id)
        return False

    logger.debug("matcher_active for %s is %s", privacy_group_id, config.matcher_active)
    return config.matcher_active

def get_privacy_group_matcher_active(privacy_group_id: str) -> bool:
    """
    Is this privacy group's matcher_active flag true? Entries in the internal
    cache are cleared every PG_CONFIG_CACHE_TIME_SECONDS seconds.

    Impl: the // is python's integer division operator. Threw me off. :)
    """
    return _get_privacy_group_matcher_active(
        privacy_group_id, time.time() // PG_CONFIG_CACHE_TIME_SECONDS
    )

@functools.lru_cache(maxsize=None)
def _get_all_matcher_active_privacy_groups(cache_buster) -> t.List[str]:
    configs = ThreatExchangeConfig.get_all()
    return list(
        map(
            lambda c: c.name,
            filter(
                lambda c: c.matcher_active,
                configs,
            ),
        )
    )

@functools.lru_cache(maxsize=None)
def _get_max_pdq_threshold_for_active_matcher_privacy_groups(
    cache_buster,
) -> int:
    active_pg_names = _get_all_matcher_active_privacy_groups(cache_buster)
    if not active_pg_names:
        return 0
    active_pdq_thresholds = [
        config.pdq_match_threshold
        for config in AdditionalMatchSettingsConfig.get_all()
        if config.name in active_pg_names
    ]
    if active_pdq_thresholds:
        return max(active_pdq_thresholds)
    # no custom threshold set for active privacy_groups
    return 0

def get_max_threshold_of_active_privacy_groups_for_signal_type(
    signal_type: t.Type[SignalType],
) -> int:
    if signal_type == PdqSignal:
        return _get_max_pdq_threshold_for_active_matcher_privacy_groups(
            time.time() // PG_CONFIG_CACHE_TIME_SECONDS
        )
    else:
        return 0

@functools.lru_cache(maxsize=128)
def _get_privacy_group_matcher_pdq_threshold(
    privacy_group_id: str, cache_buster
) -> int:
    config = AdditionalMatchSettingsConfig.get(privacy_group_id)
    if not config:
        logger.debug(
            "Privacy group %s does not have custom pdq_match_threshold. Using default defined in PDQ_CONFIDENT_MATCH_THRESHOLD",
            privacy_group_id,
        )
        return PdqSignal.PDQ_CONFIDENT_MATCH_THRESHOLD

    logger.debug(
        "pdq_match_threshold for %s is %s", privacy_group_id, config.pdq_match_threshold
    )
    return config.pdq_match_threshold

def get_privacy_group_matcher_pdq_threshold(
    privacy_group_id: str,
) -> int:
    """
    Does this privacy group's have a custom pdq threshold; if so what is it?
    otherwise return the default PDQ_CONFIDENT_MATCH_THRESHOLD.

    Entries in the internal cache are cleared
    every PG_CONFIG_CACHE_TIME_SECONDS seconds.

    ToDo this should be refactored into a signal angostic interface eventaully
        especially before we have another similarity based signal type in HMA

    Impl: the // is python's integer division operator. Threw me off. :)
    """
    return _get_privacy_group_matcher_pdq_threshold(
        privacy_group_id, time.time() // PG_CONFIG_CACHE_TIME_SECONDS
    )

class Matcher:
    """
    Match against any signal type defined on threatexchange and stored in s3.

    Once created, indexes used by this are cached on the index. Do not create
    multiple Matcher instances in the same python runtime for the same
    signal_type. This would take up more RAM than necessary.

    Indexes are pulled from S3 on first call for a signal_type.
    """

    def __init__(
        self,
        index_bucket_name: str,
        supported_signal_types: t.List[t.Type[SignalType]],
    ):
        self.index_bucket_name = index_bucket_name
        self.supported_signal_types = supported_signal_types
        self._cached_indexes: t.Dict[t.Type[SignalType], SignalTypeIndex] = {}

    def match(
        self, signal_type: t.Type[SignalType], signal_value: str
    ) -> t.List[IndexMatch[BaseIndexMetadata]]:
        """
        Returns MatchMessage which can be directly published to a queue.

        Note, this also filters out matches that are from datasets that have
        been de-activated.
        """
        index = self.get_index(signal_type)

        with metrics.timer(metrics.names.indexer.search_index):
            match_results: t.List[IndexMatch] = index.query(signal_value)

        if not match_results:
            # No matches found in the index
            return []

        return self.filter_match_results(match_results, signal_type)

    def filter_match_results(
        self, results: t.List[IndexMatch], signal_type: t.Type[SignalType]
    ) -> t.List[IndexMatch]:
        """
        For ThreatExchange, use the privacy group's matcher_active flag to
        filter out match results that should not be returned.

        If implementing a matcher for something other than threat exchange,
        consider extending this class and implementing your own.
        """

        # results is a list of match object references that live in any index
        # this method should not edit those objects directly as they could effect
        # subsequent calls made while the index is still in memory
        matches = results.copy()

        filtered_results = []
        for match in matches:
            match.metadata = [
                metadata_obj
                for metadata_obj in match.metadata
                if metadata_obj.get_source() == BANKS_SOURCE_SHORT_CODE
                or (
                    # If metadata obj is from threatexchange (one privacy group
                    # per metadata obj), check that it is active AND
                    # if the signal_type is PdqSignal that its
                    # distance is lesser than the optional pdq_threshold set for
                    # that pg.
                    metadata_obj.get_source() == THREAT_EXCHANGE_SOURCE_SHORT_CODE
                    and get_privacy_group_matcher_active(metadata_obj.privacy_group)
                    and (
                        signal_type != PdqSignal
                        or match.distance
                        <= get_privacy_group_matcher_pdq_threshold(
                            str(metadata_obj.privacy_group)
                        )
                    )
                )
            ]

            if len(match.metadata) != 0:
                filtered_results.append(match)

        return filtered_results

    def write_match_record_for_result(
        self,
        table: Table,
        signal_type: t.Type[SignalType],
        content_hash: str,
        content_id: str,
        match: IndexMatch[BaseIndexMetadata],
    ):
        """
        Write a match record to dynamodb. The content_id is not important to the
        matcher. So, the calling lambda is expected to pass on the content_id
        for match record calls.
        """
        # Write one record for TE and one for banks.. I am sure the logic can be
        # simplified. We can do filters instead of iterating with flags. But
        # umm, we can fix that later?
        bank_record_written = False
        te_record_written = False

        for metadata_obj in match.metadata:
            match_record_attributes = {
                "content_id": content_id,
                "signal_type": signal_type,
                "content_hash": content_hash,
                "updated_at": datetime.datetime.now(),
                "signal_source": metadata_obj.get_source(),
                "match_distance": int(match.distance),
            }

            if (
                metadata_obj.get_source() == THREAT_EXCHANGE_SOURCE_SHORT_CODE
                and not te_record_written
            ):
                match_record_attributes.update(
                    signal_id=metadata_obj.indicator_id,
                    signal_hash=metadata_obj.signal_value,
                )
                te_record_written = True

            elif (
                metadata_obj.get_source() == BANKS_SOURCE_SHORT_CODE
                and not bank_record_written
            ):
                match_record_attributes.update(
                    signal_id=metadata_obj.signal_id,
                    signal_hash=metadata_obj.signal_value,
                )
                bank_record_written = True

            MatchRecord(**match_record_attributes).write_to_table(table)

    @classmethod
    def write_signal_if_not_found(
        cls,
        table: Table,
        signal_type: t.Type[SignalType],
        match: IndexMatch,
    ):
        """
        Write the signal to the datastore. Only signals that have matched are
        written to the DB. The fetcher takes care of updating the signal with
        opinions or updates from the source.

        TODO: Move this out of matchers.

        This is not matcher specific functionality. Signals could benefit from
        their own store. Perhaps the API could be useful when building local
        banks. Who knows! :)
        """
        for signal in cls.get_te_metadata_objects_from_match(signal_type, match):
            if hasattr(signal, "write_to_table_if_not_found"):
                # only ThreatExchangeSignalMetadata has this method.
                # mypy not smart enough to auto cast.
                signal.write_to_table_if_not_found(table)  # type: ignore

    @classmethod
    def get_te_metadata_objects_from_match(
        cls,
        signal_type: t.Type[SignalType],
        match: IndexMatch[BaseIndexMetadata],
    ) -> t.List[ThreatExchangeSignalMetadata]:
        """
        See docstring of `write_signal_if_not_found` we will likely want to move
        this outside of Matcher. However while the MD5 expansion is still on going
        better to have it all in once place.
        Note: changes made here will have an effect on api.matches.get_match_for_hash
        """
        return [
            ThreatExchangeSignalMetadata(
                signal_id=str(match_object.indicator_id),
                privacy_group_id=match_object.privacy_group,
                updated_at=datetime.datetime.now(),
                signal_type=signal_type,
                signal_hash=match_object.signal_value,
                tags=list(match_object.tags),
            )
            for match_object in match.metadata
            if match_object.get_source() == THREAT_EXCHANGE_SOURCE_SHORT_CODE
        ]

    def get_index(self, signal_type: t.Type[SignalType]) -> SignalTypeIndex:
        """
        If cached, return an index instance for the signal_type. If not, build
        one, cache and return.
        """

        max_custom_threshold = (
            get_max_threshold_of_active_privacy_groups_for_signal_type(signal_type)
        )
        index_cls = self._get_index_for_signal_type_matching(
            signal_type, max_custom_threshold
        )

        # Check for signal_type in cache AND confirm said index class type is
        # still correct for the given [optional] max_custom_threshold
        if not signal_type in self._cached_indexes or not isinstance(
            self._cached_indexes[signal_type], index_cls
        ):
            with metrics.timer(metrics.names.indexer.download_index):
                self._cached_indexes[signal_type] = index_cls.load(
                    bucket_name=self.index_bucket_name
                )

        return self._cached_indexes[signal_type]

    @classmethod
    def _get_index_for_signal_type_matching(
        cls, signal_type: t.Type[SignalType], max_custom_threshold: int
    ):
        indexes = INDEX_MAPPING[signal_type]
        # disallow empty list
        assert indexes
        if len(indexes) == 1:
            # if we only have one option just return
            return indexes[0]

        indexes.sort(key=lambda i: i.get_index_max_distance())

        for index_cls in indexes:
            if max_custom_threshold <= index_cls.get_index_max_distance():
                return index_cls

        # if we don't have an index that supports max threshold
        # just return the one if the highest possible max distance
        return indexes[-1]

    def publish_match_message(
        self,
        content_id: str,
        content_hash: str,
        matches: t.List[IndexMatch],
        sns_client: SNSClient,
        topic_arn: str,
    ):
        """
        Creates banked signal objects and publishes one message for a list of
        matches to SNS.
        """
        banked_signals = []

        for match in matches:
            for metadata_obj in match.metadata:
                if metadata_obj.get_source() == THREAT_EXCHANGE_SOURCE_SHORT_CODE:
                    banked_signal = BankedSignal(
                        str(metadata_obj.indicator_id),
                        str(metadata_obj.privacy_group),
                        str(metadata_obj.get_source()),
                    )
                    for tag in metadata_obj.tags:
                        banked_signal.add_classification(tag)

                    banked_signals.append(banked_signal)
                elif metadata_obj.get_source() == BANKS_SOURCE_SHORT_CODE:
                    banked_signal = BankedSignal(
                        metadata_obj.signal_id,
                        metadata_obj.bank_member_id,
                        metadata_obj.get_source(),
                    )

                    banked_signals.append(banked_signal)

        match_message = MatchMessage(
            content_key=content_id,
            content_hash=content_hash,
            matching_banked_signals=banked_signals,
        )

        sns_client.publish(TopicArn=topic_arn, Message=match_message.to_aws_json())


# -*- coding: utf-8 -*-
"""

Two part script:
### 1. Calculates the proportion of mice that had graduated to stage 2 on each day of the experiment
### 2. Plots graduation day vs expression graphs

"""

# import packages and functions
from Functions_CoreFunctions_0100 import adjust_spines, makelegend2
import numpy as np
import matplotlib.pyplot as plt
import math
from scipy import stats

# Load data containing day of graduation to stage 2 for each mouse
marray = np.loadtxt('Data_Input/Behaviour_SummaryData/Stagegraduation_0100.txt',delimiter = '\t')

# specify mouse/mice and day/s to analyse
days = ['Day' + str(int(x)) for x in np.arange(1,22.1)]
mice = ['M' + str(int(x)) for x in np.arange(1,16.1)]# choose specific day/s
tetmice = ['M' + str(int(x)) for x in [2,5,6,7,8,9,11]]
conmice = ['M' + str(int(x)) for x in [1,3,4,10]]

#plot legend
def makelegend3(fig,ax):
    handles, labels = ax.get_legend_handles_labels()
    leg = fig.legend(handles,labels, loc="right", bbox_to_anchor=(0.85, 0.58), fontsize = "small", markerscale = 1)
    for l in leg.get_lines():l.set_linewidth(2)
    frame  = leg.get_frame()
    frame.set_edgecolor('w')
    frame.set_alpha(0.2)

# calculate proportion of mice that have graduated for each training day
proportionday_store = np.zeros((len(days), 4))
n_groups = 2

for daycount,day in enumerate(marray,0):
    allmice = marray[daycount,:]
    tetmice = np.hstack((marray[daycount,1],marray[daycount,4],marray[daycount,5],marray[daycount,6],marray[daycount,7],marray[daycount,8],marray[daycount,10],marray[daycount,11],marray[daycount,12],marray[daycount,13]))
    conmice = np.hstack((marray[daycount,0],marray[daycount,2],marray[daycount,3],marray[daycount,9],marray[daycount,14],marray[daycount,15]))
    htetmice = np.hstack((marray[daycount,1],marray[daycount,5],marray[daycount,6],marray[daycount,8]))
    ltetmice = np.hstack((marray[daycount,4],marray[daycount,7],marray[daycount,10],marray[daycount,11],marray[daycount,12],marray[daycount,13]))
    p_tet = (np.nansum(tetmice))/10
    p_htet = (np.nansum(htetmice))/4
    p_ltet = (np.nansum(ltetmice))/6
    p_con = (np.nansum(conmice))/6
    proportionday_store[daycount,0] = p_tet
    proportionday_store[daycount,1] = p_con
    proportionday_store[daycount,2] = p_htet
    proportionday_store[daycount,3] = p_ltet

# plot graph of proportion of graduation vs training day

daybins = np.arange(0,21.1,1) #number of days array
bar_width = 0.4
index = np.arange(n_groups)

"""
# plot proportion for GFP and TeLC mice

fig = plt.figure(figsize = (5,3))
ax = fig.add_subplot(111)
ax.plot(daybins,proportionday_store[:,0], color = '0.3', label = 'AAV-FLEX-TeLC', linewidth = 1.5)
ax.plot(daybins,proportionday_store[:,1], color = 'PaleVioletRed', label = 'AAV-FLEX-GFP',linewidth = 1.5)
makelegend3(fig,ax)
adjust_spines(ax, ['left','bottom'])
ax.tick_params(axis='x', pad = 10, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =12)
ax.tick_params(axis='y', pad = 10, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =12)
ax.set_ylabel('Proportion', fontsize=13, labelpad = 16)
ax.set_xlabel('Training day', fontsize=14, labelpad = 16)
ax.set_ylim(-0.2,1.05)
ax.set_xlim(1,20)
ax.axvline(1,linewidth=3, color="black")
ax.axhline(-0.2,linewidth=3, color="black")
plt.yticks(index, ('0','1'))
plt.yticks(rotation=70)
ax = plt.locator_params(axis = 'x', nbins = 4)
plt.subplots_adjust(hspace = 1, wspace = 0.5,  bottom = 0.3, left = 0.15, right = 0.72, top = .85)
fig.savefig('Plots/Figure5/Proportion_graduation_0200' +'.png', dpi = 200)
plt.close()
"""

# plot proportion for GFP and lTeLC and hTeLC mice

fig = plt.figure(figsize = (5,4))
ax = fig.add_subplot(111)
ax.plot(daybins,proportionday_store[:,3], color = 'blue', label = 'lTeLC', linewidth = 1.5)
ax.plot(daybins,proportionday_store[:,2], color = 'red', label = 'hTeLC', linewidth = 1.5)
ax.plot(daybins,proportionday_store[:,1], color = 'Black', label = 'GFP',linewidth = 1.5)
makelegend3(fig,ax)
adjust_spines(ax, ['left','bottom'])
ax.tick_params(axis='x', pad = 10, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =12)
ax.tick_params(axis='y', pad = 10, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =12)
ax.set_ylabel('Proportion', fontsize=13, labelpad = 16)
ax.set_xlabel('Training day', fontsize=14, labelpad = 16)
ax.set_ylim(-0.15,1.15)
ax.set_xlim(0,21)
ax.axvline(0,linewidth=3, color="black")
ax.axhline(-0.15,linewidth=3, color="black")
plt.yticks(index, ('0','1'))
plt.yticks(rotation=70)
ax = plt.locator_params(axis = 'x', nbins = 4)
plt.subplots_adjust(hspace = 1, wspace = 0.5,  bottom = 0.3, left = 0.15, right = 0.72, top = .85)
fig.savefig('Plots/Figure5/Proportion_graduation_0200' +'.png', dpi = 200)
plt.close()

# Store data

genotype = ("TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","GFP","GFP","GFP" ,"GFP" ,"GFP" ,"GFP" )
mouse = np.arange(1,16.1,1)
days = np.arange(1,22.1,1)

lTeLC = proportionday_store[:,3]
hTeLC = proportionday_store[:,2]
GFP = proportionday_store[:,1]

data = np.vstack((days, lTeLC, hTeLC, GFP))
data =  np.transpose(data)

np.savetxt('Data_Output/Figure5/Figure5_A_right_0100.csv', data,fmt =  '%s', delimiter = ',', header = 'Day, lTeLC, hTeLC, GFP ')

## Graduation days for all mice (taken from spreadsheet for data input reasons)

m1 = 7
m2 = 20
m3 = 9
m4 = 7
m5 = 15
m6 = 20
m7 = 20
m8 = 7
m9 = 20
m10 = 7
m11 = 8
m13 = 12
m14 = 11
m15 = 10
m16 = 7
m17 = 7

# stack days for tetanus and control mice
tet = np.hstack((m2,m5,m6,m7,m8,m9,m11,m13,m14,m15))
con = np.hstack((m1,m3,m4,m10,m16,m17))

htet = np.vstack((m2,m6,m7,m9))
ltet = np.vstack((m5,m8,m11,m13,m14,m15))

# average number of mice graduated per day
avg_tet = np.nanmean(tet)
sd_tet = np.nanstd(tet)/math.sqrt(10)
avg_con = np.nanmean(con)
sd_con = np.nanstd(con)/math.sqrt(6)

tt_total = np.hstack((avg_con,avg_tet))
tt_total_sd = np.hstack((sd_con,sd_tet))

hTeLC = np.nanmean(np.hstack((m2,m6,m7,m9)))
lTeLC = np.nanmean(np.hstack((m5,m8,m11,m13,m14,m15)))

# Plot vs expression

mice = [str(int(x)) for x in np.arange(1,16.1)]
expression = np.loadtxt('Data_Input/ExpressionQuantification/T15_FinalQuantification_0300.txt')

expressionstore = np.zeros((len(mice), 3))
expressionstore_d = np.zeros((len(mice), 3))
expressionstore_v = np.zeros((len(mice), 3))
sdexpressionstore = np.zeros((len(mice), 3))
sdexpressionstore_d = np.zeros((len(mice), 3))
sdexpressionstore_v = np.zeros((len(mice), 3))

for mcount,mouse in enumerate(mice):
    marraybym = expression[expression[:,0]==mcount+1,:] # get data for each mouse
    rh = marraybym[:,1:4]
    lh = marraybym[:,4:]
    avg_RH = np.nanmean(rh[:,1])
    sd_RH = stats.sem(rh[:,1])
    avg_LH = np.nanmean(lh[:,1])
    sd_LH =stats.sem(lh[:,1])
    avg = np.nanmean(np.hstack((avg_RH,avg_LH)))
    sd = stats.sem(np.hstack((avg_RH,avg_LH)))
    expressionstore[mcount, 0] = avg
    expressionstore[mcount, 1] = avg_RH
    expressionstore[mcount, 2] = avg_LH
    sdexpressionstore[mcount, 0] = sd
    sdexpressionstore[mcount, 1] = sd_RH
    sdexpressionstore[mcount, 2] = sd_LH

    #for dorsal, mid and ventral
    rh_d = np.delete(rh, np.where(rh[:,2] != 1),0)
    rh_v = np.delete(rh, np.where(rh[:,2] != 2),0)
    lh_d = np.delete(lh, np.where(lh[:,2] != 1),0)
    lh_v = np.delete(lh, np.where(lh[:,2] != 2),0)

    avg_RH_d = np.nanmean(rh_d[:,1])
    avg_LH_d = np.nanmean(lh_d[:,1])
    avg_d = np.nanmean(np.hstack((avg_RH_d,avg_LH_d)))
    avg_RH_v = np.nanmean(rh_v[:,1])
    avg_LH_v = np.nanmean(lh_v[:,1])
    avg_v = np.nanmean(np.hstack((avg_RH_v,avg_LH_v)))

    sd_RH_d = stats.sem(rh_d[:,1])
    sd_LH_d = stats.sem(lh_d[:,1])
    sd_d = stats.sem(np.hstack((avg_RH_d,avg_LH_d)))
    sd_RH_v = stats.sem(rh_v[:,1])
    sd_LH_v = stats.sem(lh_v[:,1])
    sd_v = stats.sem(np.hstack((avg_RH_v,avg_LH_v)))

    expressionstore_d[mcount, 0] = avg_d
    expressionstore_d[mcount, 1] = avg_RH_d
    expressionstore_d[mcount, 2] = avg_LH_d
    expressionstore_v[mcount, 0] = avg_v
    expressionstore_v[mcount, 1] = avg_RH_v
    expressionstore_v[mcount, 2] = avg_LH_v

    sdexpressionstore_d[mcount, 0] = sd_d
    sdexpressionstore_d[mcount, 1] = sd_RH_d
    sdexpressionstore_d[mcount, 2] = sd_LH_d
    sdexpressionstore_v[mcount, 0] = sd_v
    sdexpressionstore_v[mcount, 1] = sd_RH_v
    sdexpressionstore_v[mcount, 2] = sd_LH_v

    mcount +=1

# Average expression

total = np.hstack((expressionstore[1,1],expressionstore[4,1],expressionstore[5,1],expressionstore[6,1],expressionstore[7,1],expressionstore[8,1],expressionstore[10,1],expressionstore[11,1],expressionstore[12,1],expressionstore[13,1]))
totalcon = np.hstack((expressionstore[0,1],expressionstore[2,1],expressionstore[3,1],expressionstore[9,1],expressionstore[14,1],expressionstore[15,1]))
total_probe = np.hstack((expressionstore[4,1],expressionstore[7,1],expressionstore[10,1],expressionstore[11,1],expressionstore[12,1],expressionstore[13,1]))
dorsal = np.hstack((expressionstore_d[1,1],expressionstore_d[4,1],expressionstore_d[5,1],expressionstore_d[6,1],expressionstore_d[7,1],expressionstore_d[8,1],expressionstore_d[10,1],expressionstore_d[11,1],expressionstore_d[12,1],expressionstore_d[13,1]))
dorsalcon = np.hstack((expressionstore_d[0,1],expressionstore_d[2,1],expressionstore_d[3,1],expressionstore_d[9,1],expressionstore_d[14,1],expressionstore_d[15,1]))
dorsal_probe = np.hstack((expressionstore_d[4,1],expressionstore_d[7,1],expressionstore_d[10,1],expressionstore_d[11,1],expressionstore_d[12,1],expressionstore_d[13,1]))
ventral = np.hstack((expressionstore_v[1,1],expressionstore_v[4,1],expressionstore_v[5,1],expressionstore_v[6,1],expressionstore_v[7,1],expressionstore_v[8,1],expressionstore_v[10,1],expressionstore_v[11,1],expressionstore_v[12,1],expressionstore_v[13,1]))
ventralcon = np.hstack((expressionstore_v[0,1],expressionstore_v[2,1],expressionstore_v[3,1],expressionstore_v[9,1],expressionstore_v[14,1],expressionstore_v[15,1]))
sdtotal = np.hstack((sdexpressionstore[1,1],sdexpressionstore[4,1],sdexpressionstore[5,1],sdexpressionstore[6,1],sdexpressionstore[7,1],sdexpressionstore[8,1],sdexpressionstore[10,1],sdexpressionstore[11,1],sdexpressionstore[12,1],sdexpressionstore[13,1]))
sdtotalcon = np.hstack((sdexpressionstore[0,1],sdexpressionstore[2,1],sdexpressionstore[3,1],sdexpressionstore[9,1],sdexpressionstore[14,1],sdexpressionstore[15,1]))
sdtotal_probe = np.hstack((sdexpressionstore[4,1],sdexpressionstore[7,1],sdexpressionstore[10,1],sdexpressionstore[11,1],sdexpressionstore[12,1],sdexpressionstore[13,1]))
sddorsal = np.hstack((sdexpressionstore_d[1,1],sdexpressionstore_d[4,1],sdexpressionstore_d[5,1],sdexpressionstore_d[6,1],sdexpressionstore_d[7,1],sdexpressionstore_d[8,1],sdexpressionstore_d[10,1],sdexpressionstore_d[11,1],sdexpressionstore_d[12,1],sdexpressionstore_d[13,1]))
sddorsalcon = np.hstack((sdexpressionstore_d[0,1],sdexpressionstore_d[2,1],sdexpressionstore_d[3,1],sdexpressionstore_d[9,1],sdexpressionstore_d[14,1],sdexpressionstore_d[15,1]))
sdventral = np.hstack((sdexpressionstore_v[1,1],sdexpressionstore_v[4,1],sdexpressionstore_v[5,1],sdexpressionstore_v[6,1],sdexpressionstore_v[7,1],sdexpressionstore_v[8,1],sdexpressionstore_v[10,1],sdexpressionstore_v[11,1],sdexpressionstore_v[12,1],sdexpressionstore_v[13,1]))
sdventralcon = np.hstack((sdexpressionstore_v[0,1],sdexpressionstore_v[2,1],sdexpressionstore_v[3,1],sdexpressionstore_v[9,1],sdexpressionstore_v[14,1],sdexpressionstore_v[15,1]))

# plot graduation day vs dorsal expression

slope,intercept,r_value, p_value, std_err = stats.linregress(dorsal,tet)
ablinevalues = []
for i in dorsal:
    ablinevalues.append(slope*i+intercept)

slope,intercept,r_value, p_value, std_err = stats.linregress(dorsalcon,con)
ablinevalues1 = []
for i in dorsalcon:
    ablinevalues1.append(slope*i+intercept)
print(dorsal)
fig = plt.figure(figsize = (5,4))
ax = fig.add_subplot(111)
ax.plot(dorsal, tet, 'o',color = 'red',markersize = 5,label = '   Low')
ax.errorbar(dorsal, tet, xerr=sddorsal, fmt='o', markersize = 5,color = 'red',ecolor = 'red')
ax.plot(dorsal,ablinevalues, '-',color = 'red', linewidth = 1)
ax.plot(dorsalcon, con, 'o',color = 'Black',markersize = 5,label = '   Low')
ax.errorbar(dorsalcon, con, xerr=sddorsalcon, fmt='o', markersize = 5,color = 'Black',ecolor = 'Black')
ax.plot(dorsalcon,ablinevalues1, '-',color = 'Black', linewidth = 1)
ax.set_xlim(-50,2000)
ax.set_ylim(-1,21)
adjust_spines(ax, ['left','bottom'])
ax.tick_params(axis='x', pad = 7, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =14)
ax.tick_params(axis='y', pad = 7, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =14)
plt.xticks(rotation=70)
plt.locator_params(axis = 'y', nbins  = 5)
plt.locator_params(axis = 'x', nbins  = 5)#plt.yticks(n_groups,('Ventral','','','Dorsal'))
ax.axhline(-1,linewidth=3, color="black")
ax.axvline(-50,linewidth=3, color="black")
ax.set_xlabel('Flourescence', fontsize=16, labelpad = 20)
ax.set_ylabel('G day', fontsize=16, labelpad = 20)
plt.subplots_adjust(hspace = 1, wspace = 0.5,  bottom = 0.3, left = 0.15, right = 0.72, top = .85)
fig.savefig('Plots/Figure5/GraduationVDorsalExpression' +'.png', dpi = 200)
plt.close()

# SAVE DATA FOR R

genotype = ("TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","GFP","GFP","GFP" ,"GFP" ,"GFP" ,"GFP" )
mouse = np.arange(1,16.1,1)
fluorescence = np.hstack((dorsal,dorsalcon))
gday = np.hstack((tet,con))
data = np.vstack((mouse,genotype,fluorescence,gday))
data = np.transpose(data)
np.savetxt('Data_Output/Graduation_dorsal.csv', data,delimiter = '\t', header = 'Mouse\tGenotype\tFlourescence\tGraduation', comments='',fmt='%s')

dfluorescence = np.hstack((dorsal,dorsalcon))
vfluorescence = np.hstack((ventral,ventralcon))

# SAVE DATA FOR MANUSCIPT

data = np.vstack((genotype, gday, dfluorescence, vfluorescence)); data = np.transpose(data)
np.savetxt('Data_Output/Figure5/Figure5_A_left_0100.csv', data,fmt =  '%s', delimiter = ',', header = 'Genotype, Graduation day, Dorsal Fluorescence, Ventral Fluorescence')


"""Geodesic Active Contours and Chan-Vese "Active Contours Without Edges"

Implementation based on morphological variant of these algorithms per:
Marquez-Neila P, Baumela L, & Alvarez L. (2014).
A morphological approach to curvature-based evolution of curves and surfaces.
IEEE Transactions on Pattern Analysis and Machine Intelligence, 36(1).

This is a much-optimized version of the demo code available here:
https://github.com/pmneila/morphsnakes
In particular, a "narrow-band" approach is used whereby only pixels at the
edges of the mask are examined / updated. This speeds the code up by at least
several orders of magnitude compared to the naive approach.

Example GAC usage:
    image # image containing a structure to segment
    bounds # boolean mask with False for any image regions from which the
           # segmented structure must be excluded. (Optional.)
    initial # boolean mask with initial segmentation

    edges = ndimage.gaussian_gradient_magnitude(image, sigma)
    strong_edges = edges > threshold
    edge_gradient = numpy.gradient(edges)
    balloon_direction = -1
    # negative balloon_direction values means "deflate" the initial mask toward
    # edges except where strong_edges == True. (NB: balloon_direction can also
    # be a mask of -1, 0, and 1 values to allow spatially-varying balloon forces.)
    gac = GACMorphology(mask=initial, advection_direction=edge_gradient,
        advection_mask=strong_edges, balloon_direction=balloon_direction,
        max_region_mask=bounds)

    stopper = StoppingCondition(gac, max_iterations=100)
    while stopper.should_continue():
        gac.balloon_force() # apply balloon force
        gac.advect() # move region edges in advection_direction
        gac.smooth() # smooth region boundaries.

Obviously, different schedules of the balloon force, advection, and smoothing
can be applied. To aid in this, each of the methods above take an 'iters'
parameter to apply that many iterations of that step in a row. Also, initial
warm-up rounds with just advection or advection and smoothing may be helpful.

The smooth() method takes a 'depth' parameter that controls the spatial
smoothness of the curve. With depth=1, only very small jaggies are smoothed, but
with larger values, the curve is smoothed along larger spatial scales (rule of
thumb: min radius of curvature after smoothing is on the order of depth*3 or so).

Another useful way to calculate the advection_direction parameter, instead of
from the edge gradient (as above) is to use the edge_direction() function, which
takes a thresholded edge mask (Canny edges work great), and returns the distance
from each pixel to the nearest edge, as well as the direction from each pixel
to the nearest edge. This latter can be passed directly as the advection_direction
parameter, which will dramatically increase the capture radius for distant edges,
so a balloon_force parameter may not be needed.

Last, for Canny edge images especially, consider the EdgeClaimingAdvection
active contour class. (See the documentation for that class.)

Example ACWE usage:
    acwe = ACWEMorphology(mask=initial, image=image, max_region_mask=bounds)

    for i in range(iterations):
        acwe.acwe_step() # make inside and outside pixel means different.
        acwe.smooth() # smooth region boundaries

There is also a BinnedACWE class that uses a histogram of the image for the
inside and outside regions, rather than just means. This can be useful if
the region to segment has very distinct brightness values from the background,
but is not overall brighter or darker on average. Note that it is best to
reduce the input image to a small number of brightness values (16-64) before
using this class. The discretize_image() function is useful for this.

Last: the ActiveContour class allows for both GAC and ACWE steps, if that's
useful. BinnedActiveContour, ActiveClaimingContour, and BinnedActiveClaimingContour
similarly combine [Binned]ACWE and [EdgeClaiming]Advection functionality.
"""

import numpy
import itertools
import collections
from scipy import ndimage

from . import neighborhood

class MaskNarrowBand:
    """Track the inside and outside edge of a masked region, while allowing
    pixels from the inside edge to be moved outside and vice-versa.

    Base-class for fast morphological operations for region growing, shrinking,
    and reshaping.
    """
    S = ndimage.generate_binary_structure(2, 2)

    def __init__(self, mask, max_region_mask=None):
        mask = mask.astype(bool)
        if max_region_mask is not None:
            mask = mask & max_region_mask
        self.max_region_mask = max_region_mask
        self.mask_neighborhood = neighborhood.make_neighborhood_view(mask > 0,
            pad_mode='constant', constant_values=0) # shape = image.shape + (3, 3)
        # make self.mask identical to mask, but be a view on the center pixels of mask_neighborhood
        self.mask = self.mask_neighborhood[:,:,1,1]
        self.indices = numpy.dstack(numpy.indices(mask.shape)).astype(numpy.int32) # shape = mask.shape + (2,)
        self.index_neighborhood = neighborhood.make_neighborhood_view(self.indices,
            pad_mode='constant', constant_values=-1) # shape = mask.shape + (3, 3, 2)
        inside_border_mask = mask ^ ndimage.binary_erosion(mask, self.S) # all True pixels with a False neighbor
        outside_border_mask = mask ^ ndimage.binary_dilation(mask, self.S) # all False pixels with a True neighbor
        self.inside_border_indices = self.indices[inside_border_mask] # shape = (inside_border_mask.sum(), 2)
        self.outside_border_indices = self.indices[outside_border_mask] # shape = (outside_border_mask.sum(), 2)
        # NB: to index a numpy array with these indices, must turn shape(num_indices, 2) array
        # into tuple of two num_indices-length arrays, a la:
        # self.mask[tuple(self.outside_border_indices.T)]
        self.changed = 0

    def _assert_invariants(self):
        """Test whether the border masks and indices are in sync, and whether the
        borders are correct for the current mask."""
        inside_border_mask = self.mask ^ ndimage.binary_erosion(self.mask, self.S)
        outside_border_mask = self.mask ^ ndimage.binary_dilation(self.mask, self.S)
        assert len(self.inside_border_indices) == inside_border_mask.sum()
        assert inside_border_mask[tuple(self.inside_border_indices.T)].sum() == len(self.inside_border_indices)
        assert len(self.outside_border_indices) == outside_border_mask.sum()
        assert outside_border_mask[tuple(self.outside_border_indices.T)].sum() == len(self.outside_border_indices)

    def _move_to_outside(self, to_outside):
        """to_outside must be a boolean mask on the inside border pixels
        (in the order defined by inside_border_indices)."""
        self.inside_border_indices, self.outside_border_indices, added_idx, removed_indices = self._change_pixels(
            to_change=to_outside,
            old_border_indices=self.inside_border_indices,
            new_value=False,
            new_border_indices=self.outside_border_indices,
            some_match_old_value=numpy.any
        )
        return added_idx, removed_indices

    def _move_to_inside(self, to_inside):
        """to_inside must be a boolean mask on the outside border pixels
        (in the order defined by outside_border_indices)."""
        self.outside_border_indices, self.inside_border_indices, added_idx, removed_indices = self._change_pixels(
            to_change=to_inside,
            old_border_indices=self.outside_border_indices,
            new_value=True,
            new_border_indices=self.inside_border_indices,
            some_match_old_value=_not_all
        )
        return added_idx, removed_indices

    def _change_pixels(self, to_change, old_border_indices, new_value, new_border_indices, some_match_old_value):
        change_indices = old_border_indices[to_change]
        # prevent changes outside of the max_region_mask
        if new_value == True and self.max_region_mask is not None:
            in_region_mask = self.max_region_mask[tuple(change_indices.T)]
            change_indices = change_indices[in_region_mask]
            to_change[to_change] = in_region_mask
        change_idx = tuple(change_indices.T)
        if len(change_indices) == 0:
            return old_border_indices, new_border_indices, change_idx, change_indices
        self.changed += len(change_indices)
        # Find out which neighbors of changing pixels have the new value.
        # If we did the below after changing the mask, we would also pick up the
        # center pixels, which have the new value.
        if new_value == True:
            new_valued_neighbors = self.mask_neighborhood[change_idx]
        else:
            new_valued_neighbors = ~self.mask_neighborhood[change_idx]
        # Now, update the mask and border indices
        # (1) Update changed pixels in the mask, and remove them from old_border_indices.
        self.mask[change_idx] = new_value
        old_border_indices = old_border_indices[~to_change]

        # (2) add old-valued neighbors of newly-changed pixels to old_border_indices,
        # and then make sure the indices don't contain duplicates.
        # (Duplicates appear both because the indices might already be in the list,
        # or because the neighborhoods overlap, so old-valued neighbors might be
        # mulitply identified from several changed pixels.)
        changed_neighborhood = self.mask_neighborhood[change_idx]
        changed_neighborhood_indices = self.index_neighborhood[change_idx]
        # Find out which neighbors of changing pixels have the old value.
        # If we did the below before changing the mask, we would also pick up the
        # center pixels, which had the old value
        if new_value == True:
            old_valued_neighbors = ~changed_neighborhood
        else:
            old_valued_neighbors = changed_neighborhood
        old_valued_neighbor_indices = changed_neighborhood_indices[old_valued_neighbors]
        if new_value == True:
            # Exclude neighbors that are actually out-of-bounds "padding" pixels.
            # Only relevant for changing values to True (i.e. moving pixels inside)
            # because the out-of-bounds mask area is "False" and will thus otherwise
            # get picked up.
            # Out of bounds indices are -1 in the index_neighborhood array.
            good_indices = ~(old_valued_neighbor_indices == -1).any(axis=1)
            old_valued_neighbor_indices = old_valued_neighbor_indices[good_indices]
        # NB: many of the old_valued_neighbors are already in the old_border_indices...
        # If we kept a mask of the old_border pixels, we could look these up and
        # exclude them, which would make _unique_indices() below a bit faster. However,
        # that requires a lot of bookkeeping, so it doesn't speed things up in every case.
        old_border_indices = numpy.concatenate([old_border_indices, old_valued_neighbor_indices])
        old_border_indices = _unique_indices(old_border_indices)

        # (3) Remove all pixels from new_border_indices that no longer have any
        # old-valued neighbors left.
        # Such pixels must be a new-valued neighbor of one of the pixels
        # that changed to the new value. We know that these pixels are necessarily
        # in the new_border already because they are next to a pixel that changed.
        new_valued_neighbors_indices = changed_neighborhood_indices[new_valued_neighbors]
        # need to unique-ify indices because neighborhoods overlap and may pick up the same pixels
        new_valued_neighbors_indices = _unique_indices(new_valued_neighbors_indices)
        neighbors_of_new_valued_neighbors = self.mask_neighborhood[tuple(new_valued_neighbors_indices.T)]
        no_old_valued_neighbors = ~some_match_old_value(neighbors_of_new_valued_neighbors, axis=(1,2))
        remove_from_new_border_indices = new_valued_neighbors_indices[no_old_valued_neighbors]
        new_border_indices = _diff_indices(new_border_indices, remove_from_new_border_indices)

        # (4) Add newly-changed pixels to new_border_indices if they have an old-valued neighbor.
        changed_with_old_neighbor = some_match_old_value(changed_neighborhood, axis=(1,2))
        add_to_new_border_indices = change_indices[changed_with_old_neighbor]
        new_border_indices = numpy.concatenate([new_border_indices, add_to_new_border_indices])
        return old_border_indices, new_border_indices, change_idx, remove_from_new_border_indices

class CurvatureMorphology(MaskNarrowBand):
    """Implement basic erosion, dilation, and curvature-smoothing morphology
    steps (the latter from Marquez-Neila et al.) using a fast narrow-band approach.
    Base class for more sophisticated region-modifying steps: main function of interest
    is smooth().

    smooth_mask, if not None, is region where smoothing may be applied.
    """
    def __init__(self, mask, smooth_mask=None, max_region_mask=None):
        super().__init__(mask, max_region_mask)
        self.smooth_mask = smooth_mask
        self._reset_smoothing()

    def _reset_smoothing(self):
        self._smooth_funcs = itertools.cycle([self._SIoIS, self._ISoSI])

    def dilate(self, iters=1):
        for _ in range(iters):
            border_mask = numpy.ones(len(self.outside_border_indices), dtype=bool)
            self._move_to_inside(border_mask)

    def erode(self, iters=1):
        for _ in range(iters):
            border_mask = numpy.ones(len(self.inside_border_indices), dtype=bool)
            self._move_to_outside(border_mask)

    def smooth(self, iters=1, depth=1):
        """Apply 'iters' iterations of edge-curvature smoothing.
        'depth' controls the spatial scale of the smoothing. With depth=1, only
        the highest-frequency edges get smoothed out. Larger depth values smooth
        lower-frequency structures."""
        for _ in range(iters):
            smoother = next(self._smooth_funcs)
            smoother(depth)

    def _SI(self):
        idx, idx_mask = _masked_idx(self.smooth_mask, self.inside_border_indices)
        inside_border = self.mask_neighborhood[idx]
        on_a_line = ((inside_border[:,0,0] & inside_border[:,2,2]) |
                     (inside_border[:,1,0] & inside_border[:,1,2]) |
                     (inside_border[:,0,1] & inside_border[:,2,1]) |
                     (inside_border[:,2,0] & inside_border[:,0,2]))
        self._move_to_outside(_unmask_idx(idx_mask, ~on_a_line))

    def _IS(self):
        idx, idx_mask = _masked_idx(self.smooth_mask, self.outside_border_indices)
        outside_border = ~self.mask_neighborhood[idx]
        on_a_line = ((outside_border[:,0,0] & outside_border[:,2,2]) |
                     (outside_border[:,1,0] & outside_border[:,1,2]) |
                     (outside_border[:,0,1] & outside_border[:,2,1]) |
                     (outside_border[:,2,0] & outside_border[:,0,2]))
        self._move_to_inside(_unmask_idx(idx_mask, ~on_a_line))

    def _SIoIS(self, depth=1):
        for i in range(depth):
            self._IS()
        for i in range(depth):
            self._SI()

    def _ISoSI(self, depth=1):
        for i in range(depth):
            self._SI()
        for i in range(depth):
            self._IS()

class ACWE(CurvatureMorphology):
    def __init__(self, mask, image, acwe_mask=None, inside_bias=1,
        smooth_mask=None, max_region_mask=None):
        """Class for Active Contours Without Edges region-growing.

        Relevant methods for region-growing are smooth() and acwe_step().

        Parameters:
            mask: mask containing the initial state of the region to evolve
            image: ndarray of same shape as mask containing image values. The
                difference in mean image value inside and outside the region will
                be maximized by acwe_step()
            acwe_mask: region where ACWE updates will be applied. (Inside /
                outside values will still be computed outside this region.)
            inside_bias: weight for comparing means of inside vs. outside pixels.
                Generally 1 works properly. Values < 1 make it "easier" to add
                pixels to the region, and values > 1 make it "easier" to remove
                pixels from the region. (Simplification of the "lambda"
                parameters from the ACWE literature.)
            smooth_mask: region in which smoothing may be applied.
            max_region_mask: mask beyond which the region may not grow. If
                provided, this mask will also represent the pixels over which
                the ACWE inside/outside means are computed.
        """
        super().__init__(mask, smooth_mask, max_region_mask)
        # do work in _setup rather than __init__ to allow for complex multiple
        # inheritance from this class that super() alone can't handle. See
        # ActiveContour class.
        self._setup(image, acwe_mask, inside_bias)

    def _setup(self, image, acwe_mask, inside_bias):
        self.image = image
        assert self.image.shape == self.mask.shape
        self.acwe_mask = acwe_mask
        if acwe_mask is not None:
            assert self.image.shape == self.acwe_mask.shape
        self.inside_bias = inside_bias
        # note: self.mask is clipped to self.max_region_mask so the below works.
        self.inside_count = self.mask.sum()
        self.outside_count = numpy.product(self.mask[self.max_region_mask].shape) - self.inside_count
        self.inside_sum = self.image[self.mask].sum()
        self.outside_sum = self.image[self.max_region_mask].sum() - self.inside_sum

    def _assert_invariants(self):
        super()._assert_invariants()
        assert self.inside_count == self.mask.sum()
        assert self.outside_count == numpy.product(self.mask[self.max_region_mask].shape) - self.inside_count
        assert self.inside_sum == self.image[self.mask].sum()
        assert self.outside_sum == self.image[self.max_region_mask].sum() - self.inside_sum
        assert numpy.allclose(self.inside_sum/self.inside_count, self.image[self.mask].mean())
        if self.max_region_mask is None:
            assert numpy.allclose(self.outside_sum/self.outside_count, self.image[~self.mask].mean())
        else:
            assert numpy.allclose(self.outside_sum/self.outside_count, self.image[self.max_region_mask & ~self.mask].mean())

    def _image_sum_count(self, changed_idx):
        count = len(changed_idx[0])
        image_sum = self.image[changed_idx].sum()
        return count, image_sum

    def _move_to_outside(self, to_outside):
        """to_outside must be a boolean mask on the inside border pixels
        (in the order defined by inside_border_indices)."""
        added_idx, removed_indices = super()._move_to_outside(to_outside)
        count, image_sum = self._image_sum_count(added_idx)
        self.inside_count -= count
        self.outside_count += count
        self.inside_sum -= image_sum
        self.outside_sum += image_sum
        return added_idx, removed_indices

    def _move_to_inside(self, to_inside):
        """to_inside must be a boolean mask on the outside border pixels
        (in the order defined by outside_border_indices)."""
        added_idx, removed_indices = super()._move_to_inside(to_inside)
        count, image_sum = self._image_sum_count(added_idx)
        self.outside_count -= count
        self.inside_count += count
        self.outside_sum -= image_sum
        self.inside_sum += image_sum
        return added_idx, removed_indices

    def acwe_step(self, iters=1):
        """Apply 'iters' iterations of the Active Contours Without Edges step,
        wherein the region inside the mask is made to have a mean value as different
        from the region outside the mask as possible."""
        for _ in range(iters):
            if self.inside_count == 0 or self.outside_count == 0:
                return
            inside_mean = self.inside_sum / self.inside_count
            outside_mean = self.outside_sum / self.outside_count
            self._acwe_step(inside_mean, outside_mean, self.inside_bias,
                self.inside_border_indices, self._move_to_outside)
            self._acwe_step(outside_mean, inside_mean, 1/self.inside_bias,
                self.outside_border_indices, self._move_to_inside)

    def _acwe_step(self, mean_from, mean_to, from_bias, border_indices, move_operation):
        idx, idx_mask = _masked_idx(self.acwe_mask, border_indices)
        border_values = self.image[idx]
        to_move = from_bias*(border_values - mean_from)**2 > (border_values - mean_to)**2
        move_operation(_unmask_idx(idx_mask, to_move))

class BinnedACWE(CurvatureMorphology):
    def __init__(self, mask, image, acwe_mask=None, inside_bias=1,
        smooth_mask=None, max_region_mask=None):
        """Class for Active Contours Without Edges region-growing, using image
        histograms rather than simply means.

        Relevant methods for region-growing are smooth(), and acwe_step().

        Note that it is best to reduce the input image to a small number of
        brightness values (e.g. 16-64) before using this class.
        The discretize_image() function is useful for this.

        Parameters:
            mask: mask containing the initial state of the region to evolve
            image: ndarray of same shape as mask containing image values. The
                difference in image histograms inside and outside the region will
                be maximized by acwe_step()
            acwe_mask: region where ACWE updates will be applied. (Inside /
                outside values will still be computed outside this region.)
            inside_bias: weight for comparing means of inside vs. outside pixels.
                Generally 1 works properly. Values < 1 make it "easier" to add
                pixels to the region, and values > 1 make it "easier" to remove
                pixels from the region. (Simplification of the "lambda"
                parameters from the ACWE literature.)
            smooth_mask: region in which smoothing may be applied.
            max_region_mask: mask beyond which the region may not grow. If
                provided, this mask will also represent the pixels over which
                the ACWE inside/outside means are computed.
        """
        super().__init__(mask, smooth_mask, max_region_mask)
        # do work in _setup rather than __init__ to allow for complex multiple
        # inheritance from this class that super() alone can't handle. See
        # ActiveContour class.
        self._setup(image, acwe_mask, inside_bias)

    def _setup(self, image, acwe_mask, inside_bias):
        if image.dtype == bool:
            image = image.astype(numpy.uint8)
        self.image = image
        assert self.image.shape == self.mask.shape
        self.acwe_mask = acwe_mask
        if acwe_mask is not None:
            assert self.image.shape == self.acwe_mask.shape
        # note: self.mask is clipped to self.max_region_mask so the below works.
        self.inside_count = self.mask.sum()
        self.inside_bias = inside_bias
        self.outside_count = numpy.product(self.mask[self.max_region_mask].shape) - self.inside_count
        bincounts = numpy.bincount(self.image[self.max_region_mask].flat)
        self.bins = len(bincounts)
        self.inside_bincounts = self._bincount(self.mask)
        self.outside_bincounts = bincounts - self.inside_bincounts

    def _bincount(self, index_exp):
        return numpy.bincount(self.image[index_exp], minlength=self.bins)

    def _assert_invariants(self):
        super()._assert_invariants()
        assert self.inside_count == self.mask.sum()
        assert self.outside_count == numpy.product(self.mask[self.max_region_mask].shape) - self.inside_count
        assert (self.inside_bincounts == self._bincount(self.mask)).all()
        assert (self.outside_bincounts == self._bincount(self.max_region_mask) - self.inside_bincounts).all()
        if self.max_region_mask is not None:
            assert (self.outside_histogram == self._bincount(self.max_region_mask & ~self.mask)).all()

    def _image_bincount_count(self, changed_idx):
        count = len(changed_idx[0])
        bincounts = self._bincount(changed_idx)
        return count, bincounts

    def _move_to_outside(self, to_outside):
        """to_outside must be a boolean mask on the inside border pixels
        (in the order defined by inside_border_indices)."""
        added_idx, removed_indices = super()._move_to_outside(to_outside)
        count, bincounts = self._image_bincount_count(added_idx)
        self.inside_count -= count
        self.outside_count += count
        self.inside_bincounts -= bincounts
        self.outside_bincounts += bincounts
        return added_idx, removed_indices

    def _move_to_inside(self, to_inside):
        """to_inside must be a boolean mask on the outside border pixels
        (in the order defined by outside_border_indices)."""
        added_idx, removed_indices = super()._move_to_inside(to_inside)
        count, bincounts = self._image_bincount_count(added_idx)
        self.outside_count -= count
        self.inside_count += count
        self.outside_bincounts -= bincounts
        self.inside_bincounts += bincounts
        return added_idx, removed_indices

    def acwe_step(self, iters=1):
        """Apply 'iters' iterations of the Active Contours Without Edges step,
        wherein the region inside the mask is made to have a mean value as different
        from the region outside the mask as possible."""
        for _ in range(iters):
            if self.inside_count == 0 or self.outside_count == 0:
                return
            inside_rate = self.inside_bincounts / self.inside_count
            outside_rate = self.outside_bincounts / self.outside_count
            # small values of inside bias mean it is easier to move/stay inside
            self._acwe_step(self.inside_bias * outside_rate > inside_rate,
                self.inside_border_indices, self._move_to_outside)
            self._acwe_step(inside_rate > self.inside_bias * outside_rate,
                self.outside_border_indices, self._move_to_inside)

    def _acwe_step(self, move_bins, border_indices, move_operation):
        idx, idx_mask = _masked_idx(self.acwe_mask, border_indices)
        border_values = self.image[idx]
        to_move = move_bins[border_values]
        move_operation(_unmask_idx(idx_mask, to_move))

class BalloonForceMorphology(CurvatureMorphology):
    """Basic morphology operations plus spatially-varying balloon-force operation.
    Base-class to add balloon forces to more complex region-growing steps;
    rarely useful directly.
    """
    def __init__(self, mask, balloon_direction, smooth_mask=None, max_region_mask=None):
        """balloon_direction: (-1, 0, 1), or ndarray with same shape as 'mask'
        containing those values."""
        super().__init__(mask, smooth_mask, max_region_mask)
        # do work in _setup rather than __init__ to allow for complex multiple
        # inheritance from this class that super() alone can't handle. See
        # ActiveContour class.
        self._setup(balloon_direction)

    def _setup(self, balloon_direction):
        if numpy.isscalar(balloon_direction):
            if balloon_direction == 0:
                self.balloon_direction = None
            else:
                self.balloon_direction = numpy.zeros(self.mask.shape, dtype=numpy.int8)
                self.balloon_direction += balloon_direction
        else:
            self.balloon_direction = balloon_direction.copy() # may get changed internally by subclasses

    def balloon_force(self, iters=1):
        """Apply 'iters' iterations of balloon force region expansion / shrinkage."""
        if self.balloon_direction is None:
            return
        for _ in range(iters):
                to_erode = self.balloon_direction[tuple(self.inside_border_indices.T)] < 0
                self._move_to_outside(to_erode)
                to_dilate = self.balloon_direction[tuple(self.outside_border_indices.T)] > 0
                self._move_to_inside(to_dilate)

class GAC(BalloonForceMorphology):
    def __init__(self, mask, advection_direction, advection_mask=None,
        balloon_direction=0, smooth_mask=None, max_region_mask=None):
        """Class for Geodesic Active Contours region-growing.

        Relevant methods for region-growing are smooth(), balloon_force(),
        and advect().

        Parameters:
            mask:  mask containing the initial state of the region to evolve
            advection_direction: list of two arrays providing the x- and y-
                coordinates of the direction that the region edge should move
                in at any given point. (Only the sign of the direction matters.)
                The gradient of an edge-magnitude image works nicely here, as
                does the result of the edge_direction() function in this module
                when applied to a thresholded edge-magnitude image.
            advection_mask: boolean mask specifying where edge advection should
                be applied (versus balloon forces -- it makes no sense to try
                to apply both in the same location). If no advection_mask is
                provided, but a balloon_direction map is given, assume that
                advection is to be applied wherever the balloon_direction is
                0. If a scalar, non-zero balloon_direction is given, and no
                advection_mask is provided, then there will be no edge
                advection.
            balloon_direction: scalar balloon force direction (-1, 0, 1) or
                image map of same values. If advection_mask is provided,
                balloon_direction will be zeroed out in regions of where
                advection is allowed.
            smooth_mask: region in which smoothing may be applied.
            max_region_mask: mask beyond which the region may not grow.
        """
        CurvatureMorphology.__init__(self, mask, smooth_mask, max_region_mask)
        # do work in _setup rather than __init__ to allow for complex multiple
        # inheritance from this class that super() alone can't handle. See
        # ActiveContour class.
        self._setup(balloon_direction, advection_direction, advection_mask)

    def _setup(self, balloon_direction, advection_direction, advection_mask):
        BalloonForceMorphology._setup(self, balloon_direction)
        self.adv_dir_x, self.adv_dir_y = advection_direction
        # None balloon direction means no balloon force was asked for.
        if advection_mask is None:
            if self.balloon_direction is not None:
                self.advection_mask = self.balloon_direction == 0
            else:
                self.advection_mask = None
        else:
            self.advection_mask = advection_mask
            if self.balloon_direction is not None:
                self.balloon_direction[advection_mask] = 0

    def advect(self, iters=1):
        """Apply 'iters' iterations of edge advection, whereby the region edges
        are moved in the direction specified by advection_direction."""
        for _ in range(iters):
            # Move pixels on the inside border to the outside if advection*gradient sum > 0 (see _advect for interpretation of sum)
            self._advect(self.inside_border_indices, numpy.greater, self._move_to_outside)
            # Move pixels on the outside border to the inside if advection*gradient sum < 0
            self._advect(self.outside_border_indices, numpy.less, self._move_to_inside)

    def _advect(self, border_indices, criterion, move_operation):
        idx, idx_mask = _masked_idx(self.advection_mask, border_indices)
        neighbors = self.mask_neighborhood[idx].astype(numpy.int8)
        dx = neighbors[:,2,1] - neighbors[:,0,1]
        dy = neighbors[:,1,2] - neighbors[:,1,0]
        adv_dx = self.adv_dir_x[idx]
        adv_dy = self.adv_dir_y[idx]
        # positive gradient => outside-to-inside in that dimension
        # positive advection direction => edge should move in the positive direction
        # So:  + gradient / + advection = move pixels outside
        # + gradient / - advection or - gradient / + advection = move pixels inside
        # Tricky case: x and y disagree = go in direction with largest abs advection direction
        # To find this, see if sum of advection and gradient in each direction is > 0
        # (move pixels outside) or > 0 (move pixels inside).
        to_move = criterion(dx * adv_dx + dy * adv_dy, 0)
        move_operation(_unmask_idx(idx_mask, to_move))

class EdgeClaimingAdvection(CurvatureMorphology):
    def __init__(self, mask, edge_mask, distance_exponent=3, force_min=0.01, smooth_mask=None, max_region_mask=None):
        """Class for growing a region toward edges, such that only the edges
        that are not already at the region border influence further region
        growth.

        Given a set of image edges (ideally one pixel wide, as produced by
        Canny filtering), a region is grown toward those edges. The direction
        of growth at a region border pixel is determined by the sum of "forces",
        where each edge pixel exerts a force in inverse proportion to its
        distance from the region border pixel. This allows the region to move
        over large distances to find an edge.

        Only region border pixels that are not atop an edge pixel are free to
        have move. Only edge pixels not "captured" by a region border are free
        to exert forces on movable region border pixels. This capturing
        mechaimsm means that gaps in the edges will not cause region borders to
        "flow in" toward other edges in the image that already have a different
        part of the border atop them. Instead, gaps in the edges will be bridged
        smoothly.

        Parameters:
            mask:  mask containing the initial state of the region to evolve.
            edge_mask: mask containing the positions of the edges to occupy.
                Thinned edges (such as from Canny edge detection) are best.
            distance_exponent: the "force" on any region border pixel exerted
                by any edge pixel is 1/d**distance_exponent, where d is the
                distance in pixels between the edge pixel and the region border
                pixel. A larger value increases the influence of nearby
                uncaptured edge pixels over more distant uncaptured edge pixels.
            force_min: forces less than this minimum value will cause region
                borders to move. Important to prevent a few distant, unclaimed
                pixels from influencing borders that are bridging gaps.
            smooth_mask: region in which smoothing may be applied.
            max_region_mask: mask beyond which the region may not grow.
        """
        super().__init__(mask, smooth_mask, max_region_mask)
        # do work in _setup rather than __init__ to allow for complex multiple
        # inheritance from this class that super() alone can't handle. See
        # ActiveContour class.
        self._setup(edge_mask, distance_exponent, force_min)

    def _setup(self, edge_mask, distance_exponent, force_min):
        # Edges are "captured" iff there is an inside-border pixel on the edge.
        # Non-captured edges can generate advection forces.
        # Inside border pixels are free to advect iff they are not atop an edge.
        # Outside border pixels are free to advect iff they are (atop an edge
        # OR not adjacent a captured edge).
        self.edge_mask = edge_mask > 0
        # make an exponent to apply to the squared distance
        self.inverse_squared_distance_exponent = -distance_exponent/2
        self.force_min = force_min
        self.inside_free = ~self.edge_mask
        self.outside_free_neighborhood = neighborhood.make_neighborhood_view(numpy.zeros_like(self.edge_mask),
            pad_mode='constant', constant_values=0) # shape = edge_mask.shape + (3, 3)
        self.outside_free = self.outside_free_neighborhood[:,:,1,1]
        self.edge_indices = self.indices[self.edge_mask]
        self.noncaptured = numpy.zeros_like(self.edge_mask)
        self.update_captured()

    def update_captured(self):
        idx_mask = self.edge_mask[tuple(self.inside_border_indices.T)]
        captured_indices = self.inside_border_indices[idx_mask]
        captured_idx = tuple(captured_indices.T)
        self.captured_idx = captured_idx
        self.outside_free[:] = True
        self.outside_free[captured_idx] = False
        neighbor_indices = self.index_neighborhood[captured_idx] # shape = (m, 3, 3, 2)
        all_neighbors = neighbor_indices.reshape((numpy.product(neighbor_indices.shape[:-1]), 2)) # shape = (m*3*3, 2)
        all_idx = tuple(all_neighbors.T)
        # now erode to get pixels non-adjacent to captured edges
        self.outside_free[all_idx] = self.outside_free_neighborhood[all_idx].all()
        # Technically, we should set outside_free[captured_idx] to True,
        # but it doesn't matter because we will never look up those locations.
        # the outside_free array is only used to look at the outside_border_indices
        # positions, which are non-overlapping with inside_border_indices.
        # Now free outside pixels atop an edge:
        self.outside_free[tuple(self.edge_indices.T)] = True
        self.noncaptured_indices = _diff_indices(self.edge_indices, captured_indices)
        self.noncaptured[:] = False
        self.noncaptured[tuple(self.noncaptured_indices.T)] = True

    def advection_forces(self, border_indices, border_idx, mask_dx, mask_dy):
        """border_indices: shape (n, 2) list of indices"""
        edge_indices = self.noncaptured_indices # shape (m, 2)
        on_edge = self.noncaptured[border_idx] # shape (n,)
        dx = numpy.subtract.outer(edge_indices[:,0], border_indices[:,0]) # shape (m, n)
        dy = numpy.subtract.outer(edge_indices[:,1], border_indices[:,1])
        square_dist = dx**2 + dy**2
        square_dist[:,on_edge] = 1
        weighting = square_dist**self.inverse_squared_distance_exponent
        fx = (dx * weighting).sum(axis=0) # shape (n,)
        fx[on_edge] = -mask_dx[on_edge]
        fx[numpy.absolute(fx) < self.force_min] = 0
        fy = (dy * weighting).sum(axis=0)
        fy[on_edge] = -mask_dy[on_edge]
        fy[numpy.absolute(fy) < self.force_min] = 0
        return fx, fy

    def advect(self, iters=1):
        """Apply 'iters' iterations of edge advection, whereby uncaptured
        region edges are moved towrad uncaptured pixels in the edge_mask."""
        for _ in range(iters):
            self.update_captured()
            # Move pixels on the inside border to the outside if advection*gradient sum > 0 (see _advect for interpretation of sum)
            # and if they are not on an edge
            self._advect(self.inside_free, self.inside_border_indices, numpy.greater, self._move_to_outside)
            # Move pixels on the outside border to the inside if advection*gradient sum < 0
            # and if they arte not adjacent to a captured edge
            self._advect(self.outside_free, self.outside_border_indices, numpy.less, self._move_to_inside)

    def _advect(self, mask, border_indices, criterion, move_operation):
        idx_mask = mask[tuple(border_indices.T)]
        border_indices = border_indices[idx_mask]
        border_idx = tuple(border_indices.T)
        neighbors = self.mask_neighborhood[border_idx].astype(numpy.int8)
        dx = neighbors[:,2,1] - neighbors[:,0,1]
        dy = neighbors[:,1,2] - neighbors[:,1,0]
        adv_dx, adv_dy = self.advection_forces(border_indices, border_idx, dx, dy)
        # positive gradient => outside-to-inside in that dimension
        # positive advection direction => edge should move in the positive direction
        # So:  + gradient / + advection = move pixels outside
        # + gradient / - advection or - gradient / + advection = move pixels inside
        # Tricky case: x and y disagree = go in direction with largest abs advection direction
        # To find this, see if sum of advection and gradient in each direction is > 0
        # (move pixels outside) or > 0 (move pixels inside).
        to_move = criterion(dx * adv_dx + dy * adv_dy, 0)
        move_operation(_unmask_idx(idx_mask, to_move))

class ActiveContour(GAC, ACWE):
    def __init__(self, mask, image, advection_direction, acwe_mask=None,
            inside_bias=1, advection_mask=None, balloon_direction=0,
            smooth_mask=None, max_region_mask=None):
        """See documentation for GAC and ACWE for parameters."""
        CurvatureMorphology.__init__(self, mask, smooth_mask, max_region_mask)
        GAC._setup(self, balloon_direction, advection_direction, advection_mask)
        ACWE._setup(self, image, acwe_mask, inside_bias)

class BinnedActiveContour(GAC, BinnedACWE):
    def __init__(self, mask, image, advection_direction, acwe_mask=None,
            inside_bias=1, advection_mask=None, balloon_direction=0,
            smooth_mask=None, max_region_mask=None):
        """See documentation for GAC and BinnedACWE for parameters."""
        CurvatureMorphology.__init__(self, mask, smooth_mask, max_region_mask)
        GAC._setup(self, balloon_direction, advection_direction, advection_mask)
        BinnedACWE._setup(self, image, acwe_mask, inside_bias)

class ActiveClaimingContour(EdgeClaimingAdvection, ACWE):
    def __init__(self, mask, image, edge_mask, force_min=0.01, acwe_mask=None,
            inside_bias=1, smooth_mask=None, max_region_mask=None):
        """See documentation for EdgeClaimingAdvection and ACWE for parameters."""
        CurvatureMorphology.__init__(self, mask, smooth_mask, max_region_mask)
        EdgeClaimingAdvection._setup(self, edge_mask, force_min)
        ACWE._setup(self, image, acwe_mask, inside_bias)

class BinnedActiveClaimingContour(EdgeClaimingAdvection, BinnedACWE):
    def __init__(self, mask, image, edge_mask, distance_exponent=3, force_min=0.01,
            acwe_mask=None, inside_bias=1, smooth_mask=None, max_region_mask=None):
        """See documentation for EdgeClaimingAdvection and BinnedACWE for parameters."""
        CurvatureMorphology.__init__(self, mask, smooth_mask, max_region_mask)
        EdgeClaimingAdvection._setup(self, edge_mask, distance_exponent, force_min)
        BinnedACWE._setup(self, image, acwe_mask, inside_bias)

class StoppingCondition:
    def __init__(self, ac, max_iterations, sample_every=5, changed_min=5, cycle_max=6):
        """Class to simplify deciding when to stop active contour iteration,
        based on total number of iterations and changes to the mask.

        Usage example:
            ac = ACWE(mask, image)
            stopper = StoppingCondition(ac, max_iterations=50)
            while stopper.should_continue():
                ac.acwe_step()
                ac.smooth()

        Parameters:
            ac: active contour object of any type
            max_iterations: total number of iterations permitted
            sample_every: period at which the mask is examined to determine if
                any pixels have changed. Doing this every iteration can be slow,
                so this parameter allows trading off between stopping sooner and
                iterating faster.
            changed_min: if fewer than this number of pixels have changed, stop
                iteration.
            cycle_max: if the number of changed pixels is constant this many
                times in a row, stop iteration.
        """
        self.ac = ac
        self.max_iterations = max_iterations
        self.sample_every = sample_every
        self.changed_min = changed_min
        self.reset()

    def reset(self):
        self.recent_changes = collections.deque(maxlen=cycle_max-1)
        self.i = 0

    def should_continue(self):
        if self.i == self.max_iterations:
            return False
        if self.i % self.sample_every == 0:
            if self.i > 0:
                changed = (self.ac.mask != self.old_mask).sum()
                if changed < self.changed_min:
                    return False
                if (len(self.recent_changes) == self.recent_changes.maxlen and
                        all(prev == changed for prev in self.recent_changes)):
                    return False
                self.recent_changes.append(changed)
            self.old_mask = self.ac.mask.copy()
        self.i += 1
        return True

def edge_direction(edge_mask):
    """Given an edge mask, return the distance from each pixel to the mask,
    and the vector from each pixel to the nearest pixel in the mask.

    Parameter:
        edge_mask: boolean array that is True where image edges are.

    Returns: distances, nearest_edge
        distances: array same shape as edge_mask, containing the distance from
            every non-edge-mask pixel to the nearest pixel in the edge mask.
        nearest_edge: arary of shape (2,)+edge_mask.shape, containing the x and
            y coordinates of the vector from each non-edge pixel to the nearest
            edge pixel.
    """
    distances, nearest_edge = ndimage.distance_transform_cdt(~edge_mask,
        return_distances=True, return_indices=True)
    nearest_edge = nearest_edge - numpy.indices(edge_mask.shape)
    nearest_edge[:, edge_mask] = 0
    return distances, nearest_edge

def discretize_image(image, n_levels):
    """Given an image, reduce the number of distinct intensity levels.

    Parameters:
        image: ndarray of any type
        n_levels: number of intensity levels in the output image. Must be < 2**16.
    """
    assert n_levels < 2**16
    max = image.max()
    min = image.min()
    left_edges = numpy.linspace(min, max, n_levels)
    discretized = numpy.digitize(image, left_edges) - 1
    if n_levels > 256:
        return discretized.astype(numpy.uint8)
    else:
        return discretized.astype(numpy.uint16)

def _diff_indices(indices, to_remove):
    """Given two arrays of shape (n,2) containing x,y indices, remove those in
    the second array from the first array.
    """
    assert indices.flags.c_contiguous
    assert to_remove.flags.c_contiguous
    assert indices.dtype == to_remove.dtype
    dtype = numpy.dtype('S'+str(indices.itemsize*2)) # treat (x,y) indices as binary data instead of pairs of ints
    remaining = numpy.setdiff1d(indices.view(dtype), to_remove.view(dtype), assume_unique=True)
    return remaining.view(indices.dtype).reshape((-1, 2))

def _unique_indices(indices):
    """Given an array of shape (n,2) containing x,y indices, return only the
    unique indices from that array.
    """
    assert indices.flags.c_contiguous
    dtype = numpy.dtype('S'+str(indices.itemsize*2)) # treat (x,y) indices as binary data instead of pairs of ints
    unique = numpy.unique(indices.view(dtype))
    return unique.view(indices.dtype).reshape((-1, 2))

def _not_all(array, axis):
    return ~numpy.all(array, axis)

def _masked_idx(mask, indices):
    """Return index expression for given indices, only if mask is also true there.

    Parameters:
        mask: True/False mask
        indices: indices into array, of shape (n, mask.ndim)
    Returns:
        idx: numpy index expression for selected indices
        idx_mask: mask of same length as indices, indicating which indices had true
            values in the mask
    """
    if mask is None:
        idx_mask = None
    else:
        idx_mask = mask[tuple(indices.T)]
        indices = indices[idx_mask]
    idx = tuple(indices.T)
    return idx, idx_mask

def _unmask_idx(idx_mask, to_change):
    """Return mask over all indices originally passed to _masked_idx, marked
    False where indices were originally masked out or where to_change parameter
    is False.

    Parameters:
        idx_mask: as returned by _masked_idx
        to_change: mask over True values in idx_mask
    """
    if idx_mask is None:
        return to_change
    else:
        idx_mask[idx_mask] = to_change
        return idx_mask


import os
import shlex
import subprocess

from . import mesaurls

def install_prerequisites(directory, ostype, cleanAfter, logfile):
        """Install the pre-requisites for MESA.

        Args:
            logfile (file): File to write the output of the installation to.
        """
        if ostype == "Linux":
            subprocess.Popen(shlex.split("sudo apt-get update"), stdout=logfile, stderr=logfile).wait()
            try:
                subprocess.Popen(shlex.split("sudo apt-get install -yq build-essential wget curl binutils make perl libx11-6 \
                    libx11-dev zlib1g zlib1g-dev tcsh"), stdout=logfile, stderr=logfile).wait()
            except:
                try:
                    subprocess.Popen(shlex.split("sudo apt-get install -yq binutils make perl libx11-6 libx11-dev zlib1g zlib1g-dev tcsh"),
                                stdout=logfile, stderr=logfile).wait()
                except:
                    try:
                        subprocess.Popen(shlex.split("sudo apt-get install -yq binutils make perl libx11 zlib tcsh glibc"),
                                    stdout=logfile, stderr=logfile).wait()
                    except:
                        pass
        if "macOS" in ostype:
            # print("[green]Installing XCode Command Line Tools...")
            subprocess.Popen(shlex.split("sudo xcode-select --install"), stdout=logfile, stderr=logfile).wait()

            if not os.path.exists("/Applications/Utilities/XQuartz.app"):
                xquartz = os.path.join(directory, mesaurls.url_xquartz.split('/')[-1])

                # print("[green]Installing XQuartz...")
                subprocess.Popen(shlex.split(f"sudo installer -pkg {xquartz} -target /"), stdout=logfile, stderr=logfile).wait()
                if cleanAfter:
                    os.remove(xquartz)


import os
import zipfile
from functools import partial
import nltk
import requests
import numpy as np
from ipywidgets import widgets
from collections import namedtuple
import logging
import tqdm
from typing import List, Optional, Dict, Tuple

def download_dataset() -> None:
    """
    This func downloads the dataset and unzip to a local dir.
    """
    folder_path = os.environ.get("IR1_DATA_PATH")
    if not folder_path:
        folder_path = "./datasets/"
    os.makedirs(folder_path)

    file_location = os.path.join(folder_path, "cacm.zip")

    if not os.path.exists(file_location):

        url = "https://surfdrive.surf.nl/files/index.php/s/M0FGJpX2p8wDwxR/download"

        with open(file_location, "wb") as handle:
            print(f"Downloading file from {url} to {file_location}")
            response = requests.get(url, stream=True)
            for data in response.iter_content():
                handle.write(data)
            print("Finished downloading file")

    if not os.path.exists(os.path.join(folder_path, "train.txt")):
        # unzip file
        with zipfile.ZipFile(file_location, "r") as zip_ref:
            zip_ref.extractall(folder_path)

def read_cacm_docs(root_folder="./datasets/"):
    """
    Reads in the CACM documents. The dataset is assumed to be in the folder "./datasets/" by default
    Returns: A list of 2-tuples: (doc_id, document), where 'document' is a single string created by
        appending the title and abstract (separated by a "\n").
        In case the record doesn't have an abstract, the document is composed only by the title
    """
    file_dir = root_folder + "cacm.all"
    keeping = False
    temp = ""
    doc_list = []

    with open(file_dir) as cacm_file:
        for line in cacm_file:
            line_begin = line[0:2]

            if line_begin == ".I":
                doc_index = line.split(" ")[1].replace("\n", "")
                temp = ""

            elif line_begin == ".T":
                keeping = True

            elif line_begin == ".W":
                keeping = True
                temp += "\n "

            elif line_begin in [".B", ".A", ".N", ".K"]:
                keeping = False

            elif line_begin not in [".I", ".T", ".W", ".B", ".A", ".N", ".X", ".K"]:
                if keeping:
                    temp += line
            else:
                doc_list.append((doc_index, temp))
                keeping = False

    return doc_list

def read_queries(root_folder="./datasets/"):
    """
    Reads in the CACM queries. The dataset is assumed to be in the folder "./datasets/" by default
    Returns: A list of 2-tuples: (query_id, query)
    """
    file_dir = root_folder + "query.text"
    keeping = False
    temp = ""
    query_list = []

    with open(file_dir, "r") as query_file:
        for line in query_file:
            line_begin = line[0:2]

            if line_begin == ".I":
                doc_index = line.split(" ")[1].replace("\n", "")
                temp = ""

            elif line_begin == ".W":
                keeping = True

            elif line_begin == ".A":
                keeping = False

            elif line_begin not in [".I", ".T", ".W", ".B", ".A", ".N", ".X", ".K"]:
                if keeping:
                    temp += line
            else:
                query_list.append((doc_index, temp))
                keeping = False

    return query_list

def load_stopwords(root_folder="./datasets/"):
    """
    Loads the stopwords. The dataset is assumed to be in the folder "./datasets/" by default

    :param root_folder:
    :return: A set of stopwords
    """
    file_dir = root_folder + "common_words"
    with open(file_dir, "r") as f:
        stopwords = [line.strip() for line in f]

    return set(stopwords)

def tokenize(text: str) -> List[str]:
    """
    Tokenizes the input text. Use the WordPunctTokenizer

    :param text: text - a string
    :return: a list of tokens
    """
    return nltk.WordPunctTokenizer().tokenize(text)

def stem_token(token):
    """
    Stems the given token using the PorterStemmer from the nltk library

    :param token: a single token
    :return: the stem of the token
    """
    ps = nltk.stem.PorterStemmer()
    return ps.stem(token)

def process_text(
    text: str,
    stem: bool = False,
    remove_stopwords: str = False,
    lowercase_text: str = False,
):
    """
    Given a string, this func tokenizes it and processes it according to the flags that you set.

    :param text: input text
    :param stem: to stem or not
    :param remove_stopwords: to remove stopwords or not
    :param lowercase_text: to lowercase the text or not
    :return:
    """
    tokens = []
    for token in tokenize(text):
        if remove_stopwords and token.lower() in stopwords:
            continue
        if stem:
            token = stem_token(token)
        if lowercase_text:
            token = token.lower()
        tokens.append(token)
    return tokens

def read_qrels(root_folder="./datasets/"):
    query_f = open(os.path.join(root_folder, "qrels.text"), "r")
    query_dic = {}

    for line in query_f:
        q_id, doc_id, _, _ = line.split()
        q_id = str(int(q_id))
        if q_id not in query_dic:
            query_dic[q_id] = []
        query_dic[q_id].append(doc_id)

    return query_dic