# -*- coding: utf-8 -*-
"""
.. _tut-overview:

Overview of MEG/EEG analysis with MNE-Python
============================================

This tutorial covers the basic EEG/MEG pipeline for event-related analysis:
loading data, epoching, averaging, plotting, and estimating cortical activity
from sensor data. It introduces the core MNE-Python data structures
`~mne.io.Raw`, `~mne.Epochs`, `~mne.Evoked`, and `~mne.SourceEstimate`, and
covers a lot of ground fairly quickly (at the expense of depth). Subsequent
tutorials address each of these topics in greater detail.

We begin by importing the necessary Python modules:
"""

import os
import numpy as np
import mne

###############################################################################
# Loading data
# ^^^^^^^^^^^^
#
# MNE-Python data structures are based around the FIF file format from
# Neuromag, but there are reader functions for :ref:`a wide variety of other
# data formats <data-formats>`. MNE-Python also has interfaces to a
# variety of :ref:`publicly available datasets <datasets>`,
# which MNE-Python can download and manage for you.
#
# We'll start this tutorial by loading one of the example datasets (called
# ":ref:`sample-dataset`"), which contains EEG and MEG data from one subject
# performing an audiovisual experiment, along with structural MRI scans for
# that subject. The `mne.datasets.sample.data_path` function will automatically
# download the dataset if it isn't found in one of the expected locations, then
# return the directory path to the dataset (see the documentation of
# `~mne.datasets.sample.data_path` for a list of places it checks before
# downloading). Note also that for this tutorial to run smoothly on our
# servers, we're using a filtered and downsampled version of the data
# (:file:`sample_audvis_filt-0-40_raw.fif`), but an unfiltered version
# (:file:`sample_audvis_raw.fif`) is also included in the sample dataset and
# could be substituted here when running the tutorial locally.

sample_data_folder = mne.datasets.sample.data_path()
sample_data_raw_file = os.path.join(sample_data_folder, 'MEG', 'sample',
                                    'sample_audvis_filt-0-40_raw.fif')
raw = mne.io.read_raw_fif(sample_data_raw_file)

###############################################################################
# By default, `~mne.io.read_raw_fif` displays some information about the file
# it's loading; for example, here it tells us that there are four "projection
# items" in the file along with the recorded data; those are :term:`SSP
# projectors <projector>` calculated to remove environmental noise from the MEG
# signals, plus a projector to mean-reference the EEG channels; these are
# discussed in the tutorial :ref:`tut-projectors-background`. In addition to
# the information displayed during loading, you can get a glimpse of the basic
# details of a `~mne.io.Raw` object by printing it; even more is available by
# printing its ``info`` attribute (a `dictionary-like object <mne.Info>` that
# is preserved across `~mne.io.Raw`, `~mne.Epochs`, and `~mne.Evoked` objects).
# The ``info`` data structure keeps track of channel locations, applied
# filters, projectors, etc. Notice especially the ``chs`` entry, showing that
# MNE-Python detects different sensor types and handles each appropriately. See
# :ref:`tut-info-class` for more on the `~mne.Info` class.

print(raw)
print(raw.info)

###############################################################################
# `~mne.io.Raw` objects also have several built-in plotting methods; here we
# show the power spectral density (PSD) for each sensor type with
# `~mne.io.Raw.plot_psd`, as well as a plot of the raw sensor traces with
# `~mne.io.Raw.plot`. In the PSD plot, we'll only plot frequencies below 50 Hz
# (since our data are low-pass filtered at 40 Hz). In interactive Python
# sessions, `~mne.io.Raw.plot` is interactive and allows scrolling, scaling,
# bad channel marking, annotation, projector toggling, etc.

raw.plot_psd(fmax=50)
raw.plot(duration=5, n_channels=30)

###############################################################################
# Preprocessing
# ^^^^^^^^^^^^^
#
# MNE-Python supports a variety of preprocessing approaches and techniques
# (maxwell filtering, signal-space projection, independent components analysis,
# filtering, downsampling, etc); see the full list of capabilities in the
# :mod:`mne.preprocessing` and :mod:`mne.filter` submodules. Here we'll clean
# up our data by performing independent components analysis
# (`~mne.preprocessing.ICA`); for brevity we'll skip the steps that helped us
# determined which components best capture the artifacts (see
# :ref:`tut-artifact-ica` for a detailed walk-through of that process).

# set up and fit the ICA
ica = mne.preprocessing.ICA(n_components=20, random_state=97, max_iter=800)
ica.fit(raw)
ica.exclude = [1, 2]  # details on how we picked these are omitted here
ica.plot_properties(raw, picks=ica.exclude)

###############################################################################
# Once we're confident about which component(s) we want to remove, we pass them
# as the ``exclude`` parameter and then apply the ICA to the raw signal. The
# `~mne.preprocessing.ICA.apply` method requires the raw data to be loaded into
# memory (by default it's only read from disk as-needed), so we'll use
# `~mne.io.Raw.load_data` first. We'll also make a copy of the `~mne.io.Raw`
# object so we can compare the signal before and after artifact removal
# side-by-side:

orig_raw = raw.copy()
raw.load_data()
ica.apply(raw)

# show some frontal channels to clearly illustrate the artifact removal
chs = ['MEG 0111', 'MEG 0121', 'MEG 0131', 'MEG 0211', 'MEG 0221', 'MEG 0231',
       'MEG 0311', 'MEG 0321', 'MEG 0331', 'MEG 1511', 'MEG 1521', 'MEG 1531',
       'EEG 001', 'EEG 002', 'EEG 003', 'EEG 004', 'EEG 005', 'EEG 006',
       'EEG 007', 'EEG 008']
chan_idxs = [raw.ch_names.index(ch) for ch in chs]
orig_raw.plot(order=chan_idxs, start=12, duration=4)
raw.plot(order=chan_idxs, start=12, duration=4)

###############################################################################
# .. _overview-tut-events-section:
#
# Detecting experimental events
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# The sample dataset includes several :term:`"STIM" channels <stim channel>`
# that recorded electrical signals sent from the stimulus delivery computer (as
# brief DC shifts / squarewave pulses). These pulses (often called "triggers")
# are used in this dataset to mark experimental events: stimulus onset,
# stimulus type, and participant response (button press). The individual STIM
# channels are combined onto a single channel, in such a way that voltage
# levels on that channel can be unambiguously decoded as a particular event
# type. On older Neuromag systems (such as that used to record the sample data)
# this summation channel was called ``STI 014``, so we can pass that channel
# name to the `mne.find_events` function to recover the timing and identity of
# the stimulus events.

events = mne.find_events(raw, stim_channel='STI 014')
print(events[:5])  # show the first 5

###############################################################################
# The resulting events array is an ordinary 3-column :class:`NumPy array
# <numpy.ndarray>`, with sample number in the first column and integer event ID
# in the last column; the middle column is usually ignored. Rather than keeping
# track of integer event IDs, we can provide an *event dictionary* that maps
# the integer IDs to experimental conditions or events. In this dataset, the
# mapping looks like this:
#
# .. _sample-data-event-dict-table:
#
# +----------+----------------------------------------------------------+
# | Event ID | Condition                                                |
# +==========+==========================================================+
# | 1        | auditory stimulus (tone) to the left ear                 |
# +----------+----------------------------------------------------------+
# | 2        | auditory stimulus (tone) to the right ear                |
# +----------+----------------------------------------------------------+
# | 3        | visual stimulus (checkerboard) to the left visual field  |
# +----------+----------------------------------------------------------+
# | 4        | visual stimulus (checkerboard) to the right visual field |
# +----------+----------------------------------------------------------+
# | 5        | smiley face (catch trial)                                |
# +----------+----------------------------------------------------------+
# | 32       | subject button press                                     |
# +----------+----------------------------------------------------------+

event_dict = {'auditory/left': 1, 'auditory/right': 2, 'visual/left': 3,
              'visual/right': 4, 'smiley': 5, 'buttonpress': 32}

###############################################################################
# Event dictionaries like this one are used when extracting epochs from
# continuous data; the ``/`` character in the dictionary keys allows pooling
# across conditions by requesting partial condition descriptors (i.e.,
# requesting ``'auditory'`` will select all epochs with Event IDs 1 and 2;
# requesting ``'left'`` will select all epochs with Event IDs 1 and 3). An
# example of this is shown in the next section. There is also a convenient
# `~mne.viz.plot_events` function for visualizing the distribution of events
# across the duration of the recording (to make sure event detection worked as
# expected). Here we'll also make use of the `~mne.Info` attribute to get the
# sampling frequency of the recording (so our x-axis will be in seconds instead
# of in samples).

fig = mne.viz.plot_events(events, event_id=event_dict, sfreq=raw.info['sfreq'],
                          first_samp=raw.first_samp)

###############################################################################
# For paradigms that are not event-related (e.g., analysis of resting-state
# data), you can extract regularly spaced (possibly overlapping) spans of data
# by creating events using `mne.make_fixed_length_events` and then proceeding
# with epoching as described in the next section.
#
#
# .. _tut-section-overview-epoching:
#
# Epoching continuous data
# ^^^^^^^^^^^^^^^^^^^^^^^^
#
# The `~mne.io.Raw` object and the events array are the bare minimum needed to
# create an `~mne.Epochs` object, which we create with the `~mne.Epochs` class
# constructor. Here we'll also specify some data quality constraints: we'll
# reject any epoch where peak-to-peak signal amplitude is beyond reasonable
# limits for that channel type. This is done with a *rejection dictionary*; you
# may include or omit thresholds for any of the channel types present in your
# data. The values given here are reasonable for this particular dataset, but
# may need to be adapted for different hardware or recording conditions. For a
# more automated approach, consider using the `autoreject package`_.

reject_criteria = dict(mag=4000e-15,     # 4000 fT
                       grad=4000e-13,    # 4000 fT/cm
                       eeg=150e-6,       # 150 µV
                       eog=250e-6)       # 250 µV

###############################################################################
# We'll also pass the event dictionary as the ``event_id`` parameter (so we can
# work with easy-to-pool event labels instead of the integer event IDs), and
# specify ``tmin`` and ``tmax`` (the time relative to each event at which to
# start and end each epoch). As mentioned above, by default `~mne.io.Raw` and
# `~mne.Epochs` data aren't loaded into memory (they're accessed from disk only
# when needed), but here we'll force loading into memory using the
# ``preload=True`` parameter so that we can see the results of the rejection
# criteria being applied:

epochs = mne.Epochs(raw, events, event_id=event_dict, tmin=-0.2, tmax=0.5,
                    reject=reject_criteria, preload=True)

###############################################################################
# Next we'll pool across left/right stimulus presentations so we can compare
# auditory versus visual responses. To avoid biasing our signals to the left or
# right, we'll use `~mne.Epochs.equalize_event_counts` first to randomly sample
# epochs from each condition to match the number of epochs present in the
# condition with the fewest good epochs.

conds_we_care_about = ['auditory/left', 'auditory/right',
                       'visual/left', 'visual/right']
epochs.equalize_event_counts(conds_we_care_about)  # this operates in-place
aud_epochs = epochs['auditory']
vis_epochs = epochs['visual']
del raw, epochs  # free up memory

