#  Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import paddle
import unittest
import numpy as np
from op_test import OpTest

def ref_logsumexp(x, axis=None, keepdim=False, reduce_all=False):
    if isinstance(axis, int):
        axis = (axis, )
    elif isinstance(axis, list):
        axis = tuple(axis)
    if reduce_all:
        axis = None
    out = np.log(np.exp(x).sum(axis=axis, keepdims=keepdim))
    return out

def logsumexp_wrapper(x, axis=None, keepdim=False, allreduce=False):
    if allreduce:
        return paddle.logsumexp(x, None, keepdim)
    return paddle.logsumexp(x, axis, keepdim)

class TestLogsumexp(OpTest):
    def setUp(self):
        self.op_type = 'logsumexp'
        self.python_api = logsumexp_wrapper
        self.shape = [2, 3, 4, 5]
        self.dtype = 'float64'
        self.axis = [-1]
        self.keepdim = False
        self.reduce_all = False
        self.set_attrs()

        np.random.seed(10)
        x = np.random.uniform(-1, 1, self.shape).astype(self.dtype)
        out = ref_logsumexp(x, self.axis, self.keepdim, self.reduce_all)

        self.inputs = {'X': x}
        self.outputs = {'Out': out}
        self.attrs = {
            'axis': self.axis,
            'keepdim': self.keepdim,
            'reduce_all': self.reduce_all
        }
        self.user_defined_grads = None
        self.user_defined_grad_outputs = None
        self.set_attrs_addition()

    def set_attrs(self):
        pass

    def set_attrs_addition(self):
        pass

    def test_check_output(self):
        self.check_output(check_eager=True)

    def test_check_grad(self):
        self.check_grad(
            ['X'], ['Out'],
            user_defined_grads=self.user_defined_grads,
            user_defined_grad_outputs=self.user_defined_grad_outputs,
            check_eager=True)

    def calc_grad(self):
        dy = np.ones(1, dtype=self.dtype)
        x = self.inputs['X']
        y = self.outputs['Out']
        return dy * np.exp(x - y)

class TestLogsumexp_shape(TestLogsumexp):
    def set_attrs(self):
        self.shape = [4, 5, 6]

class TestLogsumexp_axis(TestLogsumexp):
    def set_attrs(self):
        self.axis = [0, -1]

class TestLogsumexp_axis_all(TestLogsumexp):
    def set_attrs(self):
        self.axis = [0, 1, 2, 3]

    def set_attrs_addition(self):
        if paddle.fluid.core.is_compiled_with_rocm():
            self.user_defined_grads = [self.calc_grad()]
            self.user_defined_grad_outputs = [np.ones(1, dtype=self.dtype)]

class TestLogsumexp_keepdim(TestLogsumexp):
    def set_attrs(self):
        self.keepdim = True

class TestLogsumexp_reduce_all(TestLogsumexp):
    def set_attrs(self):
        self.reduce_all = True

    def set_attrs_addition(self):
        if paddle.fluid.core.is_compiled_with_rocm():
            self.user_defined_grads = [self.calc_grad()]
            self.user_defined_grad_outputs = [np.ones(1, dtype=self.dtype)]

class TestLogsumexpError(unittest.TestCase):
    def test_errors(self):
        with paddle.static.program_guard(paddle.static.Program()):
            self.assertRaises(TypeError, paddle.logsumexp, 1)
            x1 = paddle.fluid.data(name='x1', shape=[120], dtype="int32")
            self.assertRaises(TypeError, paddle.logsumexp, x1)

class TestLogsumexpAPI(unittest.TestCase):
    def setUp(self):
        self.shape = [2, 3, 4, 5]
        self.x = np.random.uniform(-1, 1, self.shape).astype(np.float32)
        self.place = paddle.CUDAPlace(0) if paddle.fluid.core.is_compiled_with_cuda() \
            else paddle.CPUPlace()

    def api_case(self, axis=None, keepdim=False):
        out_ref = ref_logsumexp(self.x, axis, keepdim)
        with paddle.static.program_guard(paddle.static.Program()):
            x = paddle.fluid.data('X', self.shape)
            out = paddle.logsumexp(x, axis, keepdim)
            exe = paddle.static.Executor(self.place)
            res = exe.run(feed={'X': self.x}, fetch_list=[out])
        self.assertTrue(np.allclose(res[0], out_ref))

        paddle.disable_static(self.place)
        x = paddle.to_tensor(self.x)
        out = paddle.logsumexp(x, axis, keepdim)
        self.assertTrue(np.allclose(out.numpy(), out_ref))
        paddle.enable_static()

    def test_api(self):
        self.api_case()
        self.api_case(2)
        self.api_case([-1])
        self.api_case([2, -3])
        self.api_case((0, 1, -1))
        self.api_case(keepdim=True)

    def test_alias(self):
        paddle.disable_static(self.place)
        x = paddle.to_tensor(self.x)
        out1 = paddle.logsumexp(x)
        out2 = paddle.tensor.logsumexp(x)
        out3 = paddle.tensor.math.logsumexp(x)
        out_ref = ref_logsumexp(self.x)
        for out in [out1, out2, out3]:
            self.assertTrue(np.allclose(out.numpy(), out_ref))
        paddle.enable_static()

if __name__ == '__main__':
    unittest.main()


"""
This file contains several helper functions to manipulate sleep staging
(hypnogram) data. The default hypnogram format in YASA is a one dimensional
integer array where:

* -2  = Unscored
* -1  = Artefact / Movement
* 0   = Wake
* 1   = N1 sleep
* 2   = N2 sleep
* 3   = N3 sleep
* 4   = REM sleep

For more details, please refer to the following references:

- Iber, C. (2007). The AASM manual for the scoring of sleep and
associated events: rules, terminology and technical specifications.
American Academy of Sleep Medicine.

- Silber, M. H., Ancoli-Israel, S., Bonnet, M. H., Chokroverty, S.,
Grigg-Damberger, M. M., Hirshkowitz, M., … Iber, C. (2007). The visual scoring
of sleep in adults. Journal of Clinical Sleep Medicine: JCSM: Official
Publication of the American Academy of Sleep Medicine, 3(2), 121–131.

- Combrisson, E., Vallat, R., Eichenlaub, J.-B., O’Reilly, C., Lajnef, T.,
Guillot, A., … Jerbi, K. (2017). Sleep: An Open-Source Python Software for
Visualization, Analysis, and Staging of Sleep Data. Frontiers in
Neuroinformatics, 11, 60. https://doi.org/10.3389/fninf.2017.00060
"""
import mne
import logging
import numpy as np
import pandas as pd
from .io import set_log_level

__all__ = ['hypno_str_to_int', 'hypno_int_to_str', 'hypno_upsample_to_sf',
           'hypno_upsample_to_data', 'load_profusion_hypno']

logger = logging.getLogger('yasa')

#############################################################################
# STR <--> INT CONVERSION
#############################################################################

def hypno_str_to_int(hypno, mapping_dict={'w': 0, 'wake': 0, 'n1': 1, 's1': 1,
                                          'n2': 2, 's2': 2, 'n3': 3, 's3': 3,
                                          's4': 3, 'r': 4, 'rem': 4, 'art': -1,
                                          'mt': -1, 'uns': -2, 'nd': -2}):
    """Convert a string hypnogram array to integer.

    ['W', 'N2', 'N2', 'N3', 'R'] ==> [0, 2, 2, 3, 4]

    .. versionadded:: 0.1.5

    Parameters
    ----------
    hypno : array_like
        The sleep staging (hypnogram) 1D array.
    mapping_dict : dict
        The mapping dictionnary, in lowercase. Note that this function is essentially a wrapper
        around :py:meth:`pandas.Series.map`.

    Returns
    --------
    hypno : array_like
        The corresponding integer hypnogram.
    """
    assert isinstance(hypno, (list, np.ndarray, pd.Series)), 'Not an array.'
    hypno = pd.Series(np.asarray(hypno, dtype=str))
    assert not hypno.str.isnumeric().any(), 'Hypno contains numeric values.'
    return hypno.str.lower().map(mapping_dict).values

def hypno_int_to_str(hypno, mapping_dict={0: 'W', 1: 'N1', 2: 'N2', 3: 'N3',
                                          4: 'R', -1: 'Art', -2: 'Uns'}):
    """Convert an integer hypnogram array to a string array.

    [0, 2, 2, 3, 4] ==> ['W', 'N2', 'N2', 'N3', 'R']

    .. versionadded:: 0.1.5

    Parameters
    ----------
    hypno : array_like
        The sleep staging (hypnogram) 1D array.
    mapping_dict : dict
        The mapping dictionnary. Note that this function is essentially a wrapper around
        :py:meth:`pandas.Series.map`.

    Returns
    --------
    hypno : array_like
        The corresponding integer hypnogram.
    """
    assert isinstance(hypno, (list, np.ndarray, pd.Series)), 'Not an array.'
    hypno = pd.Series(np.asarray(hypno, dtype=int))
    return hypno.map(mapping_dict).values

#############################################################################
# UPSAMPLING
#############################################################################

def hypno_upsample_to_sf(hypno, sf_hypno, sf_data):
    """Upsample the hypnogram to a given sampling frequency.

    .. versionadded:: 0.1.5

    Parameters
    ----------
    hypno : array_like
        The sleep staging (hypnogram) 1D array.
    sf_hypno : float
        The current sampling frequency of the hypnogram, in Hz, e.g.

        * 1/30 = 1 value per each 30 seconds of EEG data,
        * 1 = 1 value per second of EEG data
    sf_data : float
        The desired sampling frequency of the hypnogram, in Hz (e.g. 100 Hz, 256 Hz, ...)

    Returns
    -------
    hypno : array_like
        The hypnogram, upsampled to ``sf_data``.
    """
    repeats = sf_data / sf_hypno
    assert sf_hypno <= sf_data, 'sf_hypno must be less than sf_data.'
    assert repeats.is_integer(), 'sf_hypno / sf_data must be a whole number.'
    assert isinstance(hypno, (list, np.ndarray, pd.Series))
    return np.repeat(np.asarray(hypno), repeats)

def hypno_fit_to_data(hypno, data, sf=None):
    """Crop or pad the hypnogram to fit the length of data.

    Hypnogram and data MUST have the SAME sampling frequency.

    This is an internal function.

    Parameters
    ----------
    hypno : array_like
        The sleep staging (hypnogram) 1D array.
    data : np.array_like or mne.io.Raw
        1D or 2D EEG data. Can also be a MNE Raw object, in which case data and sf will be
        automatically extracted.
    sf : float, optional
        The sampling frequency of data AND the hypnogram.

    Returns
    -------
    hypno : array_like
        Hypnogram, with the same number of samples as data.
    """
    # Check if data is an MNE raw object
    if isinstance(data, mne.io.BaseRaw):
        sf = data.info['sfreq']
        data = data.times  # 1D array and does not require to preload data
    data = np.asarray(data)
    hypno = np.asarray(hypno)
    assert hypno.ndim == 1, 'Hypno must be 1D.'
    npts_hyp = hypno.size
    npts_data = max(data.shape)  # Support for 2D data
    if npts_hyp < npts_data:
        # Hypnogram is shorter than data
        npts_diff = npts_data - npts_hyp
        if sf is not None:
            dur_diff = npts_diff / sf
            logger.warning('Hypnogram is SHORTER than data by %.2f seconds. '
                           'Padding hypnogram with last value to match data.size.' % dur_diff)
        else:
            logger.warning('Hypnogram is SHORTER than data by %i samples. '
                           'Padding hypnogram with last value to match data.size.' % npts_diff)
        hypno = np.pad(hypno, (0, npts_diff), mode='edge')
    elif npts_hyp > npts_data:
        # Hypnogram is longer than data
        npts_diff = npts_hyp - npts_data
        if sf is not None:
            dur_diff = npts_diff / sf
            logger.warning('Hypnogram is LONGER than data by %.2f seconds. '
                           'Cropping hypnogram to match data.size.' % dur_diff)
        else:
            logger.warning('Hypnogram is LONGER than data by %i samples. '
                           'Cropping hypnogram to match data.size.' % npts_diff)
        hypno = hypno[0:npts_data]
    return hypno

