"""Factors model"""
__docformat__ = "numpy"

from urllib.request import urlopen
from zipfile import ZipFile
from io import BytesIO

import statsmodels.api as sm
import yfinance as yf
import pandas as pd

def get_fama_raw():
    """Gets base Fama French data to calculate risk"""
    with urlopen(
        "http://mba.tuck.dartmouth.edu/pages/faculty/ken.french/ftp/F-F_Research_Data_Factors_CSV.zip"
    ) as url:

        # Download Zipfile and create pandas DataFrame
        with ZipFile(BytesIO(url.read())) as zipfile:
            with zipfile.open("F-F_Research_Data_Factors.CSV") as zip_open:
                df = pd.read_csv(
                    zip_open,
                    header=0,
                    names=["Date", "MKT-RF", "SMB", "HML", "RF"],
                    skiprows=3,
                )

    df = df[df["Date"].apply(lambda x: len(str(x).strip()) == 6)]
    df["Date"] = df["Date"].astype(str) + "01"
    df["Date"] = pd.to_datetime(df["Date"], format="%Y%m%d")
    df["MKT-RF"] = pd.to_numeric(df["MKT-RF"], downcast="float")
    df["SMB"] = pd.to_numeric(df["SMB"], downcast="float")
    df["HML"] = pd.to_numeric(df["HML"], downcast="float")
    df["RF"] = pd.to_numeric(df["RF"], downcast="float")
    df["MKT-RF"] = df["MKT-RF"] / 100
    df["SMB"] = df["SMB"] / 100
    df["HML"] = df["HML"] / 100
    df["RF"] = df["RF"] / 100
    df = df.set_index("Date")
    return df

def get_historical_5(ticker: str):
    """Get 5 year monthly historical performance for a ticker with dividends filtered"""
    tick = yf.Ticker(ticker)
    df = tick.history(period="5y", interval="1mo")
    df = df[df.index.to_series().apply(lambda x: x.day == 1)]
    df = df.drop(["Dividends", "Stock Splits"], axis=1)
    df = df.dropna()
    return df

def capm_information(ticker):
    """Provides information that relates to the CAPM model"""
    df_f = get_fama_raw()
    df_h = get_historical_5(ticker)
    df = df_h.join(df_f)
    df = df.dropna()
    df["Monthly Return"] = df["Close"].pct_change()
    df["Excess Monthly Return"] = df["Monthly Return"] - df["RF"]
    df["Excess MKT-RF"] = df["MKT-RF"] - df["RF"]
    df = df.dropna()
    y = df[["Excess Monthly Return"]]
    x = df["Excess MKT-RF"]
    x = sm.add_constant(x)
    model = sm.OLS(y, x).fit()
    beta = model.params["Excess MKT-RF"]
    sy = model.rsquared
    return beta, sy


"""
This function aims to simulate Thermal diffusion in the Earth’s crust
with a given initial conditions. We will be using FTCS method
to calculate the temperatures changes.
"""
import numpy as np
import matplotlib.pyplot as plt

# Constants
L = 20        # depth (m)
D = 0.1       # m^2 day^-1
A = 10        # Temperature constant (C)
B = 12        # Temperature coeffiecient (C)
N = 100       # Number of divisions in grid
a = L/N       # Grid spacing
h = 0.01      # Time-step
year = 365    # days in a year
c = h*D/(a*a) # FTCS coeffiecient
depth=np.arange(0,20+a,a)

def meanTemp(t):
    return A+(B*np.sin(2*np.pi*t/year))

# Create the arrays
T = np.empty(N+1,float)
T[N] = 11           # constant temperature at 20m
T[0] = meanTemp(0)   # temperature at the surface at t = 0 day
T[1:N] = 10         # initial temperature everywhere else

Tp = np.empty(N+1,float)
Tp[N] = 11           # constant temperature at 20m
Tp[0] = meanTemp(0)   # temperature at the surface at t = 0 day

# Create time intervals for part b
march = 9*year + year//4
june = 9*year + 2 * year//4
sep = 9*year + 3 * year//4
dec = 10*year

def FTCS(T, Tp, start, end):
    """
    This function basically performs FTCS method using the T and Tp arrays
    and updating from start time to end time.
    """
    time = start
    while time < end:
        # Update or maintain end points values
        T[0] = meanTemp(time)
        T[N] = 11
        Tp[1:N] = T[1:N] + c*(T[0:N-1] + T[2:N+1] - 2 * T[1:N])
        T,Tp = Tp,T
        time += h
    return T, Tp

# PART A
nineYears, Tp1 = FTCS(T,Tp, 0, 9*year)
# plotting
# NOTE: I adjusted the xtick to match with how the array is set up
fig, graph = plt.subplots()
graph.plot(depth,nineYears)
graph.set(xlabel='Depth (m)', ylabel='Temperature (C)',
       title="Temperature at difference depth\nof Earth\'s crust after 9 years")#,
#       xticklabels=['', 0, 4, 8, 12, 16, 20])
graph.grid()
fig.savefig("q2a.png")
plt.show()

# PART B
fig, graph = plt.subplots()

# calculating temperature profile every 3 months and plotting
tenth_march, Tp2 = FTCS(nineYears, Tp1, 9*year, march)
graph.plot(depth,tenth_march, label="1st quarter")

tenth_june, Tp3 = FTCS(tenth_march, Tp2, march, june)
graph.plot(depth,tenth_june, label="2nd quarter")

tenth_sep, Tp4 = FTCS(tenth_june, Tp3, june, sep)
graph.plot(depth,tenth_sep, label="3rd quarter")

tenth_dec, Tp5 = FTCS(tenth_sep, Tp4, sep, dec)
graph.plot(depth,tenth_dec, label="4th quarter")

graph.set(xlabel='Depth (m)', ylabel='Temperature (C)',
       title="Temperature at difference depth of Earth\'s\ncrust in quarters during the 10th year"
       )
graph.grid()
graph.legend()
fig.savefig("q2b.png")
plt.show()


"""
Construct a filter that selects extreme values in a list
"""

import numpy as np

def selectExtremes(values, max_std, min_size=1):
  """"
  Returns a selector (boolean list) that selects values
  that are most extreme according to a criteria.
  :param list-of-numbers values:
  :param float max_std: maxium standard deviation for the group,
  subject to constraints on the group size.
  :param int min_size: minimum size for the list of values
  :return list-of-bool: extreme value if bool is true
  """
  def dfm(mean, value):
    return abs(mean-value)

  if len(values) < 2:
    filter = [False for x in values]
    return filter
  cur_std = np.std(values, ddof=1)
  length = len(values)
  initial_select = [False for n in range(length)]
  paired_list = zip(initial_select, values)
  iterations = 0
  while iterations < length-min_size:
    new_values = [x for (b, x) in paired_list if not b]
    if len(new_values) < 2:
      break
    if np.std(new_values) <= max_std:
      break
  # Eliminate the next largest value
    mean = np.mean(new_values)
    max_idx = 0
    for idx in range(length):
      (b, x) = paired_list[idx]
      max_val = paired_list[max_idx][1]
      if not b and dfm(x, mean) > dfm(max_val, mean):
        max_idx = idx
    (b, x) = paired_list[max_idx]
    paired_list[max_idx] = (True, x)
    iterations += 1
  (flter, values) = zip(*paired_list)
  return list(flter)


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

import numpy as np
import pandas as pd
import multiprocessing as mp

from math import sqrt
from functools import partial
#from dask import delayed, compute
from nbeats.contrib.utils.pytorch.losses import divide_no_nan

AVAILABLE_METRICS = ['mse', 'rmse', 'mape', 'smape', 'mase', 'rmsse',
                     'mini_owa', 'pinball_loss']

######################################################################
# METRICS
######################################################################

def mse(y, y_hat):
    """Calculates Mean Squared Error.

    MSE measures the prediction accuracy of a
    forecasting method by calculating the squared deviation
    of the prediction and the true value at a given time and
    averages these devations over the length of the series.

    Parameters
    ----------
    y: numpy array
        actual test values
    y_hat: numpy array
        predicted values

    Return
    ------
    scalar: MSE
    """
    mse = np.mean(np.square(y - y_hat))

    return mse

def rmse(y, y_hat):
    """Calculates Root Mean Squared Error.

    RMSE measures the prediction accuracy of a
    forecasting method by calculating the squared deviation
    of the prediction and the true value at a given time and
    averages these devations over the length of the series.
    Finally the RMSE will be in the same scale
    as the original time series so its comparison with other
    series is possible only if they share a common scale.

    Parameters
    ----------
    y: numpy array
      actual test values
    y_hat: numpy array
      predicted values

    Return
    ------
    scalar: RMSE
    """
    rmse = sqrt(np.mean(np.square(y - y_hat)))

    return rmse

def mape(y, y_hat):
    """Calculates Mean Absolute Percentage Error.

    MAPE measures the relative prediction accuracy of a
    forecasting method by calculating the percentual deviation
    of the prediction and the true value at a given time and
    averages these devations over the length of the series.

    Parameters
    ----------
    y: numpy array
      actual test values
    y_hat: numpy array
      predicted values

    Return
    ------
    scalar: MAPE
    """
    mape = np.mean(np.abs(y - y_hat) / np.abs(y))
    mape = 100 * mape

    return mape

def smape(y, y_hat):
    """Calculates Symmetric Mean Absolute Percentage Error.

    SMAPE measures the relative prediction accuracy of a
    forecasting method by calculating the relative deviation
    of the prediction and the true value scaled by the sum of the
    absolute values for the prediction and true value at a
    given time, then averages these devations over the length
    of the series. This allows the SMAPE to have bounds between
    0% and 200% which is desireble compared to normal MAPE that
    may be undetermined.

    Parameters
    ----------
    y: numpy array
      actual test values
    y_hat: numpy array
      predicted values

    Return
    ------
    scalar: SMAPE
    """
    delta_y = np.abs(y - y_hat)
    scale = np.abs(y) + np.abs(y_hat)
    smape = divide_no_nan(delta_y, scale)
    smape = 200 * np.mean(smape)
    assert smape <= 200, 'SMAPE should be lower than 200'

    return smape

def mase(y, y_hat, y_train, seasonality=1):
    """Calculates the M4 Mean Absolute Scaled Error.

    MASE measures the relative prediction accuracy of a
    forecasting method by comparinng the mean absolute errors
    of the prediction and the true value against the mean
    absolute errors of the seasonal naive model.

    Parameters
    ----------
    y: numpy array
      actual test values
    y_hat: numpy array
      predicted values
    y_train: numpy array
      actual train values for Naive1 predictions
    seasonality: int
      main frequency of the time series
      Hourly 24,  Daily 7, Weekly 52,
      Monthly 12, Quarterly 4, Yearly 1

    Return
    ------
    scalar: MASE
    """
    scale = np.mean(abs(y_train[seasonality:] - y_train[:-seasonality]))
    mase = np.mean(abs(y - y_hat)) / scale
    mase = 100 * mase

    return mase

def rmsse(y, y_hat, y_train, seasonality=1):
    """Calculates the M5 Root Mean Squared Scaled Error.

    Parameters
    ----------
    y: numpy array
      actual test values
    y_hat: numpy array of len h (forecasting horizon)
      predicted values
    seasonality: int
      main frequency of the time series
      Hourly 24,  Daily 7, Weekly 52,
      Monthly 12, Quarterly 4, Yearly 1

    Return
    ------
    scalar: RMSSE
    """
    scale = np.mean(np.square(y_train[seasonality:] - y_train[:-seasonality]))
    rmsse = sqrt(mse(y, y_hat) / scale)
    rmsse = 100 * rmsse

    return rmsse

def mini_owa(y, y_hat, y_train, seasonality, y_bench):
    """Calculates the Overall Weighted Average for a single series.

    MASE, sMAPE for Naive2 and current model
    then calculatess Overall Weighted Average.

    Parameters
    ----------
    y: numpy array
        actual test values
    y_hat: numpy array of len h (forecasting horizon)
        predicted values
    seasonality: int
        main frequency of the time series
        Hourly 24,  Daily 7, Weekly 52,
        Monthly 12, Quarterly 4, Yearly 1
    y_train: numpy array
        insample values of the series for scale
    y_bench: numpy array of len h (forecasting horizon)
        predicted values of the benchmark model

    Return
    ------
    return: mini_OWA
    """
    mase_y = mase(y, y_hat, y_train, seasonality)
    mase_bench = mase(y, y_bench, y_train, seasonality)

    smape_y = smape(y, y_hat)
    smape_bench = smape(y, y_bench)

    mini_owa = ((mase_y/mase_bench) + (smape_y/smape_bench))/2

    return mini_owa

def pinball_loss(y, y_hat, tau=0.5):
    """Calculates the Pinball Loss.

    The Pinball loss measures the deviation of a quantile forecast.
    By weighting the absolute deviation in a non symmetric way, the
    loss pays more attention to under or over estimation.
    A common value for tau is 0.5 for the deviation from the median.

    Parameters
    ----------
    y: numpy array
      actual test values
    y_hat: numpy array of len h (forecasting horizon)
      predicted values
    tau: float
      Fixes the quantile against which the predictions are compared.

    Return
    ------
    return: pinball_loss
    """
    delta_y = y - y_hat
    pinball = np.maximum(tau * delta_y, (tau-1) * delta_y)
    pinball = pinball.mean()

    return pinball

