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

# # 14 - Advanced topics - Cement Pavers albedo example
#
# This journal creates a paver underneath the single-axis trackers, and evaluates the improvement for one day -- June 17th with and without the pavers for a location in Davis, CA.
#
# ![Paver](../images_wiki/AdvancedJournals/Pavers.PNG)
#
# Measurements:
# ![Paver](../images_wiki/AdvancedJournals/Pavers_Geometry.PNG)

# In[1]:

import os
from pathlib import Path
import pandas as pd

testfolder = str(Path().resolve().parent.parent / 'bifacial_radiance' / 'TEMP' / 'Tutorial_14')
if not os.path.exists(testfolder):
    os.makedirs(testfolder)

print ("Your simulation will be stored in %s" % testfolder)

# In[2]:

from bifacial_radiance import *
import numpy as np

# In[3]:

simulationname = 'tutorial_14'

#Location:
lat = 38.5449 # Davis, CA
lon = -121.7405 # Davis, CA
# MakeModule Parameters
moduletype='test-module'
numpanels = 1  # AgriPV site has 3 modules along the y direction (N-S since we are facing it to the south) .
x = 0.95
y = 1.838
xgap = 0.02# Leaving 2 centimeters between modules on x direction
ygap = 0.0 # 1 - up
zgap = 0.06 # gap between modules and torquetube.

# Other default values:

# TorqueTube Parameters
axisofrotationTorqueTube=True
torqueTube = False
cellLevelModule = True

numcellsx = 6
numcellsy = 10
xcell = 0.156
ycell = 0.158
xcellgap = 0.015
ycellgap = 0.015

sensorsy = numcellsy   # one sensor per cell

cellLevelModuleParams = {'numcellsx': numcellsx, 'numcellsy':numcellsy,
                         'xcell': xcell, 'ycell': ycell, 'xcellgap': xcellgap, 'ycellgap': ycellgap}

# SceneDict Parameters
gcr = 0.33 # m
albedo = 0.2  #'grass'     # ground albedo
hub_height = 1.237 # m
nMods = 20 # six modules per row.
nRows = 3  # 3 row

azimuth_ang = 90 # Facing east

demo = RadianceObj(simulationname,path = testfolder)  # Create a RadianceObj 'object'
demo.setGround(albedo) #
epwfile = demo.getEPW(lat, lon)
metdata = demo.readWeatherFile(epwfile, coerce_year=2021) # read in the EPW weather data from above
mymodule=demo.makeModule(name=moduletype,x=x,y=y,numpanels = numpanels, xgap=xgap, ygap=ygap)
mymodule.addCellModule(numcellsx=numcellsx, numcellsy=numcellsy,
                       xcell=xcell, ycell=ycell, xcellgap=xcellgap, ycellgap=ycellgap)

# In[4]:

description = 'Sherman Williams "Chantilly White" acrylic paint'
materialpav = 'sw_chantillywhite'
Rrefl = 0.5
Grefl = 0.5
Brefl = 0.5
demo.addMaterial(material=materialpav, Rrefl=Rrefl, Grefl=Grefl, Brefl=Brefl, comment=description)

# ### Simulation without Pavers

# In[5]:

timeindex = metdata.datetime.index(pd.to_datetime('2021-06-17 12:0:0 -8'))  # Davis, CA is TZ -8
demo.gendaylit(timeindex)

tilt = demo.getSingleTimestampTrackerAngle(metdata, timeindex=timeindex, gcr=gcr,
                                   azimuth=180, axis_tilt=0,
                                   limit_angle=60, backtrack=True)
# create a scene with all the variables
sceneDict = {'tilt':tilt,'gcr': gcr,'hub_height':hub_height,'azimuth':azimuth_ang, 'module_type':moduletype, 'nMods': nMods, 'nRows': nRows}
scene = demo.makeScene(module=mymodule, sceneDict=sceneDict) #makeScene creates a .rad file with 20 modules per row, 7 rows.
octfile = demo.makeOct(demo.getfilelist())  # makeOct combines all of the ground, sky and object fil|es into a .oct file.

# In[6]:

analysis = AnalysisObj(octfile, demo.name)  # return an analysis object including the scan dimensions for back irradiance
frontscan, backscan = analysis.moduleAnalysis(scene, sensorsy=sensorsy)
analysis.analysis(octfile, simulationname+"_noPavers", frontscan, backscan)  # compare the back vs front irradiance
print("Simulation without Pavers Finished")

# ## Looping on the day

# In[7]:

j=0
starttimeindex = metdata.datetime.index(pd.to_datetime('2021-06-17 7:0:0 -8'))
endtimeindex = metdata.datetime.index(pd.to_datetime('2021-06-17 19:0:0 -8'))
for timess in range (starttimeindex, endtimeindex):
    j+=1
    demo.gendaylit(timess)
    tilt = demo.getSingleTimestampTrackerAngle(metdata, timeindex=timess, gcr=gcr,
                                       azimuth=180, axis_tilt=0,
                                       limit_angle=60, backtrack=True)
    # create a scene with all the variables
    sceneDict = {'tilt':tilt,'gcr': gcr,'hub_height':hub_height,'azimuth':azimuth_ang, 'module_type':moduletype, 'nMods': nMods, 'nRows': nRows}
    scene = demo.makeScene(module=mymodule, sceneDict=sceneDict) #makeScene creates a .rad file with 20 modules per row, 7 rows.
    octfile = demo.makeOct(demo.getfilelist())  # makeOct combines all of the ground, sky and object fil|es into a .oct file
    frontscan, backscan = analysis.moduleAnalysis(scene, sensorsy=sensorsy)
    analysis.analysis(octfile, simulationname+"_noPavers_"+str(j), frontscan, backscan)  # compare the back vs front irradiance

# ### Simulation With Pavers

# In[8]:

demo.gendaylit(timeindex)
tilt = demo.getSingleTimestampTrackerAngle(metdata, timeindex=timeindex, gcr=gcr,
                                   azimuth=180, axis_tilt=0,
                                   limit_angle=60, backtrack=True)
# create a scene with all the variables
sceneDict = {'tilt':tilt,'gcr': gcr,'hub_height':hub_height,'azimuth':azimuth_ang, 'module_type':moduletype, 'nMods': nMods, 'nRows': nRows}
scene = demo.makeScene(module=mymodule, sceneDict=sceneDict) #makeScene creates a .rad file with 20 modules per row, 7 rows.

# In[9]:

torquetubelength = demo.module.scenex*(nMods)
pitch = demo.module.sceney/gcr
startpitch = -pitch * (nRows-1)/2
p_w = 0.947 # m
p_h = 0.092 # m
p_w2 = 0.187 # m
p_h2 = 0.184 # m
offset_w1y = -(p_w/2)+(p_w2/2)
offset_w2y = (p_w/2)-(p_w2/2)

customObjects = []
for i in range (0, nRows):
    name='PAVER'+str(i)
    text='! genbox {} paver{} {} {} {} | xform -t {} {} 0 | xform -t {} 0 0'.format(materialpav, i,
                                    p_w, torquetubelength, p_h,
                                    -p_w/2, (-torquetubelength+demo.module.sceney)/2.0,
                                    startpitch+pitch*i)
    text += '\r\n! genbox {} paverS1{} {} {} {} | xform -t {} {} 0 | xform -t {} 0 0'.format(materialpav, i,
                                    p_w2, torquetubelength, p_h2,
                                    -p_w2/2+offset_w1y, (-torquetubelength+demo.module.sceney)/2.0,
                                    startpitch+pitch*i)
    text += '\r\n! genbox {} paverS2{} {} {} {} | xform -t {} {} 0 | xform -t {} 0 0'.format(materialpav, i,
                                    p_w2, torquetubelength, p_h2,
                                    -p_w2/2+offset_w2y, (-torquetubelength+demo.module.sceney)/2.0,
                                    startpitch+pitch*i)

    customObject = demo.makeCustomObject(name,text)
    customObjects.append(customObject)
    demo.appendtoScene(radfile=scene.radfiles, customObject=customObject, text="!xform -rz 0")

# In[10]:

demo.makeOct()

# You can view the geometry generated in the terminal with:
#
# ### rvu -vf views\front.vp -e .01 -pe 0.01 -vp -5 -14 1 -vd 0 0.9946 -0.1040 tutorial_14.oct

# In[11]:

## Comment the ! line below to run rvu from the Jupyter notebook instead of your terminal.
## Simulation will stop until you close the rvu window

#!rvu -vf views\front.vp -e .01 -pe 0.01 -vp -5 -14 1 -vd 0 0.9946 -0.1040 tutorial_14.oct

# In[12]:

analysis = AnalysisObj(octfile, demo.name)  # return an analysis object including the scan dimensions for back irradiance
frontscan, backscan = analysis.moduleAnalysis(scene, sensorsy=sensorsy)
analysis.analysis(octfile, simulationname+"_WITHPavers", frontscan, backscan)  # compare the back vs front irradiance
print("Simulation WITH Pavers Finished")

# ## LOOP WITH PAVERS

# In[13]:

j=0
for timess in range (starttimeindex, endtimeindex):
    j+=1
    demo.gendaylit(timess)
    tilt = demo.getSingleTimestampTrackerAngle(metdata, timeindex=timess, gcr=gcr,
                                       azimuth=180, axis_tilt=0,
                                       limit_angle=60, backtrack=True)
    # create a scene with all the variables
    sceneDict = {'tilt':tilt,'gcr': gcr,'hub_height':hub_height,'azimuth':azimuth_ang, 'module_type':moduletype, 'nMods': nMods, 'nRows': nRows}
    scene = demo.makeScene(mymodule, sceneDict=sceneDict) #makeScene creates a .rad file with 20 modules per row, 7 rows.
    # Appending Pavers here
    demo.appendtoScene(radfile=scene.radfiles, customObject=customObjects[0], text="!xform -rz 0")
    demo.appendtoScene(radfile=scene.radfiles, customObject=customObjects[1], text="!xform -rz 0")
    demo.appendtoScene(radfile=scene.radfiles, customObject=customObjects[2], text="!xform -rz 0")
    octfile = demo.makeOct(demo.getfilelist())  # makeOct combines all of the ground, sky and object fil|es into a .oct file
    frontscan, backscan = analysis.moduleAnalysis(scene, sensorsy=sensorsy)
    analysis.analysis(octfile, simulationname+"_WITHPavers_"+str(j), frontscan, backscan)  # compare the back vs front irradiance

# # RESULTS ANALYSIS NOON

# In[14]:

df_0 = load.read1Result(os.path.join(testfolder, 'results', 'irr_tutorial_14_noPavers.csv'))
df_w = load.read1Result(os.path.join(testfolder, 'results', 'irr_tutorial_14_WITHPavers.csv'))

# In[15]:

df_0

# In[16]:

df_w

# ## Improvement in Rear Irradiance

# In[17]:

round((df_w['Wm2Back'].mean()-df_0['Wm2Back'].mean())*100/df_0['Wm2Back'].mean(),1)

# # RESULT ANALYSIS DAY

# In[18]:

df_0 = load.read1Result(os.path.join(testfolder, 'results', 'irr_tutorial_14_noPavers_1.csv'))
df_w = load.read1Result(os.path.join(testfolder, 'results', 'irr_tutorial_14_WITHPavers_1.csv'))

# In[19]:

df_w

# In[20]:

df_0

# In[21]:

round((df_w['Wm2Back'].mean()-df_0['Wm2Back'].mean())*100/df_0['Wm2Back'].mean(),1)

# In[22]:

average_back_d0=[]
average_back_dw=[]
average_front = []
hourly_rearirradiance_comparison = []

timessimulated = endtimeindex-starttimeindex

for i in range (1, timessimulated+1):
    df_0 = load.read1Result(os.path.join(testfolder, 'results', 'irr_tutorial_14_noPavers_'+str(i)+'.csv'))
    df_w = load.read1Result(os.path.join(testfolder, 'results', 'irr_tutorial_14_WITHPavers_'+str(i)+'.csv'))
    print(round((df_w['Wm2Back'].mean()-df_0['Wm2Back'].mean())*100/df_0['Wm2Back'].mean(),1))
    hourly_rearirradiance_comparison.append(round((df_w['Wm2Back'].mean()-df_0['Wm2Back'].mean())*100/df_0['Wm2Back'].mean(),1))
    average_back_d0.append(df_0['Wm2Back'].mean())
    average_back_dw.append(df_w['Wm2Back'].mean())
    average_front.append(df_0['Wm2Front'].mean())

# In[23]:

print("Increase in rear irradiance: ", round((sum(average_back_dw)-sum(average_back_d0))*100/sum(average_back_d0),1))

# In[24]:

print("BG no Pavers: ", round(sum(average_back_d0)*100/sum(average_front),1))
print("BG with Pavers: ", round(sum(average_back_dw)*100/sum(average_front),1))

# In[27]:

import matplotlib.pyplot as plt

#metdata.datetime[starttime].hour # 7
#metdata.datetime[endtimeindex].hour # 17
xax= [7, 8, 9, 10, 11, 12,13,14,15,16,17,18]  # Lazy way to get the x axis...

# In[28]:

plt.plot(xax,hourly_rearirradiance_comparison)
plt.ylabel('$\Delta$ in G$_{rear}$ [%] \n(G$_{rear-with}$ - G$_{rear-without}$ / G$_{rear-without}$)')
plt.xlabel('Hour')

docs/tutorials/14 - Advanced topics - Cement Pavers albedo example.py


from typing import Optional

import hypothesis.extra.numpy as hnp
import hypothesis.strategies as st
import numpy as np
import pytest
from hypothesis import assume, given, settings
from numpy.testing import assert_array_equal, assert_equal
from pytest import raises

import mygrad as mg
from mygrad import Tensor
from mygrad.errors import InvalidBackprop
from mygrad.math.misc.ops import MatMul
from mygrad.math.arithmetic.ops import Add, Divide, Multiply, Negative, Power, Subtract
from mygrad.operation_base import Operation
from tests.custom_strategies import tensors, valid_constant_arg
from tests.utils.errors import does_not_raise

def test_simple_default_constant_behavior():
    assert Tensor(1).constant is True
    assert Tensor(1.0).constant is False

@pytest.mark.parametrize(
    "data",
    [
        None,
        np.array(None, dtype="O"),
        np.array([[0], [0, 0]], dtype="O"),
        np.array(1, dtype="O"),
        0j,
        np.array(1, dtype=complex),
    ],
)
@given(constant=st.booleans(), creator=st.none() | st.just(MatMul()))
def test_input_type_checking(data, constant, creator):
    with raises(TypeError):
        Tensor(data, constant=constant, _creator=creator)

@pytest.mark.parametrize("constant", [1, (1,), np.array(1.0), "true", 0])
def test_input_constant_checking(constant):
    with raises(TypeError):
        Tensor(1.0, constant=constant)

@given(
    x=hnp.arrays(
        shape=hnp.array_shapes(min_dims=0, min_side=0), dtype=hnp.floating_dtypes()
    ),
    dtype=hnp.floating_dtypes(),
    copy=st.booleans(),
    ndmin=st.integers(0, 10),
)
def test_ndmin(x: np.ndarray, copy: bool, dtype, ndmin: int):
    """Ensure Tensor(..., ndmin=<val>) mirrors numpy behavior and
    produces appropriate view behavior"""
    arr = np.array(x, copy=copy, ndmin=ndmin, dtype=dtype)
    tensor = Tensor(x, copy=copy, ndmin=ndmin, dtype=dtype)
    assert_equal(tensor, arr)

    # Specifying array(... , ndmin=val) can create an internal
    # base that the array points to. We want to be sure that
    # this behavior doesnt manifest in tensor, which should
    # bot behave as if it is a view

    # check base behavior
    assert tensor.data.flags.writeable
    assert tensor.base is None
    assert tensor[...].base is tensor

    # check mem-lock behavior
    z = +tensor
    assert not tensor.data.flags.writeable
    del z
    assert tensor.data.flags.writeable

@given(
    x=hnp.arrays(
        shape=hnp.array_shapes(min_dims=0, min_side=0), dtype=hnp.floating_dtypes()
    ),
    constant=st.booleans(),
    data=st.data(),
)
def test_basic_backward(x: np.ndarray, constant: bool, data: st.DataObject):
    """Ensure Tensor.backward() sets the expected gradient for general array-shape/dtype"""
    grad = data.draw(hnp.arrays(shape=x.shape, dtype=x.dtype) | st.none(), label="grad")
    tensor = Tensor(x, constant=constant)
    tensor.backward(grad)
    if tensor.constant:
        assert tensor.grad is None
    else:
        assert_array_equal(
            tensor.grad, np.ones_like(tensor.data) if grad is None else grad
        )

    if tensor.grad is not None:
        assert tensor.dtype == tensor.grad.dtype

@given(
    data=hnp.arrays(shape=hnp.array_shapes(), dtype=hnp.floating_dtypes()),
    constant=st.booleans(),
    set_constant=st.booleans() | st.none(),
    invoke_backward=st.booleans(),
)
def test_copy(data, constant, set_constant, invoke_backward):
    x = Tensor(data, constant=constant)
    y = +x
    if invoke_backward:
        y.backward()
    y_copy = y.copy(constant=set_constant)

    if not invoke_backward:
        assert y.creator is not None
    assert y.dtype == y_copy.dtype
    assert y_copy.constant is (constant if set_constant is None else set_constant)
    if y.grad is None:
        assert y_copy.grad is None
    else:
        assert_array_equal(y.grad, y_copy.grad)
    assert_array_equal(y.data, y_copy.data)

@pytest.mark.parametrize("constant", [True, False])
def test_cant_set_constant(constant):
    tensor = Tensor([1.0], constant=constant)
    with pytest.raises(AttributeError):
        tensor.constant = constant

def test_to_scalar():
    nd_tensor = Tensor([1, 2])
    with raises(TypeError):
        float(nd_tensor)

    with raises(TypeError):
        int(nd_tensor)

    with raises(ValueError):
        nd_tensor.item()

    for size1_tensor in (Tensor(1), Tensor([[1]])):
        assert float(size1_tensor) == 1.0
        assert int(size1_tensor) == 1
        assert size1_tensor.item() == 1.0

@pytest.mark.parametrize(
    ("tensor", "repr_"),
    [
        (Tensor(1), "Tensor(1)"),
        (Tensor([1]), "Tensor([1])"),
        (Tensor([1, 2]), "Tensor([1, 2])"),
        (
            mg.arange(9).reshape((3, 3)),
            "Tensor([[0, 1, 2],\n        [3, 4, 5],\n        [6, 7, 8]])",
        ),
    ],
)
def test_repr(tensor, repr_):
    assert repr(tensor) == repr_

@pytest.mark.parametrize("bad_grads", ["bad", 1.0j])
@given(constant=st.booleans())
def test_invalid_gradient_raises(constant: bool, bad_grads):
    x = Tensor(3.0, constant=constant) * 2
    with (pytest.raises((TypeError, ValueError)) if not constant else does_not_raise()):
        x.backward(bad_grads)

@pytest.mark.parametrize("element", (0, [0, 1, 2]))
def test_contains(element):
    t = Tensor([[0, 1, 2], [3, 4, 5]])
    assert (element in t) is (element in t.data)

@given(
    a=hnp.arrays(
        shape=hnp.array_shapes(max_side=3, max_dims=5),
        dtype=float,
        elements=st.floats(-100, 100),
    ),
    constant=st.booleans(),
    creator=st.booleans(),
)
def test_properties(a, constant, creator):
    array = np.asarray(a)
    if creator:
        ref = Add()
        tensor = Tensor(a, constant=constant, _creator=ref)
    else:
        ref = None
        tensor = Tensor(a, constant=constant)

    assert tensor.ndim == array.ndim
    assert tensor.shape == array.shape
    assert tensor.size == array.size
    assert len(tensor) == len(array)
    assert tensor.dtype == array.dtype
    assert_equal(actual=tensor.data, desired=a)
    assert (not creator) or tensor.creator is ref

def test_init_data():
    for data in [0, [], (0, 0), ((0, 0), (0, 0)), np.random.rand(3, 4, 2)]:
        assert_equal(
            actual=Tensor(data).data,
            desired=np.asarray(data),
            err_msg="Initialization with non-tensor failed",
        )
        assert_equal(
            actual=Tensor(Tensor(data)).data,
            desired=np.asarray(data),
            err_msg="Initialization with tensor failed",
        )

@given(
    arr=hnp.arrays(
        dtype=hnp.floating_dtypes(), shape=hnp.array_shapes(min_side=1, min_dims=1)
    ),
    as_tensor=st.booleans(),
    constant=st.booleans(),
)
def test_tensor_copy_on_init_mirrors_array(
    arr: np.ndarray, as_tensor: bool, constant: bool
):
    x = Tensor(arr, copy=False) if as_tensor else arr
    tensor = Tensor(x, constant=constant)
    array = np.array(arr)
    assert np.shares_memory(tensor, arr) is np.shares_memory(array, arr)

@given(
    x=hnp.arrays(
        dtype=hnp.floating_dtypes(), shape=hnp.array_shapes(min_side=0, min_dims=0)
    )
)
def test_init_data_rand(x: np.ndarray):
    assert_equal(actual=Tensor(x).data, desired=x)

@given(
    x=hnp.arrays(
        dtype=float,
        shape=hnp.array_shapes(min_dims=0, min_side=0),
        elements=st.floats(allow_infinity=False, allow_nan=False),
    )
    | st.floats(allow_infinity=False, allow_nan=False)
    | st.integers(-100, 100)
)
def test_items(x):
    """ verify that tensor.item() mirrors array.item()"""
    tensor = Tensor(x)
    try:
        value = np.asarray(x).item()
        assert_array_equal(value, tensor.item())
    except ValueError:
        with raises(ValueError):
            tensor.item()

op = Add()
dtype_strat = st.sampled_from(
    (
        None,
        int,
        float,
        np.int8,
        np.int16,
        np.int32,
        np.int64,
        np.float16,
        np.float32,
        np.float64,
    )
)
dtype_strat_numpy = st.sampled_from(
    (np.int8, np.int16, np.int32, np.int64, np.float16, np.float32, np.float64)
)

@given(
    data=st.data(),
    creator=st.sampled_from((None, op)),
    dtype=dtype_strat,
    numpy_dtype=dtype_strat_numpy,
    ndmin=st.integers(0, 10),
    copy=st.none() | st.booleans(),
)
def test_init_params(
    data,
    creator,
    dtype,
    numpy_dtype,
    ndmin: int,
    copy: Optional[bool],
):
    """Check for bad combinations of init parameters leading to unexpected behavior"""
    elements = (
        (lambda x, y: st.floats(x, y, width=8 * np.dtype(numpy_dtype).itemsize))
        if np.issubdtype(numpy_dtype, np.floating)
        else st.integers
    )
    array_strat_args = dict(
        shape=hnp.array_shapes(max_side=3, max_dims=5),
        dtype=numpy_dtype,
        elements=elements(-100, 100),
    )
    a = data.draw(
        hnp.arrays(**array_strat_args) | tensors(**array_strat_args),
        label="a",
    )

    arr = np.array(a, dtype=dtype, ndmin=ndmin)

    constant = data.draw(valid_constant_arg(arr.dtype), label="constant")

    tensor = Tensor(
        a,
        _creator=creator,
        constant=constant,
        dtype=dtype,
        ndmin=ndmin,
        copy=copy,
    )

    if constant is None:
        constant = issubclass(tensor.dtype.type, np.integer)

    assert tensor.creator is creator
    assert tensor.constant is constant
    assert tensor.dtype is arr.dtype
    assert_equal(tensor.data, arr)
    assert tensor.grad is None
    assert tensor.base is None
    assert tensor.ndim >= ndmin

@pytest.mark.parametrize(
    ("op_name", "op"),
    [
        ("add", Add),
        ("sub", Subtract),
        ("mul", Multiply),
        ("truediv", Divide),
        ("pow", Power),
        ("matmul", MatMul),
    ],
)
@pytest.mark.parametrize("right_op", [True, False])
@given(constant_x=st.booleans(), constant_y=st.booleans())
def test_special_methods(
    op_name: str, op: Operation, constant_x: bool, constant_y: bool, right_op: bool
):
    if right_op:
        op_name = "r" + op_name
    op_name = "__" + op_name + "__"
    x = Tensor([2.0, 8.0, 5.0], constant=constant_x)
    y = Tensor([1.0, 3.0, 2.0], constant=constant_y)

    constant = constant_x and constant_y
    assert hasattr(Tensor, op_name)
    tensor_out = getattr(Tensor, op_name)(x, y)
    numpy_out = getattr(np.ndarray, op_name)(x.data, y.data)
    assert isinstance(tensor_out, Tensor)
    assert tensor_out.constant is constant
    assert_equal(tensor_out.data, numpy_out)
    assert isinstance(tensor_out.creator, op)

    if not right_op:
        assert tensor_out.creator.variables[0] is x
        assert tensor_out.creator.variables[1] is y
    else:
        assert tensor_out.creator.variables[0] is y
        assert tensor_out.creator.variables[1] is x