def hypno_upsample_to_data(hypno, sf_hypno, data, sf_data=None, verbose=True):
    """Upsample an hypnogram to a given sampling frequency and fit the
    resulting hypnogram to corresponding EEG data, such that the hypnogram
    and EEG data have the exact same number of samples.

    .. versionadded:: 0.1.5

    Parameters
    ----------
    hypno : array_like
        The sleep staging (hypnogram) 1D array.
    sf_hypno : float
        The current sampling frequency of the hypnogram, in Hz, e.g.

        * 1/30 = 1 value per each 30 seconds of EEG data,
        * 1 = 1 value per second of EEG data
    data : array_like or :py:class:`mne.io.BaseRaw`
        1D or 2D EEG data. Can also be a :py:class:`mne.io.BaseRaw`, in which
        case ``data`` and ``sf_data`` will be automatically extracted.
    sf_data : float
        The sampling frequency of ``data``, in Hz (e.g. 100 Hz, 256 Hz, ...).
        Can be omitted if ``data`` is a :py:class:`mne.io.BaseRaw`.
    verbose : bool or str
        Verbose level. Default (False) will only print warning and error
        messages. The logging levels are 'debug', 'info', 'warning', 'error',
        and 'critical'. For most users the choice is between 'info'
        (or ``verbose=True``) and warning (``verbose=False``).

    Returns
    -------
    hypno : array_like
        The hypnogram, upsampled to ``sf_data`` and cropped/padded to ``max(data.shape)``.

    Warns
    -----
    UserWarning
        If the upsampled ``hypno`` is shorter / longer than ``max(data.shape)``
        and therefore needs to be padded/cropped respectively. This output can be disabled by
        passing ``verbose='ERROR'``.
    """
    set_log_level(verbose)
    if isinstance(data, mne.io.BaseRaw):
        sf_data = data.info['sfreq']
        data = data.times
    hypno_up = hypno_upsample_to_sf(hypno=hypno, sf_hypno=sf_hypno, sf_data=sf_data)
    return hypno_fit_to_data(hypno=hypno_up, data=data, sf=sf_data)

#############################################################################
# HYPNO LOADING
#############################################################################

def load_profusion_hypno(fname, replace=True):  # pragma: no cover
    """
    Load a Compumedics Profusion hypnogram (.xml).

    The Compumedics Profusion hypnogram format is one of the two hypnogram
    formats found in the `National Sleep Research Resource (NSRR)
    <https://sleepdata.org/>`_ website. For more details on the format,
    please refer to
    https://github.com/nsrr/edf-editor-translator/wiki/Compumedics-Annotation-Format

    Parameters
    ----------
    fname : str
        Filename with full path.
    replace : bool
        If True, the integer values will be mapped to YASA default, i.e.
        0 for Wake, 1 for N1, 2 for N2, 3 for N3 / S4 and 4 for REM.
        Note that the native profusion format is identical except for REM
        sleep which is marked as 5.

    Returns
    -------
    hypno : 1D array (n_epochs, )
        Hypnogram, with one value per 30 second epochs.
    sf_hyp : float
        Sampling frequency of the hypnogram. Default is 1 / 30 Hz.
    """
    # Note that an alternative is to use the `xmltodict` library:
    # >>> with open(fname) as in_file:
    # >>>   xml = in_file.read()
    # >>> epoch_length = xml['EpochLength']
    # >>> hypno = np.array(xml['SleepStages']['SleepStage'], dtype='int')
    # >>> xml = xmltodict.parse(xml, process_namespaces=True)['CMPStudyConfig']
    # >>> annotations = pd.DataFrame(xml['ScoredEvents']['ScoredEvent'])
    # >>> annotations["Start"] = annotations["Start"].astype(float)
    # >>> annotations["Duration"] = annotations["Duration"].astype(float)
    import xml.etree.ElementTree as ET
    tree = ET.parse(fname)
    root = tree.getroot()
    epoch_length = float(root[0].text)
    sf_hyp = 1 / epoch_length
    hypno = []
    for s in root[4]:
        hypno.append(s.text)
    hypno = np.array(hypno).astype(int)
    if replace:
        # Stage 4 --> 3 and REM --> 4
        hypno = pd.Series(hypno).replace({4: 3, 5: 4}).to_numpy()
    return hypno, sf_hyp


# -*- coding: utf-8 -*-
import numpy as np
import cv2

'''
• img: 你想 绘制图形的 幅图像。
• color: 形状的颜色。以RGB为例  需要传入一个元组BGR 例如 255,0,0
   代表蓝色，第一个是蓝色通道，第二个是绿色通道，第三个是红色通道。对于灰度图只需要传入灰度值。
• thickness 线条的粗细。如果给一个闭合图形 置为 -1  那么这个图形
就会被填充。 默认值是 1.
• linetype 线条的类型， 8 连接，抗锯齿等。  默认情况是8 连接。cv2.LINE_AA
   为抗锯齿  这样看起来会非常平滑。

'''

# Create a black image
img = np.zeros((512, 512, 3), np.uint8)

# Draw a diagonal blue line with thickness of 5 px
cv2.line(img, pt1=(0, 0), pt2=(511, 511), color=(255, 0, 0), thickness=5)  # pt1, pt2, color, thickness=
# cv2.polylines() 可以 用来画很多条线。只需要把想 画的线放在一 个列表中， 将 列表传给函数就可以了。每条线 会被独立绘制。 这会比用 cv2.line() 一条一条的绘制 要快一些。
# cv2.polylines(img, pts, isClosed, color, thickness=None, lineType=None, shift=None)
cv2.arrowedLine(img,pt1=(21, 13), pt2=(151, 401), color=(255, 0, 0), thickness=5)

cv2.rectangle(img, (384, 0), (510, 128), (0, 255, 0), 3)

cv2.circle(img, center=(447, 63), radius=63, color=(0, 0, 255), thickness=-1)  # center, radius, color, thickness=None

# 一个参数是中心点的位置坐标。 下一个参数是长轴和短轴的长度。椭圆沿逆时针方向旋转的角度。
# 椭圆弧演顺时针方向起始的角度和结束角度 如果是 0 很 360 就是整个椭圆
cv2.ellipse(img, center=(256, 256), axes=(100, 50), angle=0, startAngle=0, endAngle=180, color=255,
            thickness=-1)  # center, axes, angle, startAngle, endAngle, color, thickness=

pts = np.array([[10, 5], [20, 30], [70, 20], [50, 10]], np.int32)
pts = pts.reshape((-1, 1, 2))
# 这里 reshape 的第一个参数为-1, 表明这一维的长度是根据后面的维度的计算出来的。
# 注意 如果第三个参数是 False 我们得到的多边形是不闭合的 ，首 尾不相  连 。

font = cv2.FONT_HERSHEY_SIMPLEX
#org :Bottom-left corner of the text string in the image.左下角
#或使用 bottomLeftOrigin=True,文字会上下颠倒
cv2.putText(img, text='bottomLeftOrigin', org=(10, 400), fontFace=font, fontScale=1, color=(255, 255, 255), thickness=1,bottomLeftOrigin=True)#text, org, fontFace, fontScale, color, thickness=
cv2.putText(img, text='OpenCV', org=(10, 500), fontFace=font, fontScale=4, color=(255, 255, 255), thickness=2)#text, org, fontFace, fontScale, color, thickness=

# 所有的绘图函数的返回值都是 None ，所以不能使用 img = cv2.line(img,(0,0),(5

winname = 'example'
cv2.namedWindow(winname, 0)
cv2.imshow(winname, img)

cv2.imwrite("example.png", img)

cv2.waitKey(0)
cv2.destroyAllWindows()


import os
import numpy as np
import nbformat as nbf
import mdutils

def ktx_to_dict(input_file, keystarter = '<'):
    """
        parsing keyed text to a python dictionary.
        把ktx数据从文件当中读出来，换成字典形式
    """
    answer = dict() # 字典

    with open(input_file, 'r+', encoding = 'utf-8') as f:
        lines = f.readlines() #按行读出

    k, val = '', ''
    for line in lines:
        if line.startswith(keystarter): # 对于每个新的key值
            k = line.replace(keystarter, '').strip() # 将key的箭头换掉
            val = '' #value值清空重新开始读入
        else:
            val += line  # 增加当前value内容

        if k: # 每次更新一下当前键值对信息
            answer.update({k: val.strip()})

    return answer

def dict_to_ktx(input_dict, output_file, keystarter = '<'):
    """
        Store a python dictionary to a keyed text
        逆向操作，把字典存储为ktx文件
        直接写文件
    """
    with open(output_file, 'w+') as f:
        for k, val in input_dict.items():
            f.write(f'{keystarter} {k}\n')
            f.write(f'{val}\n\n')
    return

HEADERS = ktx_to_dict(os.path.join('source', 'headers.ktx'))
QHA = ktx_to_dict(os.path.join('source', 'exercises100.ktx'))

'''
    ---------------------
    以上是预设模块直接运行
    以下是提供支持的模块
    ---------------------
'''

def create_jupyter_notebook(destination_filename = '100_Numpy_exercises.ipynb', answer_type = 0):
    """
        Programmatically create jupyter notebook with the questions
        saved under source files
        参数化生成Notebook
        在同文件夹下方生成
    """

    # Create cells sequence
    # notebook的 cell序列
    nb = nbf.v4.new_notebook()

    nb['cells'] = []

    # - Add header:
    nb['cells'].append(nbf.v4.new_markdown_cell(HEADERS["header"]))
    nb['cells'].append(nbf.v4.new_markdown_cell(HEADERS["sub_header"]))
    nb['cells'].append(nbf.v4.new_markdown_cell(HEADERS["jupyter_instruction"]))

    nb['cells'].append(nbf.v4.new_code_cell('import generators as ge'))

    # - 增加问题MD块, 提示代码块，作业代码块，答案代码块
    for n in range(1, 101):
        nb['cells'].append(nbf.v4.new_markdown_cell(f'### {n}. ' + QHA[f'q{n}']))
        nb['cells'].append(nbf.v4.new_code_cell(f'ge.hint({n})'))
        nb['cells'].append(nbf.v4.new_code_cell(''))
        if answer_type == 0:
            nb['cells'].append(nbf.v4.new_code_cell(f'ge.answer({n})'))
        else:
            nb['cells'].append(nbf.v4.new_code_cell(QHA[f'a{n}']))

    # Delete file if one with the same name is found
    # 删除同名文件
    if os.path.exists(destination_filename):
        os.remove(destination_filename)

    # Write sequence to file
    # 写入序列
    nbf.write(nb, destination_filename)
    return

'''
    以下是随机生成问题模块
'''