def panel_mape(y_hat):
    y_hat = y_hat.copy()
    y_hat['mape'] = np.abs(y_hat['y_hat']-y_hat['y'])/np.abs(y_hat['y'])
    y_hat_grouped = y_hat.groupby('unique_id').mean().reset_index()
    mape = np.mean(y_hat_grouped['mape'])
    return mape

def panel_smape(y_hat):
    y_hat = y_hat.copy()
    y_hat['smape'] = np.abs(y_hat['y_hat']-y_hat['y'])/(np.abs(y_hat['y'])+np.abs(y_hat['y_hat']))
    y_hat_grouped = y_hat.groupby('unique_id').mean().reset_index()
    smape = 2 * np.mean(y_hat_grouped['smape'])
    return smape


from quantfin.simulation import random_covariance
from quantfin.statistics import marchenko_pastur, targeted_shirinkage
from quantfin.portfolio import HRP
from matplotlib import rcParams
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

n_assets = 25
n_time = 5000
n_factors = 3
random_seed = 666

file_path = r'/Users/gustavoamarante/Dropbox/CQF/Final Project'  # Mac
# file_path = r'/Users/gusamarante/Dropbox/CQF/Final Project'  # Macbook

# ===== generate data =====
mu = np.zeros(n_assets)
cov = random_covariance(size=n_assets, n_factors=n_factors, random_seed=random_seed)
data = pd.DataFrame(data=np.random.multivariate_normal(mu, cov, n_time),
                    columns=[f'Asset {i + 1}' for i in range(n_assets)])

# ===== plot correlation matrix =====
hrp = HRP(data)
hrp.plot_corr_matrix(cmap='mako', figsize=(10, 10), show_chart=True,
                     save_path=file_path + r'/figures/Simulated Correlation Matrix.pdf')

# ===== compute different correlations matrices =====
corr = data.corr().values  # empirical correlation
corr_denoised, _, _ = marchenko_pastur(corr_matrix=corr, T=data.shape[0], N=data.shape[1])  # denoised correlation
corr_ts, _, _ = targeted_shirinkage(corr_matrix=corr, T=data.shape[0], N=data.shape[1], ts_alpha=1)  # target-shrinkage

# ===== sorted eigenvalues of each method =====
eig_empirical = np.sort(np.linalg.eig(corr)[0])[::-1]
eig_denoised = np.sort(np.linalg.eig(corr_denoised)[0])[::-1]
eig_ts = np.sort(np.linalg.eig(corr_ts)[0])[::-1]
df_eig = pd.DataFrame(data={'Empirical Eigenvalues': eig_empirical,
                            'Denoised Eigenvalues': eig_denoised,
                            'Targeted Shirinkage': eig_ts},
                      index=[i + 1 for i in range(n_assets)])

# ===== chart =====
# time series Chart
MyFont = {'fontname': 'Century Gothic'}
rcParams['font.family'] = 'sans-serif'
rcParams['font.sans-serif'] = ['Century Gothic']

plt.figure(figsize=(10, 4))
plt.plot(df_eig, linewidth=2)
plt.tick_params(axis='y', which='both', right=False, left=False, labelleft=True)
plt.tick_params(axis='x', which='both', top=False, bottom=False, labelbottom=True)
plt.legend(df_eig.columns, loc='upper right', ncol=1, frameon=True)
ax = plt.gca()
ax.set(yscale='log', xlabel='Sorted Eigenvalues', ylabel='Log-Scale')
ax.yaxis.grid(color='grey', linestyle='-', linewidth=0.5, alpha=0.5)
ax.xaxis.grid(color='grey', linestyle='-', linewidth=0.5, alpha=0.5)
plt.tight_layout()
plt.savefig(file_path + f'/figures/Simulated Sorted Eigenvalues.pdf', pad_inches=0)
plt.show()

# ===== get characteristics =====
df_char = pd.DataFrame()
corr_mat_dict = {'Empirical': corr,
                 'Denoised': corr_denoised,
                 'Targeted Shirinkage': corr_ts}
for mat in corr_mat_dict.keys():
    eigvals = np.linalg.eig(corr_mat_dict[mat])[0]
    df_char.loc['Maximum Eigenvalue', mat] = np.max(eigvals)
    df_char.loc['Minimum Eigenvalue', mat] = np.min(eigvals)
    df_char.loc['Determinant', mat] = np.linalg.det(corr_mat_dict[mat])
    df_char.loc['Max Norm', mat] = np.max(np.abs(corr_mat_dict[mat] - corr))
    df_char.loc['Frobenius Norm', mat] = np.linalg.norm(corr_mat_dict[mat] - corr, 'fro')

df_char.to_clipboard()


import scipy.sparse as sp
import numpy as np
import tensorflow as tf

def eigen_pe_sp(edges, num_nodes, pos_enc_dim):
    if isinstance(edges, tf.Tensor):
        edges = edges.numpy()
    if isinstance(num_nodes, tf.Tensor):
        num_nodes = num_nodes.numpy()
    if isinstance(pos_enc_dim, tf.Tensor):
        pos_enc_dim = pos_enc_dim.numpy()

    rows_mat = edges[:,0]
    cols_mat = edges[:,1]

    data_mat = np.ones(rows_mat.shape, dtype='float32')

    A = sp.csr_matrix((data_mat, (rows_mat,cols_mat)),
                      shape=(num_nodes,num_nodes),
                      dtype='float32')
    N = sp.diags(A.sum(axis=1).A.squeeze().clip(1) ** -0.5, dtype=float)
    L = sp.eye(num_nodes) - N * A * N

    EigVal, EigVec = sp.linalg.eigs(L, k=pos_enc_dim+1, which='SR', tol=1e-2)
    EigVec = EigVec[:, EigVal.argsort()]

    pos_enc = np.real(EigVec[:,1:pos_enc_dim+1]).astype('float32')

    return pos_enc

def eigen_pe_np(edges, num_nodes, pos_enc_dim):
    if isinstance(edges, tf.Tensor):
        edges = edges.numpy()
    if isinstance(num_nodes, tf.Tensor):
        num_nodes = num_nodes.numpy()
    if isinstance(pos_enc_dim, tf.Tensor):
        pos_enc_dim = pos_enc_dim.numpy()

    rows_mat = edges[:,0]
    cols_mat = edges[:,1]

    data_mat = np.ones(rows_mat.shape, dtype='float32')

    A = sp.csr_matrix((data_mat, (rows_mat,cols_mat)),
                      shape=(num_nodes,num_nodes),
                      dtype='float32')
    N = sp.diags(A.sum(axis=1).A.squeeze().clip(1) ** -0.5, dtype=float)
    L = sp.eye(num_nodes) - N * A * N

    EigVal, EigVec = np.linalg.eig(L.toarray())
    EigVec = np.real(EigVec[:, EigVal.argsort()])
    pos_enc = EigVec[:,1:pos_enc_dim+1].astype('float32')

    return pos_enc

def eigen_pe(edges, num_nodes, pos_enc_dim, sparse=True):
    if sparse:
        out = tf.py_function(eigen_pe_sp, [edges, num_nodes, pos_enc_dim], tf.float32)
    else:
        out = tf.py_function(eigen_pe_np, [edges, num_nodes, pos_enc_dim], tf.float32)

    if not isinstance(num_nodes, tf.Tensor):
        out.set_shape([num_nodes, None])
    if not isinstance(pos_enc_dim, tf.Tensor):
        out.set_shape([None, pos_enc_dim])
    return out

class EigenFeatures:
    def __init__(self                                     ,
                 num_features        = 8                  ,
                 sparse              = True               ,
                 edges_input_key     = 'edges'            ,
                 num_nodes_input_key = 'num_nodes'        ,
                 vector_key          = 'eigen_vectors'    ):

        self.num_features        = num_features
        self.sparse              = sparse
        self.edges_input_key     = edges_input_key
        self.num_nodes_input_key = num_nodes_input_key
        self.vector_key          = vector_key

    def __call__(self, inputs):
        edges = inputs[self.edges_input_key]
        num_nodes = inputs[self.num_nodes_input_key]

        U = eigen_pe(edges, num_nodes, self.num_features, self.sparse)

        return_dict = {
            **inputs           ,
            self.vector_key : U,
        }

        return return_dict


# Generates samples from the heart-shaped density and writes to heart.npy
import numpy as np
import torch
import random
import math
from tqdm import tqdm

def potential(z):
    z = z.view(-1, 2)
    x, y = z[:, 0], z[:, 1]
    u = 0.8 * x ** 2 + (y - ((x**2)**(1/3)))**2
    u = u / 2**2
    return u

def log_Q(potential, z_prime, z, step):
    z.requires_grad_()
    grad = torch.autograd.grad(potential(z).mean(), z)[0]
    return -(torch.norm(z_prime - z + step * grad, p=2, dim=1) ** 2) / (4 * step)

def metropolis_adjusted_langevin_algorithm(potential, n_samples=100000, step=0.1):
    burn_in = 10000
    Z0 = torch.randn(1, 2)
    Zi = Z0
    samples = []
    pbar = tqdm(range(n_samples + burn_in))
    for i in pbar:
        Zi.requires_grad_()
        u = potential(Zi).mean()
        grad = torch.autograd.grad(u, Zi)[0]
        prop_Zi = Zi.detach() - step * grad + np.sqrt(2 * step) * torch.randn(1, 2)
        log_ratio = -potential(prop_Zi).mean() + potential(Zi).mean() +\
                    log_Q(potential, Zi, prop_Zi, step) - log_Q(potential, prop_Zi, Zi, step)
        if torch.rand(1) < torch.exp(log_ratio):
            Zi = prop_Zi
        samples.append(Zi.detach().numpy())
    return np.concatenate(samples, 0)[burn_in:]

samples = metropolis_adjusted_langevin_algorithm(potential)
np.save('heart-mala.npy', samples)


import numpy as np
import matplotlib.pyplot as plt

# deriv computes 1D derivative, dF/dr, using central difference method
def deriv(F,r):
	# get the length of array F (assume same length for r)
    L = F.size

	# create empty array to store results
    result=  np.empty(L)

	# use central diff method for all interior points (we will build in tutorial)
    for i in range(L-2):
        result[i+1] = (F[i+2] - F[i]) / (r[i+2] - r[i])

    result[0] = (F[1] - F[0]) / (r[1] - r[0])
    result[L-1] = (F[L-1] - F[L-2]) / (r[L-1] - r[L-2])

    return result

# read in the files "vFieldX.csv" and "vFieldY.csv"
vFieldX= np.loadtxt( 'vFieldX.csv', delimiter = ',' )
vFieldY = np.loadtxt( 'vFieldY.csv', delimiter = ',' )

# Create a 2D grid of x, y points using numpy's meshgrid function (see Exercise 1)
nx, ny = 100,100
x = np.linspace(-5,5,nx)
y = np.linspace(-5,5,ny)
X, Y = np.meshgrid(x,y)

# Divergence
divX = np.empty(X.shape)

for j in range(ny):
    divX[j,:] = deriv(vFieldX[j,:],x)

divY = np.empty(Y.shape)

for i in range(nx):
    divY[:,i] = deriv(vFieldY[:,i], y)

totalDiv = divX + divY

# Curl
curlX = np.empty(X.shape)

for j in range(ny):
    curlX[j,:] = deriv(vFieldY[j,:], x)

curlY = np.empty(Y.shape)

for i in range(nx):
    curlY[:,i] = deriv(vFieldX[:,i], y)

totalCurl = curlX - curlY

# Plotting the Divergence and Curl using subplots
lines = 10**np.linspace(10, 12, 11)
lines = sorted(list(-lines)+list(lines))
fig, (ax1, ax2) = plt.subplots( nrows = 1, ncols = 2, sharex = False, sharey = False )

ax1.contourf(X, Y, totalDiv) #levels = Lines, colors = 'k', linewidths = 1)
CS = ax2.contour(x, y, totalCurl, ) #levels = Lines, colors = 'k', linewidths = 1)

ax1.set_title('Divergence of a Vector Field', fontweight = 'bold' )
ax2.set_title('Curl of a Vector Field', fontweight = 'bold' )

ax1.set(xlabel = "X", ylabel = "Y")
ax2.set(xlabel = "X", ylabel = "Y")

ax2.clabel(CS, inline = 1, fontsize = 8)

fig.set_size_inches(9, 5)

plt.savefig('divCurlPlot.png', dpi=300)


import numpy as np
import matplotlib.pyplot as plt

def softmax(inputs):
    """
    Calculate the softmax for the give inputs (array)
    :param inputs:
    :return:
    """
    return np.exp(inputs) / float(sum(np.exp(inputs)))

def line_graph(x, y, x_title, y_title):
    """
    Draw line graph with x and y values
    :param x:
    :param y:
    :param x_title:
    :param y_title:
    :return:
    """
    plt.plot(x, y)
    plt.xlabel(x_title)
    plt.ylabel(y_title)
    plt.show()

softmax_inputs = [2, 3, 5, 6]

print("Softmax Function Output :: {}".format(softmax(softmax_inputs)))

graph_x = range(0, 21)
graph_y = softmax(graph_x)

print("Graph X readings: {}".format(graph_x))
print("Graph Y readings: {}".format(graph_y))

line_graph(graph_x, graph_y, "Inputs", "Softmax Scores")


import math
from functools import partial
import numpy as np
import matplotlib.pyplot as plt
from scipy.misc import derivative

def _function(fun: str, x: int or float):
    """
    Transform the string to executable form.

    :param fun: a one parameter function, it must be a string.
    :param x: the parameter of the function..
    :return: the executable form of the string function.
    """
    return eval(fun)