@given(
    x=hnp.arrays(shape=hnp.array_shapes(), dtype=hnp.floating_dtypes()),
    constant=st.booleans(),
)
def test_pos(x: np.ndarray, constant: bool):
    assume(np.all(np.isfinite(x)))
    x = Tensor(x, constant=constant)
    y = +x
    assert y.creator.variables[0] is x
    assert_array_equal(y.data, x.data)
    assert y.constant is x.constant

@given(x=hnp.arrays(shape=hnp.array_shapes(), dtype=hnp.floating_dtypes()))
def test_neg(x):
    assume(np.all(np.isfinite(x)))
    x = Tensor(x)
    op_name = "__neg__"
    assert hasattr(Tensor, op_name)
    tensor_out = getattr(Tensor, "__neg__")(x)
    numpy_out = getattr(np.ndarray, "__neg__")(x.data)
    assert isinstance(tensor_out, Tensor)
    assert_equal(tensor_out.data, numpy_out)
    assert isinstance(tensor_out.creator, Negative)
    assert tensor_out.creator.variables[0] is x

@pytest.mark.parametrize(
    "op", ("__lt__", "__le__", "__gt__", "__ge__", "__eq__", "__ne__")
)
@given(
    x=hnp.arrays(
        shape=hnp.array_shapes(),
        dtype=hnp.floating_dtypes(),
        elements=st.floats(-10, 10, width=16),
    ),
    x_constant=st.booleans(),
    y_constant=st.booleans(),
    data=st.data(),
)
def test_comparison_ops(
    op: str, x: np.ndarray, x_constant: bool, y_constant: bool, data: st.SearchStrategy
):
    y = data.draw(
        hnp.arrays(shape=x.shape, dtype=x.dtype, elements=st.floats(-10, 10, width=16))
    )
    x = Tensor(x, constant=x_constant)
    y = Tensor(y, constant=y_constant)
    assert hasattr(Tensor, op), "`Tensor` is missing the attribute {}".format(op)
    tensor_out = getattr(Tensor, op)(x, y)
    array_out = getattr(np.ndarray, op)(x.data, y.data)
    assert_equal(actual=tensor_out, desired=array_out)

@pytest.mark.parametrize(
    "attr",
    (
        "sum",
        "prod",
        "cumprod",
        "cumsum",
        "mean",
        "std",
        "var",
        "max",
        "min",
        "transpose",
        "squeeze",
        "ravel",
    ),
)
@given(constant=st.booleans())
def test_math_methods(attr: str, constant: bool):
    x = Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], constant=constant)

    assert hasattr(x, attr)
    method_out = getattr(x, attr).__call__()
    function_out = getattr(mg, attr).__call__(x)
    assert_equal(method_out.data, function_out.data)
    assert method_out.constant is constant
    assert type(method_out.creator) is type(function_out.creator)

@pytest.mark.parametrize(
    "attr",
    (
        "argmax",
        "argmin",
        "any",
    ),
)
def test_numpy_math_methods(attr: str):
    x = Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])

    assert hasattr(x, attr)
    method_out = getattr(x, attr).__call__()
    function_out = getattr(np, attr).__call__(x.data)
    assert_equal(method_out, function_out)
    assert_equal(method_out, function_out)

@pytest.mark.parametrize("op", ("moveaxis", "swapaxes"))
@given(constant=st.booleans())
def test_axis_interchange_methods(op: str, constant: bool):
    x = Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], constant=constant)
    method_out = getattr(x, op)(0, -1)
    function_out = getattr(mg, op)(x, 0, -1)
    assert_equal(method_out.data, function_out.data)
    assert method_out.constant is constant
    assert type(method_out.creator) is type(function_out.creator)

@given(x=tensors(include_grad=st.booleans()), clear_graph=st.booleans())
def test_null_gradients(x: Tensor, clear_graph: bool):
    with pytest.warns(FutureWarning):
        x.null_gradients(clear_graph=clear_graph)

@given(x=tensors(include_grad=st.booleans()))
def test_null_grad(x: Tensor):
    y = x.null_grad()
    assert x.grad is None
    assert y is x

@settings(deadline=None)
@given(
    x=st.floats(min_value=-1e-6, max_value=1e6),
    y=st.floats(min_value=-1e-6, max_value=1e6),
    z=st.floats(min_value=-1e-6, max_value=1e6),
)
def test_clear_graph(x, y, z):
    x = Tensor(x)
    y = Tensor(y)
    z = Tensor(z)

    f = x * y + z
    g = x + z * f * f

    # check side effects
    unused = 2 * g - f
    w = 1 * f
    assert unused is not None

    w.clear_graph()

    assert len(g._ops) > 0
    assert g.creator is not None
    assert len(x._ops) == 0
    assert len(y._ops) == 0
    assert len(z._ops) == 0
    assert len(f._ops) == 0
    assert x.creator is None
    assert y.creator is None
    assert z.creator is None
    assert f.creator is None

    with raises(InvalidBackprop):
        g.backward()

# Tensor has its `__eq__` but not its `__hash__` overridden which leads to subtle
# problems if it ends up being used in a hashable context. See
# https://hynek.me/articles/hashes-and-equality/
# for more details. This checks to make sure that anyone who does so will get an
# error. See also https://github.com/rsokl/MyGrad/pull/276
def test_no_hash():
    try:
        {Tensor(3): "this should not work"}
    except TypeError as e:
        assert str(e) == "unhashable type: 'Tensor'"


import numpy as np

def convert_to_color_(arr_2d, palette=None):
    """Convert an array of labels to RGB color-encoded image.

    Args:
        arr_2d: int 2D array of labels
        palette: dict of colors used (label number -> RGB tuple)

    Returns:
        arr_3d: int 2D images of color-encoded labels in RGB format

    """
    arr_3d = np.zeros((arr_2d.shape[0], arr_2d.shape[1], 3), dtype=np.uint8)
    if palette is None:
        raise Exception("Unknown color palette")

    for c, i in palette.items():
        m = arr_2d == c
        arr_3d[m] = i

    return arr_3d

def convert_to_one_hot(Y, C):
    Y = np.eye(C)[Y.reshape(-1)].T
    return Y


__usage__ = """
Run:
  python return_integer.py [<f2py options>]
Examples:
  python return_integer.py --fcompiler=Gnu --no-wrap-functions
  python return_integer.py --quiet
"""

import numpy.f2py as f2py2e
from numpy import array

def build(f2py_opts):
    try:
        import f77_ext_return_integer
    except ImportError:
        assert not f2py2e.compile('''\
       function t0(value)
         integer value
         integer t0
         t0 = value
       end
       function t1(value)
         integer*1 value
         integer*1 t1
         t1 = value
       end
       function t2(value)
         integer*2 value
         integer*2 t2
         t2 = value
       end
       function t4(value)
         integer*4 value
         integer*4 t4
         t4 = value
       end
       function t8(value)
         integer*8 value
         integer*8 t8
         t8 = value
       end

       subroutine s0(t0,value)
         integer value
         integer t0
cf2py    intent(out) t0
         t0 = value
       end
       subroutine s1(t1,value)
         integer*1 value
         integer*1 t1
cf2py    intent(out) t1
         t1 = value
       end
       subroutine s2(t2,value)
         integer*2 value
         integer*2 t2
cf2py    intent(out) t2
         t2 = value
       end
       subroutine s4(t4,value)
         integer*4 value
         integer*4 t4
cf2py    intent(out) t4
         t4 = value
       end
       subroutine s8(t8,value)
         integer*8 value
         integer*8 t8
cf2py    intent(out) t8
         t8 = value
       end

''','f77_ext_return_integer',f2py_opts,source_fn='f77_ret_int.f')

    from f77_ext_return_integer import t0,t1,t2,t4,t8,s0,s1,s2,s4,s8
    test_functions = [t0,t1,t2,t4,t8,s0,s1,s2,s4,s8]
    return test_functions

def runtest(t):
    import sys
    assert t(123)==123,`t(123)`
    assert t(123.6)==123
    assert t(123l)==123
    if sys.version[:3]<'2.3':
        assert t(123.6+3j)==123
    assert t('123')==123
    assert t(-123)==-123
    assert t([123])==123
    assert t((123,))==123
    assert t(array(123))==123
    assert t(array([123]))==123
    assert t(array([[123]]))==123
    assert t(array([123],'b'))==123
    assert t(array([123],'h'))==123
    assert t(array([123],'i'))==123
    assert t(array([123],'l'))==123
    assert t(array([123],'B'))==123
    assert t(array([123],'f'))==123
    assert t(array([123],'d'))==123
    if sys.version[:3]<'2.3':
        assert t(array([123+3j],'F'))==123
        assert t(array([123],'D'))==123

    try: raise RuntimeError,`t(array([123],'c'))`
    except ValueError: pass
    try: raise RuntimeError,`t('abc')`
    except ValueError: pass

    try: raise RuntimeError,`t([])`
    except IndexError: pass
    try: raise RuntimeError,`t(())`
    except IndexError: pass

    try: raise RuntimeError,`t(t)`
    except TypeError: pass
    try: raise RuntimeError,`t({})`
    except TypeError: pass

    if t.__doc__.split()[0] in ['t8','s8']:
        try: raise RuntimeError,`t(100000000000000000000000l)`
        except OverflowError: pass
        try: raise RuntimeError,`t(10000000011111111111111.23)`
        except OverflowError: pass
    else:
        if sys.version[:3]<'2.3':
            try: raise RuntimeError,`t(10000000000000l)`
            except OverflowError: pass
            try: raise RuntimeError,`t(10000000000.23)`
            except OverflowError: pass

if __name__=='__main__':
    #import libwadpy
    status = 1
    try:
        repeat,f2py_opts = f2py2e.f2py_testing.cmdline()
        test_functions = build(f2py_opts)
        f2py2e.f2py_testing.run(runtest,test_functions,repeat)
        print 'ok'
        status = 0
    finally:
        if status:
            print '*'*20
            print 'Running f2py2e.diagnose'
            import numpy.f2py.diagnose as diagnose
            #diagnose.run()


import numpy as np

def scramble_array(a: np.ndarray) -> None:
    """In-place scrambling of an array along first dimension."""
    # FIXME: This should actually use the Matousek-Affine-Owen scrambling algorithm
    # Matousek, J. “On the L2-Discrepancy for Anchored Boxes.” Journal of Complexity.
    # Vol. 14, Number 4, 1998, pp. 527–556.
    # for now, use a Naive scrambling method
    for i in range(a.shape[1]):
        np.random.shuffle(a[:, i])


# %% [markdown]
# # Exporatory analysis of features extracted from swine audios
# This dataset contains the means of the extracted features, for each audio
# Many descriptions of features were taken from the following books:
# - https://books.google.com/books?hl=en&lr=&id=AF30yR41GIAC&oi=fnd&pg=PP9&dq=Signal+Processing+Methods+for+Music+Transcription&ots=Ooq77iPfKD&sig=WFyZTyVVFtagCBGCRLWpV8rY-Tk  # noqa: E501

# %% [markdown]
# # Imports

# %%
import matplotlib.pyplot as plt
import librosa as lr
import librosa.display as lrdisp
import numpy as np
import pandas as pd
import seaborn as sns

FILE_PREFIX = 'kdd'

# %% [markdown]
# # Loading dataset

# %%
df = pd.read_csv('features_means.csv', index_col=0, verbose=True)
df.index = pd.to_datetime(df.index)
df['rac'] = False
df.loc['2020-09-22':, 'rac'] = True  # type: ignore

# %% [markdown]
# ## Checking for and dropping duplicates

# %%
# Resetting index for duplicate analysis
df.reset_index(inplace=True)
print("Duplicates by filename:",
      df.duplicated(subset=['file_name']).value_counts(),
      sep='\n')
df.drop_duplicates(subset=['file_name'], inplace=True)
print("Duplicates by (datetime, ala, grupo):",
      df.duplicated(subset=['datetime', 'ala', 'grupo']).value_counts(),
      sep='\n')
df.drop_duplicates(subset=['datetime', 'ala', 'grupo'], inplace=True)
# Rebuilding dataframe index
df.set_index('datetime', inplace=True)

# %%
print(df.shape)

# %% [markdown]
# ## Visualizing distribution of sample dates

# %%
df_tmp = pd.DataFrame(df['file_name'].resample('1D').count())
df_tmp['count'] = df_tmp['file_name']
del df_tmp['file_name']
df_tmp['rac'] = False
df_tmp.loc['2020-09-22':, 'rac'] = True  # type: ignore

fig = plt.figure(figsize=(10, 10))
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.barplot(y=df_tmp.index, x=df_tmp['count'], hue=df_tmp['rac'])
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_files_byday.png',
            bbox_inches='tight',
            transparent=True)
plt.draw()