def question(n = 1): # 生成问题
    print(f'{n}. ' + QHA[f'q{n}'])
    return

def hint(n = 1): # 生成提示
    print(QHA[f'h{n}'])
    return

def answer(n = 1): # 生成答案
    print(QHA[f'a{n}'])
    return

def pick(): # 随机生成
    id = np.random.randint(1, 100)
    question(id)
    return id

# def create_markdown(destination_filename = '100_Numpy_exercises', with_hints = False, with_solutions = False):
#     '''
#         生成MarkDown版本的内容
#         没有做什么改动，代码非常直观
#     '''
#     if with_hints:
#         destination_filename += '_with_hints'
#     if with_solutions:
#         destination_filename += '_with_solutions'

#     # Initialise file
#     mdfile = mdutils.MdUtils(file_name = destination_filename)

#     # Add headers
#     mdfile.write(HEADERS["header"] + '\n')
#     mdfile.write(HEADERS["sub_header"] + '\n')

#     # Add questions (and hint or answers if required)
#     for n in range(1, 101):
#         mdfile.new_header(title = f"{n}. {QHA[f'q{n}']}", level = 4)
#         if with_hints:
#             mdfile.write(f"`{QHA[f'h{n}']}`")
#         if with_solutions:
#             mdfile.insert_code(QHA[f'a{n}'], language = 'python')

#     # Delete file if one with the same name is found
#     if os.path.exists(destination_filename):
#         os.remove(destination_filename)

#     # Write sequence to file
#     mdfile.create_md_file()
#     return

# def create_rst(destination_filename, with_ints = False, with_answers = False):
#     # TODO: use rstdoc python library.
#     #  also see possible integrations with https://github.com/rougier/numpy-100/pull/38
#     pass


import random

import torch
import numpy as np

def set_random_seeds(random_seed=666):
    r"""Sets the seed for generating random numbers.

    Args:
        random_seed: Desired random seed.
    """
    torch.manual_seed(random_seed)
    # torch.cuda.manual_seed(random_seed)
    # torch.backends.cudnn.deterministic = True
    # # todo remove this improves performance
    # torch.backends.cudnn.benchmark = False
    np.random.seed(random_seed)
    random.seed(random_seed)

"""
From Pytorch lightning's seed_everything
    log.info(f"Global seed set to {seed}")
    os.environ["PL_GLOBAL_SEED"] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
"""


# coding: utf-8

# # Additional tips and tricks for designing networks
#
# This tutorial assumes that you have read
# the `network_design` tutorial,
# and have designed a network or two.
# Here, we will give a few advanced tips and tricks
# for designing networks that can be reused flexibly.
# In particular, these tips will use the
# `config` system, so we will also assume that
# you have gone over the `config` tutorial.
#
# Briefly, the general principles covered
# in this tutorial are
#
# 0. Accept a network argument
# 0. Accept a config argument for groups of parameters
#
# We will demonstrate these principles
# using the two examples from the `network_design` tutorial.

# In[ ]:

import numpy as np
import matplotlib.pyplot as plt

import nengo
from nengo.dists import Choice
from nengo.utils.functions import piecewise
from nengo.utils.ipython import hide_input
def test_integrators(net):
    with net:
        piecewise_f = piecewise({0: 0, 0.2: 0.5, 1: 0, 2: -1, 3: 0, 4: 1, 5: 0})
        piecewise_inp = nengo.Node(piecewise_f)
        nengo.Connection(piecewise_inp, net.pre_integrator.input)
        input_probe = nengo.Probe(piecewise_inp)
        pre_probe = nengo.Probe(net.pre_integrator.ensemble, synapse=0.01)
        post_probe = nengo.Probe(net.post_integrator.ensemble, synapse=0.01)
    with nengo.Simulator(net) as sim:
        sim.run(6)
    plt.plot(sim.trange(), sim.data[input_probe], color='k')
    plt.plot(sim.trange(), sim.data[pre_probe], color='b')
    plt.plot(sim.trange(), sim.data[post_probe], color='g')
hide_input()

# ## 1. Accept a network argument
#
# Typically, a network-creation function
# will take a set of arguments,
# which affect some important parts
# of the network.
# When testing the network,
# it's common to change several parameters
# to see how they affect the network.
# One way to do this is to add more and more
# arguments to your function;
# this quickly gets out of hand.
# Instead, use the `config` system,
# which  enables us to set
# network-level defaults for all Nengo objects.
#
# You can either do this by creating your network
# in the context of some other network,
# or you can modify your function
# to optionally take in a network instance,
# which you will build your objects into.
#
# In the example below,
# we change both integrators to use `LIFRate` neurons
# by changing the `config` object
# in the network both integrators
# are build within.
# We also change the `post_integrator`
# to use a very small radius
# by passing in a network
# which has its default radius modified.

# In[ ]:

def Integrator(n_neurons, dimensions, tau=0.1, net=None):
    if net is None:
        net = nengo.Network()
    with net:
        net.input = nengo.Node(size_in=dimensions)
        net.ensemble = nengo.Ensemble(n_neurons, dimensions=dimensions)
        nengo.Connection(net.ensemble, net.ensemble, synapse=tau)
        nengo.Connection(net.input, net.ensemble,
                         synapse=None, transform=tau)
    return net

net = nengo.Network(label="Two integrators")
with net:
    # Make both integrators use LIFRate neurons
    net.config[nengo.Ensemble].neuron_type = nengo.LIFRate()
    net.pre_integrator = Integrator(50, 1)
    # Lower the radius of the post_integrator
    net.post_integrator = nengo.Network()
    net.post_integrator.config[nengo.Ensemble].radius = 0.2
    Integrator(50, 1, net=net.post_integrator)
    nengo.Connection(net.pre_integrator.ensemble, net.post_integrator.input)
test_integrators(net)

# ## 2. Accept a config argument for groups of parameters
#
# Often, you will not want to use the
# network-level defaults for all of your objects.
# Some objects need certain things overwritten,
# while others need other values overwritten.
# Again, it is possible to deal with this issue
# by adding more and more parameters,
# but this quickly gets out of hand.
# Instead, add a small number of arguments
# that optionally accept a `config` object,
# which allows for setting multiple parameters at once.
#
# In the coupled integrator network example,
# we make two connections.
# We have to be careful changing the defaults
# for those connections, as they are wildly different;
# one is a recurrent connection from an ensemble to itself,
# while the other is a connection from a node to an ensemble.
# We will accept a `config` object for the recurrent connection
# to make this easier.

# In[ ]:

def Integrator(n_neurons, dimensions, recurrent_config=None, net=None):
    if net is None:
        net = nengo.Network()
    if recurrent_config is None:
        recurrent_config = nengo.Config(nengo.Connection)
        recurrent_config[nengo.Connection].synapse = nengo.Lowpass(0.1)
    with net:
        net.input = nengo.Node(size_in=dimensions)
        net.ensemble = nengo.Ensemble(n_neurons, dimensions=dimensions)
        with recurrent_config:
            nengo.Connection(net.ensemble, net.ensemble)
            tau = nengo.Config.default(nengo.Connection, 'synapse').tau
        nengo.Connection(net.input, net.ensemble,
                         synapse=None, transform=tau)
    return net

net = nengo.Network(label="Two integrators")
with net:
    # Make both integrators use LIFRate neurons
    net.config[nengo.Ensemble].neuron_type = nengo.LIFRate()
    net.pre_integrator = Integrator(50, 1)
    # Give the post_integrator a shorter tau (should make integration fail)
    recurrent_config = nengo.Config(nengo.Connection)
    recurrent_config[nengo.Connection].synapse = nengo.Lowpass(0.01)
    net.post_integrator = Integrator(50, 1, recurrent_config=recurrent_config)
    nengo.Connection(net.pre_integrator.ensemble, net.post_integrator.input)
test_integrators(net)

# ## Longer example: double integrator network
#
# Recall in the previous tutorial that
# we created a model
# that released a lever 0.6 to 1.0 seconds
# after pressing a lever.
# Let's use the above principles,
# and the `config` system in general,
# to improve the code constructing this model.

# In[ ]:

def controlled_integrator(n_neurons, dimensions, recurrent_config=None, net=None):
    if net is None:
        net = nengo.Network()
    if recurrent_config is None:
        recurrent_config = nengo.Config(nengo.Connection)
        recurrent_config[nengo.Connection].synapse = nengo.Lowpass(0.1)
    with net:
        net.ensemble = nengo.Ensemble(n_neurons, dimensions=dimensions + 1)
        with recurrent_config:
            nengo.Connection(net.ensemble, net.ensemble[:dimensions],
                             function=lambda x: x[:-1] * (1.0 - x[-1]))
    return net

def medial_pfc(coupling_strength, n_neurons_per_integrator=200, recurrent_config=None, tau=0.1, net=None):
    if net is None:
        net = nengo.Network()
    with net:
        recurrent_config = nengo.Config(nengo.Connection)
        recurrent_config[nengo.Connection].synapse = nengo.Lowpass(tau)
        net.pre = controlled_integrator(n_neurons_per_integrator, 1, recurrent_config)
        net.post = controlled_integrator(n_neurons_per_integrator, 1, recurrent_config)
        nengo.Connection(net.pre.ensemble[0], net.post.ensemble[0],
                         transform=coupling_strength)
    return net

def motor_cortex(command_threshold, n_neurons_per_command=30, ens_config=None, net=None):
    if net is None:
        net = nengo.Network()
    if ens_config is None:
        ens_config = nengo.Config(nengo.Ensemble)
        ens_config[nengo.Ensemble].encoders = Choice([[1]])
        ens_config[nengo.Ensemble].intercepts = Choice([command_threshold])
    with net:
        with ens_config:
            net.press = nengo.Ensemble(n_neurons_per_command, dimensions=1)
            net.release = nengo.Ensemble(n_neurons_per_command, dimensions=1)
    return net

def double_integrator(mpfc_coupling_strength,
                      command_threshold,
                      press_to_pre_gain=3,
                      press_to_post_control=-6,
                      recurrent_tau=0.1,
                      net=None):
    if net is None:
        net = nengo.Network()
    with net:
        net.mpfc = medial_pfc(mpfc_coupling_strength)
        net.motor = motor_cortex(command_threshold)
        nengo.Connection(net.motor.press, net.mpfc.pre.ensemble[0],
                         transform=recurrent_tau * press_to_pre_gain)
        nengo.Connection(net.motor.press, net.mpfc.post.ensemble[1],
                         transform=press_to_post_control)
        nengo.Connection(net.mpfc.post.ensemble[0], net.motor.release)
    return net

def test_doubleintegrator(net):
    # Provide input and probe outside of network construction, for more flexibility
    with net:
        nengo.Connection(nengo.Node(lambda t: 1 if t < 0.2 else 0), net.motor.press)
        pr_press = nengo.Probe(net.motor.press, synapse=0.01)
        pr_release = nengo.Probe(net.motor.release, synapse=0.01)
        pr_pre_int = nengo.Probe(net.mpfc.pre.ensemble[0], synapse=0.01)
        pr_post_int = nengo.Probe(net.mpfc.post.ensemble[0], synapse=0.01)
    with nengo.Simulator(net) as sim:
        sim.run(1.4)
    t = sim.trange()
    plt.figure()
    plt.subplot(2, 1, 1)
    plt.plot(t, sim.data[pr_press], c='b', label="Press")
    plt.plot(t, sim.data[pr_release], c='g', label="Release")
    plt.axvspan(0, 0.2, color='b', alpha=0.3)
    plt.axvspan(0.8, 1.2, color='g', alpha=0.3)
    plt.xlim(right=1.4)
    plt.legend(loc="best")
    plt.subplot(2, 1, 2)
    plt.plot(t, sim.data[pr_pre_int], label="Pre Integrator")
    plt.plot(t, sim.data[pr_post_int], label="Post Integrator")
    plt.xlim(right=1.4)
    plt.legend(loc="best")