def BM(
        fun: str,
        a: int or float,
        b: int or float,
        eps1: float,
        eps2: float,
        max_time: int
) -> tuple:
    """
    BM algorithm.

    :param fun: the function which will be solved by BM, it must be a string.
    :param a: the left boundary you assign.
    :param b: the right boundary you assign.
    :param eps1: the minimum boundary of the solution.
    :param eps2: the minimum boundary of the function's value of solution.
    :param max_time: the max iteration time.
    :return: a tuple with two elements, the first element is a list which
        includes the single solution of objective equation and the second
        element is the iteration time.
    """
    f = partial(_function, fun)
    # The partial function, it can fix the parameters for a objective function.

    time = 1
    while time < max_time:
        x = (a + b) / 2
        xe = (b - a) / 2
        print(f"第{time}次迭代:{x}")

        if abs(xe) < eps2 and abs(f(x)) < eps1:
            return x, time
        else:
            if f(a) * f(x) < 0:
                b = x
            else:
                a = x
        time += 1

def BMZ(
        fun: str,
        a: int or float,
        eps1: float,
        eps2: float,
        k: int,
        dx: int or float,
        max_time: int
) -> tuple:
    """
    BMZ algorithm.

    :param fun: the function which will be solved by BM, it must be a string.
    :param a: the left boundary you assign.
    :param eps1: the minimum boundary of the solution.
    :param eps2: the minimum boundary of the function's value of solution.
    :param k: the number of all solutions.
    :param dx: the step length.
    :param max_time: the max iteration time.
    :return: a tuple with two elements, the first element is a list which
        includes the single solution of objective equation and the second
        element is the iteration time.
    """
    f = partial(_function, fun)
    # The partial function, it can fix the parameters for a objective function.

    all_solution = []

    time = 1
    while time < max_time:
        x = a + time * dx

        if abs(f(x) - 0) < eps1:
            all_solution.append(x)

        elif f(x) * f(x + dx) > 0:
            # The judgement condition in book is "f(x) * f(x + 1) > 0", but
            # while I debugged here, I discovered it will ignore one solution,
            # so I had to modify it.

            time += 1
            pass
        else:
            try:
                all_solution.append(
                    BM(fun, x, a + (time + 1) * dx, eps1, eps2, max_time)[0]
                )
                time += 1
            except Exception:
                time += 1
                pass

        if len(all_solution) == k:
            break

    return all_solution, time

def false_position(
        fun: str,
        a: int or float,
        b: int or float,
        eps1: float,
        eps2: float,
        max_time: int
) -> tuple:
    """
    Using false position algorithm to solve the function.

    :param fun: the function which will be solved by false position
        algorithm, it must be a string.
    :param a: the left boundary.
    :param b: the right boundary.
    :param eps1: the minimum boundary of the solution.
    :param eps2: the minimum boundary of the function's value of solution.
    :param max_time: the max iteration time.
    :return: the single solution of the function.
    """

    def _getX1(fun, x_left, x_right):
        x = x_left - (
                (x_right - x_left) / (fun(x_right) - fun(x_left))
            ) * fun(x_left)
        return x

    f = partial(_function, fun)
    # The partial function, it can fix the parameters for a objective function.

    getX1 = partial(_getX1, f)
    x1 = getX1(a, b)
    x = a
    xe = b

    time = 1
    while time < max_time:
        if x * x1 < 0:
            xe = x1
            x1 = getX1(x, xe)
        else:
            x = x1
            x1 = getX1(x, xe)

        if abs((xe - x) / 2) < eps1 or abs(f(x)) < eps2:
            break
        time += 1

    return x, time

def simple_iteration(
        fun: str,
        initial_value: int or float,
        eps1: int or float,
        eps2: int or float,
        max_time: int
) -> tuple:
    """
    Using simple iteration algorithm to calculate the solution of nonlinear
    equation which you assign. There is a little inconvenience, the parameter
    "fun" must be the iteration form but not the objective function. Be example,
    if you want to get the solution of "x ** 2", you must make the "fun" is
    iteration form like "x ** 2 + x" or other iteration forms.

    :param fun: The iteration form of your objective function, it must be str
        type.
    :param initial_value: the initial value.
    :param eps1: the minimum boundary of the solution.
    :param eps2: the minimum boundary of the function's value of solution.
    :param max_time: the max iteration time.
    :return: a tuple with two elements, the first element is a list which include
        all solutions of the objective equation and the second element is the
        iteration time.
    """

    f = partial(_function, fun)
    x = initial_value
    solution_sequence = [x]

    time = 1
    while time < max_time:
        x = f(x)
        solution_sequence.append(x)

        if abs(solution_sequence[-1] - solution_sequence[-2]) < eps1 \
                or abs(f(solution_sequence[-1])) < eps2:
            break

        time += 1
    return solution_sequence, time

def steffensen(
        fun: str,
        initial_value: float or int,
        eps1: float or int,
        eps2: float or int,
        max_time: int
) -> tuple:
    """
    This function implements Steffensen algorithm which can solve a nonlinear
    equation, compare to simple iteration algorithm, it has more fast convergence
    rate. And, same as simple iteration, you must transform the objective equation
    to iteration form.

    :param fun: the iteration form of the objective equation.
    :param initial_value: the initial value.
    :param eps1: the minimum boundary of the solution.
    :param eps2: the minimum boundary of the function's value of solution.
    :param max_time: the max iteration time.
    :return: a tuple with two elements, the first element is the solution and the
        second element is the iteration time.
    """
    f = partial(_function, fun)
    x = [initial_value]
    print(f"第1次：{x[-1]}， f(x) = {f(x[-1]) - x[-1]}")

    time = 1
    while time < max_time:
        x1 = f(x[-1])
        x2 = f(x1)
        x.append(x[-1] - ((x1 - x[-1]) ** 2) / (x2 - 2 * x1 + x[-1]))

        if abs(x[-1] - x[-2]) < eps1 or abs(f(x[-1])) < eps2:
            break

        print(f"第{time}次：{x[-1]}， f(x) = {f(x[-1]) - x[-1]}")

        time += 1
    print(f"最后一次：{x[-1]}，f(x) = {f(x[-1]) - x[-1]}")

    return x, time

def newton_iteration(
        fun: str,
        initial_value: int or float,
        eps1: int or float,
        eps2: int or float,
        max_time: int
) -> tuple:
    """
    This function is a implementation of Newton iteration algorithm. The
    parameter "fun" is the iteration form of the objective equation.

    :param fun: the iteration form of the objective equation.
    :param initial_value: the initial value.
    :param eps1: the minimum boundary of the solution.
    :param eps2: the minimum boundary of the function's value of solution.
    :param max_time: the max iteration time.
    :return: a tuple with two elements, the first element is the solution and the
        second element is the iteration time.
    """
    f = partial(_function, fun)
    x = initial_value
    print(f"第1次迭代结果：{x}，f(x) = {f(x)}")

    time = 1
    while time < max_time:
        xOld = x

        fd = (f(x + 0.0000001) - f(x)) / 0.0000001
        # Using the definition to calculate the value of derivatives of the
        # objective equation.

        x -= f(x) / fd

        if abs(x - xOld) < eps1 and abs(f(x)) < eps2:
            break
        print(f"第{time}次迭代结果：{x}，f(x) = {f(x)}")

        time += 1

    return x, time

def multiple_solutionNewton(
        fun: str,
        initial_value: int or float,
        eps1: int or float,
        eps2: int or float,
        max_time: int
) -> tuple:
    """
    The algorithm of this function implementation is multiple root acceleration.

    :param fun: the iteration form of the objective equation.
    :param initial_value: the initial value.
    :param eps1: the minimum boundary of the solution.
    :param eps2: the minimum boundary of the function's value of solution.
    :param max_time: the max iteration time.
    :return: a tuple with two elements, the first element is the solution and the
        second element is the iteration time.
    """
    f = partial(_function, fun)
    x = initial_value

    time = 1
    while time < max_time:
        xOld = x

        fd = derivative(f, x, dx=0.0001, n=1)
        fdd = derivative(f, x, dx=0.00001, n=2)
        # In order to calculate the value of two derivative of the objective
        # equation, I have to use the function "derivative" in "scipy". The
        # variable "fd" is the value of one derivative of the objective and
        # the "fdd" is the value of two derivative.

        x -= (f(x) * fd) / (fd ** 2 - f(x) * fdd)

        if abs(x - xOld) < eps1 and abs(f(x) < eps2):
            break

        time += 1

    return x, time

def secant(
        fun: str,
        initial_value: int or float,
        eps1: int or float,
        eps2: int or float,
        max_time: int
) -> tuple:
    """
    The algorithm this function implementation is the improvement of Newton
    iteration algorithm, this algorithm use difference quotient to replace
    differential quotient.

    :param fun: the iteration form of the objective equation.
    :param initial_value: the initial value.
    :param eps1: the minimum boundary of the solution.
    :param eps2: the minimum boundary of the function's value of solution.
    :param max_time: the max iteration time.
    :return: a tuple with two elements, the first element is the solution and the
        second element is the iteration time.
    """
    f = partial(_function, fun)
    x = [initial_value]
    x.append(x[0] - f(x[0]) / (((f(x[0]) + 0.0000001) - f(x[0])) / 0.0000001))

    time = 1
    while time < max_time:
        xNew = x[-1] - (f(x[-1]) / (f(x[-1]) - f(x[-2]))) * (x[-1] - x[-2])
        x.append(xNew)

        if abs(x[-1] - x[-2]) < eps1 and abs(f(x[-1])) < eps2:
            break

        time += 1

    return x[-1], time

def newton_descent(
        fun: str,
        initial_value: int or float,
        eps1: int or float,
        eps2: int or float,
        eps3: int or float,
        eps4: int or float,
        max_time: int
) -> tuple:
    """
    This function implements the improvement of Newton iteration, it has more fast
    convergence speed.

    :param fun: the iteration form of the objective equation.
    :param initial_value: the initial value.
    :param eps1: the minimum boundary of the solution.
    :param eps2: the minimum boundary of the function's value of solution.
    :param eps3: the minimum value of t.
    :param eps4: the correction of initial value.
    :param max_time: the max iteration time.
    :return: a tuple with two elements, the first element is the solution and the
        second element is the iteration time.
    """
    f = partial(_function, fun)
    x = [initial_value]
    t = 1

    time = 1
    while time < max_time:
        fd = (f(x[-1] + 0.0000001) - f(x[-1])) / 0.0000001
        x_ = x[-1] - t * (f(x[-1]) / fd)

        if abs(f(x_) - f(x[-1])) < eps1 or abs(f(x_)) < eps2:
            return x_, time

        if abs(f(x_)) < abs(f(x[-1])):
            x.append(x_)
            print(f"第{time}次迭代结果：{x[-1]}")
            time += 1
        else:
            if t > eps3:
                t /= 2
                time += 1
                continue
            else:
                x.append(x_ + eps4)
                print(f"第{time}次迭代结果：{x[-1]}")
                t /= 2
                time += 1

    return x[-1], time

if __name__ == '__main__':
    testBIn = "1 - x - math.sin(x)"
    print("Test of BM algorithm:")
    bm_result = BM(testBIn, 0, 1, 0.001, 0.005, 10000)
    print(f"最终结果：{bm_result[0]}")
    print(f"迭代次数：{bm_result[1]}")

    print("Test of false position algorithm:")
    false_position_re = false_position(testBIn, 0, 1, 0.01, 0.005, 10000)
    print(f"最终结果：{false_position_re[0]}")
    print(f"迭代次数：{false_position_re[1]}")

    testBIn1 = "6 * x ** 4 - 40 * x ** 2 + 9"
    print("Test of BMZ algorithm:")
    bmz_result = BMZ(testBIn1, -5, 0.001, 0.1, 4, 0.1, 100)
    print(f"所有根：{bmz_result[0]}")
    print(f"迭代次数：{bmz_result[1]}")

    testIter = "math.e ** (-x)"
    print("Test of simple iteration algorithm:")
    simpleIteration_result = simple_iteration(testIter, 0.5, 0.0001, 0.0001, 10000)
    print(f"最终结果：{simpleIteration_result[0]}")
    print(f"迭代次数：{simpleIteration_result[1]}")
    print("Test of steffensen algorithm:")
    steffensen_result = steffensen(testIter, 0.5, 0.0001, 0.0001, 10000)
    print(f"最终结果：{steffensen_result[0]}")
    print(f"迭代次数：{steffensen_result[1]}")
    plt.plot(simpleIteration_result[0], label="simple iteration")
    plt.plot(steffensen_result[0], label="steffensen")
    plt.show()

    # testNewton = "x - (1 / 2) * (x + 3 / x)"  # 3 ^ (1 / 2) iteration form.
    testNewton = "(x ** 2) ** (1 / 3)"
    # test = "(x - 1.5) ** 2"
    print("Test of Newton iteration algorithm:")
    newtonSolution = newton_iteration(testNewton, 1.5, 0.00001, 0.00001, 10000)
    print(f"最终结果：{newtonSolution[0]}")
    print(f"迭代次数：{newtonSolution[1]}")
    print("Test of multiple solution algorithm:")
    multi = multiple_solutionNewton(testNewton, 1, 0.00001, 0.00001, 10000)
    print(f"最终结果：{multi[0]}")
    print(f"迭代次数：{multi[1]}")

    print("Test of secant algorithm:")
    secantSolution = secant(testNewton, 1, 0.00001, 0.00001, 10000)
    print(f"最终结果：{secantSolution[0]}")
    print(f"迭代次数：{secantSolution[1]}")
    print("Test of Newton descent algorithm:")
    des = newton_descent(testNewton, 1.5, 0.00001, 0.00001, 0.001, 0.001, 10000)
    print(f"最终结果：{des[0]}")
    print(f"迭代次数：{des[1]}")
    #
    # t = "(x ** 2) ** (1 / 3)"
    # x = np.linspace(-10, 10, 10000)
    # plt.plot(x, list(map(partial(_function, t), x)))
    # plt.show()

    # def e(x):
    #     return math.e ** -x
    # def f(x):
    #     ff = x - (e(x) - x) ** 2 / (e(e(x)) - 2 * e(x) + x)
    #     return ff
    # print(derivative(f, 0.5671190400572149, dx=0.0001, n=1))
    # print(derivative(f, 0.5671190400572149, dx=0.0001, n=2))
    # def ff(x):
    #     return abs(derivative(e, x, dx=0.0001, n=1))
    #
    # xx = np.linspace(0, 1, 1000)
    #
    # yy = max(list(map(ff, xx)))
    # print(yy)

    # x = np.linspace(0, 1, 10000)
    # plt.plot(x, list(map(f, x)))
    # plt.show()
    # print(e(0), e(1))
    # print(f(1))