df_tmp = pd.DataFrame(df['file_name'].resample('1H').count())
df_tmp['count'] = df_tmp['file_name']
del df_tmp['file_name']
df_tmp['rac'] = False
df_tmp.loc['2020-09-22':, 'rac'] = True  # type: ignore
df_tmp = df_tmp.reset_index()
df_tmp['hour'] = df_tmp['datetime'].dt.hour

fig = plt.figure(figsize=(10, 10))
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.barplot(y=df_tmp['hour'], x=df_tmp['count'], hue=df_tmp['rac'], orient='h')
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_files_byhour.png',
            bbox_inches='tight',
            transparent=True)
plt.draw()

# %% [markdown]
# There are some data gaps, especially in the pre-ractopamine subset.
#
# There needs to be a resampling to rebalance classes, prioritizing samples
# that are closer to the start of the RAC provision.

# %% [markdown]
# ## Visualizing waveform and spectrogram of random sample

# %%
file_name = np.random.choice(df['file_name'])
x, sr = lr.load(file_name)

fig, ax = plt.subplots(2, 1, sharex=False, sharey=False, figsize=(16, 9))
lrdisp.waveplot(
    x,
    sr=sr,
    ax=ax[0],
)
img = lrdisp.specshow(lr.amplitude_to_db(np.abs(lr.stft(x)), ref=np.max),
                      y_axis='log',
                      x_axis='time',
                      ax=ax[1])
fig.colorbar(img, ax=ax[1], format="%+2.0f dB")
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_audio_sample.png',
            bbox_inches='tight',
            transparent=True)
plt.draw()

# %% [markdown]
# # Visualizing basic dataset properties

# %% [markdown]
# ### Shape

# %%
print("Shape:", df.shape)
print(df.loc[:, 'rac'].groupby(df.loc[:, 'rac']).count())
print(df.info())

# %% [markdown]
# Due to the high amount of extracted features, there must be a selection
# of the most relevant ones for further analysis with advanced methods
# (e.g., SOM)

# %% [markdown]
# ### Class distribution

# %%
df_melt = pd.melt(df, value_vars=['rac'], value_name='ractopamine')
fig = plt.figure(figsize=(10, 10))
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
ax = sns.countplot(data=df_melt, x='ractopamine', hue='ractopamine')

for p in ax.patches:
    ax.annotate(f'\n{p.get_height()}', (p.get_x() + 0.2, p.get_height()),
                ha='center',
                va='top',
                color='white',
                size=18)
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_rac_distribution.png',
            bbox_inches='tight',
            transparent=True)
plt.draw()

# %% [markdown]
# There's a substantial imbalance between the subsets.
# It's not an obstacle for this analysis. However, for the next steps,
# it may be desirable to perform an upsampling/downsampling on the dataset.

# %% [markdown]
# ### Zero Crossings
# The Zero crossing rate (ZCR) measures the number of times that the time
# domain signal changes its sign. Even though it is computed in the time
# domain, it describes the amount of high-frequency energy in the signal
# (i.e., 'brightness') and correlates strongly with the spectral centroid
# ZCR has also proven to be quite discriminative for classes of percussion
# instruments.

# %%
print(df.loc[:, ['zero_crossing_rate', 'zero_crossings']].describe())
df_melt = pd.melt(df, id_vars=['rac'], value_vars=['zero_crossing_rate'])
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.displot(
    data=df_melt,
    x='value',
    hue='rac',
    kde=True,
    rug=True,
    row='variable',
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_zero_crossing_rac.png',
            bbox_inches='tight',
            transparent=True)