for coupling_strength in (0.11, 0.16, 0.21):
    net = nengo.Network(seed=0)  # Set seed here instead
    # Try the same network with LIFRate neurons
    net.config[nengo.Ensemble].neuron_type = nengo.LIFRate()
    net = double_integrator(mpfc_coupling_strength=coupling_strength,
                            command_threshold=0.85,
                            net = net)
    test_doubleintegrator(net)


"""
Tests the 'read_fwf' function in parsers.py. This
test suite is independent of the others because the
engine is set to 'python-fwf' internally.
"""

from datetime import datetime
from io import (
    BytesIO,
    StringIO,
)
from pathlib import Path

import numpy as np
import pytest

from pandas.errors import EmptyDataError

from pandas import (
    DataFrame,
    DatetimeIndex,
)
import pandas._testing as tm

from pandas.io.parsers import (
    read_csv,
    read_fwf,
)

def test_basic():
    data = """\
A         B            C            D
201158    360.242940   149.910199   11950.7
201159    444.953632   166.985655   11788.4
201160    364.136849   183.628767   11806.2
201161    413.836124   184.375703   11916.8
201162    502.953953   173.237159   12468.3
"""
    result = read_fwf(StringIO(data))
    expected = DataFrame(
        [
            [201158, 360.242940, 149.910199, 11950.7],
            [201159, 444.953632, 166.985655, 11788.4],
            [201160, 364.136849, 183.628767, 11806.2],
            [201161, 413.836124, 184.375703, 11916.8],
            [201162, 502.953953, 173.237159, 12468.3],
        ],
        columns=["A", "B", "C", "D"],
    )
    tm.assert_frame_equal(result, expected)

def test_colspecs():
    data = """\
A   B     C            D            E
201158    360.242940   149.910199   11950.7
201159    444.953632   166.985655   11788.4
201160    364.136849   183.628767   11806.2
201161    413.836124   184.375703   11916.8
201162    502.953953   173.237159   12468.3
"""
    colspecs = [(0, 4), (4, 8), (8, 20), (21, 33), (34, 43)]
    result = read_fwf(StringIO(data), colspecs=colspecs)

    expected = DataFrame(
        [
            [2011, 58, 360.242940, 149.910199, 11950.7],
            [2011, 59, 444.953632, 166.985655, 11788.4],
            [2011, 60, 364.136849, 183.628767, 11806.2],
            [2011, 61, 413.836124, 184.375703, 11916.8],
            [2011, 62, 502.953953, 173.237159, 12468.3],
        ],
        columns=["A", "B", "C", "D", "E"],
    )
    tm.assert_frame_equal(result, expected)

def test_widths():
    data = """\
A    B    C            D            E
2011 58   360.242940   149.910199   11950.7
2011 59   444.953632   166.985655   11788.4
2011 60   364.136849   183.628767   11806.2
2011 61   413.836124   184.375703   11916.8
2011 62   502.953953   173.237159   12468.3
"""
    result = read_fwf(StringIO(data), widths=[5, 5, 13, 13, 7])

    expected = DataFrame(
        [
            [2011, 58, 360.242940, 149.910199, 11950.7],
            [2011, 59, 444.953632, 166.985655, 11788.4],
            [2011, 60, 364.136849, 183.628767, 11806.2],
            [2011, 61, 413.836124, 184.375703, 11916.8],
            [2011, 62, 502.953953, 173.237159, 12468.3],
        ],
        columns=["A", "B", "C", "D", "E"],
    )
    tm.assert_frame_equal(result, expected)

def test_non_space_filler():
    # From Thomas Kluyver:
    #
    # Apparently, some non-space filler characters can be seen, this is
    # supported by specifying the 'delimiter' character:
    #
    # http://publib.boulder.ibm.com/infocenter/dmndhelp/v6r1mx/index.jsp?topic=/com.ibm.wbit.612.help.config.doc/topics/rfixwidth.html
    data = """\
A~~~~B~~~~C~~~~~~~~~~~~D~~~~~~~~~~~~E
201158~~~~360.242940~~~149.910199~~~11950.7
201159~~~~444.953632~~~166.985655~~~11788.4
201160~~~~364.136849~~~183.628767~~~11806.2
201161~~~~413.836124~~~184.375703~~~11916.8
201162~~~~502.953953~~~173.237159~~~12468.3
"""
    colspecs = [(0, 4), (4, 8), (8, 20), (21, 33), (34, 43)]
    result = read_fwf(StringIO(data), colspecs=colspecs, delimiter="~")

    expected = DataFrame(
        [
            [2011, 58, 360.242940, 149.910199, 11950.7],
            [2011, 59, 444.953632, 166.985655, 11788.4],
            [2011, 60, 364.136849, 183.628767, 11806.2],
            [2011, 61, 413.836124, 184.375703, 11916.8],
            [2011, 62, 502.953953, 173.237159, 12468.3],
        ],
        columns=["A", "B", "C", "D", "E"],
    )
    tm.assert_frame_equal(result, expected)

def test_over_specified():
    data = """\
A   B     C            D            E
201158    360.242940   149.910199   11950.7
201159    444.953632   166.985655   11788.4
201160    364.136849   183.628767   11806.2
201161    413.836124   184.375703   11916.8
201162    502.953953   173.237159   12468.3
"""
    colspecs = [(0, 4), (4, 8), (8, 20), (21, 33), (34, 43)]

    with pytest.raises(ValueError, match="must specify only one of"):
        read_fwf(StringIO(data), colspecs=colspecs, widths=[6, 10, 10, 7])

def test_under_specified():
    data = """\
A   B     C            D            E
201158    360.242940   149.910199   11950.7
201159    444.953632   166.985655   11788.4
201160    364.136849   183.628767   11806.2
201161    413.836124   184.375703   11916.8
201162    502.953953   173.237159   12468.3
"""
    with pytest.raises(ValueError, match="Must specify either"):
        read_fwf(StringIO(data), colspecs=None, widths=None)

def test_read_csv_compat():
    csv_data = """\
A,B,C,D,E
2011,58,360.242940,149.910199,11950.7
2011,59,444.953632,166.985655,11788.4
2011,60,364.136849,183.628767,11806.2
2011,61,413.836124,184.375703,11916.8
2011,62,502.953953,173.237159,12468.3
"""
    expected = read_csv(StringIO(csv_data), engine="python")

    fwf_data = """\
A   B     C            D            E
201158    360.242940   149.910199   11950.7
201159    444.953632   166.985655   11788.4
201160    364.136849   183.628767   11806.2
201161    413.836124   184.375703   11916.8
201162    502.953953   173.237159   12468.3
"""
    colspecs = [(0, 4), (4, 8), (8, 20), (21, 33), (34, 43)]
    result = read_fwf(StringIO(fwf_data), colspecs=colspecs)
    tm.assert_frame_equal(result, expected)

def test_bytes_io_input():
    result = read_fwf(BytesIO("שלום\nשלום".encode()), widths=[2, 2], encoding="utf8")
    expected = DataFrame([["של", "ום"]], columns=["של", "ום"])
    tm.assert_frame_equal(result, expected)

def test_fwf_colspecs_is_list_or_tuple():
    data = """index,A,B,C,D
foo,2,3,4,5
bar,7,8,9,10
baz,12,13,14,15
qux,12,13,14,15
foo2,12,13,14,15
bar2,12,13,14,15
"""

    msg = "column specifications must be a list or tuple.+"

    with pytest.raises(TypeError, match=msg):
        read_fwf(StringIO(data), colspecs={"a": 1}, delimiter=",")

def test_fwf_colspecs_is_list_or_tuple_of_two_element_tuples():
    data = """index,A,B,C,D
foo,2,3,4,5
bar,7,8,9,10
baz,12,13,14,15
qux,12,13,14,15
foo2,12,13,14,15
bar2,12,13,14,15
"""

    msg = "Each column specification must be.+"

    with pytest.raises(TypeError, match=msg):
        read_fwf(StringIO(data), [("a", 1)])

@pytest.mark.parametrize(
    "colspecs,exp_data",
    [
        ([(0, 3), (3, None)], [[123, 456], [456, 789]]),
        ([(None, 3), (3, 6)], [[123, 456], [456, 789]]),
        ([(0, None), (3, None)], [[123456, 456], [456789, 789]]),
        ([(None, None), (3, 6)], [[123456, 456], [456789, 789]]),
    ],
)
def test_fwf_colspecs_none(colspecs, exp_data):
    # see gh-7079
    data = """\
123456
456789
"""
    expected = DataFrame(exp_data)

    result = read_fwf(StringIO(data), colspecs=colspecs, header=None)
    tm.assert_frame_equal(result, expected)

@pytest.mark.parametrize(
    "infer_nrows,exp_data",
    [
        # infer_nrows --> colspec == [(2, 3), (5, 6)]
        (1, [[1, 2], [3, 8]]),
        # infer_nrows > number of rows
        (10, [[1, 2], [123, 98]]),
    ],
)
def test_fwf_colspecs_infer_nrows(infer_nrows, exp_data):
    # see gh-15138
    data = """\
  1  2
123 98
"""
    expected = DataFrame(exp_data)

    result = read_fwf(StringIO(data), infer_nrows=infer_nrows, header=None)
    tm.assert_frame_equal(result, expected)

def test_fwf_regression():
    # see gh-3594
    #
    # Turns out "T060" is parsable as a datetime slice!
    tz_list = [1, 10, 20, 30, 60, 80, 100]
    widths = [16] + [8] * len(tz_list)
    names = ["SST"] + [f"T{z:03d}" for z in tz_list[1:]]

    data = """  2009164202000   9.5403  9.4105  8.6571  7.8372  6.0612  5.8843  5.5192
2009164203000   9.5435  9.2010  8.6167  7.8176  6.0804  5.8728  5.4869
2009164204000   9.5873  9.1326  8.4694  7.5889  6.0422  5.8526  5.4657
2009164205000   9.5810  9.0896  8.4009  7.4652  6.0322  5.8189  5.4379
2009164210000   9.6034  9.0897  8.3822  7.4905  6.0908  5.7904  5.4039
"""

    result = read_fwf(
        StringIO(data),
        index_col=0,
        header=None,
        names=names,
        widths=widths,
        parse_dates=True,
        date_parser=lambda s: datetime.strptime(s, "%Y%j%H%M%S"),
    )
    expected = DataFrame(
        [
            [9.5403, 9.4105, 8.6571, 7.8372, 6.0612, 5.8843, 5.5192],
            [9.5435, 9.2010, 8.6167, 7.8176, 6.0804, 5.8728, 5.4869],
            [9.5873, 9.1326, 8.4694, 7.5889, 6.0422, 5.8526, 5.4657],
            [9.5810, 9.0896, 8.4009, 7.4652, 6.0322, 5.8189, 5.4379],
            [9.6034, 9.0897, 8.3822, 7.4905, 6.0908, 5.7904, 5.4039],
        ],
        index=DatetimeIndex(
            [
                "2009-06-13 20:20:00",
                "2009-06-13 20:30:00",
                "2009-06-13 20:40:00",
                "2009-06-13 20:50:00",
                "2009-06-13 21:00:00",
            ]
        ),
        columns=["SST", "T010", "T020", "T030", "T060", "T080", "T100"],
    )
    tm.assert_frame_equal(result, expected)