"""
Filename: gth_solve.py

Author: Daisuke Oyama

Routine to compute the stationary distribution of an irreducible Markov
chain by the Grassmann-Taksar-Heyman (GTH) algorithm.

"""
import numpy as np

try:
    xrange
except:  # python3
    xrange = range

def gth_solve(A, overwrite=False):
    r"""
    This routine computes the stationary distribution of an irreducible
    Markov transition matrix (stochastic matrix) or transition rate
    matrix (generator matrix) `A`.

    More generally, given a Metzler matrix (square matrix whose
    off-diagonal entries are all nonnegative) `A`, this routine solves
    for a nonzero solution `x` to `x (A - D) = 0`, where `D` is the
    diagonal matrix for which the rows of `A - D` sum to zero (i.e.,
    :math:`D_{ii} = \sum_j A_{ij}` for all :math:`i`). One (and only
    one, up to normalization) nonzero solution exists corresponding to
    each reccurent class of `A`, and in particular, if `A` is
    irreducible, there is a unique solution; when there are more than
    one solution, the routine returns the solution that contains in its
    support the first index `i` such that no path connects `i` to any
    index larger than `i`. The solution is normalized so that its 1-norm
    equals one. This routine implements the Grassmann-Taksar-Heyman
    (GTH) algorithm [1]_, a numerically stable variant of Gaussian
    elimination, where only the off-diagonal entries of `A` are used as
    the input data. For a nice exposition of the algorithm, see Stewart
    [2]_, Chapter 10.

    Parameters
    ----------
    A : array_like(float, ndim=2)
        Stochastic matrix or generator matrix. Must be of shape n x n.
    overwrite : bool, optional(default=False)
        Whether to overwrite `A`; may improve performance.

    Returns
    -------
    x : numpy.ndarray(float, ndim=1)
        Stationary distribution of `A`.

    References
    ----------
    .. [1] W. K. Grassmann, M. I. Taksar and D. P. Heyman, "Regenerative
       Analysis and Steady State Distributions for Markov Chains,"
       Operations Research (1985), 1107-1116.

    .. [2] W. J. Stewart, Probability, Markov Chains, Queues, and
       Simulation, Princeton University Press, 2009.

    """
    A1 = np.array(A, dtype=float, copy=not overwrite)

    if len(A1.shape) != 2 or A1.shape[0] != A1.shape[1]:
        raise ValueError('matrix must be square')

    n = A1.shape[0]

    x = np.zeros(n)

    # === Reduction === #
    for i in xrange(n-1):
        scale = np.sum(A1[i, i+1:n])
        if scale <= 0:
            # There is one (and only one) recurrent class contained in
            # {0, ..., i};
            # compute the solution associated with that recurrent class.
            n = i+1
            break
        A1[i+1:n, i] /= scale

        A1[i+1:n, i+1:n] += np.dot(A1[i+1:n, i:i+1], A1[i:i+1, i+1:n])

    # === Backward substitution === #
    x[n-1] = 1
    for i in xrange(n-2, -1, -1):
        x[i] = np.dot(x[i+1:n], A1[i+1:n, i])

    # === Normalization === #
    x /= np.sum(x)

    return x


import numpy as np

def p_r_f(predictions, Y):
    '''
    Returns precision, recall, and F1 score
    '''
    tp, fp, tn, fn = 0, 0, 0, 0
    for pred, y in zip(predictions, Y):
        if pred == y:
            if pred == 1:
                tp += 1
            else:
                tn += 1
        else:
            if pred == 1:
                fp += 1
            else:
                fn += 1

    try:
        p = tp / (tp + fp)
        r = tp / (tp + fn)
        f = 2*p*r / (p + r)
    except:
        # if undefined, just default to all 0
        p, r, f = 0, 0, 0
    return p,r,f

def mean_squared_error(predictions, Y):
    squared_errors = ((pred-y)**2 for pred, y in zip(predictions, Y))
    return sum(squared_errors)/len(Y)

def KFolds(X, Y, num_folds=5):
    # get a list representing indices into X/Y, then shuffle the indices
    indices = np.arange(len(Y))
    np.random.shuffle(indices)

    for i in range(num_folds):
        # For each fold of cross val, grab the next subset of data
        # to be the hold-out portion. Since the indices were pre-shuffled,
        # we can just go in order
        start_index = int(i*len(Y)/num_folds)
        #print(f'start_index = {start_index}')
        end_index = int(start_index + len(Y)/num_folds)
        #print(f'end_index = {end_index}')
        indices_for_holdout = indices[start_index: end_index]
        #print(f'indices_for_holdout = {indices_for_holdout}')
        first_chunk = indices[0:start_index]
        second_chunk = indices[end_index:]
        indices_for_training = np.concatenate((first_chunk, second_chunk))

        training_X, training_Y = X[indices_for_training],  Y[indices_for_training]
        holdout_X, holdout_Y = X[indices_for_holdout],  Y[indices_for_holdout]

        yield training_X, training_Y, holdout_X, holdout_Y

def cross_validation(model, X, Y, task_type, num_folds=5):
    # given a model that can completely transform and predict the given data, X
    # and the expected output, Y, returns a dictionary of evaluation metrics
    # for evaluating num_folds of cross validation

    assert(task_type in {'classification', 'regression'})
    predictions = None
    truth = None
    for training_X, training_Y, holdout_X, holdout_Y in KFolds(X, Y, num_folds=num_folds):
        print(f"Fold!")
        model.fit(training_X, training_Y)
        #model.print_tree()
        current_predictions = model.predict(holdout_X)
        #print('preds')
        #print(predictions)
        #print('current preds')
        #print(current_predictions)
        if predictions is None:
            predictions = current_predictions
        else:
            predictions = np.concatenate((predictions, current_predictions))
        if truth is None:
            truth = holdout_Y
        else:
            truth = np.concatenate((truth, holdout_Y))

    #print(f'predictions = {predictions}')
    #print(f'truth = {truth}')

    if task_type == 'classification':
        p,r,f = p_r_f(predictions, truth)
        print(f'p={p}, r={r}, f={f}')
    else:
        mse= mean_squared_error(predictions, truth)
        print(f'mean squared error = {mse}')


# ============================================================================ #
#                             LINEAR REGRESSION                                #
# ============================================================================ #
# Renders stacked line charts showing convex and non-convex obj function
#%%
import pandas as pd
import plotly.offline as po
import plotly.graph_objs as go
import numpy as np
from sklearn import datasets
from sklearn.linear_model import SGDRegressor
from ml_studio.supervised_learning.regression import LinearRegression
from ml_studio.supervised_learning.training.metrics import R2
from ml_studio.utils.data_manager import StandardScaler, data_split
from ml_studio.visual.animations import SingleModelFit2D, SingleModelSearch3D
directory = "./content/figures/"
# ---------------------------------------------------------------------------- #
#                                   DATA                                       #
# ---------------------------------------------------------------------------- #
#%%
# Data
X, y, coef = datasets.make_regression(n_samples=1000, n_features=1, bias=10,
                                noise=40, coef=True, random_state=50)
X_train, X_test, y_train, y_test = data_split(X,y, test_size=0.3, seed=50)
# ---------------------------------------------------------------------------- #
#                                SCATTER PLOT                                  #
# ---------------------------------------------------------------------------- #
#%%
# Linear Regression Scatterplot
data = go.Scatter(
    x=X_train.flatten(),
    y=y_train,
    mode='markers',
    marker=dict(color='steelblue')
)
layout = go.Layout(title='Simulated Data',
                   height=400,
                   width=800,
                   showlegend=False,
                   xaxis_title="X",
                   yaxis_title="Y",
                   margin=dict(l=10,r=10,t=40,b=10),
                   template='plotly_white'
                   )
fig = go.Figure(data=data, layout=layout)
fig.show()
po.plot(fig, filename = "./content/figures/simulated_training_data.html", auto_open=False)
# ---------------------------------------------------------------------------- #
#                             LINEAR REGRESSION                                #
# ---------------------------------------------------------------------------- #
#%%
# Linear Regression
scaler = StandardScaler()
scaler.fit(X_train)
X_train = scaler.transform(X_train)
lr = LinearRegression(epochs=1000, learning_rate=0.01, val_size=0.2, patience=40,
                      early_stop=True, metric='r2', verbose=True, checkpoint=100)
lr.fit(X_train,y_train)
print(lr.intercept_)
print(lr.coef_.shape)
# ---------------------------------------------------------------------------- #
#                                ANIMATIONS                                    #
# ---------------------------------------------------------------------------- #
#%%
# Animations
plot = SingleModelSearch3D()
plot.search(lr, directory=directory, filename="linear_regression_search_test.gif")
plot = SingleModelFit2D()
plot.fit(lr, directory=directory, filename="linear_regression_fit_test.gif")
#%%
# ---------------------------------------------------------------------------- #
#                                  TEST                                        #
# ---------------------------------------------------------------------------- #
scaler.fit(X_test)
X_test = scaler.transform(X_test)

# %%
# ---------------------------------------------------------------------------- #
#                   PREDICTED VS ACTUAL ON TEST DATA                           #
# ---------------------------------------------------------------------------- #
def f_actual(x):
    y = (10 + np.dot(coef, x)).flatten()
    return y
def f_pred(x):
    y = (lr.intercept_ + X.dot(lr.coef_)).flatten()
    return y
x = np.linspace(start=np.min(X_test), stop=np.max(X_test), num=100)
y_actual = f_actual(x)
y_pred = f_pred(x)
# ---------------------------------------------------------------------------- #
#%%
# Print actual vs predicted fit
data = [go.Scatter(x=X_test.flatten(),y=y_test,
                mode='markers',
                name='Test Data',
                marker=dict(color='steelblue')),
        go.Scatter(x=x, y=y_actual,
                mode='lines',
                name='Actual Fit'),
        go.Scatter(x=x, y=y_pred,
                mode='lines',
                name='Predicted Fit')]

layout = go.Layout(title='Actual vs Predicted Simulated Test Data',
                   height=400,
                   width=800,
                   showlegend=True,
                   xaxis_title="X",
                   yaxis_title="Y",
                   margin=dict(l=10,r=10,t=40,b=10),
                   template='plotly_white'
                   )
fig = go.Figure(data=data, layout=layout)
fig.show()
po.plot(fig, filename = "./content/figures/linear_regression_test_actual_v_predicted_.html", auto_open=False)

# ---------------------------------------------------------------------------- #
#                   PREDICTED VS ACTUAL R2 ON TEST DATA                        #
# ---------------------------------------------------------------------------- #
# %%
r2_pred = lr.score(X_test, y_test)
y_test_pred = f_actual(X_test)
r2_actual = R2()(y_test, y_test_pred)
# Print actual vs predicted R2
x = ['Actual Model', 'Learned Model']
y = [r2_actual, r2_pred]
data = [go.Bar(x=x,y=y,
                text=np.round(y,4),
                textposition='auto',
                showlegend=False,
                marker_color='rgb(55, 83, 109)')]

layout = go.Layout(title='R2 for Actual vs Predicted Simulated Test Data',
                   height=400,
                   width=800,
                   showlegend=True,
                   xaxis_title="X",
                   yaxis_title="Y",
                   margin=dict(l=10,r=10,t=40,b=10),
                   template='plotly_white'
                   )
fig = go.Figure(data=data, layout=layout)
fig.show()
po.plot(fig, filename = "./content/figures/linear_regression_test_actual_v_predicted_r2.html", auto_open=False)

# %%


from numpy.testing import assert_array_equal, run_module_suite
import numpy as np
const0 = np.zeros(20)
const1 = np.ones(20)
constNeg = np.ones(20)*(-1)
constF = np.ones(20) * 2.5
lin = np.arange(20)
lin0 = np.linspace(-10,10,20)
f = 5
sample = 1000
x = np.arange(0, sample, 1)
Fs = 1000
wave = np.sin(2 * np.pi * f * x / Fs)
np.random.seed(seed=10)
noiseWave = wave + np.random.normal(0,0.1,1000)
offsetWave = wave + 2
from TSFEL import*