###############################################################################
# Like `~mne.io.Raw` objects, `~mne.Epochs` objects also have a number of
# built-in plotting methods. One is `~mne.Epochs.plot_image`, which shows each
# epoch as one row of an image map, with color representing signal magnitude;
# the average evoked response and the sensor location are shown below the
# image:

aud_epochs.plot_image(picks=['MEG 1332', 'EEG 021'])

##############################################################################
# .. note::
#
#     Both `~mne.io.Raw` and `~mne.Epochs` objects have `~mne.Epochs.get_data`
#     methods that return the underlying data as a
#     :class:`NumPy array <numpy.ndarray>`. Both methods have a ``picks``
#     parameter for subselecting which channel(s) to return; ``raw.get_data()``
#     has additional parameters for restricting the time domain. The resulting
#     matrices have dimension ``(n_channels, n_times)`` for `~mne.io.Raw` and
#     ``(n_epochs, n_channels, n_times)`` for `~mne.Epochs`.
#
# Time-frequency analysis
# ^^^^^^^^^^^^^^^^^^^^^^^
#
# The :mod:`mne.time_frequency` submodule provides implementations of several
# algorithms to compute time-frequency representations, power spectral density,
# and cross-spectral density. Here, for example, we'll compute for the auditory
# epochs the induced power at different frequencies and times, using Morlet
# wavelets. On this dataset the result is not especially informative (it just
# shows the evoked "auditory N100" response); see :ref:`here
# <inter-trial-coherence>` for a more extended example on a dataset with richer
# frequency content.

frequencies = np.arange(7, 30, 3)
power = mne.time_frequency.tfr_morlet(aud_epochs, n_cycles=2, return_itc=False,
                                      freqs=frequencies, decim=3)
power.plot(['MEG 1332'])

###############################################################################
# Estimating evoked responses
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# Now that we have our conditions in ``aud_epochs`` and ``vis_epochs``, we can
# get an estimate of evoked responses to auditory versus visual stimuli by
# averaging together the epochs in each condition. This is as simple as calling
# the `~mne.Epochs.average` method on the `~mne.Epochs` object, and then using
# a function from the :mod:`mne.viz` module to compare the global field power
# for each sensor type of the two `~mne.Evoked` objects:

aud_evoked = aud_epochs.average()
vis_evoked = vis_epochs.average()

mne.viz.plot_compare_evokeds(dict(auditory=aud_evoked, visual=vis_evoked),
                             legend='upper left', show_sensors='upper right')

###############################################################################
# We can also get a more detailed view of each `~mne.Evoked` object using other
# plotting methods such as `~mne.Evoked.plot_joint` or
# `~mne.Evoked.plot_topomap`. Here we'll examine just the EEG channels, and see
# the classic auditory evoked N100-P200 pattern over dorso-frontal electrodes,
# then plot scalp topographies at some additional arbitrary times:

# sphinx_gallery_thumbnail_number = 13
aud_evoked.plot_joint(picks='eeg')
aud_evoked.plot_topomap(times=[0., 0.08, 0.1, 0.12, 0.2], ch_type='eeg')

##############################################################################
# Evoked objects can also be combined to show contrasts between conditions,
# using the `mne.combine_evoked` function. A simple difference can be
# generated by passing ``weights=[1, -1]``. We'll then plot the difference wave
# at each sensor using `~mne.Evoked.plot_topo`:

evoked_diff = mne.combine_evoked([aud_evoked, vis_evoked], weights=[1, -1])
evoked_diff.pick_types(meg='mag').plot_topo(color='r', legend=False)

##############################################################################
# Inverse modeling
# ^^^^^^^^^^^^^^^^
#
# Finally, we can estimate the origins of the evoked activity by projecting the
# sensor data into this subject's :term:`source space` (a set of points either
# on the cortical surface or within the cortical volume of that subject, as
# estimated by structural MRI scans). MNE-Python supports lots of ways of doing
# this (dynamic statistical parametric mapping, dipole fitting, beamformers,
# etc.); here we'll use minimum-norm estimation (MNE) to generate a continuous
# map of activation constrained to the cortical surface. MNE uses a linear
# :term:`inverse operator` to project EEG+MEG sensor measurements into the
# source space. The inverse operator is computed from the
# :term:`forward solution` for this subject and an estimate of :ref:`the
# covariance of sensor measurements <tut-compute-covariance>`. For this
# tutorial we'll skip those computational steps and load a pre-computed inverse
# operator from disk (it's included with the :ref:`sample data
# <sample-dataset>`). Because this "inverse problem" is underdetermined (there
# is no unique solution), here we further constrain the solution by providing a
# regularization parameter specifying the relative smoothness of the current
# estimates in terms of a signal-to-noise ratio (where "noise" here is akin to
# baseline activity level across all of cortex).

# load inverse operator
inverse_operator_file = os.path.join(sample_data_folder, 'MEG', 'sample',
                                     'sample_audvis-meg-oct-6-meg-inv.fif')
inv_operator = mne.minimum_norm.read_inverse_operator(inverse_operator_file)
# set signal-to-noise ratio (SNR) to compute regularization parameter (λ²)
snr = 3.
lambda2 = 1. / snr ** 2
# generate the source time course (STC)
stc = mne.minimum_norm.apply_inverse(vis_evoked, inv_operator,
                                     lambda2=lambda2,
                                     method='MNE')  # or dSPM, sLORETA, eLORETA

##############################################################################
# Finally, in order to plot the source estimate on the subject's cortical
# surface we'll also need the path to the sample subject's structural MRI files
# (the ``subjects_dir``):

# path to subjects' MRI files
subjects_dir = os.path.join(sample_data_folder, 'subjects')
# plot the STC
stc.plot(initial_time=0.1, hemi='split', views=['lat', 'med'],
         subjects_dir=subjects_dir)

##############################################################################
# The remaining tutorials have *much more detail* on each of these topics (as
# well as many other capabilities of MNE-Python not mentioned here:
# connectivity analysis, encoding/decoding models, lots more visualization
# options, etc). Read on to learn more!
#
# .. LINKS
#
# .. _`autoreject package`: http://autoreject.github.io/


import numpy as np

def test(test_cases, target):
    success = 0
    for test_case in test_cases:
        try:
            if test_case['name'] == "datatype_check":
                assert isinstance(target(*test_case['input']),
                                  test_case["expected"])
                success += 1
            if test_case['name'] == "equation_output_check":
                assert np.allclose(test_case["expected"],
                                   target(*test_case['input']))
                success += 1
            if test_case['name'] == "shape_check":
                assert test_case['expected'].shape == target(*test_case['input']).shape
                success += 1
        except:
            print("Error: " + test_case['error'])

    if success == len(test_cases):
        print("\033[92m All tests passed.")
    else:
        print('\033[92m', success," Tests passed")
        print('\033[91m', len(test_cases) - success, " Tests failed")
        raise AssertionError("Not all tests were passed for {}. Check your equations and avoid using global variables inside the function.".format(target.__name__))


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

Source alignment and coordinate frames
======================================

The aim of this tutorial is to show how to visually assess that the data are
well aligned in space for computing the forward solution, and understand
the different coordinate frames involved in this process.

.. contents:: Topics
   :local:
   :depth: 2