df_tmp = df.reset_index()
df_tmp['hour'] = df_tmp['datetime'].dt.hour
df_tmp['day_quarter'] = df_tmp['hour'] // 6
df_melt = pd.melt(df_tmp,
                  id_vars=['day_quarter'],
                  value_vars=['zero_crossing_rate'])
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.displot(
    data=df_melt,
    x='value',
    hue='day_quarter',
    palette='vlag',
    kde=True,
    rug=True,
    row='variable',
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_zero_crossing_day.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# The distribution of the zero crossing seems slightly different after the
# ractopamine provision.
#
# As noted by Ricardo, this may be related to an increase in the animals'
# vocal activity.

# %% [markdown]
# ### Spectrogram and Mel-Spectrogram
# There exists only one Fourier transform of a given signal.
# However, there is an infinite number of time-frequency representations.
# The most popular one is the spectrogram, defined as the Fourier transform of
# successive signal frames. Frames are widely used in audio processing
# algorithms. They are portions of the signal with given time localizations.
#
# Standard window shapes are Gaussian, Hamming, Hanning, or rectangular and
# typical frame durations are from 20 ms to 100 ms in audio processing. As a
# rule of thumb, rectangular windows should not be used in practice, except
# under special circumstances. From frames, it is easy to build short time
# Fourier transforms (STFTs) as the FT of successive frames.
#
# Spectrograms are energy representations and they are defined as the squared
# modulus of the STFT.
#
# The mel-spectrogram is an spectrogram, converted to the Mel scale

# %%
print(df.loc[:, ['spectrogram', 'mel_spectrogram']].describe())
df_melt = pd.melt(df,
                  id_vars=['rac'],
                  value_vars=['spectrogram', 'mel_spectrogram'])
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.displot(
    data=df_melt,
    x='value',
    hue='rac',
    col='variable',
    kde=True,
    rug=True,
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_spectrogram_rac.png',
            bbox_inches='tight',
            transparent=True)

df_tmp = df.reset_index()
df_tmp['hour'] = df_tmp['datetime'].dt.hour
df_tmp['day_quarter'] = df_tmp['hour'] // 6
df_melt = pd.melt(df_tmp,
                  id_vars=['day_quarter'],
                  value_vars=['spectrogram', 'mel_spectrogram'])
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.displot(
    data=df_melt,
    x='value',
    hue='day_quarter',
    palette='vlag',
    kde=True,
    rug=True,
    col='variable',
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_spectrogram_day.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# No relevant differences in the distributions were observed.
#
# These features are discard candidates.
#
# Further analysis required with correlation heatmap.

# %% [markdown]
# ### Harmonics and Perceptual Shock Wave
# Harmonics are characteristichs that represent the sound color
# Perceptrual shock wave represents the sound rhythm and emotion

# %%
print(df.loc[:, ['harmonics', 'perceptual_shock_wave']].describe())
df_melt = pd.melt(df,
                  id_vars=['rac'],
                  value_vars=['harmonics', 'perceptual_shock_wave'])
sns.displot(
    data=df_melt,
    x='value',
    hue='rac',
    col='variable',
    kde=True,
    rug=True,
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_harmonics_rac.png',
            bbox_inches='tight',
            transparent=True)

df_tmp = df.reset_index()
df_tmp['hour'] = df_tmp['datetime'].dt.hour
df_tmp['day_quarter'] = df_tmp['hour'] // 6
df_melt = pd.melt(df_tmp,
                  id_vars=['day_quarter'],
                  value_vars=['harmonics', 'perceptual_shock_wave'])
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.displot(
    data=df_melt,
    x='value',
    hue='day_quarter',
    palette='vlag',
    kde=True,
    rug=True,
    col='variable',
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_harmonics_day.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# No relevant differences in the distributions were observed.
#
# These features are discard candidates.
#
# Further analysis required with correlation heatmap.

# %% [markdown]
# ### Spectral Centroids
# In addition to the rough spectrum described by the MFCCs and bandwise
# energy descriptors, more simple spectral shape features are also useful.
# These include the first four moments of the spectrum, i.e.,
# spectral centroid, spectralspread, spectral skewness, and spectral kurtosis.
#
# The spectral centroid is the "center-of-mass" of the analysed audio file.

# %%
print(df.loc[:, ['spectral_centroids']].describe())
df_melt = pd.melt(df, id_vars=['rac'], value_vars=['spectral_centroids'])
sns.displot(
    data=df_melt,
    x='value',
    hue='rac',
    col='variable',
    kde=True,
    rug=True,
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_spectral_centroids_rac.png',
            bbox_inches='tight',
            transparent=True)

df_tmp = df.reset_index()
df_tmp['hour'] = df_tmp['datetime'].dt.hour
df_tmp['day_quarter'] = df_tmp['hour'] // 6
df_melt = pd.melt(df_tmp,
                  id_vars=['day_quarter'],
                  value_vars=['spectral_centroids'])
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.displot(
    data=df_melt,
    x='value',
    hue='day_quarter',
    palette='vlag',
    kde=True,
    rug=True,
    col='variable',
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_spectral_centroids_day.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# The distribution of the spectral centroids seems to have a displacement to
# higher values, after the provision of ractopamine to the animals.
#
# This seems to suggest that the most intense peaks of audio were at higher
# frequencies, with rac=True.

# %% [markdown]
# ### Spectral Centroids (Delta and Accelerate)

# %%

print(df.loc[:, ['spectral_centroids_delta', 'spectral_centroids_accelerate']].
      describe())
df_melt = pd.melt(
    df,
    id_vars=['rac'],
    value_vars=['spectral_centroids_delta', 'spectral_centroids_accelerate'])
sns.displot(
    data=df_melt,
    x='value',
    hue='rac',
    col='variable',
    kde=True,
    rug=True,
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(
    f'./output_{FILE_PREFIX}/{FILE_PREFIX}'
    '_spectral_centroids_derivatives_rac.png',
    bbox_inches='tight',
    transparent=True)

df_tmp = df.reset_index()
df_tmp['hour'] = df_tmp['datetime'].dt.hour
df_tmp['day_quarter'] = df_tmp['hour'] // 6
df_melt = pd.melt(
    df_tmp,
    id_vars=['day_quarter'],
    value_vars=['spectral_centroids_delta', 'spectral_centroids_accelerate'])
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.displot(
    data=df_melt,
    x='value',
    hue='day_quarter',
    palette='vlag',
    kde=True,
    rug=True,
    col='variable',
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(
    f'./output_{FILE_PREFIX}/{FILE_PREFIX}'
    '_spectral_centroids_derivatives_day.png',
    bbox_inches='tight',
    transparent=True)

# %% [markdown]
# No relevant differences in the distributions were observed.
#
# These features are discard candidates.
#
# Further analysis required with correlation heatmap.

# %% [markdown]
# ### Chroma features
# The chroma vector is a perceptually motivated feature vector using the
# concept of chroma in Shepard's helix representation of musical pitch
# perception. According to Shepard [584], the perception of pitch with respect
# to a musical context can be graphically represented by using a continually
# cyclic helix that has two dimensions, chroma and height. Chroma refers to the
# position of a musical pitch within an octave that corresponds to a cycle of
# the helix.

# %%

# Visualizing single-file chromagram
chroma = lr.feature.chroma_stft(y=x, sr=sr, n_chroma=12, n_fft=2048)
fig, ax = plt.subplots(figsize=(16, 9))
img = lr.display.specshow(chroma, y_axis='chroma', x_axis='time', ax=ax)
fig.colorbar(img, ax=ax)
ax.set(title='Chromagram')
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_chromagram.png',
            bbox_inches='tight',
            transparent=True)

chroma_cols = [
    'chroma1',
    'chroma2',
    'chroma3',
    'chroma4',
    'chroma5',
    'chroma6',
    'chroma7',
    'chroma8',
    'chroma9',
    'chroma10',
    'chroma11',
    'chroma12',
]
print(df.loc[:, chroma_cols].describe())
df_melt = pd.melt(df, id_vars=['rac'], value_vars=chroma_cols)
sns.displot(
    data=df_melt,
    x='value',
    hue='rac',
    row='variable',
    kde=True,
    rug=True,
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_chroma_features_rac.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# Chroma features may be very useful for distinguishing musical tones.
# However, related research papers didn't seem to use it for pig audio
# analysis. With this, they are discard candidates.
#
# Still, there were some interesting changes in distribution for chroma7,
# chroma8, chroma9, chroma11, and chroma12, which should be further analyzed.

# %%
plt.figure(figsize=(12, 12))
sns.heatmap(df.loc[:, chroma_cols + ['rac']].corr(),
            vmin=-1.0,
            vmax=1.0,
            cmap="coolwarm",
            center=0.0,
            robust=True,
            annot=True,
            fmt='.1f')
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_chroma_features_corr.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# As expected, the chroma features are highly correlated with their neighbors.
#
# When comparing to our class (rac), the hightest absolute correlation between
# it and a chroma feature is 0.15 which is not very substantial.

# %% [markdown]
# ### Tempo (BMP)
# Estimate mean of the tempo (beats per minute), as extracted by librosa.

# %%
print(df.loc[:, ['tempo_bpm']].describe())
df_melt = pd.melt(df, id_vars=['rac'], value_vars=['tempo_bpm'])
sns.displot(
    data=df_melt,
    x='value',
    hue='rac',
    col='variable',
    kde=True,
    rug=True,
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_tempo_rac.png',
            bbox_inches='tight',
            transparent=True)

df_tmp = df.reset_index()
df_tmp['hour'] = df_tmp['datetime'].dt.hour
df_tmp['day_quarter'] = df_tmp['hour'] // 6
df_melt = pd.melt(df_tmp, id_vars=['day_quarter'], value_vars=['tempo_bpm'])
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.displot(
    data=df_melt,
    x='value',
    hue='day_quarter',
    palette='vlag',
    kde=True,
    rug=True,
    col='variable',
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_tempo_day.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# This feature, similarly to the chroma features, is mostly used with music
# analysis. No relevant differences in the distributions were observed.
#
# This feature is a discard candidate.
#
# Further analysis required with correlation heatmap.

# %% [markdown]
# ### Spectral Rolloff (85%)
# Spectral roll-off is defined as the frequency index R below which a certain
# fraction (85%) of the spectral energy resides.

# %%
print(df.loc[:, ['spectral_rolloff']].describe())
df_melt = pd.melt(df, id_vars=['rac'], value_vars=['spectral_rolloff'])
sns.displot(
    data=df_melt,
    x='value',
    hue='rac',
    col='variable',
    kde=True,
    rug=True,
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_spectral_rolloff_rac.png',
            bbox_inches='tight',
            transparent=True)

df_tmp = df.reset_index()
df_tmp['hour'] = df_tmp['datetime'].dt.hour
df_tmp['day_quarter'] = df_tmp['hour'] // 6
df_melt = pd.melt(df_tmp,
                  id_vars=['day_quarter'],
                  value_vars=['spectral_rolloff'])
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.displot(
    data=df_melt,
    x='value',
    hue='day_quarter',
    palette='vlag',
    kde=True,
    rug=True,
    col='variable',
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_spectral_rolloff_day.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# A slight displacement can be seem between the two groups of files.
# This could reinforce the hypothesis drawn from the spectral centroid
# analysis, as the 85% threshold is at a higher frequency after RAC provision.

# %% [markdown]
# ### Spectral Flux
# The spectral flux, also known as the delta spectrum magnitude, is a measure
# of local spectral change. It is defined as the squared norm of the
# frame-to-frame spectral difference.

# %%
print(df.loc[:, ['spectral_flux']].describe())
df_melt = pd.melt(df, id_vars=['rac'], value_vars=['spectral_flux'])
sns.displot(
    data=df_melt,
    x='value',
    hue='rac',
    col='variable',
    kde=True,
    rug=True,
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_spectral_flux_rac.png',
            bbox_inches='tight',
            transparent=True)

df_tmp = df.reset_index()
df_tmp['hour'] = df_tmp['datetime'].dt.hour
df_tmp['day_quarter'] = df_tmp['hour'] // 6
df_melt = pd.melt(df_tmp,
                  id_vars=['day_quarter'],
                  value_vars=['spectral_flux'])
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.displot(
    data=df_melt,
    x='value',
    hue='day_quarter',
    palette='vlag',
    kde=True,
    rug=True,
    col='variable',
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_spectral_flux_day.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# There has been no noticeable major changes in the distribution.
# This suggests that the number of "peaks"/onsets is similar after the rac
# provision.
# This feature is a discard candidate.

# %% [markdown]
# ### Spectral Bandwidth
# The bandwidth of the spectrum is described by spectral spread
#
# The spectral bandwidth 1 at frame t is computed by:
#
# (sum_k S[k, t] * (freq[k, t] - centroid[t])**p)**(1/p)

# %%
print(df.loc[:, [
    'spectral_bandwidth_2', 'spectral_bandwidth_3', 'spectral_bandwidth_4'
]].describe())
df_melt = pd.melt(df,
                  id_vars=['rac'],
                  value_vars=[
                      'spectral_bandwidth_2', 'spectral_bandwidth_3',
                      'spectral_bandwidth_4'
                  ])
sns.displot(
    data=df_melt,
    x='value',
    hue='rac',
    col='variable',
    kde=True,
    rug=True,
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(
    f'./output_{FILE_PREFIX}/{FILE_PREFIX}_spectral_bandwidths_rac.png',
    bbox_inches='tight',
    transparent=True)

# %% [markdown]
# No relevant differences in the distributions were observed.
#
# These features are discard candidates.
#
# Further analysis required with correlation heatmap.

# %% [markdown]
# ### Correlation matrix of subset of features

# %%
plt.figure(figsize=(12, 12))
sns.heatmap(df.loc[:, [
    'zero_crossing_rate', 'spectrogram', 'mel_spectrogram', 'harmonics',
    'perceptual_shock_wave', 'spectral_centroids', 'tempo_bpm',
    'spectral_rolloff', 'spectral_flux', 'spectral_bandwidth_2',
    'spectral_bandwidth_3', 'spectral_bandwidth_4', 'rac'
]].corr(),
            vmin=-1.0,
            vmax=1.0,
            cmap="coolwarm",
            center=0.0,
            robust=True,
            annot=True,
            fmt='.1f')
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_corr.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# The spectral centroids feature had the highest correlation with 'rac',
# between the analyzed features.

# %% [markdown]
# ### MFCC
# Mel-frequency cepstral coefficients (MFCCs) describe the rough shape of
# the signal spectrum and are widely used in speech recognition [536].
# Similarly, they are often encountered in percussion transcription algorithms.
# Usually the coefficients are calculated in short (about 20 ms) partially
# overlapping frames over the analysed segment, and 5 to 15 coefficients are
# retained in each frame. Instead of using these directly as features,
# typically the mean and variance of each coefficient over the segment are
# used. In addition, the first- and second-order temporal differences of the
# coefficients, and the means and variances of these, are commonly used as
# features.

# %%

# Visualizing single-file mfccs
mfccs = lr.feature.mfcc(y=x, sr=sr, n_mfcc=20)
fig, ax = plt.subplots(figsize=(16, 9))
img = lr.display.specshow(mfccs, x_axis='time', ax=ax)
fig.colorbar(img, ax=ax)
ax.set(title='MFCC')
plt.draw()
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_mfcc.png')

mfcc_cols = [
    'mfcc0',
    'mfcc1',
    'mfcc2',
    'mfcc3',
    'mfcc4',
    'mfcc5',
    'mfcc6',
    'mfcc7',
    'mfcc8',
    'mfcc9',
    'mfcc10',
    'mfcc11',
    'mfcc12',
    'mfcc13',
    'mfcc14',
    'mfcc15',
    'mfcc16',
    'mfcc17',
    'mfcc18',
    'mfcc19',
]
print(df.loc[:, mfcc_cols].describe())
df_melt = pd.melt(df, id_vars=['rac'], value_vars=mfcc_cols)
sns.displot(
    data=df_melt,
    x='value',
    hue='rac',
    row='variable',
    kde=True,
    rug=True,
    height=9,
    aspect=1,
)
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_mfccs_rac.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# From related papers, the MFCCs are some of the most representative features.
# The main question related the MFCCs is the number of coefficients that will
# be used. For most use cases, 12 or 13 coefficients seem to be enough.
#
# TODO Study and write detailed description of MFCC in monograph.
#
# mfcc14 through mfcc19 are discard candidates.

# %%
plt.figure(figsize=(12, 12))
sns.heatmap(df.loc[:, mfcc_cols + ['rac']].corr(),
            vmin=-1.0,
            vmax=1.0,
            cmap="coolwarm",
            center=0.0,
            robust=True,
            annot=True,
            fmt='.1f')
plt.draw()
plt.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_mfccs_corr.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# Initially the highest absolute correlation with 'rac' is onm mfcc9, mfcc15,
# and mfcc16.
#
# The coefficients by themselfs seem to be relatively independent of each
# other, taking into account only the correlation matrix.
# This may be an indicator of high representativity of the dataset.

# %% [markdown]
# ### Visualizing discard candidates

# %%
discard_features = [
    'zero_crossings',
    'spectrogram',
    'mel_spectrogram',
    'harmonics',
    'perceptual_shock_wave',
    'tempo_bpm',
    'spectral_flux',
    'spectral_bandwidth_2',
    'spectral_bandwidth_3',
    'spectral_bandwidth_4',
] + chroma_cols + [
    'mfcc14',
    'mfcc15',
    'mfcc16',
    'mfcc17',
    'mfcc18',
    'mfcc19',
    'mfcc14_delta',
    'mfcc15_delta',
    'mfcc16_delta',
    'mfcc17_delta',
    'mfcc18_delta',
    'mfcc19_delta',
    'mfcc14_accelerate',
    'mfcc15_accelerate',
    'mfcc16_accelerate',
    'mfcc17_accelerate',
    'mfcc18_accelerate',
    'mfcc19_accelerate',
]

print('# of discarded features:', len(discard_features))
print(
    '# of remaining features:',
    df.loc[:, 'zero_crossing_rate':].shape[1] -  # type: ignore
    len(discard_features))

# %% [markdown]
# # Final comments
# After analyzing the extracted features for our dataset, the following
# features are discard candidates:
#
# - spectrogram';
# - mel-spectrogram';
# - harmonics';
# - perceptual shock wave;
# - chroma features;
# - tempo (bpm);
# - spectral flux;
# - spectral bandwidth;
# - mfcc14, ..., mfcc19.
#
# For now, the remaining features shall be used for further analysis and
# data visualization methods, such as SOMs.
#
# If necessary, more features, such as deltas and accelerates, can be
# discarded, for performance reasons. The relevance of these features should
# be further analyzed.

# %%


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

# ---
# title: ERP Core N400 in MNE-Python: part III
# date: 2021-04-23
# image:
#   preview_only: true
# tags:
# - Python
# - EEG
# - Preprocessing
# categories:
# - Python
# - EEG
# - English
# summary: "Adapting ICA artifact rejection."
# copyright_license:
#   enable: true
#
# ---

# In the [last post](https://msilvestrin.me/post/n400_2/) on adapting ERP CORE's N400 pipeline to MNE-Python, we did the ICA-related steps to correct blinking and horizontal eye-movement artfacts from the signal. Here we do some final steps in cleaning the data: interpolation of bad channels, epoching and rejection of epochs with possible artifacts. These steps are decipted in the  _Elist_Bin_Epoch_ script of ERP CORE resources (script 5).
#
# Contrary to previous posts, I won't make the whole process in a single participant and show the loop for all participants at the end. That is because the result for some af the steps are saved as intermediate stages of the preprocessing pipeline. I should also tell you up-front that I wasn't able to find implemented equivalents for some of the procedures ERP CORE uses to define artifacts. That forced me to implement a python version of what they do. Therefore, this post is a little longer the the other ones, so buckle up! On the bright side, at the end we will have the data ready to start looking for those N400 ERPs.
#

# ## Interpolation
# Our first step here is to interpolate bad channels. Bad channels for each participant are listed in the file "Interpolate_Channels_N400.xlsx" provided by ERP CORE. Below, I load the information in that file, make a list of bad channels for each participant and add that list to their `raw.info['bads']`.
#
# Load the information:

# In[3]:

import os
os.chdir('D:\\EEGdata\\Erp CORE\\N400')

import mne
# import some methods directly so we can call them by name
from mne.io import read_raw
from mne.epochs import read_epochs

import numpy as np
import pandas as pd

get_ipython().run_line_magic('matplotlib', 'qt')
get_ipython().run_line_magic('gui', 'qt')

# In[8]:

# load as DataFrame
interpolate_path = "D:\\EEGdata\\Erp CORE\\N400\\EEG_ERP_Processing\\Interpolate_Channels_N400.xlsx"
chans_interpolate = pd.read_excel(interpolate_path)
# Turn strings to a list of strings for each participant
chans_interpolate = chans_interpolate.loc[:,"Name of Interpolated Channel"].str.split(" ")

# Now I'll deal with some nuisences in the file: there is an asterisk at the beginning of the channel list for each participant, and Fp channels are in upper-case. We remove the asterisk and correct the captalization.

# In[11]:

# make function to remove asterisk and correct
def remove_star(item):
    # if item is a list
    if isinstance(item, list):
        # if list is longer than 1
        if len(item) > 1:
            # remove asterisk
            clean_names = [item[0].lstrip("*")]
            # add channel to a list clean_names
            clean_names.extend(item[1:])
            # look for Fp channels and correct captalization
            for n, elec in enumerate(clean_names):
                if elec =='FP1' or elec =='FP2':
                    clean_names[n] = elec.capitalize()
        else:
            clean_names = [item[0].lstrip("*")]
            if clean_names[0] =='FP1' or clean_names[0] =='FP2':
                clean_names[0] = clean_names[0].capitalize()
    else:
        clean_names = np.nan
    return clean_names

# apply function
chans_interpolate = chans_interpolate.apply(remove_star)
chans_interpolate.loc[chans_interpolate.notnull()].head()

# We have already indentified the channels to be interpolated, however if we try to do the interpolation now MNE will throw us an error because we have more EEG channels than mapped on the 3D digitation. The channels HEOG_left, HEOG_right and VEOG_lower are currently identified as EEG channels, but they are actually extra channels we used to make the proper HEOG and VEOG channels. We will drop them, since they already served their purpose.
#
# Let's do it for all participants and save an intermediate ("interpol") raw file:

# In[ ]:

for subj in range(1,41):
    raw_correct_path = "D:\\EEGdata\\Erp CORE\\N400\\mne\\raw_corrected\\N400_ERP_CORE_{}_correct-raw.fif".format(subj)
    raw_correct = read_raw(raw_correct_path, preload = True)
    #only do interpolation if there is at least one channel to interpolate
    if  not chans_interpolate.isna()[subj-1]:
        raw_correct.info['bads'].extend(chans_interpolate[subj-1])
        # drop extra channels
        raw_correct.drop_channels(['HEOG_left', 'HEOG_right', 'VEOG_lower'])
        raw_correct.interpolate_bads()
    else:
        raw_correct.drop_channels(['HEOG_left', 'HEOG_right', 'VEOG_lower'])

    interpol_path = "D:\\EEGdata\\Erp CORE\\N400\\mne\\raw_interpol\\N400_ERP_CORE_{}_interpol-raw.fif".format(subj)
    raw_correct.save(interpol_path, overwrite=True)

# ## Event list for epochs
# Now let's make our event list and event_ids for epoching.
#
# As we did at the back at beginning of preprocessing [(in this post)](https://msilvestrin.me/post/n400_1/), we will get the events from the `annotations` in the data. Until now, we have two types of annotation descriptions in the data: "BAD_seg" and numbers corresponding to triggers used in the experiment. The ERP CORE file "N400_Event_Code_Scheme.xlsx" has the meaning for each code. For instance, the code 121 refers to a Priming stimuli which is unrelated to the target stimulus and comes from the second list of words; while the code 201 means that the participant gave a correct answer.
#
# We are going to change these codes to a descriptive account of each event. A nice feature of MNE is that we can make _hierarchical ids_. We will follow the hierarchy shown in the N400_Event_Code_Scheme file: Word Type > Relatedness > Word List.
#
# First, we will create a dictionary with the correspondence between codes and events to properly identify them later. We will be able to use these ids to select epochs afterwards.

# In[ ]:

# load data
subj = 1
interpol_path = "D:\\EEGdata\\Erp CORE\\N400\\mne\\raw_interpol\\N400_ERP_CORE_{}_interpol-raw.fif".format(subj)
raw_interpol = read_raw(interpol_path)
# make dict with correspondences between codes and hierarchical ids
description_dict = {'111' : 'Prime/Related/L1',
             '112' : 'Prime/Related/L2',
             '121' : 'Prime/Unrelated/L1',
             '122' : 'Prime/Unrelated/L2',
             '211' : 'Target/Related/L1',
             '212' : 'Target/Related/L2',
             '221' : 'Target/Unrelated/L1',
             '222' : 'Target/Unrelated/L2',
             '201' : 'Hit',
             '202' : 'Miss',
             'BAD_seg': 'BAD_seg'
                   }

# We will change the description list with some [Pandas](https://pandas.pydata.org/) functionality. It is beyond this tutorial to teach Pandas' skills, but it is nice for you to know that, although most of the backbone of MNE is in [Numpy](https://numpy.org/), it also allows exportation to Pandas' data structures.
#
# In summary, what we do in the one-liner below is make a copy of the annotations descriptions while transforming it to a Pandas series, then we make the substituitions according to our `description dictionary` and turn it back to a numpy array.

# In[ ]:

raw_interpol.annotations.description = pd.Series(raw_interpol.annotations.description).map(description_dict).to_numpy()
# substitute annotations description
set(raw_interpol.annotations.description)

# Now we are almost ready to make the `events` numpy array we will need to epoch the data. We will use the `events_from_annotations` function. One last step before that: create a dictionary with correspondences between hierarchical ids and trigger numbers.
#
# You may be wondering why we would do that when we just made the reverse dictionary on our last step. The thing is that MNE's event list is a _numeric array_ composed of columns _instant_, _blank column_ and _event code_, and the default behavior of `events_from_annotations` is to assign the event codes as \[0,1,2,3...\]. However, I want to keep the original codes from ERP CORE, so I need to use the `events_id` input with a dict showing the correspondences. I should say that for the rest of our MNE operations the numeric codes themselves are actually rather irrelevant, what really matters are the correspondences in `event_ids`. I am just being picky to keep as close as possible to the original ERP CORE material.

# In[ ]:

event_ids = {'Prime/Related/L1': 111,
             'Prime/Related/L2': 112,
             'Prime/Unrelated/L1' : 121,
             'Prime/Unrelated/L2' : 122,
             'Target/Related/L1' : 211,
             'Target/Related/L2' : 212,
             'Target/Unrelated/L1' : 221,
             'Target/Unrelated/L2' : 222,
             'Hit' : 201,
             'Miss' : 202}
events, event_ids = mne.events_from_annotations(raw_interpol, event_ids)

# So, our events look like this (_instant_, _blank column_, _event code_):

# In[7]:

events

# ## Epoching and artfact rejection
# The ERP CORE pipeline is __very throrough__ in its epoch dropping due to ocular artifact detection. There are different thresholds for epoch dropping for: (1) EEG channels, (2) the VEOG channel and (3 and 4) for both variations of the HEOG channel (with and without ICA applied). In MNE, we can set epoch rejection thresholds for each channel type by setting the `reject` input when we create the epochs.
#
# However, as far as I understand from the documentation, MNE uses a simple peak-to-peak algorithm to do its automatic rejection. ERP CORE pipeline uses such an algorithm only for rejection on the scalp EEG channels. For VEOG and uncorrected HEOG channels it uses a windowed peak-to-peak algorithm and for ICA corrected HEOG it uses a step algorithm (both as implemented in the ERP LAB toolbox).
#
# Since ERP CORE gives us the parameters, and the algorithms themselves are quite straightforward, I have implemented them below. I make use of the nifty feature of MNE that allows us to loop through epochs. On each epoch I run the algorithms on the corresponding channels and, if some amplitude is above the threshold, it is marked as bad with a `MW_Reject` annotation. The value compared to the treshold on each algorithm are the following:
#
# * Step algorithm: $\bmod (mean(window1)-mean(window2))$
# * Peak-to-peak algorithm: $\bmod (max(window1)-min(window2))$
#
# Below, we (1) load the the information we need for the artfact rejection procedure, (2) create the dictionaries we need for marking the events in the experiment (as shown above); then, in the loop, we (3) add the annotations for epoching, (4) do the epoching with `mne.Epochs()`, (5) run the artfact rejection algorithm and (6) save the epochs.

# In[ ]:

#get eog channels numbers
subj = 1
interpol_path = "D:\\EEGdata\\Erp CORE\\N400\\mne\\raw_interpol\\N400_ERP_CORE_{}_interpol-raw.fif".format(subj)
raw_interpol = read_raw(interpol_path)
eog_nums = [np.where(np.array(raw_interpol.ch_names) == 'VEOG')[0].item(0),
            np.where(np.array(raw_interpol.ch_names) == 'HEOG')[0].item(0),
            np.where(np.array(raw_interpol.ch_names) == 'HEOG_ICA')[0].item(0)]

#get veog moving windows parameters from files
veog_path = "D:\\EEGdata\\Erp CORE\\N400\\EEG_ERP_Processing\\AR_Parameters_for_MW_CRAP_N400.xlsx"
veog_mw_params = pd.read_excel(veog_path)
#convert time columns to seconds
veog_mw_params.iloc[:,3:] = veog_mw_params.iloc[:,3:].apply(lambda x: x*1e-3)
#convert threshold to volts
veog_mw_params.Threshold = veog_mw_params.Threshold*1e-4

# the same for non ICA corrected heog moving windows
heog_path = "D:\\EEGdata\\Erp CORE\\N400\\EEG_ERP_Processing\\AR_Parameters_for_MW_Blinks_N400.xlsx"
heog_mw_params = pd.read_excel(heog_path)
heog_mw_params.iloc[:,3:] = heog_mw_params.iloc[:,3:].apply(lambda x: x*1e-3)
heog_mw_params.Threshold = heog_mw_params.Threshold*1e-4

#... and for ICA corrected HEOG moving windows
heog_ica_path = "D:\\EEGdata\\Erp CORE\\N400\\EEG_ERP_Processing\\AR_Parameters_for_SL_HEOG_N400.xlsx"
heog_ica_mw_params = pd.read_excel(heog_ica_path)
heog_ica_mw_params.iloc[:,3:] = heog_ica_mw_params.iloc[:,3:].apply(lambda x: x*1e-3)
heog_ica_mw_params.Threshold = heog_ica_mw_params.Threshold*1e-4

#put moving windows dfs in a list
mw_params = [veog_mw_params, heog_mw_params, heog_ica_mw_params]

# make dict with correspondences between codes and hierarchical ids
description_dict = {'111' : 'Prime/Related/L1',
                    '112' : 'Prime/Related/L2',
                    '121' : 'Prime/Unrelated/L1',
                    '122' : 'Prime/Unrelated/L2',
                    '211' : 'Target/Related/L1',
                    '212' : 'Target/Related/L2',
                    '221' : 'Target/Unrelated/L1',
                    '222' : 'Target/Unrelated/L2',
                    '201' : 'Hit',
                    '202' : 'Miss',
                    'BAD_seg': 'BAD_seg'
               }

 #create correspondence between description and trigger number
event_ids = {'Prime/Related/L1': 111,
             'Prime/Related/L2': 112,
             'Prime/Unrelated/L1' : 121,
             'Prime/Unrelated/L2' : 122,
             'Target/Related/L1' : 211,
             'Target/Related/L2' : 212,
             'Target/Unrelated/L1' : 221,
             'Target/Unrelated/L2' : 222,
             'Hit' : 201,
             'Miss' : 202}

# loop
for subj in range(1,41):
    interpol_path = "D:\\EEGdata\\Erp CORE\\N400\\mne\\raw_interpol\\N400_ERP_CORE_{}_interpol-raw.fif".format(subj)
    raw_interpol = read_raw(interpol_path)

    #add event descriptions to annotations
    raw_interpol.annotations.description = pd.Series(raw_interpol.annotations.description).map(description_dict).to_numpy()
    set(raw_interpol.annotations.description)

    #create events and event ids
    events, event_ids = mne.events_from_annotations(raw_interpol, event_ids)

    #change heog channel types
    raw_interpol.set_channel_types({'HEOG_ICA':'emg', 'HEOG':'misc'})
    epochs = mne.Epochs(raw_interpol,
                        events,
                        event_ids,
                        baseline = (-.2,0),
                        tmax = .8,
                        reject = dict(eeg  = 200e-6)
                       )

    epos_to_drop = []
    # run through every epoch
    for i in range(len(epochs.events)):
        # for each channel of interest
        for eog_chan_num, mw_df in zip(eog_nums, mw_params): #
            # get moving windows parameters
                w_size = mw_df['Window Size'].iloc[subj-1]
                w_start = mw_df['Time Window Minimum'].iloc[subj-1].round(2)
                w_start_last = np.round(mw_df['Time Window Maximum'].iloc[subj-1] - mw_df['Window Size'].iloc[subj-1],2)
                w_stop = mw_df['Time Window Maximum'].iloc[subj-1].round(2)
                w_stop_first = np.round(w_start + w_size,2)
               # create an array with starting and stopping times of every window
                mw_starts = np.arange(w_start,
                                      w_start_last,
                                      w_size)
                mw_stops = np.arange(w_stop_first,
                                     w_stop,
                                     w_size)
        # only run if channel is not empty, e.g. already marked as bad
            if epochs[i].get_data()[:,eog_chan_num,:].shape[0] != 0:
                #get channel data
                chan_data = epochs[i].get_data()[:,eog_chan_num,:].flatten()
                #make list with data for each window
                mws = [ chan_data[epochs.time_as_index(start).item(0) : epochs.time_as_index(stop).item(0)] for start, stop in zip(mw_starts, mw_stops)]

               # if channel is veog or uncorrected heog do peak-to-peak artfact detection
                if eog_chan_num < 32:
                    mws = pd.DataFrame(mws).T
                    peak2peak = mws.apply(lambda x: np.abs(x.max() - x.min()))
                    if any(peak2peak > mw_df['Threshold'].iloc[subj-1]):
                        epos_to_drop.append(i)
                # if channel is heog_ica do step artfact detection
                else:
                    steps = [True for i in range(len(mws)-1) if np.abs(mws[i].mean()-mws[i+1].mean()) > mw_df['Threshold'].iloc[subj-1]]
                    if steps:
                        epos_to_drop.append(i)
    # if any epoch was set as bad add it to drop list with label MW Reject
    if epos_to_drop:
        epochs.drop(epos_to_drop,'MW_Reject')
    #Save epochs
    epochs_path = "D:\\EEGdata\\Erp CORE\\N400\\mne\\epochs\\N400_ERP_CORE_{}-epo.fif".format(subj)
    print('Subject {}'.format(subj)) # line to keep up with analysis as it runs
    epochs.save(epochs_path, overwrite=True)

# Lastly, let's check the percentage of dropped epochs for each participant.

# In[ ]:

percent_dropped = pd.DataFrame(data = {'Subj': np.repeat(np.nan, 40),
                                       'Dropped': np.repeat(np.nan, 40)})
for subj in range(1,41):
    epochs_path = "D:\\EEGdata\\Erp CORE\\N400\\mne\\epochs\\N400_ERP_CORE_{}-epo.fif".format(subj)
    epochs = read_epochs(epochs_path)

    percent_dropped['Subj'].iloc[subj-1] = subj
    percent_dropped['Dropped'].iloc[subj-1] = np.round(epochs.drop_log_stats(),2)

# In[6]:

percent_dropped['Dropped'].describe().round(2)

# Dropped epochs per participant range from roughly 4% to 85%. Here, we will set consider participants with more than 30% of epochs dropped as participants with too many artfacts to be used. The ones that do not pass this filter are below:

# In[7]:

bad_subjs = percent_dropped.loc[percent_dropped.Dropped >30].Subj
bad_subjs.to_csv('bad_subjs.csv')
bad_subjs.count()

# We have 5 unsuable participants. This is different from ERP CORE's original preprocessing results, where they lose a single participant. However, I wasn't able to exactly reproduce the ICA step in the pipeline, so that is probably where we are falling a little behind. Anyway, 35 participants are more than enough for us to grasp the N400 effects in the following analyses. So, enough preprocesing, let's get to those ERPs in the next post!
#
# Let's finish contemplating some hard earned epochs:

# In[12]:

epochs.plot(n_epochs = 4)


#   ***********************************************************************
#
#   FILE         plots.py
#
#   AUTHOR       Dr. Vishal Sharma
#
#   VERSION      1.0.0-alpha5
#
#   WEBSITE      https://github.com/vxsharma-14/project-NAnPack
#
#   NAnPack Learner's Edition is distributed under the MIT License.
#
#   Copyright (c) 2022 Vishal Sharma
#
#   Permission is hereby granted, free of charge, to any person
#   obtaining a copy of this software and associated documentation
#   files (the "Software"), to deal in the Software without restriction,
#   including without limitation the rights to use, copy, modify, merge,
#   publish, distribute, sublicense, and/or sell copies of the Software,
#   and to permit persons to whom the Software is furnished to do so,
#   subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be
#   included in all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
#   EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
#   OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
#   NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
#   BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
#   ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
#   CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#   SOFTWARE.
#
#   You should have received a copy of the MIT License along with
#   NAnPack Learner's Edition.
#
#   ***********************************************************************
import matplotlib.pyplot as plt
import numpy as np

def Plot1D(dataFiles, uAxis, legend, markers, useFileCol,
          title, xlbl, ylbl):
    """Plot results along 1D axis."""
    print("Preparing data to plot results...")
    countFiles = len(dataFiles)

    # Assign font family and create axis for plotting
    plt.rc('font', family='sans-serif', size=10)
    fig, ax = plt.subplots(dpi=150)

    data = {}  # Initialize a dictionary

    # Use dictionary to create multiple line plots from saved files
    for i in range(countFiles):
        data[f"x{i}"] = np.loadtxt(dataFiles[i], unpack=True, skiprows=3,
                                   usecols=0)
        data[f"u{i}"] = np.loadtxt(dataFiles[i], unpack=True, skiprows=3,
                                   usecols=useFileCol)

    # Start plotting
    print("Plotting 1D results")
    for i in range(countFiles):
        if markers is not None:
            mrkr = markers[i]
        else:
            mrkr = None
        if uAxis == "Y":
            ax.plot(data[f"x{i}"], data[f"u{i}"], linewidth=0.5,
                    label=legend[i], marker=mrkr, markersize=3,
                    markevery=3)
        elif uAxis == "X":
            ax.plot(data[f"u{i}"], data[f"x{i}"], linewidth=0.5,
                    label=legend[i], marker=mrkr, markersize=3,
                    markevery=3)
    # Format and customize plot
    plt.grid(which='both', axis='both', color='lightgrey', linestyle=':',
             linewidth=0.5)
    plt.xlabel(xlbl)
    plt.ylabel(ylbl)
    plt.title(title, fontsize=10)
    ax.legend(fontsize=9)
    # plt.tight_layout()
    plt.show()

def Plot2D(dataFile, title, xlbl, ylbl, cbarlbl, ptype, cMap, shade,
          clevel, alpha):
    """Plot results within 2D domain."""
    print("Preparing data to plot results...")

    with open(dataFile) as f:
        grid = f.readline().split(",")
    iM = int((grid[0].split("= "))[1])
    jM = int((grid[1].split("= "))[1])

    # Assign font family and create axis for plotting
    plt.rc('font', family='sans-serif', size=10)
    fig, ax = plt.subplots(dpi=150)

    # Read data from saved file
    x, y, u = np.loadtxt(dataFile, unpack=True, skiprows=3)

    # Reshape data based on iM and jM
    X = np.reshape(x, (iM, jM))
    Y = np.reshape(y, (iM, jM))
    U = np.reshape(u, (iM, jM))

    if ptype == "pcolormesh":
        # pcolormesh plot
        plt.pcolormesh(X, Y, U, cmap=cMap, shading=shade)
    elif ptype == "contourf":
        # filled contour plots
        plt.contourf(X, Y, U, cmap=cMap, levels=clevel)
    elif ptype == "contour":
        # contour plots
        plt.contour(X, Y, U, cmap=cMap, levels=clevel)
    elif ptype == "imshow":
        # image plots
        plt.imshow(U, origin="lower", cmap=cMap, interpolation="bilinear")
    elif ptype == "pcolor":
        #
        plt.pcolor(X, Y, U, alpha=alpha, cmap=cMap)

    # Format and customize plot
    plt.xlabel(xlbl, size=10)
    plt.ylabel(ylbl)
    plt.title(title, fontsize=10)
    if cbarlbl is not None:
        cbar = plt.colorbar(format='%.0e')
        cbar.set_label(cbarlbl, rotation=270, size=10, labelpad=15)

    plt.tight_layout()
    plt.gca().set_aspect('equal', adjustable='box')
    plt.show()

def MultiPlot2D(nplots, nrow, ncol, dataFiles, title,
                xlbl, ylbl, cbarlbl, ptype, cMap, shade, clevel, alpha):
    """Plot multiple results in a single image."""
    print("Preparing data to plot results...")

    # Assign font family and create axis for plotting
    plt.rc('font', family='sans-serif', size=10)
    fig, ax = plt.subplots(dpi=150)

    data = {}  # Initialize a dictionary
    DATA = {}
    iM = []
    jM = []
    for i in range(nplots):
        with open(dataFiles[i]) as f:
            grid = f.readline().split(",")
        iM.append(int((grid[0].split("= "))[1]))
        jM.append(int((grid[1].split("= "))[1]))

    # Use dictionary to create multiple plots
    for i in range(nplots):
        data[f"x{i}"], data[f"y{i}"], data[f"u{i}"] = (
            np.loadtxt(dataFiles[i], unpack=True, skiprows=3))

        # Reshape data based on iM and jM
        DATA[f"X{i}"] = np.reshape(data[f"x{i}"], (iM[i], jM[i]))
        DATA[f"Y{i}"] = np.reshape(data[f"y{i}"], (iM[i], jM[i]))
        DATA[f"U{i}"] = np.reshape(data[f"u{i}"], (iM[i], jM[i]))

    for i in range(nplots):
        plt.subplot(nrow, ncol, i+1)

        if ptype == "pcolormesh":
            # pcolormesh plot
            img = plt.pcolormesh(DATA[f"X{i}"], DATA[f"Y{i}"],
                                 DATA[f"U{i}"],
                                 cmap=cMap, shading=shade)
        elif ptype == "contourf":
            # filled contour plots
            img = plt.contourf(DATA[f"X{i}"], DATA[f"Y{i}"], DATA[f"U{i}"],
                               cmap=cMap, levels=clevel)
        elif ptype == "contour":
            # contour plots
            img = plt.contour(DATA[f"X{i}"], DATA[f"Y{i}"], DATA[f"U{i}"],
                              cmap=cMap, levels=clevel)
        elif ptype == "imshow":
            # image plots
            img = plt.imshow(DATA[f"U{i}"], origin="lower", cmap=cMap,
                             interpolation="bilinear")
        elif ptype == "pcolor":
            # pcolor plots
            img = plt.pcolor(DATA[f"X{i}"], DATA[f"Y{i}"], DATA[f"U{i}"],
                             alpha=alpha, cmap=cMap)

        cbar_axismap(img, cbarlbl)

        # Format and customize plot

        plt.xlabel(str(xlbl), size=10)
        plt.ylabel(str(ylbl), size=10)
        plt.title(str(title[i]), size=10)
        plt.tight_layout()
        plt.gca().set_aspect("equal", adjustable="box")
    plt.show()

def cbar_axismap(mappable, ctitle):
    """Control the size of colorbar (aspect ratio) in multi-plot axis."""
    from mpl_toolkits.axes_grid1 import make_axes_locatable

    last_axes = plt.gca()
    ax = mappable.axes
    fig = ax.figure
    divider = make_axes_locatable(ax)
    cax = divider.append_axes('right', size='5%', pad=0.05)
    cbar = fig.colorbar(mappable, cax=cax)
    if ctitle is not None:
        cbar.set_label(ctitle, rotation=270, size=12, labelpad=15)
    plt.sca(last_axes)

    return cbar


""" Tests the NumPy introduction exercises

:Author: Jonathan Karr <jonrkarr@gmail.com>
:Date: 2017-07-11
:Copyright: 2017, Karr Lab
:License: MIT
"""

from intro_to_wc_modeling.concepts_skills.software_engineering import numpy_exercises
import unittest

class TestNumPyExercises(unittest.TestCase):

    def test(self):
        numpy_exercises.main()


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

# In[2]:

import os
project_name = "reco-tut-mlh"; branch = "main"; account = "sparsh-ai"
project_path = os.path.join('/content', project_name)

# In[3]:

if not os.path.exists(project_path):
    get_ipython().system(u'cp /content/drive/MyDrive/mykeys.py /content')
    import mykeys
    get_ipython().system(u'rm /content/mykeys.py')
    path = "/content/" + project_name;
    get_ipython().system(u'mkdir "{path}"')
    get_ipython().magic(u'cd "{path}"')
    import sys; sys.path.append(path)
    get_ipython().system(u'git config --global user.email "recotut@recohut.com"')
    get_ipython().system(u'git config --global user.name  "reco-tut"')
    get_ipython().system(u'git init')
    get_ipython().system(u'git remote add origin https://"{mykeys.git_token}":x-oauth-basic@github.com/"{account}"/"{project_name}".git')
    get_ipython().system(u'git pull origin "{branch}"')
    get_ipython().system(u'git checkout main')
else:
    get_ipython().magic(u'cd "{project_path}"')

# ---

# # Exploratory Data Analysis
#
# In this notebook we explore the MovieLens 100k dataset.
#
#
# *   Find missing/null values
# *   Examine the distribution of ratings
# *   Examine movies and users with most reviews
# *   Examine correlation between time and reviews
#
#

# # Imports

# In[13]:

import math
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import re
import requests
import seaborn as sns
from scipy.stats.stats import pearsonr
from tqdm import tqdm

# # Prepare data

# In[16]:

# Load reviews.
fp = os.path.join('./data/bronze', 'u.data')
raw_data = pd.read_csv(fp, sep='\t', names=['userId', 'movieId', 'rating', 'timestamp'])
raw_data.head()

# In[17]:

# Load movie titles.
fp = os.path.join('./data/bronze', 'u.item')
movie_titles = pd.read_csv(fp, sep='|', names=['movieId', 'title'], usecols = range(2), encoding='iso-8859-1')
movie_titles.head()

# In[18]:

# Merge dataframes.
raw_data = raw_data.merge(movie_titles, how='left', on='movieId')
raw_data.head()

# In[19]:

# Change timestamp to datatime.
raw_data.timestamp = pd.to_datetime(raw_data.timestamp, unit='s')
raw_data.head()

# # Exploration

# ## Unique and null values

# We first see that there are 100k observations in our dataset. There are 943 unique users and 1682 unique movies, and the rating system is out of 5. We then check to see if there are any missing data points in the set, which we find there are none.

# In[20]:

print(f'Shape: {raw_data.shape}')
raw_data.sample(5, random_state=123)

# In[21]:

raw_data.nunique()

# In[22]:

raw_data.info()

# In[23]:

print(f'Shape: {movie_titles.shape}')
movie_titles.sample(5, random_state=123)

# ## Summary Stats

# ### Ratings
#
# Next, we look at the summary statistics of each feature in the dataset. We notice that the mean rating of the movies is 3.5 and that the minimum and maximum rating is 1 and 5 respectivle, and that the ratings are discrete (no in-between values). The most common rating is 4, with the second most common being 3. There are very few reviews with a 1 rating (about 6000/100,000). In fact looking at our boxplots, reviews where the movie is rated 1 might even be considered an outlier.

# In[24]:

raw_data.describe()

# In[25]:

plt.figure(figsize=(7,5))
sns.histplot(raw_data.rating)
plt.show()

# In[26]:

plt.figure(figsize=(10,6))
sns.boxplot(x = raw_data.rating)
plt.show()

# ### Time
#
# Actual reviews were made starting from September 20, 1997 to April 22, 1998, about 7 months of data.
#
# Actual movies reviewed were released from 1922 to 1998, with 4 years missing in that timespan. There are also a couple of movies with no year given. We assigned these movies to year 0.

# In[28]:

raw_data.timestamp.describe(datetime_is_numeric=True)

# In[29]:

def get_year(title):
    year=re.search(r'\(\d{4}\)', title)
    if year:
        year=year.group()
        return int(year[1:5])
    else:
        return 0

# In[30]:

raw_data['year'] = raw_data.title.apply(get_year)
raw_data.year.sort_values().unique()

# In[31]:

raw_data[['year']].nunique()

# In[32]:

sns.histplot(raw_data['year'][raw_data['year'] != 0])
plt.show()

# ## Users with most reviews
#
# The most movies single user has reviewed is 737 reviews. The minimum number of reviews a user has reviewed in the dataset is 20. This is good since when creating recommendation systems, you want users with lots or reviews, allowing for us to test our recomendations. We also notice that most users reviewed less than 65 movies.

# In[33]:

users_count = raw_data.groupby('userId')['rating'].count().sort_values(ascending=False).reset_index()
users_count

# In[34]:

# Plot how many movies a user reviewed
plt.figure(figsize=(10, 6))
fig = sns.histplot(users_count['rating'])
plt.show()

# In[35]:

users_count['rating'].median()

# ## Movies with most reviews
#
# As we can expect, popular movies such as 'Star Wars' and 'Toy Story' have the most reviews. The highest number of reviews is 583 while the lowest number of reviews is 1.

# In[36]:

movies_count = raw_data.groupby('title')['rating'].count().sort_values(ascending=False).reset_index()
movies_count

# In[37]:

# Plot 50 most reviewed movies.
plt.figure(figsize=(15,10))
fig = sns.barplot(x=movies_count.head(50)['title'], y=movies_count.head(50)['rating'])
fig.set_xticklabels(fig.get_xticklabels(), rotation=45, horizontalalignment='right')
plt.tight_layout()
plt.show()

# ## Time correlation
#
# Lastly we will examine if there is a correlation between then the movie was made and the rating given.

# ## Year movie released vs rating

# With a correlation coefficient of -0.1050, there is a tiny inverse relationship between when a movie was released and the rating given to it. The p-value is also much lower than 0.05 meaning that we can conclude that the correlation is statistically significant. Older movies were rating more generously than newer movies.
#
# This could be because older movies do not have as many ratings as the newer movies. People who would actually watch and rate old movies from the 20s and 30s would typically be film enthusiasts and thus have a bias towards older movies.

# In[38]:

plt.figure(figsize=(10, 6))
mean_rating = raw_data.groupby('year')['rating'].mean().reset_index()
mean_rating = mean_rating[mean_rating.year != 0]
sns.lineplot(x=mean_rating.year, y=mean_rating.rating)
plt.ylabel('avg_rating')
plt.show()

# In[39]:

pearsonr(raw_data.year, raw_data.rating)


import torch
import numpy as np
import random
seed = 0 #1
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
np.random.seed(seed)  # Numpy module.
random.seed(seed)  # Python random module.
torch.manual_seed(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True


# -*- coding: utf-8 -*-
"""
.. _tut-viz-stcs:

====================================
Visualize source time courses (stcs)
====================================

This tutorial focuses on visualization of :term:`source estimates <STC>`.

Surface Source Estimates
------------------------
First, we get the paths for the evoked data and the source time courses (stcs).
"""

# %%

import os.path as op

import numpy as np
import matplotlib.pyplot as plt

import mne
from mne.datasets import sample, fetch_hcp_mmp_parcellation
from mne.minimum_norm import apply_inverse, read_inverse_operator
from mne import read_evokeds

data_path = sample.data_path()
meg_path = data_path / 'MEG' / 'sample'
subjects_dir = data_path / 'subjects'

fname_evoked = meg_path / 'sample_audvis-ave.fif'
fname_stc = meg_path / 'sample_audvis-meg'
fetch_hcp_mmp_parcellation(subjects_dir)

# %%
# Then, we read the stc from file.
stc = mne.read_source_estimate(fname_stc, subject='sample')

# %%
# This is a :class:`SourceEstimate <mne.SourceEstimate>` object.
print(stc)

# %%
# The SourceEstimate object is in fact a *surface* source estimate. MNE also
# supports volume-based source estimates but more on that later.
#
# We can plot the source estimate using the
# :func:`stc.plot <mne.SourceEstimate.plot>` just as in other MNE
# objects. Note that for this visualization to work, you must have ``PyVista``
# installed on your machine.
initial_time = 0.1
brain = stc.plot(subjects_dir=subjects_dir, initial_time=initial_time,
                 clim=dict(kind='value', lims=[3, 6, 9]),
                 smoothing_steps=7)

# %%
# You can also morph it to fsaverage and visualize it using a flatmap.

# sphinx_gallery_thumbnail_number = 3
stc_fs = mne.compute_source_morph(stc, 'sample', 'fsaverage', subjects_dir,
                                  smooth=5, verbose='error').apply(stc)
brain = stc_fs.plot(subjects_dir=subjects_dir, initial_time=initial_time,
                    clim=dict(kind='value', lims=[3, 6, 9]),
                    surface='flat', hemi='both', size=(1000, 500),
                    smoothing_steps=5, time_viewer=False,
                    add_data_kwargs=dict(
                        colorbar_kwargs=dict(label_font_size=10)))

# to help orient us, let's add a parcellation (red=auditory, green=motor,
# blue=visual)
brain.add_annotation('HCPMMP1_combined', borders=2)

# You can save a movie like the one on our documentation website with:
# brain.save_movie(time_dilation=20, tmin=0.05, tmax=0.16,
#                  interpolation='linear', framerate=10)

# %%
# Note that here we used ``initial_time=0.1``, but we can also browse through
# time using ``time_viewer=True``.
#
# In case ``PyVista`` is not available, we also offer a ``matplotlib``
# backend. Here we use verbose='error' to ignore a warning that not all
# vertices were used in plotting.
mpl_fig = stc.plot(subjects_dir=subjects_dir, initial_time=initial_time,
                   backend='matplotlib', verbose='error', smoothing_steps=7)

# %%
#
# Volume Source Estimates
# -----------------------
# We can also visualize volume source estimates (used for deep structures).
#
# Let us load the sensor-level evoked data. We select the MEG channels
# to keep things simple.
evoked = read_evokeds(fname_evoked, condition=0, baseline=(None, 0))
evoked.pick_types(meg=True, eeg=False).crop(0.05, 0.15)
# this risks aliasing, but these data are very smooth
evoked.decimate(10, verbose='error')

# %%
# Then, we can load the precomputed inverse operator from a file.
fname_inv = meg_path / 'sample_audvis-meg-vol-7-meg-inv.fif'
inv = read_inverse_operator(fname_inv)
src = inv['src']
mri_head_t = inv['mri_head_t']

# %%
# The source estimate is computed using the inverse operator and the
# sensor-space data.
snr = 3.0
lambda2 = 1.0 / snr ** 2
method = "dSPM"  # use dSPM method (could also be MNE or sLORETA)
stc = apply_inverse(evoked, inv, lambda2, method)
del inv

# %%
# This time, we have a different container
# (:class:`VolSourceEstimate <mne.VolSourceEstimate>`) for the source time
# course.
print(stc)

# %%
# This too comes with a convenient plot method.
stc.plot(src, subject='sample', subjects_dir=subjects_dir)

# %%
# For this visualization, ``nilearn`` must be installed.
# This visualization is interactive. Click on any of the anatomical slices
# to explore the time series. Clicking on any time point will bring up the
# corresponding anatomical map.
#
# We could visualize the source estimate on a glass brain. Unlike the previous
# visualization, a glass brain does not show us one slice but what we would
# see if the brain was transparent like glass, and
# :term:`maximum intensity projection`) is used:
stc.plot(src, subject='sample', subjects_dir=subjects_dir, mode='glass_brain')

# %%
# You can also extract label time courses using volumetric atlases. Here we'll
# use the built-in ``aparc+aseg.mgz``:

fname_aseg = op.join(subjects_dir, 'sample', 'mri', 'aparc+aseg.mgz')
label_names = mne.get_volume_labels_from_aseg(fname_aseg)
label_tc = stc.extract_label_time_course(fname_aseg, src=src)

lidx, tidx = np.unravel_index(np.argmax(label_tc), label_tc.shape)
fig, ax = plt.subplots(1)
ax.plot(stc.times, label_tc.T, 'k', lw=1., alpha=0.5)
xy = np.array([stc.times[tidx], label_tc[lidx, tidx]])
xytext = xy + [0.01, 1]
ax.annotate(
    label_names[lidx], xy, xytext, arrowprops=dict(arrowstyle='->'), color='r')
ax.set(xlim=stc.times[[0, -1]], xlabel='Time (s)', ylabel='Activation')
for key in ('right', 'top'):
    ax.spines[key].set_visible(False)
fig.tight_layout()

# %%
# We can plot several labels with the most activation in their time course
# for a more fine-grained view of the anatomical loci of activation.
labels = [label_names[idx] for idx in np.argsort(label_tc.max(axis=1))[:7]
          if 'unknown' not in label_names[idx].lower()]  # remove catch-all
brain = mne.viz.Brain('sample', hemi='both', surf='pial', alpha=0.5,
                      cortex='low_contrast', subjects_dir=subjects_dir)
brain.add_volume_labels(aseg='aparc+aseg', labels=labels)
brain.show_view(azimuth=250, elevation=40, distance=400)

# %%
# And we can project these label time courses back to their original
# locations and see how the plot has been smoothed:

stc_back = mne.labels_to_stc(fname_aseg, label_tc, src=src)
stc_back.plot(src, subjects_dir=subjects_dir, mode='glass_brain')

# %%
# Vector Source Estimates
# -----------------------
# If we choose to use ``pick_ori='vector'`` in
# :func:`apply_inverse <mne.minimum_norm.apply_inverse>`
fname_inv = op.join(data_path, 'MEG', 'sample',
                    'sample_audvis-meg-oct-6-meg-inv.fif')
inv = read_inverse_operator(fname_inv)
stc = apply_inverse(evoked, inv, lambda2, 'dSPM', pick_ori='vector')
brain = stc.plot(subject='sample', subjects_dir=subjects_dir,
                 initial_time=initial_time, brain_kwargs=dict(
                     silhouette=True), smoothing_steps=7)

# %%
# Dipole fits
# -----------
# For computing a dipole fit, we need to load the noise covariance, the BEM
# solution, and the coregistration transformation files. Note that for the
# other methods, these were already used to generate the inverse operator.
fname_cov = meg_path / 'sample_audvis-cov.fif'
fname_bem = subjects_dir / 'sample' / 'bem' / 'sample-5120-bem-sol.fif'
fname_trans = meg_path / 'sample_audvis_raw-trans.fif'

##############################################################################
# Dipoles are fit independently for each time point, so let us crop our time
# series to visualize the dipole fit for the time point of interest.
evoked.crop(0.1, 0.1)
dip = mne.fit_dipole(evoked, fname_cov, fname_bem, fname_trans)[0]

##############################################################################
# Finally, we can visualize the dipole.

dip.plot_locations(fname_trans, 'sample', subjects_dir)