#### STATISTICAL FEATURES ####
def test_skew():
    np.testing.assert_equal(skew(const0), 0.0)
    np.testing.assert_equal(skew(const1), 0.0)
    np.testing.assert_equal(skew(constNeg), 0.0)
    np.testing.assert_equal(skew(constF), 0.0)
    np.testing.assert_equal(skew(lin), 0)
    np.testing.assert_almost_equal(skew(lin0), -1.0167718723297815e-16, decimal=5)
    np.testing.assert_almost_equal(skew(wave), -2.009718347115232e-17, decimal=5)
    np.testing.assert_almost_equal(skew(offsetWave), 9.043732562018544e-16, decimal=5)
    np.testing.assert_almost_equal(skew(noiseWave), -0.0004854111290521465, decimal=5)

def test_kurtosis():
    np.testing.assert_equal(kurtosis(const0), -3)
    np.testing.assert_equal(kurtosis(const1), -3)
    np.testing.assert_equal(kurtosis(constNeg), -3)
    np.testing.assert_equal(kurtosis(constF), -3.0)
    np.testing.assert_almost_equal(kurtosis(lin), -1.206015037593985, decimal=2)
    np.testing.assert_almost_equal(kurtosis(lin0), -1.2060150375939847, decimal=2)
    np.testing.assert_almost_equal(kurtosis(wave), -1.501494077162359, decimal=2)
    np.testing.assert_almost_equal(kurtosis(offsetWave), -1.5014940771623597, decimal=2)
    np.testing.assert_almost_equal(kurtosis(noiseWave), -1.4606204906023366, decimal=2)

def test_mean():
    np.testing.assert_equal(np.mean(const0), 0.0)
    np.testing.assert_equal(np.mean(const1), 1.0)
    np.testing.assert_equal(np.mean(constNeg), -1.0)
    np.testing.assert_equal(np.mean(constF), 2.5)
    np.testing.assert_equal(np.mean(lin), 9.5)
    np.testing.assert_almost_equal(np.mean(lin0), -3.552713678800501e-16, decimal=5)
    np.testing.assert_almost_equal(np.mean(wave), 7.105427357601002e-18, decimal=5)
    np.testing.assert_almost_equal(np.mean(offsetWave), 2.0, decimal=5)
    np.testing.assert_almost_equal(np.mean(noiseWave), -0.0014556635615470554, decimal=5)

def test_median():
    np.testing.assert_equal(np.median(const0), 0.0)
    np.testing.assert_equal(np.median(const1), 1.0)
    np.testing.assert_equal(np.median(constNeg), -1.0)
    np.testing.assert_equal(np.median(constF), 2.5)
    np.testing.assert_equal(np.median(lin), 9.5)
    np.testing.assert_almost_equal(np.median(lin0), -3.552713678800501e-16, decimal=5)
    np.testing.assert_almost_equal(np.median(wave), 7.105427357601002e-18, decimal=5)
    np.testing.assert_almost_equal(np.median(offsetWave), 2.0, decimal=5)
    np.testing.assert_almost_equal(np.median(noiseWave), 0.013846093997438328, decimal=5)

def test_hist():
    np.testing.assert_equal(hist(const0, 10, 5), (0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0))
    np.testing.assert_equal(hist(const1, 10, 5), (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0))
    np.testing.assert_equal(hist(constNeg, 10, 5), (0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0))
    np.testing.assert_equal(hist(constF, 10, 5),  (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0))
    np.testing.assert_equal(hist(lin, 10, 5), (0.0, 0.0, 0.0, 0.0, 0.0, 0.16666666666666666, 0.16666666666666666, 0.16666666666666666, 0.16666666666666666, 0.3333333333333333))
    np.testing.assert_almost_equal(hist(wave, 10, 5), (0.   , 0.   , 0.   , 0.   , 0.499, 0.496, 0.005, 0.   , 0.   ,
       0.   ), decimal=5)
    np.testing.assert_almost_equal(hist(offsetWave, 10, 5), (0.   , 0.   , 0.   , 0.   , 0.   , 0.   , 0.499, 0.496, 0.005,
       0.  ), decimal=5)
    np.testing.assert_almost_equal(hist(noiseWave, 10, 5), (0.   , 0.   , 0.   , 0.048, 0.446, 0.45 , 0.056, 0.   , 0.   ,
       0.   ), decimal=5)

def test_max():
    np.testing.assert_equal(np.max(const0), 0.0)
    np.testing.assert_equal(np.max(const1), 1.0)
    np.testing.assert_equal(np.max(constNeg), -1.0)
    np.testing.assert_equal(np.max(constF), 2.5)
    np.testing.assert_equal(np.max(lin), 19)
    np.testing.assert_almost_equal(np.max(lin0), 10.0, decimal=5)
    np.testing.assert_almost_equal(np.max(wave), 1.0, decimal=5)
    np.testing.assert_almost_equal(np.max(noiseWave), 1.221757617217142, decimal=5)
    np.testing.assert_almost_equal(np.max(offsetWave), 3.0, decimal=5)

def test_min():
    np.testing.assert_equal(np.min(const0), 0.0)
    np.testing.assert_equal(np.min(const1), 1.0)
    np.testing.assert_equal(np.min(constNeg), -1.0)
    np.testing.assert_equal(np.min(constF), 2.5)
    np.testing.assert_equal(np.min(lin), 0)
    np.testing.assert_almost_equal(np.min(lin0), -10.0, decimal=5)
    np.testing.assert_almost_equal(np.min(wave), -1.0, decimal=5)
    np.testing.assert_almost_equal(np.min(noiseWave), -1.2582533627830566, decimal=5)
    np.testing.assert_almost_equal(np.min(offsetWave), 1.0, decimal=5)

def test_variance():
    np.testing.assert_equal(np.var(const0), 0.0)
    np.testing.assert_equal(np.var(const1), 0.0)
    np.testing.assert_equal(np.var(constNeg), 0.0)
    np.testing.assert_equal(np.var(constF), 0.0)
    np.testing.assert_equal(np.var(lin), 33.25)
    np.testing.assert_almost_equal(np.var(lin0), 36.84210526315789, decimal=5)
    np.testing.assert_almost_equal(np.var(wave), 0.5, decimal=5)
    np.testing.assert_almost_equal(np.var(offsetWave), 0.5, decimal=5)
    np.testing.assert_almost_equal(np.var(noiseWave), 0.5081167177369529, decimal=5)

def test_std():
    np.testing.assert_equal(np.std(const0), 0.0)
    np.testing.assert_equal(np.std(const1), 0.0)
    np.testing.assert_equal(np.std(constNeg), 0.0)
    np.testing.assert_equal(np.std(constF), 0.0)
    np.testing.assert_equal(np.std(lin), 5.766281297335398)
    np.testing.assert_almost_equal(np.std(lin0), 6.069769786668839, decimal=5)
    np.testing.assert_almost_equal(np.std(wave), 0.7071067811865476, decimal=5)
    np.testing.assert_almost_equal(np.std(offsetWave), 0.7071067811865476, decimal=5)
    np.testing.assert_almost_equal(np.std(noiseWave), 0.7128230620125536, decimal=5)

def test_int_range():
    np.testing.assert_equal(interq_range(const0), 0.0)
    np.testing.assert_equal(interq_range(const1), 0.0)
    np.testing.assert_equal(interq_range(constNeg), 0.0)
    np.testing.assert_equal(interq_range(constF), 0.0)
    np.testing.assert_equal(interq_range(lin), 9.5)
    np.testing.assert_almost_equal(interq_range(lin0), 10.0, decimal=5)
    np.testing.assert_almost_equal(interq_range(wave), 1.414213562373095, decimal=5)
    np.testing.assert_almost_equal(interq_range(offsetWave), 1.414213562373095, decimal=5)
    np.testing.assert_almost_equal(interq_range(noiseWave), 1.4277110228590328, decimal=5)

def test_calc_meanad():
    np.testing.assert_equal(calc_meanad(const0), 0.0)
    np.testing.assert_equal(calc_meanad(const1), 0.0)
    np.testing.assert_equal(calc_meanad(constNeg), 0.0)
    np.testing.assert_equal(calc_meanad(constF), 0.0)
    np.testing.assert_equal(calc_meanad(lin), 5.0)
    np.testing.assert_almost_equal(calc_meanad(lin0), 5.263157894736842, decimal=5)
    np.testing.assert_almost_equal(calc_meanad(wave), 0.6365674116287159, decimal=5)
    np.testing.assert_almost_equal(calc_meanad(noiseWave), 0.6392749078483896, decimal=5)
    np.testing.assert_almost_equal(calc_meanad(offsetWave), 0.6365674116287157, decimal=5)

def test_calc_medad():
    np.testing.assert_equal(calc_medad(const0), 0.0)
    np.testing.assert_equal(calc_medad(const1), 0.0)
    np.testing.assert_equal(calc_medad(constNeg), 0.0)
    np.testing.assert_equal(calc_medad(constF), 0.0)
    np.testing.assert_equal(calc_medad(lin), 5.0)
    np.testing.assert_almost_equal(calc_medad(lin0), 5.2631578947368425, decimal=5)
    np.testing.assert_almost_equal(calc_medad(wave), 0.7071067811865475, decimal=5)
    np.testing.assert_almost_equal(calc_medad(offsetWave), 0.7071067811865475, decimal=5)
    np.testing.assert_almost_equal(calc_medad(noiseWave), 0.7068117164205888, decimal=5)

def test_rms():
    np.testing.assert_equal(rms(const0), 0.0)
    np.testing.assert_equal(rms(const1), 1.0)
    np.testing.assert_equal(rms(constNeg), 1.0)
    np.testing.assert_equal(rms(constF), 2.5)
    np.testing.assert_equal(rms(lin), 11.113055385446435)
    np.testing.assert_almost_equal(rms(lin0), 6.06976978666884, decimal=5)
    np.testing.assert_almost_equal(rms(wave), 0.7071067811865476, decimal=5)
    np.testing.assert_almost_equal(rms(offsetWave), 2.1213203435596424, decimal=5)
    np.testing.assert_almost_equal(rms(noiseWave), 0.7128245483240299, decimal=5)

#### TEMPORAL FEATURES ####
def test_distance():
    np.testing.assert_equal(distance(const0), 19.0)
    np.testing.assert_equal(distance(const1), 19.0)
    np.testing.assert_equal(distance(constNeg), 19.0)
    np.testing.assert_equal(distance(constF), 19.0)
    np.testing.assert_equal(distance(lin), 26.87005768508881)
    np.testing.assert_almost_equal(distance(lin0), 27.586228448267438, decimal=5)
    np.testing.assert_almost_equal(distance(wave), 999.2461809866238, decimal=5)
    np.testing.assert_almost_equal(distance(offsetWave), 999.2461809866238, decimal=5)
    np.testing.assert_almost_equal(distance(noiseWave), 1007.8711901383033, decimal=5)

def test_minpeaks():
    np.testing.assert_equal(minpeaks(const0), 0.0)
    np.testing.assert_equal(minpeaks(const1), 0.0)
    np.testing.assert_equal(minpeaks(constNeg), 0.0)
    np.testing.assert_equal(minpeaks(constF), 0.0)
    np.testing.assert_equal(minpeaks(lin), 0.0)
    np.testing.assert_almost_equal(minpeaks(lin0), 0.0, decimal=5)
    np.testing.assert_almost_equal(minpeaks(wave), 5, decimal=5)
    np.testing.assert_almost_equal(minpeaks(offsetWave), 5, decimal=5)
    np.testing.assert_almost_equal(minpeaks(noiseWave), 323, decimal=5)

def test_maxpeaks():
    np.testing.assert_equal(maxpeaks(const0), 0.0)
    np.testing.assert_equal(maxpeaks(const1), 0.0)
    np.testing.assert_equal(maxpeaks(constNeg), 0.0)
    np.testing.assert_equal(maxpeaks(constF), 0.0)
    np.testing.assert_equal(maxpeaks(lin), 0.0)
    np.testing.assert_almost_equal(maxpeaks(lin0), 0.0, decimal=5)
    np.testing.assert_almost_equal(maxpeaks(wave), 5, decimal=5)
    np.testing.assert_almost_equal(maxpeaks(offsetWave), 5, decimal=5)
    np.testing.assert_almost_equal(maxpeaks(noiseWave), 322, decimal=5)

def test_centroid():
    np.testing.assert_equal(centroid(const0, Fs), 0.0)
    np.testing.assert_equal(centroid(const1, Fs), 0.009)
    np.testing.assert_equal(centroid(constNeg, Fs), 0.009)
    np.testing.assert_equal(centroid(constF, Fs), 0.009)
    np.testing.assert_equal(centroid(lin, Fs), 0.013571428571428573)
    np.testing.assert_almost_equal(centroid(lin0, Fs), 0.00708955223880597, decimal=5)
    np.testing.assert_almost_equal(centroid(wave, Fs), 0.45, decimal=5)
    np.testing.assert_almost_equal(centroid(offsetWave, Fs), 0.44999999999999996, decimal=5)
    np.testing.assert_almost_equal(centroid(noiseWave, Fs), 0.4494737253547235, decimal=5)