Let's start out by loading some data.
"""
import os.path as op

import numpy as np
from mayavi import mlab

import mne
from mne.datasets import sample

print(__doc__)

data_path = sample.data_path()
subjects_dir = op.join(data_path, 'subjects')
raw_fname = op.join(data_path, 'MEG', 'sample', 'sample_audvis_raw.fif')
trans_fname = op.join(data_path, 'MEG', 'sample',
                      'sample_audvis_raw-trans.fif')
raw = mne.io.read_raw_fif(raw_fname)
trans = mne.read_trans(trans_fname)
src = mne.read_source_spaces(op.join(subjects_dir, 'sample', 'bem',
                                     'sample-oct-6-src.fif'))

###############################################################################
# Understanding coordinate frames
# -------------------------------
# For M/EEG source imaging, there are three **coordinate frames** that we must
# bring into alignment using two 3D
# `transformation matrices <trans_matrices_>`_
# that define how to rotate and translate points in one coordinate frame
# to their equivalent locations in another.
#
# :func:`mne.viz.plot_alignment` is a very useful function for inspecting
# these transformations, and the resulting alignment of EEG sensors, MEG
# sensors, brain sources, and conductor models. If the ``subjects_dir`` and
# ``subject`` parameters are provided, the function automatically looks for the
# Freesurfer MRI surfaces to show from the subject's folder.
#
# We can use the ``show_axes`` argument to see the various coordinate frames
# given our transformation matrices. These are shown by axis arrows for each
# coordinate frame:
#
# * shortest arrow is (**R**)ight/X
# * medium is forward/(**A**)nterior/Y
# * longest is up/(**S**)uperior/Z
#
# i.e., a **RAS** coordinate system in each case. We can also set
# the ``coord_frame`` argument to choose which coordinate
# frame the camera should initially be aligned with.
#
# Let's take a look:

mne.viz.plot_alignment(raw.info, trans=trans, subject='sample',
                       subjects_dir=subjects_dir, surfaces='head-dense',
                       show_axes=True, dig=True, eeg=[], meg='sensors',
                       coord_frame='meg')
mlab.view(45, 90, distance=0.6, focalpoint=(0., 0., 0.))
print('Distance from head origin to MEG origin: %0.1f mm'
      % (1000 * np.linalg.norm(raw.info['dev_head_t']['trans'][:3, 3])))
print('Distance from head origin to MRI origin: %0.1f mm'
      % (1000 * np.linalg.norm(trans['trans'][:3, 3])))

###############################################################################
# Coordinate frame definitions
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# .. raw:: html
#
#    <style>
#    .pink {color:DarkSalmon; font-weight:bold}
#    .blue {color:DeepSkyBlue; font-weight:bold}
#    .gray {color:Gray; font-weight:bold}
#    .magenta {color:Magenta; font-weight:bold}
#    .purple {color:Indigo; font-weight:bold}
#    .green {color:LimeGreen; font-weight:bold}
#    .red {color:Red; font-weight:bold}
#    </style>
#
# .. role:: pink
# .. role:: blue
# .. role:: gray
# .. role:: magenta
# .. role:: purple
# .. role:: green
# .. role:: red
#
# 1. Neuromag head coordinate frame ("head", :pink:`pink axes`)
#      Defined by the intersection of 1) the line between the LPA
#      (:red:`red sphere`) and RPA (:purple:`purple sphere`), and
#      2) the line perpendicular to this LPA-RPA line one that goes through
#      the Nasion (:green:`green sphere`).
#      The axes are oriented as **X** origin→RPA, **Y** origin→Nasion,
#      **Z** origin→upward (orthogonal to X and Y).
#
#      .. note:: This gets defined during the head digitization stage during
#                acquisition, often by use of a Polhemus or other digitizer.
#
# 2. MEG device coordinate frame ("meg", :blue:`blue axes`)
#      This is defined by the MEG manufacturers. From the Elekta user manual:
#
#          The origin of the device coordinate system is located at the center
#          of the posterior spherical section of the helmet with axis going
#          from left to right and axis pointing front. The axis is, again
#          normal to the plane with positive direction up.
#
#      .. note:: The device is coregistered with the head coordinate frame
#                during acquisition via emission of sinusoidal currents in
#                head position indicator (HPI) coils
#                (:magenta:`magenta spheres`) at the beginning of the
#                recording. This is stored in ``raw.info['dev_head_t']``.
#
# 3. MRI coordinate frame ("mri", :gray:`gray axes`)
#      Defined by Freesurfer, the MRI (surface RAS) origin is at the
#      center of a 256×256×256 1mm anisotropic volume (may not be in the center
#      of the head).
#
#      .. note:: This is aligned to the head coordinate frame that we
#                typically refer to in MNE as ``trans``.
#
# A bad example
# -------------
# Let's try using ``trans=None``, which (incorrectly!) equates the MRI
# and head coordinate frames.

mne.viz.plot_alignment(raw.info, trans=None, subject='sample', src=src,
                       subjects_dir=subjects_dir, dig=True,
                       surfaces=['head-dense', 'white'], coord_frame='meg')

###############################################################################
# It is quite clear that the MRI surfaces (head, brain) are not well aligned
# to the head digitization points (dots).
#
# A good example
# --------------
# Here is the same plot, this time with the ``trans`` properly defined
# (using a precomputed matrix).

mne.viz.plot_alignment(raw.info, trans=trans, subject='sample',
                       src=src, subjects_dir=subjects_dir, dig=True,
                       surfaces=['head-dense', 'white'], coord_frame='meg')

###############################################################################
# Defining the head↔MRI ``trans`` using the GUI
# ---------------------------------------------
# You can try creating the head↔MRI transform yourself using
# :func:`mne.gui.coregistration`.
#
# * First you must load the digitization data from the raw file
#   (``Head Shape Source``). The MRI data is already loaded if you provide the
#   ``subject`` and ``subjects_dir``. Toggle ``Always Show Head Points`` to see
#   the digitization points.
# * To set the landmarks, toggle ``Edit`` radio button in ``MRI Fiducials``.
# * Set the landmarks by clicking the radio button (LPA, Nasion, RPA) and then
#   clicking the corresponding point in the image.
# * After doing this for all the landmarks, toggle ``Lock`` radio button. You
#   can omit outlier points, so that they don't interfere with the finetuning.
#
#   .. note:: You can save the fiducials to a file and pass
#             ``mri_fiducials=True`` to plot them in
#             :func:`mne.viz.plot_alignment`. The fiducials are saved to the
#             subject's bem folder by default.
# * Click ``Fit Head Shape``. This will align the digitization points to the
#   head surface. Sometimes the fitting algorithm doesn't find the correct
#   alignment immediately. You can try first fitting using LPA/RPA or fiducials
#   and then align according to the digitization. You can also finetune
#   manually with the controls on the right side of the panel.
# * Click ``Save As...`` (lower right corner of the panel), set the filename
#   and read it with :func:`mne.read_trans`.
#
# For more information, see step by step instructions
# `in these slides
# <https://www.slideshare.net/mne-python/mnepython-coregistration>`_.
# Uncomment the following line to align the data yourself.

# mne.gui.coregistration(subject='sample', subjects_dir=subjects_dir)

###############################################################################
# .. _plot_source_alignment_without_mri:
#
# Alignment without MRI
# ---------------------
# The surface alignments above are possible if you have the surfaces available
# from Freesurfer. :func:`mne.viz.plot_alignment` automatically searches for
# the correct surfaces from the provided ``subjects_dir``. Another option is
# to use a :ref:`spherical conductor model <ch_forward_spherical_model>`. It is
# passed through ``bem`` parameter.

sphere = mne.make_sphere_model(info=raw.info, r0='auto', head_radius='auto')
src = mne.setup_volume_source_space(sphere=sphere, pos=10.)
mne.viz.plot_alignment(
    raw.info, eeg='projected', bem=sphere, src=src, dig=True,
    surfaces=['brain', 'outer_skin'], coord_frame='meg', show_axes=True)

###############################################################################
# It is also possible to use :func:`mne.gui.coregistration`
# to warp a subject (usually ``fsaverage``) to subject digitization data, see
# `these slides
# <https://www.slideshare.net/mne-python/mnepython-scale-mri>`_.
#
# .. _trans_matrices: https://en.wikipedia.org/wiki/Transformation_matrix


#!/usr/bin/env python
# coding: utf-8

# ## Rover Lab Notebook
# This notebook contains the functions from the lesson and provides the scaffolding you need to test out your mapping methods.  The steps you need to complete in this notebook for the project are the following:
#
# * First just run each of the cells in the notebook, examine the code and the results of each.
#
# **Note: For the online lab, data has been collected and provided for you. If you would like to try locally please do so! Please continue instructions from the continue point.**
# * Run the simulator in "Training Mode" and record some data. Note: the simulator may crash if you try to record a large (longer than a few minutes) dataset, but you don't need a ton of data, just some example images to work with.
# * Change the data directory path (2 cells below) to be the directory where you saved data
# * Test out the functions provided on your data
#
# **Continue Point**
# * Write new functions (or modify existing ones) to report and map out detections of obstacles and rock samples (yellow rocks)
# * Populate the `process_image()` function with the appropriate steps/functions to go from a raw image to a worldmap.
# * Run the cell that calls `process_image()` using `moviepy` functions to create video output
# * Once you have mapping working, move on to modifying `perception.py` and `decision.py` in the project to allow your rover to navigate and map in autonomous mode!
#
# **Note: If, at any point, you encounter frozen display windows or other confounding issues, you can always start again with a clean slate by going to the "Kernel" menu above and selecting "Restart & Clear Output".**
#
# **Run the next cell to get code highlighting in the markdown cells.**

# In[31]:

get_ipython().run_cell_magic('HTML', '', '<style> code {background-color : orange !important;} </style>')

# In[32]:

get_ipython().run_line_magic('matplotlib', 'inline')
#%matplotlib qt # Choose %matplotlib qt to plot to an interactive window (note it may show up behind your browser)
# Make some of the relevant imports
import cv2 # OpenCV for perspective transform
import numpy as np
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import scipy.misc # For saving images as needed
import glob  # For reading in a list of images from a folder

# ## Quick Look at the Data
# There's some example data provided in the `test_dataset` folder.  This basic dataset is enough to get you up and running but if you want to hone your methods more carefully you should record some data of your own to sample various scenarios in the simulator.
#
# Next, read in and display a random image from the `test_dataset` folder

# In[33]:

path = './test_dataset/IMG/*'
img_list = glob.glob(path)
# Grab a random image and display it
idx = np.random.randint(0, len(img_list)-1)
image = mpimg.imread(img_list[idx])
plt.imshow(image)

# ## Calibration Data
# Read in and display example grid and rock sample calibration images.  You'll use the grid for perspective transform and the rock image for creating a new color selection that identifies these samples of interest.

# In[34]:

# In the simulator you can toggle on a grid on the ground for calibration
# You can also toggle on the rock samples with the 0 (zero) key.
# Here's an example of the grid and one of the rocks
example_grid = './calibration_images/example_grid1.jpg'
example_rock = './calibration_images/example_rock1.jpg'
grid_img = mpimg.imread(example_grid)
rock_img = mpimg.imread(example_rock)

fig = plt.figure(figsize=(12,3))
plt.subplot(121)
plt.imshow(grid_img)
plt.subplot(122)
plt.imshow(rock_img)

# ## Perspective Transform
#
# Define the perspective transform function from the lesson and test it on an image.

# In[35]:

# Define a function to perform a perspective transform
# I've used the example grid image above to choose source points for the
# grid cell in front of the rover (each grid cell is 1 square meter in the sim)
# Define a function to perform a perspective transform
def perspect_transform(img, src, dst):
    M = cv2.getPerspectiveTransform(src, dst)
    warped = cv2.warpPerspective(img, M, (img.shape[1], img.shape[0]))# keep same size as input image
    mask = cv2.warpPerspective(np.ones_like(img[:,:,0]), M, (img.shape[1], img.shape[0]))
    return warped, mask

# Define calibration box in source (actual) and destination (desired) coordinates
# These source and destination points are defined to warp the image
# to a grid where each 10x10 pixel square represents 1 square meter
# The destination box will be 2*dst_size on each side
dst_size = 5
# Set a bottom offset to account for the fact that the bottom of the image
# is not the position of the rover but a bit in front of it
# this is just a rough guess, feel free to change it!
bottom_offset = 6
source = np.float32([[14, 140], [301 ,140],[200, 96], [118, 96]])
destination = np.float32([[image.shape[1]/2 - dst_size, image.shape[0] - bottom_offset],
                  [image.shape[1]/2 + dst_size, image.shape[0] - bottom_offset],
                  [image.shape[1]/2 + dst_size, image.shape[0] - 2*dst_size - bottom_offset],
                  [image.shape[1]/2 - dst_size, image.shape[0] - 2*dst_size - bottom_offset],
                  ])
warped, mask = perspect_transform(image, source, destination)
fig = plt.figure(figsize=(12,3))
plt.subplot(121)
plt.imshow(warped)
plt.subplot(122)
plt.imshow(mask)
#scipy.misc.imsave('../output/warped_example.jpg', warped)

# ## Color Thresholding
# Define the color thresholding function from the lesson and apply it to the warped image
#
# **TODO:** Ultimately, you want your map to not just include navigable terrain but also obstacles and the positions of the rock samples you're searching for.  Modify this function or write a new function that returns the pixel locations of obstacles (areas below the threshold) and rock samples (yellow rocks in calibration images), such that you can map these areas into world coordinates as well.
# **Suggestion:** Think about imposing a lower and upper boundary in your color selection to be more specific about choosing colors.  Feel free to get creative and even bring in functions from other libraries.  Here's an example of [color selection](http://opencv-python-tutroals.readthedocs.io/en/latest/py_tutorials/py_imgproc/py_colorspaces/py_colorspaces.html) using OpenCV.
# **Beware:** if you start manipulating images with OpenCV, keep in mind that it defaults to `BGR` instead of `RGB` color space when reading/writing images, so things can get confusing.

# In[36]:

# Identify pixels above the threshold
# Threshold of RGB > 160 does a nice job of identifying ground pixels only
def color_thresh(img, rgb_thresh=(160, 160, 160)):
    # Create an array of zeros same xy size as img, but single channel
    color_select = np.zeros_like(img[:,:,0])
    # Require that each pixel be above all three threshold values in RGB
    # above_thresh will now contain a boolean array with "True"
    # where threshold was met
    above_thresh = (img[:,:,0] > rgb_thresh[0])                 & (img[:,:,1] > rgb_thresh[1])                 & (img[:,:,2] > rgb_thresh[2])
    # Index the array of zeros with the boolean array and set to 1
    color_select[above_thresh] = 1

    # Return the binary image
    return color_select

def yellow_rock_thresh(img):
    # identify if the image has a rock in it
    # define a threshold for the rock object
    color_select_yellow = np.zeros_like(img[:,:,0])
    yellow_thresh = (img[:,:,0] > 50) & (img[:,:,0] < 258)                  & (img[:,:,1] > 115) & (img[:,:,1] < 256)                  & (img[:,:,2] > -1) &  (img[:,:,2] < 100)

    color_select_yellow[yellow_thresh] = 1

    is_yellow_rock = False
    for nested_list in color_select_yellow:
        for item in nested_list:
            if item:
                is_yellow_rock = True
                break
    return color_select_yellow, is_yellow_rock

threshed = color_thresh(warped)
yellow_rock_img, is_yellow_rock = yellow_rock_thresh(warped)
fig = plt.figure(figsize=(12,3))
plt.subplot(121)
plt.imshow(threshed, cmap='gray')
plt.subplot(122)
plt.imshow(yellow_rock_img, cmap='gray')
print (is_yellow_rock)
#scipy.misc.imsave('../output/warped_threshed.jpg', threshed*255)

# ## Coordinate Transformations
# Define the functions used to do coordinate transforms and apply them to an image.

# In[37]:

def rover_coords(binary_img):
    # Identify nonzero pixels
    ypos, xpos = binary_img.nonzero()
    # Calculate pixel positions with reference to the rover position being at the
    # center bottom of the image.
    x_pixel = np.absolute(ypos - binary_img.shape[0]).astype(np.float)
    y_pixel = -(xpos - binary_img.shape[0]).astype(np.float)
    return x_pixel, y_pixel

# Define a function to convert to radial coords in rover space
def to_polar_coords(x_pixel, y_pixel):
    # Convert (x_pixel, y_pixel) to (distance, angle)
    # in polar coordinates in rover space
    # Calculate distance to each pixel
    dist = np.sqrt(x_pixel**2 + y_pixel**2)
    # Calculate angle away from vertical for each pixel
    angles = np.arctan2(y_pixel, x_pixel)
    return dist, angles

# Define a function to apply a rotation to pixel positions
def rotate_pix(xpix, ypix, yaw):
    # TODO:
    # Convert yaw to radians
    yaw_rad = yaw * np.pi / 180

    # Apply a rotation
    x_rotated = xpix * np.cos(yaw_rad) - ypix * np.sin(yaw_rad)
    y_rotated = xpix * np.sin(yaw_rad) + ypix * np.cos(yaw_rad)
    # Return the result
    return x_rotated, y_rotated

# Define a function to perform a translation
def translate_pix(xpix_rot, ypix_rot, xpos, ypos, scale):
    # TODO:
    # Apply a scaling and a translation
    # Assume a scale factor of 10 between world space pixels and rover space pixels
    # Perform translation and convert to integer since pixel values can't be float
    x_world = np.int_(xpos + (xpix_rot / scale))
    y_world = np.int_(ypos + (ypix_rot / scale))

    # Return the result
    return x_world, y_world

# Define a function to apply rotation and translation (and clipping)
# Once you define the two functions above this function should work
def pix_to_world(xpix, ypix, xpos, ypos, yaw, world_size, scale):
    # Apply rotation
    xpix_rot, ypix_rot = rotate_pix(xpix, ypix, yaw)
    # Apply translation
    xpix_tran, ypix_tran = translate_pix(xpix_rot, ypix_rot, xpos, ypos, scale)
    # Perform rotation, translation and clipping all at once
    x_pix_world = np.clip(np.int_(xpix_tran), 0, world_size - 1)
    y_pix_world = np.clip(np.int_(ypix_tran), 0, world_size - 1)
    # Return the result
    return x_pix_world, y_pix_world

# Grab another random image
idx = np.random.randint(0, len(img_list)-1)
#image = mpimg.imread(img_list[idx])
image = rock_img
warped, mask = perspect_transform(image, source, destination)
threshed = color_thresh(warped)  # navigable terrain pixels
obs_map = np.absolute(np.float32(threshed)-1)*mask  #obstacles pixels

# Calculate pixel values in rover-centric coords and distance/angle to all pixels
xpix, ypix = rover_coords(threshed)
dist, angles = to_polar_coords(xpix, ypix)
mean_dir = np.mean(angles)

# Do some plotting
fig = plt.figure(figsize=(12,9))
plt.subplot(221)
plt.imshow(image)
plt.subplot(222)
plt.imshow(warped)
plt.subplot(223)
plt.imshow(threshed, cmap='gray')
plt.subplot(224)
plt.plot(xpix, ypix, '.')
plt.ylim(-160, 160)
plt.xlim(0, 160)
arrow_length = 100
x_arrow = arrow_length * np.cos(mean_dir)
y_arrow = arrow_length * np.sin(mean_dir)
plt.arrow(0, 0, x_arrow, y_arrow, color='red', zorder=2, head_width=10, width=2)

# ## Read in saved data and ground truth map of the world
# The next cell is all setup to read your saved data into a `pandas` dataframe.  Here you'll also read in a "ground truth" map of the world, where white pixels (pixel value = 1) represent navigable terrain.
#
# After that, we'll define a class to store telemetry data and pathnames to images.  When you instantiate this class (`data = Databucket()`) you'll have a global variable called `data` that you can refer to for telemetry and map data within the `process_image()` function in the following cell.
#

# In[38]:

# Import pandas and read in csv file as a dataframe
import pandas as pd
# Change this path to your data directory
df = pd.read_csv('./test_dataset/robot_log.csv')
csv_img_list = df["Path"].tolist() # Create list of image pathnames
# Read in ground truth map and create a 3-channel image with it
ground_truth = mpimg.imread('./calibration_images/map_bw.png')
ground_truth_3d = np.dstack((ground_truth*0, ground_truth*255, ground_truth*0)).astype(np.float)

# Creating a class to be the data container
# Will read in saved data from csv file and populate this object
# Worldmap is instantiated as 200 x 200 grids corresponding
# to a 200m x 200m space (same size as the ground truth map: 200 x 200 pixels)
# This encompasses the full range of output position values in x and y from the sim
class Databucket():
    def __init__(self):
        self.images = csv_img_list
        self.xpos = df["X_Position"].values
        self.ypos = df["Y_Position"].values
        self.yaw = df["Yaw"].values
        self.count = -1 # This will be a running index, setting to -1 is a hack
                        # because moviepy (below) seems to run one extra iteration
        self.worldmap = np.zeros((200, 200, 3)).astype(np.float)
        self.ground_truth = ground_truth_3d # Ground truth worldmap

# Instantiate a Databucket().. this will be a global variable/object
# that you can refer to in the process_image() function below
data = Databucket()

# ## Write a function to process stored images
#
# Modify the `process_image()` function below by adding in the perception step processes (functions defined above) to perform image analysis and mapping.  The following cell is all set up to use this `process_image()` function in conjunction with the `moviepy` video processing package to create a video from the images you saved taking data in the simulator.
#
# In short, you will be passing individual images into `process_image()` and building up an image called `output_image` that will be stored as one frame of video.  You can make a mosaic of the various steps of your analysis process and add text as you like (example provided below).
#
#
#
# To start with, you can simply run the next three cells to see what happens, but then go ahead and modify them such that the output video demonstrates your mapping process.  Feel free to get creative!

# In[43]:

# Define a function to pass stored images to
# reading rover position and yaw angle from csv file
# This function will be used by moviepy to create an output video
def process_image(img):
    # Example of how to use the Databucket() object defined above
    # to print the current x, y and yaw values
    # print(data.xpos[data.count], data.ypos[data.count], data.yaw[data.count])

    # TODO:
    # 1) Define source and destination points for perspective transform
    source = np.float32([
                [14, 140], [301, 140], # bottom left to right
                [200, 96], [118, 96]]) # top rigt to left
    dst_size = 5
    bottom_offset = 6
    destination = np.float32([[image.shape[1]/2 - dst_size, image.shape[0] - bottom_offset],
                  [image.shape[1]/2 + dst_size, image.shape[0] - bottom_offset],
                  [image.shape[1]/2 + dst_size, image.shape[0] - 2*dst_size - bottom_offset],
                  [image.shape[1]/2 - dst_size, image.shape[0] - 2*dst_size - bottom_offset],
                  ])

    # 2) Apply perspective transform
    warped, mask = perspect_transform(img, source, destination)
    # 3) Apply color threshold to identify navigable terrain/obstacles/rock samples

    # def color_thresh
    threshed = color_thresh(warped) # navigable pixels
    yellow_rock_threshed, is_yellow_rock = yellow_rock_thresh(warped)  # rock pixels
    obs_map = np.absolute(np.float32(threshed)-1)*mask  #obstacles pixels

    # 4) Convert thresholded image pixel values to rover-centric coords
    xpix, ypix = rover_coords(threshed)
    xpix_obs, ypix_obs = rover_coords(obs_map)

    # 5) Convert rover-centric pixel values to world coords
    scale = 10
    x_world, y_world = pix_to_world(xpix, ypix, data.xpos[data.count], data.ypos[data.count], \
        data.yaw[data.count], data.worldmap.shape[0], scale)
    x_obs, y_obs = pix_to_world(xpix_obs, ypix_obs, data.xpos[data.count], \
        data.ypos[data.count], data.yaw[data.count], data.worldmap.shape[0], scale)

    # 6) Update worldmap (to be displayed on right side of screen)
        # Example: data.worldmap[obstacle_y_world, obstacle_x_world, 0] += 1
        #          data.worldmap[rock_y_world, rock_x_world, 1] += 1
        #          data.worldmap[navigable_y_world, navigable_x_world, 2] += 1
    data.worldmap[y_world, x_world, 2] += 255  # blue nav terrain
    data.worldmap[y_obs, x_obs, 0] += 255 # red obstacles
    nav_pix = data.worldmap[:,:,2] > 0
    data.worldmap[nav_pix, 0] = 0

    if is_yellow_rock:
        xpix_rock, ypix_rock = rover_coords(yellow_rock_threshed)
        x_world_rock, y_world_rock = pix_to_world(xpix_rock, ypix_rock, data.xpos[data.count], \
            data.ypos[data.count], data.yaw[data.count], data.worldmap.shape[0], scale)
        data.worldmap[y_world_rock, x_world_rock, :] = 255

    # 7) Make a mosaic image, below is some example code
        # First create a blank image (can be whatever shape you like)
    output_image = np.zeros((img.shape[0] + data.worldmap.shape[0], img.shape[1]*2, 3))
        # Next you can populate regions of the image with various output
        # Here I'm putting the original image in the upper left hand corner
    #left corner
    output_image[0:img.shape[0], 0:img.shape[1]] = img

    # first a warped image
    # upper right
    output_image[0:img.shape[0], img.shape[1]:] = warped

    # Overlay worldmap with ground truth map
    map_add = cv2.addWeighted(data.worldmap, 1, data.ground_truth, 0.5, 0)
        # Flip map overlay so y-axis points upward and add to output_image
    #lower left
    output_image[img.shape[0]:, 0:data.worldmap.shape[1]] = np.flipud(map_add)
    #lower right ??????

        # Then putting some text over the image
    if is_yellow_rock:
        cv2.putText(output_image,"Rock Found!!!", (20, 20),
                cv2.FONT_HERSHEY_COMPLEX, 0.4, (255, 255, 255), 1)
    else:
         cv2.putText(output_image,"No Rock", (20, 20),
                cv2.FONT_HERSHEY_COMPLEX, 0.4, (255, 255, 255), 1)
    data.count += 1 # Keep track of the index in the Databucket()

    return output_image

# ## Make a video from processed image data
# Use the [moviepy](https://zulko.github.io/moviepy/) library to process images and create a video.
#

# In[44]:

# Import everything needed to edit/save/watch video clips
from moviepy.editor import VideoFileClip
from moviepy.editor import ImageSequenceClip

# Define pathname to save the output video
output = './output/test_mapping.mp4'
data = Databucket() # Re-initialize data in case you're running this cell multiple times
clip = ImageSequenceClip(data.images, fps=60) # Note: output video will be sped up because
                                          # recording rate in simulator is fps=25
new_clip = clip.fl_image(process_image) #NOTE: this function expects color images!!
get_ipython().run_line_magic('time', 'new_clip.write_videofile(output, audio=False)')

# ### This next cell should function as an inline video player
# If this fails to render the video, try running the following cell (alternative video rendering method).  You can also simply have a look at the saved mp4 in your `/output` folder

# In[45]:

output = './output/test_mapping.mp4'
from IPython.display import HTML
HTML("""
<video width="960" height="540" controls>
  <source src="{0}">