def test_fwf_for_uint8():
    data = """1421302965.213420    PRI=3 PGN=0xef00      DST=0x17 SRC=0x28    04 154 00 00 00 00 00 127
1421302964.226776    PRI=6 PGN=0xf002               SRC=0x47    243 00 00 255 247 00 00 71"""  # noqa
    df = read_fwf(
        StringIO(data),
        colspecs=[(0, 17), (25, 26), (33, 37), (49, 51), (58, 62), (63, 1000)],
        names=["time", "pri", "pgn", "dst", "src", "data"],
        converters={
            "pgn": lambda x: int(x, 16),
            "src": lambda x: int(x, 16),
            "dst": lambda x: int(x, 16),
            "data": lambda x: len(x.split(" ")),
        },
    )

    expected = DataFrame(
        [
            [1421302965.213420, 3, 61184, 23, 40, 8],
            [1421302964.226776, 6, 61442, None, 71, 8],
        ],
        columns=["time", "pri", "pgn", "dst", "src", "data"],
    )
    expected["dst"] = expected["dst"].astype(object)
    tm.assert_frame_equal(df, expected)

@pytest.mark.parametrize("comment", ["#", "~", "!"])
def test_fwf_comment(comment):
    data = """\
  1   2.   4  #hello world
  5  NaN  10.0
"""
    data = data.replace("#", comment)

    colspecs = [(0, 3), (4, 9), (9, 25)]
    expected = DataFrame([[1, 2.0, 4], [5, np.nan, 10.0]])

    result = read_fwf(StringIO(data), colspecs=colspecs, header=None, comment=comment)
    tm.assert_almost_equal(result, expected)

def test_fwf_skip_blank_lines():
    data = """

A         B            C            D

201158    360.242940   149.910199   11950.7
201159    444.953632   166.985655   11788.4

201162    502.953953   173.237159   12468.3

"""
    result = read_fwf(StringIO(data), skip_blank_lines=True)
    expected = DataFrame(
        [
            [201158, 360.242940, 149.910199, 11950.7],
            [201159, 444.953632, 166.985655, 11788.4],
            [201162, 502.953953, 173.237159, 12468.3],
        ],
        columns=["A", "B", "C", "D"],
    )
    tm.assert_frame_equal(result, expected)

    data = """\
A         B            C            D
201158    360.242940   149.910199   11950.7
201159    444.953632   166.985655   11788.4

201162    502.953953   173.237159   12468.3
"""
    result = read_fwf(StringIO(data), skip_blank_lines=False)
    expected = DataFrame(
        [
            [201158, 360.242940, 149.910199, 11950.7],
            [201159, 444.953632, 166.985655, 11788.4],
            [np.nan, np.nan, np.nan, np.nan],
            [np.nan, np.nan, np.nan, np.nan],
            [201162, 502.953953, 173.237159, 12468.3],
        ],
        columns=["A", "B", "C", "D"],
    )
    tm.assert_frame_equal(result, expected)

@pytest.mark.parametrize("thousands", [",", "#", "~"])
def test_fwf_thousands(thousands):
    data = """\
 1 2,334.0    5
10   13     10.
"""
    data = data.replace(",", thousands)

    colspecs = [(0, 3), (3, 11), (12, 16)]
    expected = DataFrame([[1, 2334.0, 5], [10, 13, 10.0]])

    result = read_fwf(
        StringIO(data), header=None, colspecs=colspecs, thousands=thousands
    )
    tm.assert_almost_equal(result, expected)

@pytest.mark.parametrize("header", [True, False])
def test_bool_header_arg(header):
    # see gh-6114
    data = """\
MyColumn
   a
   b
   a
   b"""

    msg = "Passing a bool to header is invalid"
    with pytest.raises(TypeError, match=msg):
        read_fwf(StringIO(data), header=header)

def test_full_file():
    # File with all values.
    test = """index                             A    B    C
2000-01-03T00:00:00  0.980268513777    3  foo
2000-01-04T00:00:00  1.04791624281    -4  bar
2000-01-05T00:00:00  0.498580885705   73  baz
2000-01-06T00:00:00  1.12020151869     1  foo
2000-01-07T00:00:00  0.487094399463    0  bar
2000-01-10T00:00:00  0.836648671666    2  baz
2000-01-11T00:00:00  0.157160753327   34  foo"""
    colspecs = ((0, 19), (21, 35), (38, 40), (42, 45))
    expected = read_fwf(StringIO(test), colspecs=colspecs)

    result = read_fwf(StringIO(test))
    tm.assert_frame_equal(result, expected)

def test_full_file_with_missing():
    # File with missing values.
    test = """index                             A    B    C
2000-01-03T00:00:00  0.980268513777    3  foo
2000-01-04T00:00:00  1.04791624281    -4  bar
                     0.498580885705   73  baz
2000-01-06T00:00:00  1.12020151869     1  foo
2000-01-07T00:00:00                    0  bar
2000-01-10T00:00:00  0.836648671666    2  baz
                                      34"""
    colspecs = ((0, 19), (21, 35), (38, 40), (42, 45))
    expected = read_fwf(StringIO(test), colspecs=colspecs)

    result = read_fwf(StringIO(test))
    tm.assert_frame_equal(result, expected)

def test_full_file_with_spaces():
    # File with spaces in columns.
    test = """
Account                 Name  Balance     CreditLimit   AccountCreated
101     Keanu Reeves          9315.45     10000.00           1/17/1998
312     Gerard Butler         90.00       1000.00             8/6/2003
868     Jennifer Love Hewitt  0           17000.00           5/25/1985
761     Jada Pinkett-Smith    49654.87    100000.00          12/5/2006
317     Bill Murray           789.65      5000.00             2/5/2007
""".strip(
        "\r\n"
    )
    colspecs = ((0, 7), (8, 28), (30, 38), (42, 53), (56, 70))
    expected = read_fwf(StringIO(test), colspecs=colspecs)

    result = read_fwf(StringIO(test))
    tm.assert_frame_equal(result, expected)

def test_full_file_with_spaces_and_missing():
    # File with spaces and missing values in columns.
    test = """
Account               Name    Balance     CreditLimit   AccountCreated
101                           10000.00                       1/17/1998
312     Gerard Butler         90.00       1000.00             8/6/2003
868                                                          5/25/1985
761     Jada Pinkett-Smith    49654.87    100000.00          12/5/2006
317     Bill Murray           789.65
""".strip(
        "\r\n"
    )
    colspecs = ((0, 7), (8, 28), (30, 38), (42, 53), (56, 70))
    expected = read_fwf(StringIO(test), colspecs=colspecs)

    result = read_fwf(StringIO(test))
    tm.assert_frame_equal(result, expected)

def test_messed_up_data():
    # Completely messed up file.
    test = """
   Account          Name             Balance     Credit Limit   Account Created
       101                           10000.00                       1/17/1998
       312     Gerard Butler         90.00       1000.00

       761     Jada Pinkett-Smith    49654.87    100000.00          12/5/2006
  317          Bill Murray           789.65
""".strip(
        "\r\n"
    )
    colspecs = ((2, 10), (15, 33), (37, 45), (49, 61), (64, 79))
    expected = read_fwf(StringIO(test), colspecs=colspecs)

    result = read_fwf(StringIO(test))
    tm.assert_frame_equal(result, expected)

def test_multiple_delimiters():
    test = r"""
col1~~~~~col2  col3++++++++++++++++++col4
~~22.....11.0+++foo~~~~~~~~~~Keanu Reeves
  33+++122.33\\\bar.........Gerard Butler
++44~~~~12.01   baz~~Jennifer Love Hewitt
~~55       11+++foo++++Jada Pinkett-Smith
..66++++++.03~~~bar           Bill Murray
""".strip(
        "\r\n"
    )
    delimiter = " +~.\\"
    colspecs = ((0, 4), (7, 13), (15, 19), (21, 41))
    expected = read_fwf(StringIO(test), colspecs=colspecs, delimiter=delimiter)

    result = read_fwf(StringIO(test), delimiter=delimiter)
    tm.assert_frame_equal(result, expected)

def test_variable_width_unicode():
    data = """
שלום שלום
ום   שלל
של   ום
""".strip(
        "\r\n"
    )
    encoding = "utf8"
    kwargs = {"header": None, "encoding": encoding}

    expected = read_fwf(
        BytesIO(data.encode(encoding)), colspecs=[(0, 4), (5, 9)], **kwargs
    )
    result = read_fwf(BytesIO(data.encode(encoding)), **kwargs)
    tm.assert_frame_equal(result, expected)

@pytest.mark.parametrize("dtype", [{}, {"a": "float64", "b": str, "c": "int32"}])
def test_dtype(dtype):
    data = """ a    b    c
1    2    3.2
3    4    5.2
"""
    colspecs = [(0, 5), (5, 10), (10, None)]
    result = read_fwf(StringIO(data), colspecs=colspecs, dtype=dtype)

    expected = DataFrame(
        {"a": [1, 3], "b": [2, 4], "c": [3.2, 5.2]}, columns=["a", "b", "c"]
    )

    for col, dt in dtype.items():
        expected[col] = expected[col].astype(dt)

    tm.assert_frame_equal(result, expected)

def test_skiprows_inference():
    # see gh-11256
    data = """
Text contained in the file header

DataCol1   DataCol2
     0.0        1.0
   101.6      956.1
""".strip()
    skiprows = 2
    expected = read_csv(StringIO(data), skiprows=skiprows, delim_whitespace=True)

    result = read_fwf(StringIO(data), skiprows=skiprows)
    tm.assert_frame_equal(result, expected)

def test_skiprows_by_index_inference():
    data = """
To be skipped
Not  To  Be  Skipped
Once more to be skipped
123  34   8      123
456  78   9      456
""".strip()
    skiprows = [0, 2]
    expected = read_csv(StringIO(data), skiprows=skiprows, delim_whitespace=True)

    result = read_fwf(StringIO(data), skiprows=skiprows)
    tm.assert_frame_equal(result, expected)

def test_skiprows_inference_empty():
    data = """
AA   BBB  C
12   345  6
78   901  2
""".strip()

    msg = "No rows from which to infer column width"
    with pytest.raises(EmptyDataError, match=msg):
        read_fwf(StringIO(data), skiprows=3)

def test_whitespace_preservation():
    # see gh-16772
    header = None
    csv_data = """
 a ,bbb
 cc,dd """

    fwf_data = """
 a bbb
 ccdd """
    result = read_fwf(
        StringIO(fwf_data), widths=[3, 3], header=header, skiprows=[0], delimiter="\n\t"
    )
    expected = read_csv(StringIO(csv_data), header=header)
    tm.assert_frame_equal(result, expected)

def test_default_delimiter():
    header = None
    csv_data = """
a,bbb
cc,dd"""

    fwf_data = """
a \tbbb
cc\tdd """
    result = read_fwf(StringIO(fwf_data), widths=[3, 3], header=header, skiprows=[0])
    expected = read_csv(StringIO(csv_data), header=header)
    tm.assert_frame_equal(result, expected)