def test_calc_meandiff():
    np.testing.assert_equal(calc_meandiff(const0), 0.0)
    np.testing.assert_equal(calc_meandiff(const1), 0.0)
    np.testing.assert_equal(calc_meandiff(constNeg), 0.0)
    np.testing.assert_equal(calc_meandiff(constF), 0.0)
    np.testing.assert_equal(calc_meandiff(lin), 1.0)
    np.testing.assert_almost_equal(calc_meandiff(lin0), 1.0526315789473684, decimal=5)
    np.testing.assert_almost_equal(calc_meandiff(wave), -3.1442201279407477e-05, decimal=5)
    np.testing.assert_almost_equal(calc_meandiff(offsetWave), -3.1442201279407036e-05, decimal=5)
    np.testing.assert_almost_equal(calc_meandiff(noiseWave), -0.00010042477181949707, decimal=5)

def test_calc_meddiff():
    np.testing.assert_equal(calc_meddiff(const0), 0.0)
    np.testing.assert_equal(calc_meddiff(const1), 0.0)
    np.testing.assert_equal(calc_meddiff(constNeg), 0.0)
    np.testing.assert_equal(calc_meddiff(constF), 0.0)
    np.testing.assert_equal(calc_meddiff(lin), 1.0)
    np.testing.assert_almost_equal(calc_meddiff(lin0), 1.0526315789473684, decimal=5)
    np.testing.assert_almost_equal(calc_meddiff(wave), -0.0004934396342684, decimal=5)
    np.testing.assert_almost_equal(calc_meddiff(offsetWave), -0.0004934396342681779, decimal=5)
    np.testing.assert_almost_equal(calc_meddiff(noiseWave), -0.004174819648320949, decimal=5)

def test_calc_meanadiff():
    np.testing.assert_equal(calc_meanadiff(const0), 0.0)
    np.testing.assert_equal(calc_meanadiff(const1), 0.0)
    np.testing.assert_equal(calc_meanadiff(constNeg), 0.0)
    np.testing.assert_equal(calc_meanadiff(constF), 0.0)
    np.testing.assert_equal(calc_meanadiff(lin), 1.0)
    np.testing.assert_almost_equal(calc_meanadiff(lin0), 1.0526315789473684, decimal=5)
    np.testing.assert_almost_equal(calc_meanadiff(wave), 0.019988577818740614, decimal=5)
    np.testing.assert_almost_equal(calc_meanadiff(offsetWave), 0.019988577818740614, decimal=5)
    np.testing.assert_almost_equal(calc_meanadiff(noiseWave), 0.10700252903161508, decimal=5)

def test_calc_medadiff():
    np.testing.assert_equal(calc_medadiff(const0), 0.0)
    np.testing.assert_equal(calc_medadiff(const1), 0.0)
    np.testing.assert_equal(calc_medadiff(constNeg), 0.0)
    np.testing.assert_equal(calc_medadiff(constF), 0.0)
    np.testing.assert_equal(calc_medadiff(lin), 1.0)
    np.testing.assert_almost_equal(calc_medadiff(lin0), 1.0526315789473681, decimal=5)
    np.testing.assert_almost_equal(calc_medadiff(wave), 0.0218618462348652, decimal=5)
    np.testing.assert_almost_equal(calc_medadiff(offsetWave), 0.021861846234865645, decimal=5)
    np.testing.assert_almost_equal(calc_medadiff(noiseWave), 0.08958750592592835, decimal=5)

def test_calc_sadiff():
    np.testing.assert_equal(calc_sadiff(const0), 0.0)
    np.testing.assert_equal(calc_sadiff(const1), 0.0)
    np.testing.assert_equal(calc_sadiff(constNeg), 0.0)
    np.testing.assert_equal(calc_sadiff(constF), 0.0)
    np.testing.assert_equal(calc_sadiff(lin), 19)
    np.testing.assert_almost_equal(calc_sadiff(lin0), 20.0, decimal=5)
    np.testing.assert_almost_equal(calc_sadiff(wave), 19.968589240921872, decimal=5)
    np.testing.assert_almost_equal(calc_sadiff(offsetWave), 19.968589240921872, decimal=5)
    np.testing.assert_almost_equal(calc_sadiff(noiseWave), 106.89552650258346, decimal=5)

def test_zeroCross():
    np.testing.assert_equal(zero_cross(const0), 0.0)
    np.testing.assert_equal(zero_cross(const1), 0.0)
    np.testing.assert_equal(zero_cross(constNeg), 0.0)
    np.testing.assert_equal(zero_cross(constF), 0.0)
    np.testing.assert_equal(zero_cross(lin), 1.0)
    np.testing.assert_almost_equal(zero_cross(lin0), 1.0, decimal=5)
    np.testing.assert_almost_equal(zero_cross(wave), 10, decimal=5)
    np.testing.assert_almost_equal(zero_cross(offsetWave), 0.0, decimal=5)
    np.testing.assert_almost_equal(zero_cross(noiseWave), 38, decimal=5)

def test_autocorr():
    np.testing.assert_equal(autocorr(const0), 0.0)
    np.testing.assert_equal(autocorr(const1), 20.0)
    np.testing.assert_equal(autocorr(constNeg), 20.0)
    np.testing.assert_equal(autocorr(constF), 125.0)
    np.testing.assert_equal(autocorr(lin), 2470.0)
    np.testing.assert_almost_equal(autocorr(lin0), 736.8421052631579, decimal=0)
    np.testing.assert_almost_equal(autocorr(wave), 500.5, decimal=0)
    np.testing.assert_almost_equal(autocorr(offsetWave), 4500.0, decimal=0)
    np.testing.assert_almost_equal(autocorr(noiseWave), 508.6149018530489, decimal=0)

#### TEMPORAL FEATURES ####
def test_max_fre():
    np.testing.assert_equal(max_frequency(const0, Fs), 55.55555555555556)
    np.testing.assert_equal(max_frequency(const1, Fs), 444.44444444444446)
    np.testing.assert_equal(max_frequency(constNeg, Fs), 444.44444444444446)
    np.testing.assert_equal(max_frequency(constF, Fs), 444.44444444444446)
    np.testing.assert_equal(max_frequency(lin, Fs), 500.0)
    np.testing.assert_almost_equal(max_frequency(lin0, Fs), 500.0, decimal=5)
    np.testing.assert_almost_equal(max_frequency(wave, Fs), 5.0100200400801596, decimal=5)
    np.testing.assert_almost_equal(max_frequency(offsetWave, Fs), 5.0100200400801596, decimal=5)
    np.testing.assert_almost_equal(max_frequency(noiseWave, Fs), 464.9298597194388, decimal=5)
    np.testing.assert_almost_equal(max_frequency(x, Fs),  390.7815631262525, decimal=1)

def test_med_fre():
    np.testing.assert_equal(median_frequency(const0, Fs), 55.55555555555556)
    np.testing.assert_equal(median_frequency(const1, Fs), 222.22222222222223)
    np.testing.assert_equal(median_frequency(constNeg, Fs), 222.22222222222223)
    np.testing.assert_equal(median_frequency(constF, Fs), 222.22222222222223)
    np.testing.assert_equal(median_frequency(lin, Fs), 166.66666666666669)
    np.testing.assert_almost_equal(median_frequency(lin0, Fs), 166.66666666666669, decimal=5)
    np.testing.assert_almost_equal(median_frequency(wave, Fs), 5.0100200400801596, decimal=5)
    np.testing.assert_almost_equal(median_frequency(offsetWave, Fs), 5.0100200400801596, decimal=5)
    np.testing.assert_almost_equal(median_frequency(noiseWave, Fs), 146.29258517034066, decimal=5)
    np.testing.assert_almost_equal(median_frequency(x, Fs), 19.03807615230461, decimal=1)

def test_fund_fre():
    np.testing.assert_equal(fundamental_frequency(const0, 1), 0.0)
    np.testing.assert_equal(fundamental_frequency(const1, 1), 0.0)
    np.testing.assert_equal(fundamental_frequency(constNeg, Fs), 0.0)
    np.testing.assert_equal(fundamental_frequency(constF, Fs), 0.0)
    np.testing.assert_equal(fundamental_frequency(lin, Fs), 0.0)
    np.testing.assert_almost_equal(fundamental_frequency(lin0, Fs), 0.0, decimal=5)
    np.testing.assert_almost_equal(fundamental_frequency(wave, Fs), 4.008016032064128, decimal=1)
    np.testing.assert_almost_equal(fundamental_frequency(offsetWave, Fs), 3.006012024048096, decimal=1)
    np.testing.assert_almost_equal(fundamental_frequency(noiseWave, Fs), 39.07815631262525, decimal=1)

def test_power_spec():
    np.testing.assert_equal(max_power_spectrum(const0, Fs), 0.0)
    np.testing.assert_equal(max_power_spectrum(const1, Fs), 0.0)
    np.testing.assert_equal(max_power_spectrum(constNeg, Fs), 0.0)
    np.testing.assert_equal(max_power_spectrum(constF, Fs), 0.0)
    np.testing.assert_equal(max_power_spectrum(lin, Fs), 0.004621506382612649)
    np.testing.assert_almost_equal(max_power_spectrum(lin0, Fs), 0.0046215063826126525, decimal=5)
    np.testing.assert_almost_equal(max_power_spectrum(wave, Fs), 0.6666666666666667, decimal=5)
    np.testing.assert_almost_equal(max_power_spectrum(offsetWave, Fs), 0.6666666666666667, decimal=5)
    np.testing.assert_almost_equal(max_power_spectrum(noiseWave, Fs),0.6570878541643916, decimal=5)

def test_total_energy():
    np.testing.assert_equal(total_energy(const0, Fs), 0.0)
    np.testing.assert_equal(total_energy(const1, Fs), 1052.6315789473686)
    np.testing.assert_equal(total_energy(constNeg, Fs), 1052.6315789473686)
    np.testing.assert_equal(total_energy(constF, Fs), 6578.9473684210525)
    np.testing.assert_equal(total_energy(lin, Fs), 130000.0)
    np.testing.assert_almost_equal(total_energy(lin0, Fs), 38781.16343490305, decimal=5)
    np.testing.assert_almost_equal(total_energy(wave, Fs), 500.5005005005005, decimal=5)
    np.testing.assert_almost_equal(total_energy(offsetWave, Fs), 4504.504504504504, decimal=5)
    np.testing.assert_almost_equal(total_energy(noiseWave, Fs), 508.6274641575148, decimal=5)

def test_spectral_centroid():
    np.testing.assert_equal(spectral_centroid(const0, Fs), 0.0)
    np.testing.assert_equal(spectral_centroid(const1, Fs), 232.87645373386948)
    np.testing.assert_equal(spectral_centroid(constNeg, Fs), 232.87645373386948)
    np.testing.assert_equal(spectral_centroid(constF, Fs), 231.25899007425215)
    np.testing.assert_equal(spectral_centroid(lin, Fs), 189.72282595943128)
    np.testing.assert_almost_equal(spectral_centroid(lin0, Fs), 189.7228259594313, decimal=5)
    np.testing.assert_almost_equal(spectral_centroid(wave, Fs), 5.010020040084022, decimal=5)
    np.testing.assert_almost_equal(spectral_centroid(offsetWave, Fs), 5.010020040084512, decimal=5)
    np.testing.assert_almost_equal(spectral_centroid(noiseWave, Fs), 181.26672828206048, decimal=5)

def test_spectral_spread():
    np.testing.assert_almost_equal(spectral_spread(const0, Fs), 0.0, decimal=5)
    np.testing.assert_almost_equal(spectral_spread(const1, Fs), 12780.633762829042, decimal=5)
    np.testing.assert_almost_equal(spectral_spread(constNeg, Fs), 12780.633762829042, decimal=5)
    np.testing.assert_almost_equal(spectral_spread(constF, Fs), 13180.136854206703, decimal=5)
    np.testing.assert_almost_equal(spectral_spread(lin, Fs), 19861.962160017465, decimal=5)
    np.testing.assert_almost_equal(spectral_spread(lin0, Fs), 19861.962160017472, decimal=5)
    np.testing.assert_almost_equal(spectral_spread(wave, Fs), 0.0, decimal=5)
    np.testing.assert_almost_equal(spectral_spread(offsetWave, Fs), 1.4368545161461419e-09, decimal=5)
    np.testing.assert_almost_equal(spectral_spread(noiseWave, Fs), 27432.31810321154, decimal=5)

def test_spectral_skewness():
    np.testing.assert_almost_equal(spectral_skewness(const0, Fs), 0.0, decimal=5)
    np.testing.assert_almost_equal(spectral_skewness(const1, Fs), 0.47000214356667996, decimal=5)
    np.testing.assert_almost_equal(spectral_skewness(constNeg, Fs), 0.47000214356667996, decimal=5)
    np.testing.assert_almost_equal(spectral_skewness(constF, Fs), 0.5161486069831996, decimal=5)
    np.testing.assert_almost_equal(spectral_skewness(lin, Fs), 0.8140329168647059, decimal=5)
    np.testing.assert_almost_equal(spectral_skewness(lin0, Fs), 0.8140329168647046, decimal=5)
    np.testing.assert_almost_equal(spectral_skewness(wave, Fs), 10624606.203489874, decimal=1)
    np.testing.assert_almost_equal(spectral_skewness(offsetWave, Fs), 9940161.672715995, decimal=1)
    np.testing.assert_almost_equal(spectral_skewness(noiseWave, Fs), 0.41151531285467985, decimal=1)