</video>
""".format(output))

# In[ ]:

# In[ ]:


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:light
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.4.2
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## Import Main Librairies

# +
# Basic Python modules
import random
import csv
import pickle
import os

# Scientific modules
import numpy as np
import scipy as scp
import matplotlib
from matplotlib import pyplot as plt

from scipy.io import loadmat
import networkx as nx

# For better looking graphs
import seaborn as sns
sns.set()

# Just a custom color palette that I use
# blue, darkblue, red, orange, green, palegreen, yellow, brokenwhite, brokengrey
colors = ["#AF3127","#6182B5","#112A3C","#D99C37","#90A954","#C5B868","#FAC764","#DAC0A6","#C4C2D1"]
sns.set_palette(sns.color_palette(colors))
# -

# ## Configure paths for data once for all

# +
#generate_images = False

#pickle_data_folder = 'new_new_pickle_data_cross/'

# Folder in which data is placed
data_folder = '../fake_data'

# Folder in which some computed objects will be saved
pickle_folder = 'cross_validation'

# Control whether we recompute data even if it has been previously
# computed and pickled
recompute_pickle = False

# Suffix at the end of each .mat file; it is specified here in order to select
# relevant files, as well as to make file name lighter during loading for
# further operations (such as printing subjects names), since it does not carry
# additional information.
suffix = '_fiber_number.mat'

# For instance here, with those setting, every ../fake_data/*_fiber_number.mat
# will be loaded

# Keys used to split data between patients and controls. Subject whose filename
# contains one of the control_keys will be affected to the control cohort, and
# similarly for patients.
control_keys = ['060', 'dep', 'dpr', 'S', 'TI']
patient_keys = ['lgp']

# By default, the code expects a "table.csv" present in data_folder, containing
# information about the patients, such as their age, the duration of the
# disease, etc.
csv_path = data_folder + "/table.csv"

# -

# ## Load data

# +
def get_matrix_file_list(data_folder, suffix):
    """
        Return the list of files in the folder data_folder ending by suffix
    """
    file_list = [f for f in os.listdir(data_folder) if f.endswith(suffix)]
    return list(map(lambda x: data_folder + '/' + x, file_list))

def load_matrix_file(file):
    """
        Return the matrix loaded from a Matlab data file. Note that the
        'Measure' key is hardcoded, so you might need to adjust it to your own
        data.
    """
    return loadmat(file)['Measure']

# Eventually create the folder in which data will be saved
if not os.path.exists(pickle_folder):
    os.mkdir(pickle_folder)

# Create a dictionnary of all the matrices, where each matrix gets associated to
# the filename of the corresponding .mat file minus the suffix.
connectivity_matrices = {}

for f in get_matrix_file_list(data_folder, suffix):
    connectivity_matrices[f.replace(suffix,'').replace(data_folder + '/','')] = load_matrix_file(f)

# Create a dictionnary of metadata for each patient, obtained from the file
# at csv_path (by default 'data_foler/table.csv')
patient_info_dict = {}

with open(csv_path, 'r') as csv_file:
    metadata = csv.DictReader(csv_file)
    # Each patient is associated to a dictionnary containing all its information
    for row in metadata:
        metadata_dict = {key:row[key] for key in metadata.fieldnames if key != 'Subject'}
        patient_info_dict[row['Subject']] = metadata_dict

print("Succesfully loaded {} matrices from {}.".format(len(connectivity_matrices), data_folder))
print("Metadata has been found in {} for {} subjects.".format(csv_path,len(patient_info_dict)))
# -

# ## Split Data into Cohorts

# +
# list of controls and patients names
controls = []
patients = []

# The following can be used to limit the number of either controls or patients
# considered for the study if they are set to some non infinite number.
controls_count = np.inf
patients_count = np.inf

current_control = 0
current_patient = 0

for key in [*connectivity_matrices]:
    # Use patients_keys and control_keys list to classify subject into cohorts
    if any(list(map(lambda x: x in key, patient_keys))) and current_patient < patients_count:
        patients.append(key)
        current_patient += 1
    elif any(list(map(lambda x: x in key, control_keys))) and current_control < controls_count:
        controls.append(key)
        current_control += 1
    else:
        print("Patient {} cannot be classified either as control or patient.".format(key))

controls_count = current_control
patients_count = current_patient

subject_count = len(patients) + len(controls)

print("Classified {} controls and {} patients (total {} subjects)".format(controls_count, patients_count, subject_count))

# -

# ## Basic network manipulation functions

# +
def plot_connectivity_matrix(*matrix, save=None):
    """
        Plot the connectivity matrix as a colormap
        It uses a logarithmic scale for the colors and it is possible to plot
        multiple matrices at once.
    """

    max_rows = 2
    line = len(matrix) / 3 + 1

    for i, mat in enumerate(matrix):
        plt.subplot(line, max_rows, i+1)
        plt.imshow(np.log(mat+1),cmap='viridis')

    plt.tight_layout()
    if save is None:
        plt.show()
    else:
        plt.savefig(save,bbox_inches='tight')

def get_network(matrix, threshold = 0):
    """
        Return the network (as a networkx data structure) defined by matrix.
        It is possible to specify a threshold that will disregard all the
        edges below this threshold when creating the network
    """
    G = nx.Graph()
    N = matrix.shape[0]
    G.add_nodes_from(list(range(N)))
    G.add_weighted_edges_from([(i,j,1.0*matrix[i][j]) for i in range(0,N) for j in range(0,i) \
                                                                   if matrix[i][j] >= threshold])
    return G

def filter_weights(matrix, ratio):
    """
        Only keep a fraction of the weights of the matrix, fraction specified
        via the threshold parameter
    """
    n = matrix.shape[0]
    filtered_matrix = np.zeros_like(matrix)
    total_weight = np.sum(matrix)
    weights_id = sorted([(matrix[i,j],i,j) for i in range(n-1) for j in range(i+1,n)], reverse=True)

    filtered_weight = 0
    for (w,i,j) in weights_id:
        filtered_weight += 2*w

        if filtered_weight > ratio*total_weight:
            break

        filtered_matrix[i,j] = w
        filtered_matrix[j,i] = w

    return filtered_matrix

# -

# ## Laplacian manipulation

# +
def get_laplacian_matrix(matrix, threshold = 0):
    """
        Return the laplacian matrix of the network defined by matrix.
        It is possible to specify a threshold that will disregard all the
        edges below this threshold when creating the network
    """

    L = np.zeros_like(matrix)

    assert(L.shape[0] == L.shape[1]), "The network should be encoded by a square matrix"

    N = L.shape[0]

    # Small one liner trick to get rid of the diagonal and apply the threshold.
    A = np.multiply(np.where(matrix >= threshold, matrix , np.zeros_like(matrix)), np.ones((N,N)) - np.eye(N))

    # Use the definition of L as L = D - A
    L = np.eye(N) * np.array([np.sum(A[i,:]) for i in range(N)]) - A

    return L

def get_stabilized_laplacian_matrix(matrix, threshold = 0, ratio=1.0, alpha = 1e-2):
    """
        Prefered function to get the laplacian matrix of a network encoded by a
        matrix. Can apply a ratio filtering of the weights of the matrix (i.e.
        only keep ratio of the total weight) or / and apply a threshold. The
        alpha parameter corrects a numerical issue.
    """
    L = get_laplacian_matrix(filter_weights(matrix, ratio = ratio), threshold)

    n = L.shape[0]

    # Add a small epsilon to the diagonal of the Laplacian matrix. It should
    # make the computation more stable. Idea from https://github.com/Hermina/GOT

    return L + alpha * np.eye(n)

# -

# Small tests to assess consistency of the different functions.
assert np.allclose(nx.linalg.laplacian_matrix(get_network(connectivity_matrices[controls[0]])).toarray(), get_laplacian_matrix(connectivity_matrices[controls[0]])), "Inconsistency in the Laplacian functions"
assert np.allclose(nx.linalg.laplacian_matrix(get_network(connectivity_matrices[patients[0]], threshold=10)).toarray(), get_stabilized_laplacian_matrix(connectivity_matrices[patients[0]], alpha=0, threshold=10)), "Inconsistency in the Laplacian functions (test 2)"

# ## Wasserstein distance functions

# +
# Those are dictionaries used to speed up computations, by avoiding multiple
# computations of the same objects.
precomp_dict_net = {}
precomp_dict_cov = {}

def wasserstein_networks(A, B, threshold = 0, threshold2 = None, ratio = 1.0):
    """
            Returns the Wasserstein distance (in the sense of the GOT paper)
            between the two networks encoded by A and B. It is possible to
            apply a threshold to the network obtained via both matrices (and
            even a different one per matrices)
    """

    if threshold2 is None:
        threshold2 = threshold

    assert(A.shape == B.shape), "compared networks should have the \
                                                        same number of nodes"

    L1dag = None
    L1dag_sqrt = None
    L2dag = None
    L2dag_sqrt = None

    if (A.tobytes(), threshold, ratio) not in precomp_dict_net:
        L1 = get_stabilized_laplacian_matrix(A, threshold = threshold, ratio = ratio)
        L1dag = np.real(scp.linalg.pinv(L1))
        L1dag_sqrt = np.real(scp.linalg.sqrtm(L1dag))

        precomp_dict_net[(A.tobytes(), threshold,ratio)] = (L1dag, L1dag_sqrt)
    else:
        L1dag, L1dag_sqrt = precomp_dict_net[(A.tobytes(), threshold, ratio)]

    if (B.tobytes(), threshold2, ratio) not in precomp_dict_net:
        L2 = get_stabilized_laplacian_matrix(B, threshold = threshold2, ratio = ratio)
        L2dag = np.real(scp.linalg.pinv(L2))
        L2dag_sqrt = np.real(scp.linalg.sqrtm(L2dag))

        precomp_dict_net[(B.tobytes(), threshold2, ratio)] = (L2dag, L2dag_sqrt)

    else :
        L2dag, L2dag_sqrt = precomp_dict_net[(B.tobytes(), threshold2, ratio)]

    W = np.trace(L1dag) + np.trace(L2dag) - 2*np.trace(np.real(scp.linalg.sqrtm(L1dag_sqrt @ L2dag @ L1dag_sqrt)))

    return W

def wasserstein_network_covariance(A, cov, threshold = 0, ratio = 1.0):
    """
            Return the Wasserstein distance (in the sense of the GOT paper)
            between the the network encoded by a matrix A, and the normal
            distribution whose covariance matrix is given by cov.
            It is possible to apply a threshold to the network obtained via
            matrix A.
    """

    assert(A.shape == cov.shape), "compared networks should have the \
                                                        same number of nodes"

    N = A.shape[0]

    L1dag = None
    L1dag_sqrt = None
    L2dag = None
    L2dag_sqrt = None

    if (A.tobytes(), threshold, ratio) not in precomp_dict_net:
        L1 = get_stabilized_laplacian_matrix(A, threshold = threshold, ratio = ratio)
        L1dag = np.real(scp.linalg.pinv(L1))
        L1dag_sqrt = np.real(scp.linalg.sqrtm(L1dag))

        precomp_dict_net[(A.tobytes(), threshold ,ratio)] = (L1dag, L1dag_sqrt)
    else:
        L1dag, L1dag_sqrt = precomp_dict_net[(A.tobytes(), threshold ,ratio)]

    if cov.tobytes() not in precomp_dict_cov:
        L2dag = cov
        L2dag_sqrt = np.real(scp.linalg.sqrtm(L2dag))

        precomp_dict_cov[cov.tobytes()] = L2dag_sqrt

    else:
        L2dag = cov
        L2dag_sqrt = precomp_dict_cov[cov.tobytes()]

    W = np.trace(L1dag) + np.trace(L2dag) - 2*np.trace(np.real(scp.linalg.sqrtm(L1dag_sqrt @ L2dag @ L1dag_sqrt)))
    return W

def clip_remove(array, threshold):
    """
        Small helper function to clip an array by removing values above a
        certain threshold
    """
    r = array[array < threshold]
    print("Kept {} out of {} ({} %)".format(r.shape[0], array.shape[0], r.shape[0]/array.shape[0] * 100))
    return r

def barycenter_iteration(cov_list, iteration_count = 10):
    """
        Return the covariance matrix of the barycenter of the distributions
        whose covariances are given as cov_list ; iteration_count controls the
        number of iteration made by the iterative algorithm.
    """
    covariance_count = len(cov_list)

    size = cov_list[0].shape[0]
    curr_cov = np.eye(size)

    for iteration in range(iteration_count):
        curr_cov_sqrt = np.real(scp.linalg.sqrtm(curr_cov))
        curr_cov_pinv_sqrt = scp.linalg.pinv(curr_cov_sqrt)
        big_sum = np.zeros_like(curr_cov)
        for sigma in cov_list:
            big_sum += 1/covariance_count * np.real(scp.linalg.sqrtm(curr_cov_sqrt @ sigma @ curr_cov_sqrt))
        big_sum = big_sum @ big_sum
        curr_cov = curr_cov_pinv_sqrt @ big_sum @ curr_cov_pinv_sqrt

    return curr_cov

# Some functions to print a barycenter covariance as an approximate connectivity matrix
def cov_to_laplacian(cov, alpha = 1e-2):
    n = cov.shape[0]
    return scp.linalg.pinv(cov) - alpha * np.eye(n)

def laplacian_to_connectivity_matrix(laplacian):
    n = laplacian.shape[0]
    return np.multiply(laplacian, -(np.ones_like(laplacian) - 2*np.eye(n)))
# -

# ## Compute a distance matrix (distance subjects to controls)

# ## Compute a barycenter with subsample of controls

def subsample_controls(controls, sample_size, ratio = 1.0):
    """
        Return a random draw of sample_size controls, whose name are taken from
        the controls array, as well as the covariance matrices of their
        associated distribution. Ratio is used to control the sparsity of the
        networks (see get_stabilized_laplacian_matrix for its exact effect)
    """
    controls_count = len(controls)

    random_controls_sample = random.sample(range(0, controls_count),sample_size)
    controls_barycenter = np.array(controls)[random_controls_sample]

    pseudo_inverses = [scp.linalg.pinv(get_stabilized_laplacian_matrix(connectivity_matrices[subject],ratio=ratio)) for subject in controls_barycenter]

    return (controls_barycenter, pseudo_inverses)

def do_some_folds(fold_count = 50, control_sample_size = 10, ratios = [1.0]):
    """
        Perform fold_count random folds. Each fold consists in: subsampling the
        control group to randomly select control_sample_size controls, whose
        barycenter is computed. Then all remaining controls + each patients are
        compared to the barycenter. The results of multiple folds are aggregated
        into distance_dict and distance_fold_count (the average distance for
        subject s to the barycenters is thus

                        distance_dict[s] / distance_fold_count[s]

        ratios can be used to perform tests on multiple sparsity ratios at a
        time
    """
    for ratio in ratios:
        true_folds = 0
        print("Folding with sparsity ratio {:.2f}".format(ratio))

        pickle_file_path = pickle_folder + "/dist_and_fold_count_r{:.2f}_k{}.pickle".format(ratio, control_sample_size)
        print("Will save or update {}".format(pickle_file_path))

        distance_dict = {s : 0.0 for s in patients + controls}
        distance_fold_count = {s:0 for s in patients + controls}

        for fold in range(fold_count):
            if fold % 10 == 0: print("Fold {}".format(fold))

            controls_sample, pinvs = subsample_controls(controls, control_sample_size, ratio = ratio)

            barycenter_cov = barycenter_iteration(pinvs)

            # list of all the other subjects (non sampled controls + patients)
            subjects = patients + [c for c in controls if c not in controls_sample]
            current_fold_distance_dict = {subject:wasserstein_network_covariance(connectivity_matrices[subject], barycenter_cov, ratio = ratio) for subject in subjects}

            for subject, d in current_fold_distance_dict.items():
                distance_fold_count[subject] += 1
                distance_dict[subject] += d

            true_folds += 1

        # create or update the pickled data
        if os.path.exists(pickle_file_path) and not recompute_pickle:
            with open(pickle_file_path, 'rb') as filehandler:
                loaded_dist_dict, loaded_dist_count, loaded_fold_count = pickle.load(filehandler)

                for s in patients + controls:
                    distance_dict[s] += loaded_dist_dict[s]
                    distance_fold_count[s] += loaded_dist_count[s]

                print("{} ({} folds) will be updated with {} more folds".format(pickle_file_path, loaded_fold_count, fold_count))
                true_folds += loaded_fold_count

        with open(pickle_file_path, 'wb') as filehandler:
            print("Saving {}".format(pickle_file_path))
            pickle.dump((distance_dict, distance_fold_count, true_folds), filehandler)

        # plot the results

        # first, obtain the true distance lists from the distance_dict and
        # distance_fold_count dictionnaries
        dist_list = []
        for s in patients:
            dist_list.append(distance_dict[s] / distance_fold_count[s])

        for c in controls:
            dist_list.append(distance_dict[c] / distance_fold_count[c])

        #print(scp.stats.ttest_ind(dist_list[:patients_count], dist_list[patients_count:], equal_var=True))

        #threshold = 0.0025
        #print("Plot thresholded at {}".format(threshold))
        plt.hist(dist_list[:patients_count], stacked=True, density=True, alpha=0.3, label='Patients')
        plt.hist(dist_list[patients_count:], stacked=True, density=True, alpha = 0.3, label='Controls')
        sns.kdeplot(dist_list[:patients_count], shade=False, color='C0', alpha=0.3);
        sns.kdeplot(dist_list[patients_count:], shade=False, color='C1', alpha=0.3);
        plt.xlabel("Distance to barycenter ({:.2f})".format(ratio))
        plt.legend()
        #import tikzplotlib
        #tikzplotlib.save("figures_generation/dist_to_barycenter_full.tex".format(ratio), wrap=False)

        plt.show()

do_some_folds()

import numpy,import scipy,from scipy,import networkx


"""
Selecting Data II - Access a Column
"""

import numpy as np
import pandas as pd
wine_reviews = pd.read_csv('../winemag-data-130k.csv')

# Return the winery column in a variable called "wineries".

# Run a command to return the type of wineries.


import numpy as np
import pytest

import taichi as ti
from tests import test_utils

def with_data_type(dt):
    val = ti.field(ti.i32)

    n = 4

    ti.root.dense(ti.i, n).place(val)

    @ti.kernel
    def test_numpy(arr: ti.ext_arr()):
        for i in range(n):
            arr[i] = arr[i]**2

    a = np.array([4, 8, 1, 24], dtype=dt)

    for i in range(n):
        a[i] = i * 2

    test_numpy(a)

    for i in range(n):
        assert a[i] == i * i * 4

@test_utils.test()
def test_numpy_f32():
    with_data_type(np.float32)

@test_utils.test(require=ti.extension.data64)
def test_numpy_f64():
    with_data_type(np.float64)

@test_utils.test()
def test_numpy_i32():
    with_data_type(np.int32)

@test_utils.test(require=ti.extension.data64)
def test_numpy_i64():
    with_data_type(np.int64)

@test_utils.test()
def test_numpy_2d():
    val = ti.field(ti.i32)

    n = 4
    m = 7

    ti.root.dense(ti.i, n).dense(ti.j, m).place(val)

    @ti.kernel
    def test_numpy(arr: ti.ext_arr()):
        for i in range(n):
            for j in range(m):
                arr[i, j] += i + j

    a = np.empty(shape=(n, m), dtype=np.int32)

    for i in range(n):
        for j in range(m):
            a[i, j] = i * j

    test_numpy(a)

    for i in range(n):
        for j in range(m):
            assert a[i, j] == i * j + i + j

@test_utils.test()
def test_numpy_2d_transpose():
    val = ti.field(ti.i32)

    n = 8
    m = 8

    ti.root.dense(ti.ij, (n, m)).place(val)

    @ti.kernel
    def test_numpy(arr: ti.ext_arr()):
        for i in ti.grouped(val):
            val[i] = arr[i]

    a = np.empty(shape=(n, m), dtype=np.int32)

    for i in range(n):
        for j in range(m):
            a[i, j] = i * j + i * 4

    test_numpy(a.transpose())

    for i in range(n):
        for j in range(m):
            assert val[i, j] == i * j + j * 4

@test_utils.test()
def test_numpy_3d():
    val = ti.field(ti.i32)

    n = 4
    m = 7
    p = 11

    ti.root.dense(ti.i, n).dense(ti.j, m).dense(ti.k, p).place(val)

    @ti.kernel
    def test_numpy(arr: ti.ext_arr()):
        for i in range(n):
            for j in range(m):
                for k in range(p):
                    arr[i, j, k] += i + j + k * 2

    a = np.empty(shape=(n, m, p), dtype=np.int32)

    for i in range(n):
        for j in range(m):
            for k in range(p):
                a[i, j, k] = i * j * (k + 1)

    test_numpy(a)

    for i in range(n):
        for j in range(m):
            for k in range(p):
                assert a[i, j, k] == i * j * (k + 1) + i + j + k * 2

@test_utils.test()
def test_numpy_3d_error():
    val = ti.field(ti.i32)

    n = 4
    m = 7
    p = 11

    ti.root.dense(ti.i, n).dense(ti.j, m).dense(ti.k, p).place(val)

    @ti.kernel
    def test_numpy(arr: ti.ext_arr()):
        for i in range(n):
            for j in range(m):
                for k in range(p):
                    arr[i, j] += i + j + k * 2

    a = np.empty(shape=(n, m, p), dtype=np.int32)

    with pytest.raises(ti.TaichiCompilationError):
        test_numpy(a)

@test_utils.test()
def test_numpy_multiple_external_arrays():

    n = 4

    @ti.kernel
    def test_numpy(a: ti.ext_arr(), b: ti.ext_arr()):
        for i in range(n):
            a[i] = a[i] * b[i]
            b[i] = a[i] + b[i]

    a = np.array([4, 8, 1, 24], dtype=np.int32)
    b = np.array([5, 6, 12, 3], dtype=np.int32)
    c = a * b
    d = c + b

    test_numpy(a, b)
    for i in range(n):
        assert a[i] == c[i]
        assert b[i] == d[i]

@test_utils.test()
def test_index_mismatch():
    with pytest.raises(AssertionError):
        val = ti.field(ti.i32, shape=(1, 2, 3))
        val[0, 0] = 1

@test_utils.test()
def test_numpy_zero():
    @ti.kernel
    def test_numpy(arr: ti.ext_arr()):
        pass

    test_numpy(np.empty(shape=(0), dtype=np.int32))
    test_numpy(np.empty(shape=(0, 5), dtype=np.int32))
    test_numpy(np.empty(shape=(5, 0), dtype=np.int32))

@test_utils.test()
def test_numpy_struct_for():
    @ti.kernel
    def func1(a: ti.any_arr()):
        for i, j in a:
            a[i, j] = i + j

    m = np.zeros((123, 456), dtype=np.int32)
    func1(m)
    for i in range(123):
        for j in range(456):
            assert m[i, j] == i + j

    @ti.kernel
    def func2(a: ti.any_arr()):
        for I in ti.grouped(a):
            a[I] = I.sum()

    n = np.zeros((98, 76, 54), dtype=np.int32)
    func2(n)
    for i, j, k in ti.ndrange(98, 76, 54):
        assert n[i, j, k] == i + j + k


"""
  Unit tests for the class NNModifier in nn_modifiers.py
  -- kandasamy@cs.cmu.edu
"""

# pylint: disable=invalid-name

from copy import deepcopy
import numpy as np
import os
import six
from shutil import rmtree
# Local imports
from nn import nn_constraint_checkers
from nn import nn_modifiers
from nn.neural_network import NeuralNetwork
from nn.nn_visualise import visualise_nn
from unittest_neural_network import generate_cnn_architectures, generate_mlp_architectures
from utils.base_test_class import BaseTestClass, execute_tests

def test_if_two_networks_are_equal(net1, net2, false_if_net1_is_net2=True):
  """ Returns true if both net1 and net2 are equal.
      If any part of net1 is copied onto net2, then the output will be false
      if false_if_net1_is_net2 is True (default).
  """
  is_true = True
  for key in net1.__dict__.keys():
    val1 = net1.__dict__[key]
    val2 = net2.__dict__[key]
    is_true = True
    if isinstance(val1, dict):
      if false_if_net1_is_net2:
        is_true = is_true and (val1 is not val2)
      for val_key in val1.keys():
        is_true = is_true and np.all(val1[val_key] == val2[val_key])
    elif hasattr(val1, '__iter__'):
      if false_if_net1_is_net2:
        is_true = is_true and (val1 is not val2)
      is_true = is_true and np.all(val1 == val2)
    else:
      is_true = is_true and val1 == val2
    if not is_true: # break here if necessary
      return is_true
  return is_true

def test_for_orig_vs_modifications(save_dir, save_prefix, old_nn,
                                   get_modifications, constraint_checker, write_result):
  """ Tests for the original network and the modifications. Also, visualises the networks.
  """
  visualise_nn(old_nn, os.path.join(save_dir, '%s_orig'%(save_prefix)))
  old_nn_copy = deepcopy(old_nn)
  # Get the modified networks.
  new_nns = get_modifications(old_nn)
  # Go through each new network.
  for new_idx, new_nn in enumerate(new_nns):
    assert isinstance(new_nn, NeuralNetwork)
    assert constraint_checker(new_nn)
    visualise_nn(new_nn, os.path.join(save_dir, '%s_%d'%(save_prefix, new_idx)))
  # Finally test if the networks have not changed.
  assert test_if_two_networks_are_equal(old_nn, old_nn_copy)
  write_result('%s (%s):: #new-networks: %d.'%(
    save_prefix, old_nn.nn_class, len(new_nns)), 'test_result')

class NNModifierTestCase(BaseTestClass):
  """ Unit tests for the NNModifier class. """

  def __init__(self, *args, **kwargs):
    """ Constructor. """
    super(NNModifierTestCase, self).__init__(*args, **kwargs)
    self.cnns = generate_cnn_architectures()
    self.mlps = generate_mlp_architectures()
    self.save_dir = '../scratch/unittest_modifier_class/'
    self.cnn_constraint_checker = nn_constraint_checkers.CNNConstraintChecker(
      50, 4, np.inf, 4.0, 5, 5, 100, 8000, 8)
    self.mlp_constraint_checker = nn_constraint_checkers.MLPConstraintChecker(
      50, 4, np.inf, 4.0, 5, 5, 100, 8000, 8)
    self.cnn_modifier = nn_modifiers.NNModifier(self.cnn_constraint_checker)
    self.mlp_modifier = nn_modifiers.NNModifier(self.mlp_constraint_checker)
    self.modifier_wo_cc = nn_modifiers.NNModifier(None)

  def _get_modifier_and_cc(self, nn):
    """ Returns modifier for the neural network nn."""
    if nn.nn_class == 'cnn':
      modifier = self.cnn_modifier
      constraint_checker = self.cnn_constraint_checker
    else:
      modifier = self.mlp_modifier
      constraint_checker = self.mlp_constraint_checker
    return modifier, constraint_checker

  def test_get_primitives(self):
    """ Test for the get_primitives_grouped_by_type method. """
    self.report('Testing get_primitives_grouped_by_type')
    test_nns = self.cnns + self.mlps
    primitives, _ = self.cnn_modifier.get_primitives_grouped_by_type(self.cnns[0])
    self.report('Types of primitives: %s'%(primitives.keys()), 'test_result')
    for idx, nn in enumerate(test_nns):
      nn_copy = deepcopy(nn)
      modifier, _ = self._get_modifier_and_cc(nn)
      primitives, _ = modifier.get_primitives_grouped_by_type(nn)
      report_str = '%d (%s n=%d,m=%d):: '%(idx, nn.nn_class, nn.num_layers,
                                           nn.get_total_num_edges())
      total_num_primitives = 0
      for _, list_or_prims in six.iteritems(primitives):
        report_str += '%d, '%(len(list_or_prims))
        total_num_primitives += len(list_or_prims)
      report_str += 'tot=%d'%(total_num_primitives)
      self.report(report_str, 'test_result')
      assert test_if_two_networks_are_equal(nn_copy, nn)

  def test_get_single_step_modifications(self):
    """ Tests single step modifications. """
    self.report('Testing single step modifications.')
    save_dir = os.path.join(self.save_dir, 'single_step')
    if os.path.exists(save_dir):
      rmtree(save_dir)
    # Now iterate through the test networks
    test_nns = self.cnns + self.mlps
    for idx, old_nn in enumerate(test_nns):
      save_prefix = str(idx)
      modifier, constraint_checker = self._get_modifier_and_cc(old_nn)
      if idx in [2, 12]:
        num_modifications = 'all'
      else:
        num_modifications = 'all'
      get_modifications = lambda arg_nn: modifier.get_single_step_modifications(
                                           arg_nn, num_modifications)
      test_for_orig_vs_modifications(save_dir, save_prefix, old_nn,
        get_modifications, constraint_checker, self.report)

  def test_multi_step_modifications(self):
    """ Tests multi step modifications. """
    num_steps = 4
    self.report('Testing %d-step modifications.'%(num_steps))
    num_modifications = 20
    save_dir = os.path.join(self.save_dir, 'multi_step_%d'%(num_steps))
    if os.path.exists(save_dir):
      rmtree(save_dir)
    # Now iterate through the test networks
    test_nns = self.cnns + self.mlps
    for idx, old_nn in enumerate(test_nns):
      save_prefix = str(idx)
      modifier, constraint_checker = self._get_modifier_and_cc(old_nn)
      get_modifications = lambda arg_nn: modifier.get_multi_step_modifications(
                                           arg_nn, num_steps, num_modifications)
      test_for_orig_vs_modifications(save_dir, save_prefix, old_nn,
        get_modifications, constraint_checker, self.report)

  def test_call(self):
    """ Tests the __call__ function with a single input of the modifier. """
    self.report('Testing the __call__ function with single input of the modifier.')
    num_modifications = 20
    num_steps_probs = [0.5, 0.25, 0.125, 0.075, 0.05]
    save_dir = os.path.join(self.save_dir, 'modifier_call_single')
    if os.path.exists(save_dir):
      rmtree(save_dir)
    test_nns = self.cnns + self.mlps
    for idx, old_nn in enumerate(test_nns):
      save_prefix = str(idx)
      modifier, constraint_checker = self._get_modifier_and_cc(old_nn)
      get_modifications = lambda arg_nn: modifier(arg_nn, num_modifications,
                                                  num_steps_probs)
      test_for_orig_vs_modifications(save_dir, save_prefix, old_nn,
        get_modifications, constraint_checker, self.report)

  def test_call_with_list(self):
    """ Tests the __call__ function with a single input of the modifier. """
    self.report('Testing the __call__ function with a list of inputs.')
    num_modifications = 40
    num_steps_probs = [0.5, 0.25, 0.125, 0.075, 0.05]
    save_dir = os.path.join(self.save_dir, 'modifier_call_list')
    if os.path.exists(save_dir):
      rmtree(save_dir)
    test_probs = [self.cnns, self.mlps, generate_mlp_architectures('class')]
    for idx, prob in enumerate(test_probs):
      save_prefix = str(idx)
      modifier = self.modifier_wo_cc
      modifications = modifier(prob, num_modifications, num_steps_probs)
      for new_idx, new_nn in enumerate(modifications):
        assert isinstance(new_nn, NeuralNetwork)
        visualise_nn(new_nn, os.path.join(save_dir, '%s_%d'%(save_prefix, new_idx)))
      self.report('With list of %d nns(%s):: #new-networks: %d.'%(
                   len(prob), prob[0].nn_class, len(modifications)), 'test_result')

if __name__ == '__main__':
  execute_tests()


"""
Memory Customization
====================

**Author**: Yi-Hsiang Lai (seanlatias@github)

In this tutorial, we demonstrate how memory customization works in HeteroCL.
"""
import heterocl as hcl
import numpy as np
##############################################################################
# Memory Customization in HeteroCL
# --------------------------------
# There are two types of memory customization in HeteroCL. The first one is
# similar to what we have seen in
# :ref:`sphx_glr_tutorials_tutorial_04_compute.py`, where we demonstrate some
# primitives that will be synthesized as pragmas. An example of such primitive
# is ``partition``. Following is an example. Note that the primitive is
# directly applied on the schedule instead of a stage. This is because we are
# modifying the property of a tensor.

hcl.init()

A = hcl.placeholder((10, 10), "A")

def kernel(A):
    return hcl.compute(A.shape, lambda x, y: A[x][y]+1, "B")

s = hcl.create_schedule(A, kernel)
s.partition(A)
print(hcl.lower(s))

##############################################################################
# In the IR, we should see a line that annotates tensor ``A`` to be
# partitioned completely.
#
# .. note::
#
#    For more information, please see
#    :obj:`heterocl.schedule.Schedule.partition`
#
# Data Reuse in HeteroCL
# ----------------------
# The other type of memory customization primitives involves the introduction
# of allocation of new memory buffers. An example is data reuse. The idea of
# data reuse is to reduce the number of accesses to a tensor by introducing
# an intermediate buffer that holds the values being reused across different
# iterations. This finally leads to better performance in hardware.
#
# Example: 2D Convolution
# -----------------------
# To demonstrate this, we use the computation of 2D convolution as an example.
# Let's see how we can define 2D convolution in HeteroCL.

hcl.init()

A = hcl.placeholder((6, 6), "A")
F = hcl.placeholder((3, 3), "F")

def kernel(A, F):
    r = hcl.reduce_axis(0, 3)
    c = hcl.reduce_axis(0, 3)
    return hcl.compute((4, 4),
            lambda y, x: hcl.sum(A[y+r, x+c]*F[r, c], axis=[r, c]), "B")

s = hcl.create_schedule([A, F], kernel)
print(hcl.lower(s))

##############################################################################
# In the above example, we convolve the input tensor ``A`` with a filter ``F``.
# Then, we store the output in tensor ``B``. Note that the output shape is
# different from the shape of the input tensor. Let's give some real inputs.

hcl_A = hcl.asarray(np.random.randint(0, 10, A.shape))
hcl_F = hcl.asarray(np.random.randint(0, 10, F.shape))
hcl_B = hcl.asarray(np.zeros((4, 4)))
f = hcl.build(s)
f(hcl_A, hcl_F, hcl_B)
print('Input:')
print(hcl_A)
print('Filter:')
print(hcl_F)
print('Output:')
print(hcl_B)

##############################################################################
# To analyze the data reuse, let's take a closer look to the generated IR.
# To begin with, we can see that in two consecutive iterations of ``x`` (i.e.,
# the inner loop), there are 6 pixels that are overlapped, as illustrated in
# the figure below. Without any optimization, **we are reading 9 values from
# the input for each iteration**.
#
# .. figure:: ../../../tutorials/moving_x/Slide1.png
#    :scale: 60 %
#
# Introduce Data Reuse: Window Buffer
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# To reuse the overlapped pixels, we can introduce a reuse buffer. Since the
# filter moves like a window, we call the buffer a window buffer ``WB``. The
# window buffers stores the reused pixels and also the new pixels that will
# be used in the current iteration. For each iteration, to update the values
# inside the window buffer, the last two columns, in this case, shift left.
# After that, the last column is replaced with the pixels read from the
# input. Now, we only **read 3 values from the input for each iteration**.
#
# .. figure:: ../../../tutorials/moving_x/Slide2.png
#    :scale: 60 %
#
# To introduce such reuse buffers in HeteroCL, we use the API ``reuse_at``.
# The first argument is the **tensor** whose values will be reused. The
# second argument is the output **stage** that reuses the values of the
# tensor. The reason why we need to specify this is because we may have
# multiple stages reusing the values from the same input tensor. The third
# argument is the desired axis to be reused. It must be the output axis.
# Finally, we can specify the name of the reuse buffer. The API returns
# a new tensor.

s_x = hcl.create_schedule([A, F], kernel)
WB = s_x.reuse_at(A, s_x[kernel.B], kernel.B.axis[1], "WB")
print(hcl.lower(s_x))

##############################################################################
# In the printed IR, you should be able to see a buffer ``WB`` with size
# (3, 3) being allocated. Moreover, in the ``produce WB`` scope, you should
# see the update algorithm described above. Now let's test the function again.

hcl_Bx = hcl.asarray(np.zeros((4, 4)))
f = hcl.build(s_x)
f(hcl_A, hcl_F, hcl_Bx)
print('Output without WB:')
print(hcl_B)
print('Output with WB:')
print(hcl_Bx)

##############################################################################
# You should see the same results with and without the window buffer.
#
# Reuse at a Different Dimension: Linebuffer
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Similarly, we can create a reuse buffer for two consecutive iterations of
# ``y``. In this case, in each iteration of ``y``, we read an entire row from
# input ``A``. Meanwhile, we update the reuse buffer by shifting up. Since it
# reads an entire line at a time, we call it a linebuffer ``LB``. The
# operation is illustrated in the figure below.
#
# .. figure:: ../../../tutorials/moving_x/Slide3.png
#    :scale: 60 %
#
# Similar to the window buffer, we can introduce the linebuffer in HeteroCL by
# using a single ``reuse_at`` API. We show the code below.

s_y = hcl.create_schedule([A, F], kernel)
LB = s_y.reuse_at(A, s_y[kernel.B], kernel.B.axis[0], "LB")
print(hcl.lower(s_y))

hcl_By = hcl.asarray(np.zeros((4, 4)))
f = hcl.build(s_y)
f(hcl_A, hcl_F, hcl_By)
print('Output without LB:')
print(hcl_B)
print('Output with LB:')
print(hcl_By)

##############################################################################
# Note that the difference between WB and LB is the we reuse at different
# axes. We can also see from the printed IR that the allocated size is larger,
# which is the same as illustrated in the figure above. In this case, we read
# 6 pixels from the input for each iteration of ``y``, which means we read 1
# pixel for each iteration of ``x`` **effectively**. Namely, this is not true
# in terms of hardware execution. Can we do even better?
#
# Combine Window Buffer and Linebuffer
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# We do not need to restrict ourselves to reuse at a single dimension. Since
# we have data reuse in both dimension, we can reuse both. In this case,
# we generate both a linebuffer and a window buffer. Let's take a look at
# the figure first.
#
# .. figure:: ../../../tutorials/moving_x/Slide4.png
#    :scale: 60 %
#
# What happens here is that, we first update the linebuffer (blue arrows),
# then we update the window buffer (purple arrows). More precisely, for each
# iteration of ``x``, we **read 1 pixel from input ``A``**. We simultaneously
# shift up the linebffer. After we update the linebuffer, we go on update the
# window buffer by reading pixels updated in the linebuffer. Then we shift the
# window buffer. To describe such behavior in HeteroCL is very easy. We only
# need to apply ``reuse_at`` twice. We just need to specify the corresponding
# reuse tensors and the reuse axes. In this case, the linebuffer reuses the
# pixels from the input ``A`` while the window buffer reuses from the
# linebuffer. Following we show the code and its IR.

s_xy = hcl.create_schedule([A, F], kernel)
LB = s_xy.reuse_at(A, s_xy[kernel.B], kernel.B.axis[0], "LB")
WB = s_xy.reuse_at(LB, s_xy[kernel.B], kernel.B.axis[1], "WB")
print(hcl.lower(s_xy))

hcl_Bxy = hcl.asarray(np.zeros((4, 4)))
f = hcl.build(s_xy)
f(hcl_A, hcl_F, hcl_Bxy)
print('Output without reuse buffers:')
print(hcl_B)
print('Output with reuse buffers:')
print(hcl_Bxy)

##############################################################################
# We can see from the IR that the allocation sizes are indeed as expected.
#
# Further Optimization
# ~~~~~~~~~~~~~~~~~~~~
# To further optimize the design, we need to think more carefully. For each
# iteration of ``x``, there are three pixels in LB that are being read/write
# simultaneously. Thus, to maximize the memory bandwidth, we need to partition
# LB in the row direction. For WB, all pixels are updated at the same time.
# Therefore, we need to partition the whole WB completely. Finally, we can
# pipeline the whole design. Don't forget that we also need to partition the
# filter ``F``.

s_final = hcl.create_schedule([A, F], kernel)
LB = s_final.reuse_at(A, s_final[kernel.B], kernel.B.axis[0], "LB")
WB = s_final.reuse_at(LB, s_final[kernel.B], kernel.B.axis[1], "WB")
s_final.partition(LB, dim=1)
s_final.partition(WB)
s_final.partition(F)
s_final[kernel.B].pipeline(kernel.B.axis[1])
print(hcl.lower(s_final))

##############################################################################
# Finally, we can generate the HLS code and see if the II is indeed 1.

f = hcl.build(s_final, target="vhls")

##############################################################################
# Following is a sample report from Vivado_HLS.
#
# .. code::
#
#    + Latency (clock cycles):
#        * Summary:
#        +-----+-----+-----+-----+---------+
#        |  Latency  |  Interval | Pipeline|
#        | min | max | min | max |   Type  |
#        +-----+-----+-----+-----+---------+
#        |   42|   42|   43|   43|   none  |
#        +-----+-----+-----+-----+---------+
#
#        + Detail:
#            * Instance:
#            N/A
#
#            * Loop:
#            +----------+-----+-----+----------+-----------+-----------+------+----------+
#            |          |  Latency  | Iteration|  Initiation Interval  | Trip |          |
#            | Loop Name| min | max |  Latency |  achieved |   target  | Count| Pipelined|
#            +----------+-----+-----+----------+-----------+-----------+------+----------+
#            |- Loop 1  |   40|   40|         6|          1|          1|    36|    yes   |
#            +----------+-----+-----+----------+-----------+-----------+------+----------+
#
# Limitations
# -----------
# Following we list the limitations of using reuse buffers in HeteroCL.
#
# 1. We do not accept non-linear index patterns, e.g., ``y*y+c``, ``y*(y+c)``
# 2. The stride is not one, e.g., ``2*y+c``
# 3. There is no overlapped pixel between two consecutive iterations of the
#    specified axis, e.g., ``[x+r, y]`` and reuse ``y``
#
# More Examples: 2D Image Blur
# ----------------------------
# HeteroCL is also able to infer reuse buffers for explicit reduction
# operations. Namely, instead of using ``hcl.sum``, we can expand the compute
# patterns. Following is an example of 2D blur.

hcl.init()
A = hcl.placeholder((10, 10), "A")

def kernel_blur(A):
    return hcl.compute((8, 8), lambda y, x: A[y, x] + A[y+1, x+1] + A[y+2, x+2], "B")

s_blur = hcl.create_schedule(A, kernel_blur)
B = kernel_blur.B
RB_y = s_blur.reuse_at(A, s_blur[B], B.axis[0], "RB_y")
RB_x = s_blur.reuse_at(RB_y, s_blur[B], B.axis[1], "RB_x")
print(hcl.lower(s_blur))


import cv2
import numpy as np
def main():

    resim=np.zeros((400,400,3),dtype="uint8")

    cv2.circle(resim,(200,200),50,(0,255,0),2)
    #ikinci parametremiz dairenin merkezi..Üçüncü parametre yarıçap istiyor
    #4. parametre BGR istiyor
    #5.parametrede kalınlık istiyor.

    cv2.imshow("DAIRE",resim)

    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()