@pytest.mark.parametrize("infer", [True, False])
def test_fwf_compression(compression_only, infer):
    data = """1111111111
    2222222222
    3333333333""".strip()

    compression = compression_only
    extension = "gz" if compression == "gzip" else compression

    kwargs = {"widths": [5, 5], "names": ["one", "two"]}
    expected = read_fwf(StringIO(data), **kwargs)

    data = bytes(data, encoding="utf-8")

    with tm.ensure_clean(filename="tmp." + extension) as path:
        tm.write_to_compressed(compression, path, data)

        if infer is not None:
            kwargs["compression"] = "infer" if infer else compression

        result = read_fwf(path, **kwargs)
        tm.assert_frame_equal(result, expected)

def test_binary_mode():
    """
    read_fwf supports opening files in binary mode.

    GH 18035.
    """
    data = """aas aas aas
bba bab b a"""
    df_reference = DataFrame(
        [["bba", "bab", "b a"]], columns=["aas", "aas.1", "aas.2"], index=[0]
    )
    with tm.ensure_clean() as path:
        Path(path).write_text(data)
        with open(path, "rb") as file:
            df = read_fwf(file)
            file.seek(0)
            tm.assert_frame_equal(df, df_reference)

@pytest.mark.parametrize("memory_map", [True, False])
def test_encoding_mmap(memory_map):
    """
    encoding should be working, even when using a memory-mapped file.

    GH 23254.
    """
    encoding = "iso8859_1"
    data = BytesIO(" 1 A Ä 2\n".encode(encoding))
    df = read_fwf(
        data,
        header=None,
        widths=[2, 2, 2, 2],
        encoding=encoding,
        memory_map=memory_map,
    )
    data.seek(0)
    df_reference = DataFrame([[1, "A", "Ä", 2]])
    tm.assert_frame_equal(df, df_reference)

@pytest.mark.parametrize(
    "colspecs, names, widths, index_col",
    [
        (
            [(0, 6), (6, 12), (12, 18), (18, None)],
            list("abcde"),
            None,
            None,
        ),
        (
            None,
            list("abcde"),
            [6] * 4,
            None,
        ),
        (
            [(0, 6), (6, 12), (12, 18), (18, None)],
            list("abcde"),
            None,
            True,
        ),
        (
            None,
            list("abcde"),
            [6] * 4,
            False,
        ),
        (
            None,
            list("abcde"),
            [6] * 4,
            True,
        ),
        (
            [(0, 6), (6, 12), (12, 18), (18, None)],
            list("abcde"),
            None,
            False,
        ),
    ],
)
def test_len_colspecs_len_names(colspecs, names, widths, index_col):
    # GH#40830
    data = """col1  col2  col3  col4
    bab   ba    2"""
    msg = "Length of colspecs must match length of names"
    with pytest.raises(ValueError, match=msg):
        read_fwf(
            StringIO(data),
            colspecs=colspecs,
            names=names,
            widths=widths,
            index_col=index_col,
        )

@pytest.mark.parametrize(
    "colspecs, names, widths, index_col, expected",
    [
        (
            [(0, 6), (6, 12), (12, 18), (18, None)],
            list("abc"),
            None,
            0,
            DataFrame(
                index=["col1", "ba"],
                columns=["a", "b", "c"],
                data=[["col2", "col3", "col4"], ["b   ba", "2", np.nan]],
            ),
        ),
        (
            [(0, 6), (6, 12), (12, 18), (18, None)],
            list("ab"),
            None,
            [0, 1],
            DataFrame(
                index=[["col1", "ba"], ["col2", "b   ba"]],
                columns=["a", "b"],
                data=[["col3", "col4"], ["2", np.nan]],
            ),
        ),
        (
            [(0, 6), (6, 12), (12, 18), (18, None)],
            list("a"),
            None,
            [0, 1, 2],
            DataFrame(
                index=[["col1", "ba"], ["col2", "b   ba"], ["col3", "2"]],
                columns=["a"],
                data=[["col4"], [np.nan]],
            ),
        ),
        (
            None,
            list("abc"),
            [6] * 4,
            0,
            DataFrame(
                index=["col1", "ba"],
                columns=["a", "b", "c"],
                data=[["col2", "col3", "col4"], ["b   ba", "2", np.nan]],
            ),
        ),
        (
            None,
            list("ab"),
            [6] * 4,
            [0, 1],
            DataFrame(
                index=[["col1", "ba"], ["col2", "b   ba"]],
                columns=["a", "b"],
                data=[["col3", "col4"], ["2", np.nan]],
            ),
        ),
        (
            None,
            list("a"),
            [6] * 4,
            [0, 1, 2],
            DataFrame(
                index=[["col1", "ba"], ["col2", "b   ba"], ["col3", "2"]],
                columns=["a"],
                data=[["col4"], [np.nan]],
            ),
        ),
    ],
)
def test_len_colspecs_len_names_with_index_col(
    colspecs, names, widths, index_col, expected
):
    # GH#40830
    data = """col1  col2  col3  col4
    bab   ba    2"""
    result = read_fwf(
        StringIO(data),
        colspecs=colspecs,
        names=names,
        widths=widths,
        index_col=index_col,
    )
    tm.assert_frame_equal(result, expected)


# Coding Syntax

## Single line comment
"""
multiline comment
"""
## -------------------------------------------------

## Operators
d = 7//2		# floor division
d = 7%2			# modulus
d = 7**2		# to the power
long_str = 5*"o"	# some are also defined for other var types, but be careful

## Printing
d = input("what's your name?")
print("hi", d)		# multiple printing
print('hi', 'duck', sep='\n')
print('hi', 'duck', sep='/') # / for file path, "," for CSV format
print('hi', 'duck', end='\r')
f"something {d}" # better string formating, PEP	498
## -------------------------------------------------

# Variables
## Collections
##### List (ordered & changeable)
list_var1 = [1, 2, 3]
list_var7 = [i+1 for i in list_var1] # list comprehension
list_var5 = [1, 2, "d", [2, 7]]
list_var2 = ["d", "u", "c"]
list_var2[0] # "d"
list_var2[-1] # "c"
list_var2[1:3] # only 1st and 2nd element, 3rd is not included
list_var2.append("k")
del list_var2[3]
list_var2.reverse()
list_var2[0] = 0
for mem in list_var1:
  print()
print(*list) # will print members of list without []

##### Tuple (ordered & unchangeable)
tuple_var1 = (1, 2, 3)
tuple_var2 = (1, "d", [1, 2])
tuple_var1 + (4, ) # append tuple

##### Dictionaries (unordered, changable & indexed)
dict_var = {"a":1, "b":2, "c":3}
dict_var["a"] # 1
dict_var["b"] # 2
dict_var("k", 7) # get value of key "k", if doesn't exist, take default value 7
dict_var.keys() # iterate over all keys
for key in dict_var.keys():
  print()
dict_var.values() # iterate over all values
for val in dict_var.values():
  print()
del dict_var["a"] # delete pair of key and val
dict_var["d"] = 7 # add key-val pair
dict_var({"b":5, "d":7}) # add & overwrite

##### Set ( unordered, changeable & unindexed)
set_var = {"a", "b", "c"}
## -------------------------------------------------

## Conditional - Control flow
x = 72
print((x <= 72 and x > 72) or x == 72) # True

if x > 72:
	print()
elif x == 27:
	print()
else:
	print()

i=27
while i < x:
	print()
	i+=1

for i in range(0,7):
  print()
  if i==1:
    continue	# prematurely skip current iteration
  elif i==2:
    break		# prematurely terminate a loop

try:
  print(duck) # duck is not yet define -> error
except:
  print("Something went wrong")
else:
  print("Nothing went wrong")
finally: # useful to close objects and clean up resources, execute regardless if there's error
  print("The 'try except' is finished")

x = "hello"
if not type(x) is int:
  raise TypeError("Only integers are allowed")
## -------------------------------------------------

## Functions
def square(x):
  return x**2
print(square(3)) # 9

def greet(name="", last_name=""):
  print("hello",name)
  return name, name + last_name

cubed = lambda x: x**3
print(cubed(2)) # 8

# args, kwargs
def print_all(*args,**kwargs):
  print("args",args)
  print("kwargs", kwargs)
print_all(1,2,3,b=7,c=4)
# args (1, 2, 3)
# kwargs {'b': 7, 'c': 4}
## -------------------------------------------------

## Classes
class Person:
  def __init__(self, name, allergies=""):
    self.name = name
    self.allergies = allergies
  def talk(self):
    print("my name is", self.name)
    if self.allergies != "":
      print("I am allergic to ", self.allergies)

p1 = Person("Antine")
p1.name = "Antoine"
p1.talk()
## -------------------------------------------------

## Common library
help(print) # If possible, just use this first

import numpy as np
a = np.array([1,2,3])
a[1]; a[:,2]; a[1:3,1]
b = np.array([[1,2],[3,4]], dtype=float)
# np.zeros, ones, linspace, random.random
# shape, ndim, size, dtype
# sum, min, max, mean, median, std
# round, ceil, floor

import pandas as pd
file_name="test.txt"
df = pd.read_csv(file_name)
df.to_csv(file_name)
df_2 = pd.read_excel(file_name)
df_2.to_excel(file_name, sheet_name="sheet")
df.head() # columns,

from sqlalchemy import create_engine
engine = create_engine('sqlite:///foo.db')
df_3 = pd.read_sql_table("tableName", engine)
df_3.to_sql("tableName",engine)

data = [["tom",10],["pete",15],["jean",30],["puff",35],["pete",5]]
df = pd.DataFrame(data=data, columns=["name","age"])
# info
df.columns
df.shape
df.info()
# filters
df[(df.age > 10) & (df.name == "pete")]
df.iloc[0] # by position (row)

# operations: sum, cumsum, min, max, mean, median
df["age"].sum()
# apply a function to cells
sum_one = lambda x: x + 1
df["new age"] = df["age"].apply(sum_one)
upper = lambda s: s[0].upper() + s[1:]
df["name"] = df["name"].apply(upper)

## Write & Open file
## File access modes: r, r+, w, w+, a, a+
file_obj = open(r"file_name", "access_mode") # add r to prevent special character, the string is consider raw
file_obj.close() # free the memory space acquired by that file
file_obj.write(str) #writelines(list_of_str)
file_obj.read # readline, readlines
file_obj.writelines(list_of_str)

import matplotlib.pyplot as plt
plt.plot() # subplot
plt.figure(figure_size()), grid, title
plt.scatter, hist
plt.show()

from sklearn import duck

import os
os.getcwd()
os.chdir()
os.system('mkdir duck')
os.walk()
os.listdir()
os.path.dirname(dir_path)
os.path.realpath(__file__)

import glob
glob.glob(name, recursive = False)
glob.escape
glob.glob('*.py')
for name in glob.glob('dir/*'):
	cmd
for name in glob.glob('dir/name?.txt'):
	cmd
# Find recursively
for files in glob.glob('dir/**/*.txt', recursive = True):
	cmd

import sys
print(sys.argv)
sys.path
sys.path.append(path)
sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))

import argparse
parser = argparse.ArgumentParser(prog = 'top',
    description = 'Show top lines from each file')
parser.add_argument('filenames', nargs='+')
parser.add_argument('-l', '--lines', type=int, default=10)
args = parser.parse_args()
print(args)