def test_spectral_kurtosis():
    np.testing.assert_almost_equal(spectral_kurtosis(const0, Fs), 0.0, decimal=5)
    np.testing.assert_almost_equal(spectral_kurtosis(const1, Fs), 2.046321609772459, decimal=5)
    np.testing.assert_almost_equal(spectral_kurtosis(constNeg, Fs), 2.046321609772459, decimal=5)
    np.testing.assert_almost_equal(spectral_kurtosis(constF, Fs), 2.0107996983500787, decimal=5)
    np.testing.assert_almost_equal(spectral_kurtosis(lin, Fs), 2.0107996983500787, decimal=0)
    np.testing.assert_almost_equal(spectral_kurtosis(lin0, Fs), 2.4060168768515413, decimal=5)
    np.testing.assert_almost_equal(spectral_kurtosis(wave, Fs),  120631070963747.72, decimal=1)
    np.testing.assert_almost_equal(spectral_kurtosis(offsetWave, Fs), 105458851126694.56, decimal=5)
    np.testing.assert_almost_equal(spectral_kurtosis(noiseWave, Fs), 1.7244872082677942, decimal=5)

def test_spectral_slope():
    np.testing.assert_equal(spectral_slope(const0, Fs), 0.0)
    np.testing.assert_equal(spectral_slope(const1, Fs), -5.721690608887972e-19)
    np.testing.assert_equal(spectral_slope(constNeg, Fs), -5.721690608887972e-19)
    np.testing.assert_equal(spectral_slope(constF, Fs), -1.330865622866306e-18)
    np.testing.assert_equal(spectral_slope(lin, Fs), -0.09209918800721191)
    np.testing.assert_almost_equal(spectral_slope(lin0, Fs), -0.09209918800721191, decimal=1)
    np.testing.assert_almost_equal(spectral_slope(wave, Fs), -0.011807228915662668, decimal=5)
    np.testing.assert_almost_equal(spectral_slope(offsetWave, Fs), -0.011807228915662668, decimal=5)
    np.testing.assert_almost_equal(spectral_slope(noiseWave, Fs), -0.011996993466635138, decimal=5)

def test_spectral_decrease():
    np.testing.assert_equal(spectral_decrease(const0, Fs), 0.0)
    np.testing.assert_equal(spectral_decrease(const1, Fs), -1.092141531796704)
    np.testing.assert_equal(spectral_decrease(constNeg, Fs), -1.092141531796704)
    np.testing.assert_equal(spectral_decrease(constF, Fs), -1.0677129498080251)
    np.testing.assert_equal(spectral_decrease(lin, Fs), -0.41194423707863714)
    np.testing.assert_almost_equal(spectral_decrease(lin0, Fs), -0.4119442370786371, decimal=5)
    np.testing.assert_almost_equal(spectral_decrease(wave, Fs), 0.333333333333328, decimal=5)
    np.testing.assert_almost_equal(spectral_decrease(offsetWave, Fs), 0.3333333333333272, decimal=5)
    np.testing.assert_almost_equal(spectral_decrease(noiseWave, Fs), 0.09171304789096776, decimal=5)

def test_spectral_roll_on():
    np.testing.assert_equal(spectral_roll_on(const0, Fs), 55.55555555555556)
    np.testing.assert_equal(spectral_roll_on(const1, Fs), 111.11111111111111)
    np.testing.assert_equal(spectral_roll_on(constNeg, Fs), 111.11111111111111)
    np.testing.assert_equal(spectral_roll_on(constF, Fs), 111.11111111111111)
    np.testing.assert_equal(spectral_roll_on(lin, Fs), 55.55555555555556)
    np.testing.assert_almost_equal(spectral_roll_on(lin0, Fs), 55.55555555555556, decimal=5)
    np.testing.assert_almost_equal(spectral_roll_on(wave, Fs), 5.0100200400801596, decimal=5)
    np.testing.assert_almost_equal(spectral_roll_on(offsetWave, Fs), 5.0100200400801596, decimal=5)
    np.testing.assert_almost_equal(spectral_roll_on(noiseWave, Fs), 5.0100200400801596, decimal=5)

def test_spectral_roll_off():
    np.testing.assert_equal(spectral_roll_off(const0, Fs), 55.55555555555556)
    np.testing.assert_equal(spectral_roll_off(const1, Fs), 444.44444444444446)
    np.testing.assert_equal(spectral_roll_off(constNeg, Fs), 444.44444444444446)
    np.testing.assert_equal(spectral_roll_off(constF, Fs), 444.44444444444446)
    np.testing.assert_equal(spectral_roll_off(lin, Fs), 500.0)
    np.testing.assert_almost_equal(spectral_roll_off(lin0, Fs), 500.0, decimal=5)
    np.testing.assert_almost_equal(spectral_roll_off(wave, Fs), 5.0100200400801596, decimal=5)
    np.testing.assert_almost_equal(spectral_roll_off(offsetWave, Fs), 5.0100200400801596, decimal=5)
    np.testing.assert_almost_equal(spectral_roll_off(noiseWave, Fs), 464.9298597194388, decimal=5)

def test_spectral_curve_distance():
    np.testing.assert_equal(curve_distance(const0, Fs), 0.0)
    np.testing.assert_equal(curve_distance(const1, Fs), -3.0871200587084623e-15)
    np.testing.assert_equal(curve_distance(constNeg, Fs), -3.0871200587084623e-15)
    np.testing.assert_equal(curve_distance(constF, Fs), -7.0852193832076304e-15)
    np.testing.assert_equal(curve_distance(lin, Fs), -403.8425293964846)
    np.testing.assert_almost_equal(curve_distance(lin0, Fs), -425.0973993647206, decimal=5)
    np.testing.assert_almost_equal(curve_distance(wave, Fs), -122750.0000000002, decimal=5)
    np.testing.assert_almost_equal(curve_distance(offsetWave, Fs), -122750.0000000002, decimal=5)
    np.testing.assert_almost_equal(curve_distance(noiseWave, Fs), -125369.50449385124, decimal=5)

def test_spect_variation():
    np.testing.assert_equal(spect_variation(const0, Fs), 1.0)
    np.testing.assert_equal(spect_variation(const1, Fs), 1.0)
    np.testing.assert_equal(spect_variation(constNeg, Fs), 1.0)
    np.testing.assert_equal(spect_variation(constF, Fs), 0.9872320416232951)
    np.testing.assert_equal(spect_variation(lin, Fs), 0.7225499818571235)
    np.testing.assert_almost_equal(spect_variation(lin0, Fs), 0.7225499818571237, decimal=5)
    np.testing.assert_almost_equal(spect_variation(wave, Fs), 1.0, decimal=5)
    np.testing.assert_almost_equal(spect_variation(offsetWave, Fs), 1.0, decimal=5)
    np.testing.assert_almost_equal(spect_variation(noiseWave, Fs), 0.9961527995674907, decimal=5)

def test_spectral_maxpeaks():
    np.testing.assert_equal(spectral_maxpeaks(const0, Fs), 0.0)
    np.testing.assert_equal(spectral_maxpeaks(const1, Fs), 4.0)
    np.testing.assert_equal(spectral_maxpeaks(constNeg, Fs), 4.0)
    np.testing.assert_equal(spectral_maxpeaks(constF, Fs), 4.0)
    np.testing.assert_equal(spectral_maxpeaks(lin, Fs), 0.0)
    np.testing.assert_almost_equal(spectral_maxpeaks(lin0, Fs), 0.0, decimal=5)
    np.testing.assert_almost_equal(spectral_maxpeaks(wave, Fs), 163, decimal=0)
    np.testing.assert_almost_equal(spectral_maxpeaks(offsetWave, Fs), 160, decimal=1)
    np.testing.assert_almost_equal(spectral_maxpeaks(noiseWave, Fs), 171.0, decimal=1)

#if __name__ == "__main__":
    #run_module_suite()
run_module_suite()


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

"""
Example of running Nestle to fit the parameters of a straight line.
"""

from __future__ import print_function, division

import os
import sys
from scipy.special import ndtri
import numpy as np

# import Nestle
import nestle

# import model and data
from createdata import *

LN2PI = np.log(2.*np.pi)
LNSIGMA = np.log(sigma)

def prior_transform(theta):
    """
    A function defining the tranform between the parameterisation in the unit hypercube
    to the true parameters.

    Args:
        theta (tuple): a tuple containing the parameters.

    Returns:
        tuple: a new tuple or array with the transformed parameters.
    """

    mprime, cprime = theta # unpack the parameters (in their unit hypercube form)

    cmin = -10.  # lower bound on uniform prior on c
    cmax = 10.   # upper bound on uniform prior on c

    mmu = 0.     # mean of Gaussian prior on m
    msigma = 10. # standard deviation of Gaussian prior on m

    m = mmu + msigma*ndtri(mprime) # convert back to m
    c = cprime*(cmax-cmin) + cmin  # convert back to c

    return (m, c)

def loglikelihood_nestle(theta):
    """
    The log-likelihood function.
    """

    m, c = theta # unpack the parameters

    # normalisation
    norm = -0.5*M*LN2PI - M*LNSIGMA

    # chi-squared (data, sigma and x are global variables defined early on in this notebook)
    chisq = np.sum(((data-straight_line(x, m, c))/sigma)**2)

    return norm - 0.5*chisq

nlive = 1024     # number of live points
method = 'multi' # use MutliNest algorithm
ndims = 2        # two parameters
tol= 0.5         # the stopping criterion (this is the nestle default, so doesn't need to be set)

res = nestle.sample(loglikelihood_nestle, prior_transform, ndims, method=method, npoints=nlive, dlogz=tol)

logZnestle = res.logz                         # value of logZ
infogainnestle = res.h                        # value of the information gain in nats
logZerrnestle = np.sqrt(infogainnestle/nlive) # estimate of the statistcal uncertainty on logZ

# output marginal likelihood
print('Marginalised evidence is {} ± {}'.format(logZnestle, logZerrnestle))

# re-scale weights to have a maximum of one
nweights = res.weights/np.max(res.weights)

# get the probability of keeping a sample from the weights
keepidx = np.where(np.random.rand(len(nweights)) < nweights)[0]

# get the posterior samples
postsamples = res.samples[keepidx,:]

print('Number of posterior samples is {}'.format(postsamples.shape[0]))

# plot posterior samples (if corner.py is installed)
try:
    import matplotlib as mpl
    mpl.use("Agg") # force Matplotlib backend to Agg
    import corner # import corner.py
except ImportError:
    sys.exit(1)

fig = corner.corner(postsamples, labels=[r"$m$", r"$c$"], truths=[m, c])
fig.savefig('Nestle.png')


import math

import scipy.stats
import torch

from uncertify.utils.custom_types import Tensor

from typing import Union

def normalize_to_0_1(tensor: Tensor, min_val: Union[float, torch.tensor] = None,
                     max_val: Union[float, torch.tensor] = None) -> Tensor:
    """Takes a pytorch tensor and normalizes the values to the range [0, 1], i.e. largest value gets 1, smallest 0."""
    if min_val is None:
        min_val = tensor.min()
    if max_val is None:
        max_val = tensor.max()
    return (tensor - min_val) / (max_val - min_val)

def get_mean_and_std(tensor: Tensor) -> Tensor:
    return torch.mean(tensor), torch.std(tensor)

def get_min_and_max(tensor: Tensor) -> Tensor:
    return torch.min(tensor), torch.max(tensor)

def print_scipy_stats_description(description_result: scipy.stats.stats.DescribeResult,
                                  name: str = None, column_spacer: int = 25) -> None:
    name_part = f'Name: {name}' if name is not None else ''
    min_part = f'min: {description_result.minmax[0]:.2f}'
    max_part = f'max: {description_result.minmax[1]:.2f}'
    mean_part = f'mean: {description_result.mean:.2f}'
    std_part = f'std: {math.sqrt(description_result.variance):.2f}'

    print(f''.join([f'{part:{column_spacer}}' for part in [name_part, min_part, max_part, mean_part, std_part]]))


import numpy as np
from sklearn.datasets import make_blobs

def _generate_2d_rotation(theta=0):
    R = np.array([[np.cos(theta), np.sin(theta)], [-np.sin(theta), np.cos(theta)]])

    return R

def generate_gaussian_parity(
    n_samples,
    centers=None,
    class_label=None,
    cluster_std=0.25,
    angle_params=None,
    random_state=None,
):
    """
    Generate 2-dimensional Gaussian XOR distribution.
    (Classic XOR problem but each point is the
    center of a Gaussian blob distribution)
    Parameters
    ----------
    n_samples : int
        Total number of points divided among the four
        clusters with equal probability.
    centers : array of shape [n_centers,2], optional (default=None)
        The coordinates of the ceneter of total n_centers blobs.
    class_label : array of shape [n_centers], optional (default=None)
        class label for each blob.
    cluster_std : float, optional (default=1)
        The standard deviation of the blobs.
    angle_params: float, optional (default=None)
        Number of radians to rotate the distribution by.
    random_state : int, RandomState instance, default=None
        Determines random number generation for dataset creation. Pass an int
        for reproducible output across multiple function calls.
    Returns
    -------
    X : array of shape [n_samples, 2]
        The generated samples.
    y : array of shape [n_samples]
        The integer labels for cluster membership of each sample.
    """

    if random_state != None:
        np.random.seed(random_state)

    if centers == None:
        centers = np.array([(-0.5, 0.5), (0.5, 0.5), (-0.5, -0.5), (0.5, -0.5)])

    if class_label == None:
        class_label = [0, 1, 1, 0]

    blob_num = len(class_label)

    # get the number of samples in each blob with equal probability
    samples_per_blob = np.random.multinomial(
        n_samples, 1 / blob_num * np.ones(blob_num)
    )

    X, y = make_blobs(
        n_samples=samples_per_blob,
        n_features=2,
        centers=centers,
        cluster_std=cluster_std,
    )

    for blob in range(blob_num):
        y[np.where(y == blob)] = class_label[blob]

    if angle_params != None:
        R = _generate_2d_rotation(angle_params)
        X = X @ R

    return X, y.astype(int)