import cv2
# cv2 axis: x - left to right, y - up to down
cv2.imread #imwrite, imshow, VideoWriter, waitKey, destroyAllWindows
cv2.VideoWriter(name, cv2.VideoWriter_fourcc(*'DIVX'), fps, size)
## The fun parts of opencv
cv2.cvtColor, flip, resize, threshold, adaptiveThreshold
cv2.bitwise_and #or, not, xor
## Contour stuffs
contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
# CHAIN_APPROX_NONE for all contour points
# check coutour retrieval mode for more infos on hierarchy
cv2.drawContours, contourArea, arcLength, approxPolyDP, isContourConvex, convexHull, boundingRect, minAreaRect
# this thickness is also stupid, becareful :)
## Basic drawings
cv2.line, polylines, circle, text
# the width is stupid, be careful when it's small
## Advance feature extraction
cv2.HoughLines, HoughLinesP, cornerHarris
cv2.norm, moments
cv2.addWeighted
cv2.getRotationMatrix2D
cv2.warpAffine
cv2.countNonZero

from tqdm import tqdm
for i in tqdm(range(1,5)): # wrap around iterator
  print(i)

import pickle, cpickle
pickle.dump(var, open('name.pickle', 'wb'))
var = pickle.load(open('file.pickle', 'rb'))
## use dumps and loads with gzip
## Python 2 use .pkl, python 3 use .pickle
with open('', '') as f:
	pickle.dump()

import gzip
var = gzip.open('file.gz', 'rb')
with gzip.open('file.txt.gz', 'wb') as f:
	f.write(content)
## Save a compressed pickle file
file = gzip.GzipFile(filename, 'wb')
file.write(pickle.dumps, protocol = 0)
file.close
## Load a compressed pickle file
file = gzip.GzipFile(filename, 'rb')
data = file.read()
object = pickle.loads(data)
file.close

from typing import NewType, Final
UserId = NewType('UserId', int)
some_id = UserId(524313)
## Final: A special typing construct to indicate to type checkers that a name cannot be re-assigned or overridden in a subclass.
MAX_SIZE: Final = 9000
MAX_SIZE += 1  # Error reported by type checker
class Connection:
    TIMEOUT: Final[int] = 10
class FastConnector(Connection):
    i=2
    # TIMEOUT = 1  # Error reported by type checker

from shapely import *
from shapely.geometry import box, LineString, LinearRing, Point
from shapely.geometry import Polygon, MultiPolygon, MultiLineString
from shapely.geometry import CAP_STYLE, JOIN_STYLE
from shapely.ops import polygonize, cascaded_union, unary_union


"""
.. _ex-publication-figure:

===================================
Make figures more publication ready
===================================

In this example, we show several use cases to take MNE plots and
customize them for a more publication-ready look.
"""

# Authors: Eric Larson <larson.eric.d@gmail.com>
#          Daniel McCloy <dan.mccloy@gmail.com>
#          Stefan Appelhoff <stefan.appelhoff@mailbox.org>
#
# License: BSD (3-clause)

###############################################################################
# Imports
# -------
# We are importing everything we need for this example:

import os.path as op

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import (make_axes_locatable, ImageGrid,
                                     inset_locator)

import mne

###############################################################################
# Evoked plot with brain activation
# ---------------------------------
#
# Suppose we want a figure with an evoked plot on top, and the brain activation
# below, with the brain subplot slightly bigger than the evoked plot. Let's
# start by loading some :ref:`example data <sample-dataset>`.

data_path = mne.datasets.sample.data_path()
subjects_dir = op.join(data_path, 'subjects')
fname_stc = op.join(data_path, 'MEG', 'sample', 'sample_audvis-meg-eeg-lh.stc')
fname_evoked = op.join(data_path, 'MEG', 'sample', 'sample_audvis-ave.fif')

evoked = mne.read_evokeds(fname_evoked, 'Left Auditory')
evoked.pick_types(meg='grad').apply_baseline((None, 0.))
max_t = evoked.get_peak()[1]

stc = mne.read_source_estimate(fname_stc)

###############################################################################
# During interactive plotting, we might see figures like this:

evoked.plot()

stc.plot(views='lat', hemi='split', size=(800, 400), subject='sample',
         subjects_dir=subjects_dir, initial_time=max_t,
         time_viewer=False, show_traces=False)

###############################################################################
# To make a publication-ready figure, first we'll re-plot the brain on a white
# background, take a screenshot of it, and then crop out the white margins.
# While we're at it, let's change the colormap, set custom colormap limits and
# remove the default colorbar (so we can add a smaller, vertical one later):

colormap = 'viridis'
clim = dict(kind='value', lims=[4, 8, 12])

# Plot the STC, get the brain image, crop it:
brain = stc.plot(views='lat', hemi='split', size=(800, 400), subject='sample',
                 subjects_dir=subjects_dir, initial_time=max_t, background='w',
                 colorbar=False, clim=clim, colormap=colormap,
                 time_viewer=False, show_traces=False)
screenshot = brain.screenshot()
brain.close()

###############################################################################
# Now let's crop out the white margins and the white gap between hemispheres.
# The screenshot has dimensions ``(h, w, 3)``, with the last axis being R, G, B
# values for each pixel, encoded as integers between ``0`` and ``255``. ``(255,
# 255, 255)`` encodes a white pixel, so we'll detect any pixels that differ
# from that:

nonwhite_pix = (screenshot != 255).any(-1)
nonwhite_row = nonwhite_pix.any(1)
nonwhite_col = nonwhite_pix.any(0)
cropped_screenshot = screenshot[nonwhite_row][:, nonwhite_col]

# before/after results
fig = plt.figure(figsize=(4, 4))
axes = ImageGrid(fig, 111, nrows_ncols=(2, 1), axes_pad=0.5)
for ax, image, title in zip(axes, [screenshot, cropped_screenshot],
                            ['Before', 'After']):
    ax.imshow(image)
    ax.set_title('{} cropping'.format(title))

###############################################################################
# A lot of figure settings can be adjusted after the figure is created, but
# many can also be adjusted in advance by updating the
# :data:`~matplotlib.rcParams` dictionary. This is especially useful when your
# script generates several figures that you want to all have the same style:

# Tweak the figure style
plt.rcParams.update({
    'ytick.labelsize': 'small',
    'xtick.labelsize': 'small',
    'axes.labelsize': 'small',
    'axes.titlesize': 'medium',
    'grid.color': '0.75',
    'grid.linestyle': ':',
})

###############################################################################
# Now let's create our custom figure. There are lots of ways to do this step.
# Here we'll create the figure and the subplot axes in one step, specifying
# overall figure size, number and arrangement of subplots, and the ratio of
# subplot heights for each row using :mod:`GridSpec keywords
# <matplotlib.gridspec>`. Other approaches (using
# :func:`~matplotlib.pyplot.subplot2grid`, or adding each axes manually) are
# shown commented out, for reference.

# sphinx_gallery_thumbnail_number = 4
# figsize unit is inches
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(4.5, 3.),
                         gridspec_kw=dict(height_ratios=[3, 4]))

# alternate way #1: using subplot2grid
# fig = plt.figure(figsize=(4.5, 3.))
# axes = [plt.subplot2grid((7, 1), (0, 0), rowspan=3),
#         plt.subplot2grid((7, 1), (3, 0), rowspan=4)]

# alternate way #2: using figure-relative coordinates
# fig = plt.figure(figsize=(4.5, 3.))
# axes = [fig.add_axes([0.125, 0.58, 0.775, 0.3]),  # left, bot., width, height
#         fig.add_axes([0.125, 0.11, 0.775, 0.4])]

# we'll put the evoked plot in the upper axes, and the brain below
evoked_idx = 0
brain_idx = 1

# plot the evoked in the desired subplot, and add a line at peak activation
evoked.plot(axes=axes[evoked_idx])
peak_line = axes[evoked_idx].axvline(max_t, color='#66CCEE', ls='--')
# custom legend
axes[evoked_idx].legend(
    [axes[evoked_idx].lines[0], peak_line], ['MEG data', 'Peak time'],
    frameon=True, columnspacing=0.1, labelspacing=0.1,
    fontsize=8, fancybox=True, handlelength=1.8)
# remove the "N_ave" annotation
axes[evoked_idx].texts = []
# Remove spines and add grid
axes[evoked_idx].grid(True)
axes[evoked_idx].set_axisbelow(True)
for key in ('top', 'right'):
    axes[evoked_idx].spines[key].set(visible=False)
# Tweak the ticks and limits
axes[evoked_idx].set(
    yticks=np.arange(-200, 201, 100), xticks=np.arange(-0.2, 0.51, 0.1))
axes[evoked_idx].set(
    ylim=[-225, 225], xlim=[-0.2, 0.5])

# now add the brain to the lower axes
axes[brain_idx].imshow(cropped_screenshot)
axes[brain_idx].axis('off')
# add a vertical colorbar with the same properties as the 3D one
divider = make_axes_locatable(axes[brain_idx])
cax = divider.append_axes('right', size='5%', pad=0.2)
cbar = mne.viz.plot_brain_colorbar(cax, clim, colormap, label='Activation (F)')

# tweak margins and spacing
fig.subplots_adjust(
    left=0.15, right=0.9, bottom=0.01, top=0.9, wspace=0.1, hspace=0.5)

# add subplot labels
for ax, label in zip(axes, 'AB'):
    ax.text(0.03, ax.get_position().ymax, label, transform=fig.transFigure,
            fontsize=12, fontweight='bold', va='top', ha='left')

###############################################################################
# Custom timecourse with montage inset
# ------------------------------------
#
# Suppose we want a figure with some mean timecourse extracted from a number of
# sensors, and we want a smaller panel within the figure to show a head outline
# with the positions of those sensors clearly marked.
# If you are familiar with MNE, you know that this is something that
# :func:`mne.viz.plot_compare_evokeds` does, see an example output in
# :ref:`ex-hf-sef-data` at the bottom.
#
# In this part of the example, we will show you how to achieve this result on
# your own figure, without having to use :func:`mne.viz.plot_compare_evokeds`!
#
# Let's start by loading some :ref:`example data <sample-dataset>`.

data_path = mne.datasets.sample.data_path()
fname_raw = op.join(data_path, "MEG", "sample", "sample_audvis_raw.fif")
raw = mne.io.read_raw_fif(fname_raw)

# For the sake of the example, we focus on EEG data
raw.pick_types(meg=False, eeg=True)

###############################################################################
# Let's make a plot.

# channels to plot:
to_plot = [f"EEG {i:03}" for i in range(1, 5)]

# get the data for plotting in a short time interval from 10 to 20 seconds
start = int(raw.info['sfreq'] * 10)
stop = int(raw.info['sfreq'] * 20)
data, times = raw.get_data(picks=to_plot,
                           start=start, stop=stop, return_times=True)

# Scale the data from the MNE internal unit V to µV
data *= 1e6
# Take the mean of the channels
mean = np.mean(data, axis=0)
# make a figure
fig, ax = plt.subplots(figsize=(4.5, 3))
# plot some EEG data
ax.plot(times, mean)

###############################################################################
# So far so good. Now let's add the smaller figure within the figure to show
# exactly, which sensors we used to make the timecourse.
# For that, we use an "inset_axes" that we plot into our existing axes.
# The head outline with the sensor positions can be plotted using the
# `~mne.io.Raw` object that is the source of our data.
# Specifically, that object already contains all the sensor positions,
# and we can plot them using the ``plot_sensors`` method.

# recreate the figure (only necessary for our documentation server)
fig, ax = plt.subplots(figsize=(4.5, 3))
ax.plot(times, mean)
axins = inset_locator.inset_axes(ax, width="30%", height="30%", loc=2)

# pick_channels() edits the raw object in place, so we'll make a copy here
# so that our raw object stays intact for potential later analysis
raw.copy().pick_channels(to_plot).plot_sensors(title="", axes=axins)

###############################################################################
# That looks nice. But the sensor dots are way too big for our taste. Luckily,
# all MNE-Python plots use Matplotlib under the hood and we can customize
# each and every facet of them.
# To make the sensor dots smaller, we need to first get a handle on them to
# then apply a ``*.set_*`` method on them.

# If we inspect our axes we find the objects contained in our plot:
print(axins.get_children())

###############################################################################
# That's quite a a lot of objects, but we know that we want to change the
# sensor dots, and those are most certainly a "PathCollection" object.
# So let's have a look at how many "collections" we have in the axes.
print(axins.collections)

###############################################################################
# There is only one! Those must be the sensor dots we were looking for.
# We finally found exactly what we needed. Sometimes this can take a bit of
# experimentation.

sensor_dots = axins.collections[0]

# Recreate the figure once more; shrink the sensor dots; add axis labels
fig, ax = plt.subplots(figsize=(4.5, 3))
ax.plot(times, mean)
axins = inset_locator.inset_axes(ax, width="30%", height="30%", loc=2)
raw.copy().pick_channels(to_plot).plot_sensors(title="", axes=axins)
sensor_dots = axins.collections[0]
sensor_dots.set_sizes([1])
# add axis labels, and adjust bottom figure margin to make room for them
ax.set(xlabel="Time (s)", ylabel="Amplitude (µV)")
fig.subplots_adjust(bottom=0.2)


import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from cycler import cycler
from IPython.display import display, set_matplotlib_formats, HTML

display(HTML(data="""
<style>
    div#notebook-container    { width: 95%; }
    div#menubar-container     { width: 65%; }
    div#maintoolbar-container { width: 99%; }
</style>
"""))

set_matplotlib_formats('pdf', 'png')
plt.rcParams['savefig.dpi'] = 300
plt.rcParams['image.cmap'] = "viridis"
plt.rcParams['image.interpolation'] = "none"
plt.rcParams['savefig.bbox'] = "tight"
plt.rcParams['lines.linewidth'] = 2
plt.rcParams['legend.numpoints'] = 1

np.set_printoptions(precision=3, suppress=True)

pd.set_option("display.max_columns", 8)
pd.set_option('precision', 2)

__all__ = ['np', 'display', 'plt', 'pd']


import numpy as np
from .data_iterator import DataIterator

class BatchIterator(DataIterator):
    """TODO: BatchIterator docs"""
    def __init__(self, batch_size, shuffle=False):
        super().__init__()

        self.batch_size = batch_size
        self.shuffle = shuffle

    def __call__(self, inputs, targets):
        starts = np.arange(0, len(inputs), self.batch_size)
        if self.shuffle:
            np.random.shuffle(starts)

        for start in starts:
            end = start + self.batch_size
            batch_inputs = inputs[start:end]
            batch_targets = targets[start:end]
            yield batch_inputs, batch_targets


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

# # MODFLOW 6: Observation packages
#

# ## Introduction to Observations
#
# Observations can be set for any package through the `package.obs` object, and
# each package.obs object has several attributes that can be set:
#
# | Attribute | Type | Description |
# | :---      | :---- | :----      |
# | package.obs.filename | str | Name of observations file to create. The default is packagename + '.obs'.|
# | package.obs.continuous | dict | A dictionary that has file names as keys and a list of observations as the dictionary values. |
# | package.obs.digits | int | Number of digits to write the observation values. Default is 10. |
# | package.obs.print_input | bool | Flag indicating whether or not observations are written to listing file. |
#
# The following code sets up a simulation used in the observation examples.

# package import
import os

import numpy as np

import flopy

# set up where simulation workspace will be stored
workspace = os.path.join("data", "mf6_working_with_data")
name = "example_1"
if not os.path.exists(workspace):
    os.makedirs(workspace)

# create the flopy simulation and tdis objects
sim = flopy.mf6.MFSimulation(
    sim_name=name, exe_name="mf6", version="mf6", sim_ws=workspace
)
tdis_rc = [(1.0, 1, 1.0), (10.0, 5, 1.0), (10.0, 5, 1.0), (10.0, 1, 1.0)]
tdis_package = flopy.mf6.modflow.mftdis.ModflowTdis(
    sim, time_units="DAYS", nper=4, perioddata=tdis_rc
)
# create the flopy groundwater flow (gwf) model object
model_nam_file = f"{name}.nam"
gwf = flopy.mf6.ModflowGwf(sim, modelname=name, model_nam_file=model_nam_file)
# create the flopy iterative model solver (ims) package object
ims = flopy.mf6.modflow.mfims.ModflowIms(sim, pname="ims", complexity="SIMPLE")
# create the discretization package
bot = np.linspace(-3.0, -50.0 / 3.0, 3)
delrow = delcol = 4.0
dis = flopy.mf6.modflow.mfgwfdis.ModflowGwfdis(
    gwf,
    pname="dis",
    nogrb=True,
    nlay=3,
    nrow=101,
    ncol=101,
    delr=delrow,
    delc=delcol,
    top=0.0,
    botm=bot,
)
# create the initial condition (ic) and node property flow (npf) packages
ic_package = flopy.mf6.modflow.mfgwfic.ModflowGwfic(gwf, strt=50.0)
npf_package = flopy.mf6.modflow.mfgwfnpf.ModflowGwfnpf(
    gwf,
    save_flows=True,
    icelltype=[1, 0, 0],
    k=[5.0, 0.1, 4.0],
    k33=[0.5, 0.005, 0.1],
)

# ## Observation Example 1
#
# One method to build the observation package is to pass a dictionary with
# the observations containing "observations" parameters of the parent package.
#
# This example uses the observation package in a `GHB` package.  First the
# stress period data for a ghb package is built.

# build ghb stress period data
ghb_spd = {}
ghb_period = []
for layer, cond in zip(range(1, 3), [15.0, 1500.0]):
    for row in range(0, 15):
        ghb_period.append(((layer, row, 9), 1.0, cond, "Estuary-L2"))
ghb_spd[0] = ghb_period

# The next step is to build the observation data in a dictionary.  The
# dictionary key is the filename of the observation output file and
# optionally a "binary" keyword to make the file binary.  When the optional
# "binary" keyword is used the dictionary key is a tuple, otherwise it is a
# string.  The dictionary value is a list of tuples containing the contents
# of the observation package's continuous block, with each tuple containing
# one line of information.

# build obs data
ghb_obs = {
    ("ghb_obs.csv", "binary"): [
        ("ghb-2-6-10", "GHB", (1, 5, 9)),
        ("ghb-3-6-10", "GHB", (2, 5, 9)),
    ],
    "ghb_flows.csv": [
        ("Estuary2", "GHB", "Estuary-L2"),
        ("Estuary3", "GHB", "Estuary-L3"),
    ],
}

# The ghb package is now constructed with observations by setting the
# `observations` parameter to `ghb_obs` on construction of the ghb package.

# build ghb package passing obs dictionary to package constructor
ghb = flopy.mf6.modflow.mfgwfghb.ModflowGwfghb(
    gwf,
    print_input=True,
    print_flows=True,
    save_flows=True,
    boundnames=True,
    observations=ghb_obs,
    pname="ghb",
    maxbound=30,
    stress_period_data=ghb_spd,
)

# Observation information such as the print_input option can then be set using
# the package's `obs` parameter.

ghb.obs.print_input = True

# clean up for next example
gwf.remove_package("ghb")

# ## Observation Example 2
#
# Alternatively, an obs package can be built by initializing obs
# through `ghb.obs.initialize`.

# First, a `GHB` package is built without defining observations.

# build ghb package
ghb = flopy.mf6.modflow.mfgwfghb.ModflowGwfghb(
    gwf,
    print_input=True,
    print_flows=True,
    save_flows=True,
    boundnames=True,
    maxbound=30,
    stress_period_data=ghb_spd,
    pname="ghb",
)

# Then the ghb observations are defined in a dictionary similar to example 1.

# build obs data
ghb_obs = {
    ("ghb_obs.csv", "binary"): [
        ("ghb-2-6-10", "GHB", (1, 5, 9)),
        ("ghb-3-6-10", "GHB", (2, 5, 9)),
    ],
    "ghb_flows.csv": [
        ("Estuary2", "GHB", "Estuary-L2"),
        ("Estuary3", "GHB", "Estuary-L3"),
    ],
}

# The observations can then be added to the ghb package using the obs
# attribute's initialize method.  The observation package's file name,
# digits, and print_input options, along with the continuous block data
# are set in the initialize method.

# initialize obs package
ghb.obs.initialize(
    filename="child_pkgs_test.ghb.obs",
    digits=9,
    print_input=True,
    continuous=ghb_obs,
)


"""The WaveBlocks Project

Function for stem-plotting functions of the type f:I -> C
with abs(f) as y-value and phase(f) as color code.
This function makes a stem plot.

@author: R. Bourquin
@copyright: Copyright (C) 2010, 2011, 2012, 2016 R. Bourquin
@license: Modified BSD License
"""

from numpy import array, zeros, real
from matplotlib.collections import LineCollection
from matplotlib.pyplot import gca

from WaveBlocksND.Plot.color_map import color_map

def stemcf(grid, phase, modulus, darken=None, axes=None, linestylep="solid", linewidthp=2, color=None, markerp="o", **kwargs):
    r"""Stemplot the modulus of a complex valued function :math:`f:I -> \mathbb{C}` together with its phase in a color coded fashion.
    Additional keyword arguments are passed to the plot function.

    :param grid: The grid nodes of the real domain grid :math:`\Gamma`
    :param phase: The phase of the complex domain result :math:`f(\Gamma)`
    :param modulus: The modulus of the complex domain result :math:`f(\Gamma)`
    :param darken: Whether to take into account the modulus of the data to darken colors.
    :param axes: The axes instance used for plotting.
    :param linestylep: The line style of the phase curve.
    :param linewidthp: The line width of the phase curve.
    :param color: The color of the stemmed markers.
    :param markerp: The shape of the stemmed markers.
    """
    # Color mapping
    rgb_colors = color_map(grid, phase=phase, modulus=modulus, darken=darken)

    # Put all the vertical line into a collection
    segments = [array([[node, 0], [node, value]]) for node, value in zip(grid, modulus)]
    line_segments = LineCollection(segments)

    # Set some properties of the lines
    rgb_colors = line_segments.to_rgba(rgb_colors)
    line_segments.set_color(rgb_colors[0])
    line_segments.set_linestyle(linestylep)
    line_segments.set_linewidth(linewidthp)

    # Plot to the given axis instance or retrieve the current one
    if axes is None:
        axes = gca()

    # Plot the phase
    axes.add_collection(line_segments)
    # Plot the modulus
    if color is None:
        # Scatter has a problem with complex data type, make sure values are purely real
        axes.scatter(grid, real(modulus), c=rgb_colors[0], **kwargs)
    else:
        axes.plot(grid, modulus, linestyle="", marker=markerp, color=color, **kwargs)
    # Plot the ground line
    axes.plot(grid, zeros(grid.shape), linestyle=linestylep, color="k", **kwargs)