import numpy as np

import openmdao.api as om

class CannonballSizeComp(om.ExplicitComponent):
    """
    Compute the reference area and mass of a cannonball with a given radius and density.

    Notes
    -----
    This component is not vectorized with 'num_nodes' as is the usual way with Dymos, but is instead
    intended to compute a scalar mass and reference area from scalar radius and density inputs. This
    component does not reside in the ODE but instead its outputs are connected to the trajectory via
    input design parameters.
    """
    def setup(self):
        self.add_input(name='radius', val=1.0, desc='cannonball radius', units='m')
        self.add_input(name='dens', val=7870., desc='cannonball density', units='kg/m**3')

        self.add_output(name='mass', shape=(1,), desc='cannonball mass', units='kg')
        self.add_output(name='S', shape=(1,), desc='aerodynamic reference area', units='m**2')

        self.declare_partials(of='mass', wrt='dens')
        self.declare_partials(of='mass', wrt='radius')

        self.declare_partials(of='S', wrt='radius')

    def compute(self, inputs, outputs):
        radius = inputs['radius']
        dens = inputs['dens']

        outputs['mass'] = (4/3.) * dens * np.pi * radius ** 3
        outputs['S'] = np.pi * radius ** 2

    def compute_partials(self, inputs, partials):
        radius = inputs['radius']
        dens = inputs['dens']

        partials['mass', 'dens'] = (4/3.) * np.pi * radius ** 3
        partials['mass', 'radius'] = 4. * dens * np.pi * radius ** 2

        partials['S', 'radius'] = 2 * np.pi * radius


import os
import numpy as np
import math
from scipy import ndimage
import nibabel as nib
from PIL import Image
from scipy import signal

'''
Function to generate derivatives of Gaussian kernels, in either 1D, 2D, or 3D.
Source code in MATLAB obtained from Qiyuan Tian, Stanford University, September 2015
Edited to work in Python by Tony.
'''

def doggen(sigma):
    halfsize = np.ceil(3 * np.max(sigma))
    x = range(np.single(-halfsize), np.single(halfsize + 1));  # Python colon is not inclusive at end, while MATLAB is.
    dim = len(sigma);

    if dim == 1:
        X = np.array(x);  # Remember that, by default, numpy arrays are elementwise multiplicative
        k = -X * np.exp(-X**2/(2 * sigma**2));

    elif dim == 2:
        [X, Y] = np.meshgrid(x, x);
        k = -X * np.exp(-X**2/(2*sigma[0]^2) * np.exp(-Y**2))

    elif dim == 3:
        [X, Y, Z] = np.meshgrid(x, x, x);
        X = X.transpose(0, 2, 1);  # Obtained through vigorous testing (see below...)
        Y = Y.transpose(2, 0, 1);
        Z = Z.transpose(2, 1, 0);

        X = X.astype(float);
        Y = Y.astype(float);
        Z = Z.astype(float);
        k = -X * np.exp(np.divide(-np.power(X, 2), 2 * np.power(sigma[0], 2))) * np.exp(np.divide(-np.power(Y,2), 2 * np.power(sigma[1],2))) * np.exp(np.divide(-np.power(Z,2), 2 * np.power(sigma[2],2)))

    else:
        print 'Only supports up to 3 dimensions'

    return np.divide(k, np.sum(np.abs(k[:])));

'''
Function to generate Gaussian kernels, in 1D, 2D and 3D.
Source code in MATLAB obtained from Qiyuan Tian, Stanford University, September 2015
Edited to work in Python by Tony.
'''

def gaussgen(sigma):
    halfsize = np.ceil(3 * max(sigma));
    x = range(np.single(-halfsize), np.single(halfsize + 1));

    dim = len(sigma);

    if dim == 1:
        x = x.astype(float);
        k = np.exp(-x**2 / (2 * sigma^2));

    elif dim == 2:
        [X, Y] = np.meshgrid(x, x);
        X = X.astype(float);
        Y = Y.astype(float);

        k = np.exp(-X**2 / (2 * sigma[0]**2)) * np.exp(-Y**2 / (2 * sigma[1]**2));

    elif dim == 3:
        [X, Y, Z] = np.meshgrid(x, x, x);
        X = X.transpose(0, 2, 1);  # Obtained through vigorous testing (see below...)
        Y = Y.transpose(2, 0, 1);
        Z = Z.transpose(2, 1, 0);

        X = X.astype(float);  # WHY PYTHON?
        Y = Y.astype(float);
        Z = Z.astype(float);
        k = np.exp(-X**2 / (2 * sigma[0]**2)) * np.exp(-Y**2 / (2 * sigma[1]**2)) * np.exp(-Z**2 / (2 * sigma[2]**2));

    else:
        print 'Only supports up to dimension 3'

    return np.divide(k, np.sum(np.abs(k)));

'''
Function takes a single image (TIFF, or other also works), and returns
the single image as a numpy array.  Called by tiff_stack_to_array.
'''
def tiff_to_array(input_path):
    im = Image.open(input_path)
    # im.show()

    imarray = np.array(im)
    # print(imarray)

    return imarray

'''
Function takes input_path, which should should lead to a directory.
Loads all TIFFs in input_path, then generates numpy arrays from the
TIFF stack by calling tiff_to_array helper function.
'''
def tiff_stack_to_array(input_path):
    im_list = [];
    for filename in os.listdir(input_path):
        if filename.endswith(".tiff"):
            # print(os.path.join(directory, filename))
            im_arr = tiff_to_array(filename)
            im_list.append(im_arr)

    s = np.stack(im_list, axis=2)
    print s.shape
    return s

'''
Function takes a numpy array (from TIFF_stack_to_array) and saves output
FSL structure tensor as filename string. Allows inputting alternate dogsigmaArr,
gausigmaArr, angleArr, although defaults to currently to parameters from MATLAB script.
Also returns tensorfsl (the tensor fsl structure) image numpy array.

## Parameters (the script loops through all parameters and saves each result automatically)
# dogsigmaArr = [1]; Sigma values for derivative of gaussian filter, recommended value: 0.6 - 1.3 (based on actual data)
# gausigmaArr = [2.3]; Sigma values for gaussian filter, recommended value: 1.3 - 2.3 (based on actual data)
# angleArr = [25]; Angle thresholds for fiber tracking, recommended value: 20 - 30.

Follows code from MATLAB CAPTURE scripts, edited to work in Python by Tony.
'''
def generate_FSL_structure_tensor(img_data, filename, dogsigmaArr=[1], gausigmaArr=[2.3], angleArr=[25]):
    for jj in range(len(dogsigmaArr)):
        dogsigma = dogsigmaArr[jj];
        print "Start DoG Sigma on " + str(dogsigma);

        # Generate dog kernels
        dogkercc = doggen([dogsigma, dogsigma, dogsigma]);
        dogkercc = np.transpose(dogkercc, (0, 2, 1));  # annoying

        #print dogkercc.shape;
        #print dogkercc[:, :, 0];

        dogkerrr = np.transpose(dogkercc, (1, 0, 2));

        #print dogkerrr[:, :, 0];
        dogkerzz = np.transpose(dogkercc, (0, 2, 1));

        #print dogkerzz[:, :, 0];

        # Compute gradients
        grr = signal.convolve(img_data, dogkerrr, 'same');

        #print grr[:, :, 0];

        gcc = signal.convolve(img_data, dogkercc, 'same');

        #print gcc[:, :, 0];

        gzz = signal.convolve(img_data, dogkerzz, 'same');

        #print gzz[:, :, 0];

        # Compute gradient products
        gprrrr = np.multiply(grr, grr);

        #print gprrrr[:, :, 0];

        gprrcc = np.multiply(grr, gcc);

        #print gprrcc[:, :, 0];

        gprrzz = np.multiply(grr, gzz);

        #print gprrzz[:, :, 0]

        gpcccc = np.multiply(gcc, gcc);
        gpcczz = np.multiply(gcc, gzz);
        gpzzzz = np.multiply(gzz, gzz);

        # Compute gradient amplitudes
        # print ga.dtype;
        ga = np.sqrt(gprrrr + gpcccc + gpzzzz);

        #print ga[:, :, 0];

        print "GA SHAPE:"
        print ga.shape;

        # Convert numpy ndarray object to Nifti data type
        gradient_amplitudes_data = nib.Nifti1Image(ga, affine=np.eye(4));

        # Save gradient amplitudes image
        nib.save(gradient_amplitudes_data, 'gradient_amplitudes.nii');

        # Compute gradient vectors
        gv = np.concatenate((grr[..., np.newaxis], gcc[..., np.newaxis], gzz[..., np.newaxis]), axis = 3);

        #print gv[:, :, 0, 0];

        gv = np.divide(gv, np.tile(ga[..., None], [1, 1, 1, 3]));
        #print gv[:, :, 0, 1];

        print "GV SHAPE:"
        print gv.shape;

        # Convert numpy ndarray object to Nifti data type
        gradient_vectors_data = nib.Nifti1Image(gv, affine=np.eye(4));

        # Save gradient vectors
        nib.save(gradient_vectors_data, 'gradient_vectors.nii');

        # Compute structure tensor
        for kk in range(len(gausigmaArr)):
            gausigma = gausigmaArr[kk];
            print "Start Gauss Sigma with gausigma = " + str(gausigma);

            print "Generating Gaussian kernel..."
            gaussker = np.single(gaussgen([gausigma, gausigma, gausigma]));

            #print gaussker[:, :, 0];

            print "Blurring gradient products..."
            gprrrrgauss = signal.convolve(gprrrr, gaussker, "same");

            #print gprrrrgauss[:, :, 0];

            gprrccgauss = signal.convolve(gprrcc, gaussker, "same");

            #print gprrccgauss[:, :, 0];

            gprrzzgauss = signal.convolve(gprrzz, gaussker, "same");
            gpccccgauss = signal.convolve(gpcccc, gaussker, "same");
            gpcczzgauss = signal.convolve(gpcczz, gaussker, "same");
            gpzzzzgauss = signal.convolve(gpzzzz, gaussker, "same");

            print "Saving a copy for this Gaussian sigma..."
            tensorfsl = np.concatenate((gprrrrgauss[..., np.newaxis], gprrccgauss[..., np.newaxis], gprrzzgauss[..., np.newaxis], gpccccgauss[..., np.newaxis], gpcczzgauss[..., np.newaxis], gpzzzzgauss[..., np.newaxis]), axis = 3);

             # Convert numpy ndarray object to Nifti data type
            tensor_fsl_data = nib.Nifti1Image(tensorfsl, affine=np.eye(4));

            nib.save(tensor_fsl_data, str(filename) + "dogsigma_" + str(jj) + "gausigma_" + str(kk) + 'tensorfsl.nii');

    print 'Completed computing on ' + str(filename) + '!'
    return tensorfsl


""" Verification tests for the multi-sample classes.

Author:

    C.M. Gosmeyer

Date:

    Mar 2018

References:

    "Introduction to Statistical Problem Solving in Geography",
    J.C. McGrew, Jr., A.J. Lembo, Jr., C.B. Monroe

"""

import numpy as np
import pytest
from stats.inferential_stats.multisample import *

class TestANOVA(object):
    """ Uses table 11.1.
    """
    def setup(self):
        means = np.array([28.780, 25.914, 30.609, 25.819])
        ns = np.array([65, 14, 90, 31])
        variances = np.array([1.544, 4.295, 3.006, 5.104])
        return ANOVA(means, ns, variances)

    def test_between_group_sum_of_squares(self):
        Ftest = self.setup()
        val = round(Ftest.SS_B, 1)
        assert val == 682.5

    def test_between_group_mean_squares(self):
        Ftest = self.setup()
        val = round(Ftest.MS_B, 1)
        assert val == 227.5

    def test_within_group_sum_of_squares(self):
        Ftest = self.setup()
        val = round(Ftest.SS_W, 0)
        assert val == 1978.0

    def test_within_group_mean_squares(self):
        Ftest = self.setup()
        val = round(Ftest.MS_W, 1)
        assert val == 10.1

    def test_F(self):
        Ftest = self.setup()
        val = round(Ftest.MS_W, 1)
        assert val == 10.1

class TestKruskalWallis(object):
    """ Uses tables 11.3 - 11.4.
    """
    def setup(self):
        rs = np.array([
            [17, 23, 5, 6, 22, 3, 12, 15, 10],
            [13.5, 19, 25.5, 33, 29, 21, 25.5, 27, 8, 30],
            [11, 2, 1, 9, 4, 18, 7, 16],
            [13.5, 24, 28, 31, 34, 32, 20]])
        return KruskalWallis(rs)

    def test_N(self):
        Htest = self.setup()
        val = Htest.N
        assert val == 34

    def test_T(self):
        Htest = self.setup()
        val = Htest.T
        assert val == 0

    def test_mean_rs1(self):
        Htest = self.setup()
        val = round(Htest.mean_rs[0], 1)
        assert val == 12.6

    def test_sum_n_r(self):
        Htest = self.setup()
        val = round(Htest.sum_n_r, 1)
        assert val == 12114.0

    def test_H(self):
        Htest = self.setup()
        val = round(Htest.H, 2)
        assert val == 17.16