import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
import sys
from sklearn import linear_model

def scale(t):
    tMin = t.min(axis=0)
    tMax = t.max(axis=0)
    return (t-tMin)/(tMax-tMin)

class Problem1Base():
    def __init__(self):
        tf.reset_default_graph()

    def create_dataset(self):
        n_samples = 30
        self.x_point = np.linspace(0,20,n_samples)

        self.y_point = 3.7 * self.x_point + 14 + 4 * np.random.randn(n_samples)

        plt.plot(self.x_point, self.y_point,'o')
        plt.savefig("lesson5.png")

class Problem1SK(Problem1Base):
    def __init__(self):
        self.RANDOM_SEED = 42
        np.random.seed(self.RANDOM_SEED)

    def compute_regression(self):
        linreg = linear_model.LinearRegression()
        self.x_point = np.reshape(self.x_point,(len(self.x_point),1))
        linreg.fit(self.x_point,self.y_point,self.y_point)
        return linreg.coef_,linreg.intercept_

class Problem1TF(Problem1Base):
    def __init__(self):
        self.graph = tf.Graph()
        self.RANDOM_SEED = 42
        tf.set_random_seed(self.RANDOM_SEED)
        np.random.seed(self.RANDOM_SEED)

    def build(self):
        self.X = tf.placeholder(tf.float32)
        self.Y = tf.placeholder(tf.float32)
        self.LR = tf.placeholder(tf.float32)
        self.slope = tf.Variable(tf.random_uniform([1],-1.0,1.0))
        self.intercept = tf.Variable(tf.zeros([1]))
        self.response = self.slope*self.X + self.intercept
        self.cost = tf.reduce_mean(tf.square(self.response - self.Y))
        self.optimizer = tf.train.GradientDescentOptimizer(self.LR).minimize(self.cost)

    def compute_regression(self,epochs,learning_rate):
        init = tf.global_variables_initializer()
        with tf.Session() as session:
            session.run(init)
            for epoch in range(epochs):
                session.run(
                        self.optimizer,feed_dict = {
                            self.LR: learning_rate,
                            self.X : self.x_point,
                            self.Y : self.y_point})
                if (epoch % (epochs/20)) == 0:
                    c = session.run(
                            self.cost, feed_dict = {
                                self.LR: learning_rate,
                                self.X : self.x_point,
                                self.Y : self.y_point})
                    print("Epoch:", epoch , "Cost:", c, "Slope:",
                            session.run(self.slope),
                            "Intercept:",
                            session.run(self.intercept))
            c = session.run(self.cost, feed_dict = {
                self.LR: learning_rate,
                self.X : self.x_point,
                self.Y : self.y_point})
            print("Epoch:", epoch , "Cost:", c, "Slope:",
                    session.run(self.slope),
                    "Intercept:",
                    session.run(self.intercept))
            return c,session.run(self.slope),session.run(self.intercept)

class Problem2Base():
    def __init__(self):
        tf.reset_default_graph()

    def create_dataset(self):
        f = open("00 kc_house_data.csv")
        f.readline()
        dataset = np.genfromtxt(fname = f, delimiter = ',',usecols=(2,3,5))
        self.predictors = dataset[:,1:3]
        predictorsMin = self.predictors.min(axis=0)
        predictorsMax = self.predictors.max(axis=0)
        self.predictors_scaled = scale(self.predictors)
        self.response = dataset[:,0:1]
        responseMin = self.response.min(axis=0)
        responseMax = self.response.max(axis=0)
        self.response_scaled = scale(self.response)

class Problem2SK(Problem2Base):
    def __init__(self):
        self.RANDOM_SEED = 42
        np.random.seed(self.RANDOM_SEED)

    def compute_regression(self):
        linreg = linear_model.LinearRegression()
        linreg.fit(self.predictors_scaled,self.response_scaled)
        return linreg.coef_,linreg.intercept_

class Problem2TF(Problem2Base):
    def __init__(self):
        self.RANDOM_SEED = 42
        tf.set_random_seed(self.RANDOM_SEED)
        np.random.seed(self.RANDOM_SEED)

    def build(self):
        self.X1 = tf.placeholder(tf.float32)
        self.X2 = tf.placeholder(tf.float32)
        self.Y = tf.placeholder(tf.float32)

        self.LR = tf.placeholder(tf.float32)

        self.slope1 = tf.Variable([0],dtype=tf.float32,name="weight1")
        self.slope2 = tf.Variable([0],dtype=tf.float32,name="weight2")
        self.intercept = tf.Variable(tf.zeros([1]),dtype=tf.float32,name="bias")
        self.response = self.slope1*self.X1 + self.slope2*self.X2 + self.intercept
        self.cost = tf.reduce_mean(tf.square(self.response - self.Y))
        self.optimizer = tf.train.GradientDescentOptimizer(self.LR).minimize(self.cost)

    def compute_regression(self,epochs,learning_rate):
        init = tf.global_variables_initializer()
        with tf.Session() as session:
            session.run(init)
            for epoch in range(epochs):
                session.run(
                        self.optimizer,feed_dict = {
                            self.LR: learning_rate,
                            self.X1:self.predictors_scaled[:,0:1],
                            self.X2:self.predictors_scaled[:,1:2],
                            self.Y:self.response_scaled})
                if not epoch % 100000:
                    c = session.run(
                            self.cost, feed_dict = {
                                self.LR:learning_rate,
                                self.X1:self.predictors_scaled[:,0:1],
                                self.X2:self.predictors_scaled[:,1:2],
                                self.Y:self.response_scaled})
                    print("Epoch:", epoch , "Cost:", c, "Slope1:",
                            session.run(self.slope1),
                            "Slope2:",
                            session.run(self.slope2),
                            "Intercept:",
                            session.run(self.intercept))
            c = session.run(
                    self.cost, feed_dict = {
                        self.LR: learning_rate,
                        self.X1 : self.predictors_scaled[:,0:1],
                        self.X2 : self.predictors_scaled[:,1:2],
                        self.Y : self.response_scaled} )
            print("Epoch:", epoch , "Cost:", c, "Slope1:",
                    session.run(self.slope1),
                    "Slope2:",
                    session.run(self.slope2),
                    "Intercept:",
                    session.run(self.intercept))

            return c,session.run(self.slope1),session.run(self.slope2),session.run(self.intercept)

def main():

    print("Problem 1 - SKLearn")
    print("====================")
    problem1 = Problem1SK()
    problem1.create_dataset()
    slope,intercept = problem1.compute_regression()
    print(slope,intercept)
    assert ("%.8f" % slope[0]) == "3.54942311"
    assert ("%.15f" % intercept) == "14.841075232789862"

    print("\nProblem 1, solution 1 - Tensorflow")
    print("====================================")
    print("   Epoch 50000, Learning rate 0.0001")
    problem1 = Problem1TF()
    problem1.create_dataset()
    problem1.build()
    cost,slope,intercept = problem1.compute_regression(50000,0.0001)
    print(cost,slope,intercept)
    assert ("%.5f" % cost) == "11.36642"
    assert ("%.7f" % slope[0]) == "3.5775077"
    assert ("%.6f" % intercept[0]) == "14.184714"

    print("\nProblem 1, solution 2 - Tensorflow")
    print("====================================")
    print("   Epoch 5000, Learning rate 0.001")
    problem1 = Problem1TF()
    problem1.create_dataset()
    problem1.build()
    cost,slope,intercept = problem1.compute_regression(5000,0.001)
    print(cost,slope,intercept)
    assert ("%.6f" % cost) == "11.365601"
    assert ("%.7f" % slope[0]) == "3.5774026"
    assert ("%.6f" % intercept[0]) == "14.186133"

    print("Problem 2 - SKLearn")
    print("====================")
    problem2 = Problem2SK()
    problem2.create_dataset()
    slope,intercept = problem2.compute_regression()
    print(slope,intercept)

    print("Problem 2 - Tensorflow")
    print("======================")
    print("   Epoch 1000000, Learning rate 0.01")
    problem2 = Problem2TF()
    problem2.create_dataset()
    problem2.build()
    cost,slope1,slope2,intercept = problem2.compute_regression(1000000,0.01)

if __name__ == "__main__":
    main()


import numpy as np;
import scipy.linalg as spla;
from gauss_elimination import gauss_elimination;

def finite_difference(yhist, h, n=1, customge=False):
    '''
    Forward, Central, Backwards finite difference calculation of derivative
    This function uses the Central Finite Difference method to compute the
    derivative for the given data vector. At the ends of the array, central
    difference doesn't work, so the forward difference method and backwards
    difference methods are used instead.

    Note that this function decides for you to use forward and bakwards
    differencing functions at either end of the dataset. This cannot be
    turned off or changed.

    @arg
    yhist    - M x N numpy.ndarray
               Function value time history, where N is the length of the
               dataset and M is the number different things to take the
               derivative of. M is usually 1. N corresponds to the number
               of samples in the time history of the dataset.
    h        - double
               Time step
    n        - double (optional)
               Order of finite difference
    customge - bool (optional)
               Specify to use the custom gauss elimination function.
               True to use the function, False use numpy's built in function.
               False by default.

    @return
    ydothist - N x M numpy.ndarray
               Finite difference derivative time history

    @dependencies
    python 3.6.0
    numpy
    scipy

    @author: Matt Marti
    @date: 2019-06-14
    '''

    # Raise Type Errors on bad input
    if type(yhist) != np.ndarray:
        raise TypeError("Argument 'yhist' is not of type numpy.ndarray");
    #
    if type(h) not in [int, float]:
        raise TypeError("Argument 'h' is not of type [int, float]");
    elif type(h) == int:
        h = h + 0.0;
    #
    if type(n) != int:
        raise TypeError("Argument 'n' is not of type int");
    #

    # Check yhist shape
    assert n > 0, "Argument 'n' is not greater than zero";
    if yhist.shape.__len__() == 1:
        M = 1;
        yhistold = yhist.copy();
        yhist = np.ndarray((1, yhistold.shape[0]));
        for i in range(0, yhistold.shape[0]):
            yhist[0,i] = yhistold[i];
        #
    else:
        M = yhist.shape[0];
    #
    N = yhist.shape[1];
    assert N >= n+1, 'Not enough data points for given order';

    # Preallocate output
    ydothist = np.zeros(yhist.shape, np.float64);
    hmat = np.zeros((n,n), np.float64);

    # Forward difference method
    for i in range(0, n):

        # Delta t matrix
        for j in range(1, n+1):
            d = 1.0;
            hj = h*j;
            for k in range(1, n+1):
                d = d*k;
                hmat[j-1,k-1] = (hj**k)/d;
            #
        #

        # Forward finite difference
        k = 0;
        yfvec = np.ndarray((M, n));
        for j in range(i+1, i+n+1):
            yfvec[:,k] = yhist[:,j] - yhist[:,i];
            k += 1;
        #

        # Compute derivative
        if customge:
            ydoti = gauss_elimination(hmat, yfvec.transpose());
        else:
            ydoti = spla.solve(hmat, yfvec.transpose());
        #

        # Assign output
        for j in range(ydothist.shape[0]):
            ydothist[j,i] = ydoti[0,j];
        #
    #

    # Central difference method
    for i in range( n, N-n ):

        # Delta t matrix
        for j in range(1, n+1):
            d = 1.0;
            hj = h*j;
            for k in range(1, n+1):
                d = d*k;
                hmat[j-1,k-1] = (hj**k)/d;
            #
        #

        # Forward finite difference
        k = 0;
        yfvec = np.ndarray((M, n));
        for j in range(i+1, i+n+1):
            yfvec[:,k] = yhist[:,j] - yhist[:,i];
            k += 1;
        #

        # Backward finite difference
        k = 0;
        ybvec = np.ndarray((M, n));
        for j in range(i-1, i-n-1, -1):
            ybvec[:,k] = yhist[:,i] - yhist[:,j];
            k += 1;
        #

        # Compute derivative
        if customge:
            ydoti_f = gauss_elimination(hmat, yfvec.transpose());
            ydoti_b = gauss_elimination(hmat, ybvec.transpose());
            ydoti = 0.5 * (ydoti_f + ydoti_b);
        else:
            ydoti_f = spla.solve(hmat, yfvec.transpose());
            ydoti_b = spla.solve(hmat, ybvec.transpose());
            ydoti = 0.5 * (ydoti_f + ydoti_b);
        #

        # Assign output
        for j in range(ydothist.shape[0]):
            ydothist[j,i] = ydoti[0,j];
        #
    #

    # Backwards difference method
    for i in range(N-n, N):

        # Delta t matrix
        for j in range(1, n+1):
            d = 1.0;
            hj = h*j;
            for k in range(1, n+1):
                d = d*k;
                hmat[j-1,k-1] = (hj**k)/d;
            #
        #

        # Backward finite difference
        k = 0;
        ybvec = np.ndarray((M, n));
        for j in range(i-1, i-n-1, -1):
            ybvec[:,k] = yhist[:,i] - yhist[:,j];
            k += 1;
        #

        # Compute derivative
        if customge:
            ydoti = gauss_elimination(hmat, ybvec.transpose());
        else:
            ydoti = spla.solve(hmat, ybvec.transpose());
        #

        # Assign output
        for j in range(ydothist.shape[0]):
            ydothist[j,i] = ydoti[0,j];
        #
    #
    return ydothist;
#


import numpy as np

def sigmoid(z: np.ndarray) -> np.ndarray:
    return 1 / (1 + np.exp(-z))

def relu(z: np.ndarray) -> np.ndarray:
    return np.maximum(0, z)

def leaky_relu(z: np.ndarray, rate: float = .01) -> np.ndarray:
    return np.maximum(z * rate, z)

# def sigmoid(z):
#     return 1 / (1 + np.exp(-z))

# def relu(z):
#     return np.maximum(0, z)

# def leaky_relu(z, rate=.01):
#     return np.maximum(z * rate, z)

# def sigmoid_prime(z):
#     return 1 / (1 + np.exp(-z))

def relu_prime(z):
    return (z > 0) * 1

def sigmoid_prime(z):
    return sigmoid(z) * (1 - sigmoid(z))

# def leaky_relu(z, rate=.01):
#     return np.maximum(z * rate, z)


# -*- coding: utf-8 -*-
"""
Created on Fri Oct 19, 2018

@author: Timofey Golubev

This contains an implementation of the Thomas algorithm for solving a tridiagonal matrix equation.
"""
import numpy as np
import time
from numba import jit

@jit(nopython=True, parallel = True)
def thomas_solve(diagonal, upper, lower, rhs):
    '''
    Solves a tridiagonal matrix equation using the Thomas algorithm [1].
    Inputs: The matrix is passed in terms of 3 NumPy arrays corresponding to the upper, lower, and main
        diagonal.
        rhs: array for the right hand side

    Reference:
        [1] https://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm
    '''

    num_elements = len(diagonal) - 1

    #start = time.clock()

    diagonal = np.copy(diagonal)  #do a copy  to ensure the original diag is not changed

    x = np.empty(num_elements+2)

    # Forward substitution
    for i in range(2, num_elements + 1): # recall range uses [ ) ]
        cdiag_ratio = lower[i-1]/diagonal[i-1]
        diagonal[i] -= cdiag_ratio * upper[i-1]
        rhs[i] -= cdiag_ratio * rhs[i-1]

    # Backward substitution
    x[num_elements] = rhs[num_elements]/diagonal[num_elements] #lin eqn. corresponding to last row
    for i in range(num_elements, 1, -1):
        x[i-1] = (rhs[i-1] - x[i]*upper[i-1])/diagonal[i-1]

    #end = time.clock()

    return x


"""
Module: libfmp.c5.c5s3_chord_rec_hmm
Author: Meinard Müller, Christof Weiss
License: The MIT license, https://opensource.org/licenses/MIT

This file is part of the FMP Notebooks (https://www.audiolabs-erlangen.de/FMP)
"""

import numpy as np
from sklearn.preprocessing import normalize
from scipy.linalg import circulant
from numba import jit
from matplotlib import pyplot as plt

import libfmp.c3
from libfmp.c5 import get_chord_labels

def generate_sequence_hmm(N, A, C, B, details=False):
    """Generate observation and state sequence from given HMM

    Notebook: C5/C5S3_HiddenMarkovModel.ipynb

    Args:
        N (int): Number of observations to be generated
        A (np.ndarray): State transition probability matrix of dimension I x I
        C (np.ndarray): Initial state distribution of dimension I
        B (np.ndarray): Output probability matrix of dimension I x K
        details (bool): If "True" then shows details (Default value = False)

    Returns:
        O (np.ndarray): Observation sequence of length N
        S (np.ndarray): State sequence of length N
    """
    assert N > 0, "N should be at least one"
    I = A.shape[1]
    K = B.shape[1]
    assert I == A.shape[0], "A should be an I-square matrix"
    assert I == C.shape[0], "Dimension of C should be I"
    assert I == B.shape[0], "Column-dimension of B should be I"

    O = np.zeros(N, int)
    S = np.zeros(N, int)
    for n in range(N):
        if n == 0:
            i = np.random.choice(np.arange(I), p=C)
        else:
            i = np.random.choice(np.arange(I), p=A[i, :])
        k = np.random.choice(np.arange(K), p=B[i, :])
        S[n] = i
        O[n] = k
        if details:
            print('n = %d, S[%d] = %d, O[%d] = %d' % (n, n, S[n], n, O[n]))
    return O, S

def estimate_hmm_from_o_s(O, S, I, K):
    """Estimate the state transition and output probability matrices from
    a given observation and state sequence

    Notebook: C5/C5S3_HiddenMarkovModel.ipynb

    Args:
        O (np.ndarray): Observation sequence of length N
        S (np.ndarray): State sequence of length N
        I (int): Number of states
        K (int): Number of observation symbols

    Returns:
        A_est (np.ndarray): State transition probability matrix of dimension I x I
        B_est (np.ndarray): Output probability matrix of dimension I x K
    """
    # Estimate A
    A_est = np.zeros([I, I])
    N = len(S)
    for n in range(N-1):
        i = S[n]
        j = S[n+1]
        A_est[i, j] += 1
    A_est = normalize(A_est, axis=1, norm='l1')

    # Estimate B
    B_est = np.zeros([I, K])
    for i in range(I):
        for k in range(K):
            B_est[i, k] = np.sum(np.logical_and(S == i, O == k))
    B_est = normalize(B_est, axis=1, norm='l1')
    return A_est, B_est

@jit(nopython=True)
def viterbi(A, C, B, O):
    """Viterbi algorithm for solving the uncovering problem

    Notebook: C5/C5S3_Viterbi.ipynb

    Args:
        A (np.ndarray): State transition probability matrix of dimension I x I
        C (np.ndarray): Initial state distribution  of dimension I
        B (np.ndarray): Output probability matrix of dimension I x K
        O (np.ndarray): Observation sequence of length N

    Returns:
        S_opt (np.ndarray): Optimal state sequence of length N
        D (np.ndarray): Accumulated probability matrix
        E (np.ndarray): Backtracking matrix
    """
    I = A.shape[0]    # Number of states
    N = len(O)  # Length of observation sequence

    # Initialize D and E matrices
    D = np.zeros((I, N))
    E = np.zeros((I, N-1)).astype(np.int32)
    D[:, 0] = np.multiply(C, B[:, O[0]])

    # Compute D and E in a nested loop
    for n in range(1, N):
        for i in range(I):
            temp_product = np.multiply(A[:, i], D[:, n-1])
            D[i, n] = np.max(temp_product) * B[i, O[n]]
            E[i, n-1] = np.argmax(temp_product)

    # Backtracking
    S_opt = np.zeros(N).astype(np.int32)
    S_opt[-1] = np.argmax(D[:, -1])
    for n in range(N-2, -1, -1):
        S_opt[n] = E[int(S_opt[n+1]), n]

    return S_opt, D, E

@jit(nopython=True)
def viterbi_log(A, C, B, O):
    """Viterbi algorithm (log variant) for solving the uncovering problem

    Notebook: C5/C5S3_Viterbi.ipynb

    Args:
        A (np.ndarray): State transition probability matrix of dimension I x I
        C (np.ndarray): Initial state distribution  of dimension I
        B (np.ndarray): Output probability matrix of dimension I x K
        O (np.ndarray): Observation sequence of length N

    Returns:
        S_opt (np.ndarray): Optimal state sequence of length N
        D_log (np.ndarray): Accumulated log probability matrix
        E (np.ndarray): Backtracking matrix
    """
    I = A.shape[0]    # Number of states
    N = len(O)  # Length of observation sequence
    tiny = np.finfo(0.).tiny
    A_log = np.log(A + tiny)
    C_log = np.log(C + tiny)
    B_log = np.log(B + tiny)

    # Initialize D and E matrices
    D_log = np.zeros((I, N))
    E = np.zeros((I, N-1)).astype(np.int32)
    D_log[:, 0] = C_log + B_log[:, O[0]]

    # Compute D and E in a nested loop
    for n in range(1, N):
        for i in range(I):
            temp_sum = A_log[:, i] + D_log[:, n-1]
            D_log[i, n] = np.max(temp_sum) + B_log[i, O[n]]
            E[i, n-1] = np.argmax(temp_sum)

    # Backtracking
    S_opt = np.zeros(N).astype(np.int32)
    S_opt[-1] = np.argmax(D_log[:, -1])
    for n in range(N-2, -1, -1):
        S_opt[n] = E[int(S_opt[n+1]), n]

    return S_opt, D_log, E

def plot_transition_matrix(A, log=True, ax=None, figsize=(6, 5), title='',
                           xlabel='State (chord label)', ylabel='State (chord label)',
                           cmap='gray_r', quadrant=False):
    """Plot a transition matrix for 24 chord models (12 major and 12 minor triads)

    Notebook: C5/C5S3_ChordRec_HMM.ipynb

    Args:
        A: Transition matrix
        log: Show log probabilities (Default value = True)
        ax: Axis (Default value = None)
        figsize: Width, height in inches (only used when ax=None) (Default value = (6, 5))
        title: Title for plot (Default value = '')
        xlabel: Label for x-axis (Default value = 'State (chord label)')
        ylabel: Label for y-axis (Default value = 'State (chord label)')
        cmap: Color map (Default value = 'gray_r')
        quadrant: Plots additional lines for C-major and C-minor quadrants (Default value = False)

    Returns:
        fig: The created matplotlib figure or None if ax was given.
        ax: The used axes.
        im: The image plot
    """
    fig = None
    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=figsize)
        ax = [ax]

    if log is True:
        A_plot = np.log(A)
        cbar_label = 'Log probability'
        clim = [-6, 0]
    else:
        A_plot = A
        cbar_label = 'Probability'
        clim = [0, 1]
    im = ax[0].imshow(A_plot, origin='lower', aspect='equal', cmap=cmap)
    im.set_clim(clim)
    plt.sca(ax[0])
    cbar = plt.colorbar(im)
    ax[0].set_xlabel(xlabel)
    ax[0].set_ylabel(ylabel)
    ax[0].set_title(title)
    cbar.ax.set_ylabel(cbar_label)

    chord_labels = get_chord_labels()
    chord_labels_squeezed = chord_labels.copy()
    for k in [1, 3, 6, 8, 10, 11, 13, 15, 17, 18, 20, 22]:
        chord_labels_squeezed[k] = ''

    ax[0].set_xticks(np.arange(24))
    ax[0].set_yticks(np.arange(24))
    ax[0].set_xticklabels(chord_labels_squeezed)
    ax[0].set_yticklabels(chord_labels)

    if quadrant is True:
        ax[0].axvline(x=11.5, ymin=0, ymax=24, linewidth=2, color='r')
        ax[0].axhline(y=11.5, xmin=0, xmax=24, linewidth=2, color='r')

    return fig, ax, im

def matrix_circular_mean(A):
    """Computes circulant matrix with mean diagonal sums

    Notebook: C5/C5S3_ChordRec_HMM.ipynb

    Args:
        A (np.ndarray): Square matrix

    Returns:
        A_mean (np.ndarray): Circulant output matrix
    """
    N = A.shape[0]
    A_shear = np.zeros((N, N))
    for n in range(N):
        A_shear[:, n] = np.roll(A[:, n], -n)
    circ_sum = np.sum(A_shear, axis=1)
    A_mean = circulant(circ_sum) / N
    return A_mean

def matrix_chord24_trans_inv(A):
    """Computes transposition-invariant matrix for transition matrix
    based 12 major chords and 12 minor chords

    Notebook: C5/C5S3_ChordRec_HMM.ipynb

    Args:
        A (np.ndarray): Input transition matrix

    Returns:
        A_ti (np.ndarray): Output transition matrix
    """
    A_ti = np.zeros(A.shape)
    A_ti[0:12, 0:12] = matrix_circular_mean(A[0:12, 0:12])
    A_ti[0:12, 12:24] = matrix_circular_mean(A[0:12, 12:24])
    A_ti[12:24, 0:12] = matrix_circular_mean(A[12:24, 0:12])
    A_ti[12:24, 12:24] = matrix_circular_mean(A[12:24, 12:24])
    return A_ti

def uniform_transition_matrix(p=0.01, N=24):
    """Computes uniform transition matrix

    Notebook: C5/C5S3_ChordRec_HMM.ipynb

    Args:
        p (float): Self transition probability (Default value = 0.01)
        N (int): Column and row dimension (Default value = 24)

    Returns:
        A (np.ndarray): Output transition matrix
    """
    off_diag_entries = (1-p) / (N-1)     # rows should sum up to 1
    A = off_diag_entries * np.ones([N, N])
    np.fill_diagonal(A, p)
    return A

@jit(nopython=True)
def viterbi_log_likelihood(A, C, B_O):
    """Viterbi algorithm (log variant) for solving the uncovering problem

    Notebook: C5/C5S3_Viterbi.ipynb

    Args:
        A (np.ndarray): State transition probability matrix of dimension I x I
        C (np.ndarray): Initial state distribution  of dimension I
        B_O (np.ndarray): Likelihood matrix of dimension I x N

    Returns:
        S_opt (np.ndarray): Optimal state sequence of length N
        S_mat (np.ndarray): Binary matrix representation of optimal state sequence
        D_log (np.ndarray): Accumulated log probability matrix
        E (np.ndarray): Backtracking matrix
    """
    I = A.shape[0]    # Number of states
    N = B_O.shape[1]  # Length of observation sequence
    tiny = np.finfo(0.).tiny
    A_log = np.log(A + tiny)
    C_log = np.log(C + tiny)
    B_O_log = np.log(B_O + tiny)

    # Initialize D and E matrices
    D_log = np.zeros((I, N))
    E = np.zeros((I, N-1)).astype(np.int32)
    D_log[:, 0] = C_log + B_O_log[:, 0]

    # Compute D and E in a nested loop
    for n in range(1, N):
        for i in range(I):
            temp_sum = A_log[:, i] + D_log[:, n-1]
            D_log[i, n] = np.max(temp_sum) + B_O_log[i, n]
            E[i, n-1] = np.argmax(temp_sum)

    # Backtracking
    S_opt = np.zeros(N).astype(np.int32)
    S_opt[-1] = np.argmax(D_log[:, -1])
    for n in range(N-2, -1, -1):
        S_opt[n] = E[int(S_opt[n+1]), n]

    # Matrix representation of result
    S_mat = np.zeros((I, N)).astype(np.int32)
    for n in range(N):
        S_mat[S_opt[n], n] = 1

    return S_mat, S_opt, D_log, E

def chord_recognition_all(X, ann_matrix, p=0.15, filt_len=None, filt_type='mean'):
    """Conduct template- and HMM-based chord recognition and evaluates the approaches

    Notebook: C5/C5S3_ChordRec_Beatles.ipynb

    Args:
        X (np.ndarray): Chromagram
        ann_matrix (np.ndarray): Reference annotation as given as time-chord binary matrix
        p (float): Self-transition probability used for HMM (Default value = 0.15)
        filt_len (int): Filter length used for prefilitering (Default value = None)
        filt_type (str): Filter type used for prefilitering (Default value = 'mean')

    Returns:
        result_Tem (tuple): Chord recogntion evaluation results ([P, R, F, TP, FP, FN]) for template-based approach
        result_HMM (tuple): Chord recogntion evaluation results ([P, R, F, TP, FP, FN]) for HMM-based approach
        chord_Tem (np.ndarray): Template-based chord recogntion result given as binary matrix
        chord_HMM (np.ndarray): HMM-based chord recogntion result given as binary matrix
        chord_sim (np.ndarray): Chord similarity matrix
    """
    if filt_len is not None:
        if filt_type == 'mean':
            X, Fs_X = libfmp.c3.smooth_downsample_feature_sequence(X, Fs=1, filt_len=filt_len, down_sampling=1)
        if filt_type == 'median':
            X, Fs_X = libfmp.c3.median_downsample_feature_sequence(X, Fs=1, filt_len=filt_len, down_sampling=1)
    # Template-based chord recogntion
    chord_sim, chord_Tem = libfmp.c5.chord_recognition_template(X, norm_sim='1')
    result_Tem = libfmp.c5.compute_eval_measures(ann_matrix, chord_Tem)
    # HMM-based chord recogntion
    A = libfmp.c5.uniform_transition_matrix(p=p)
    C = 1 / 24 * np.ones((1, 24))
    B_O = chord_sim
    chord_HMM, _, _, _ = libfmp.c5.viterbi_log_likelihood(A, C, B_O)
    result_HMM = libfmp.c5.compute_eval_measures(ann_matrix, chord_HMM)
    return result_Tem, result_HMM, chord_Tem, chord_HMM, chord_sim


import numpy as np

def swap_rows(A, a, b):
    """Rows two rows in a matrix, switch row a with row b

    args:
        A: matrix to perform row swaps on
        a: row index of matrix
        b: row index of matrix

    returns: nothing

    side effects:
    changes A to rows a and b swapped
    """
    assert (a>=0) and (b>=0)
    N = A.shape[0] #number of rows
    assert (a<N) and (b<N) #less than because 0-based indexing
    temp = A[a,:].copy()
    A[a,:] = A[b,:].copy()
    A[b,:] = temp.copy()

# LU_factor function from chapter 8
def LU_factor(A,LOUD=True):
    """Factor in place A in L*U=A. The lower triangular parts of A
    are the L matrix. The L has implied ones on the diagonal.

    Args:
        A: N by N array
    Returns:
        a vector holding the order of the rows, relative to the original order
    Side Effects:
        A is factored in place.
    """
    [Nrow, Ncol] = A.shape
    assert Nrow == Ncol
    N = Nrow
    #create scale factors
    s = np.zeros(N)
    count = 0
    row_order = np.arange(N)
    for row in A:
        s[count] = np.max(np.fabs(row))
        count += 1
    if LOUD:
        print("s =",s)
    if LOUD:
        print("Original Matrix is\n",A)
    for column in range(0,N):
        #swap rows if needed
        largest_pos = np.argmax(np.fabs(A[column:N,column]/s[column])) + column
        if (largest_pos != column):
            if (LOUD):
                print("Swapping row",column,"with row",largest_pos)
                print("Pre swap\n",A)
            swap_rows(A,column,largest_pos)
            #keep track of changes to RHS
            tmp = row_order[column]
            row_order[column] = row_order[largest_pos]
            row_order[largest_pos] = tmp
            #re-order s
            tmp = s[column]
            s[column] = s[largest_pos]
            s[largest_pos] = tmp
            if (LOUD):
                print("A =\n",A)
        for row in range(column+1,N):
            mod_row = A[row]
            factor = mod_row[column]/A[column,column]
            mod_row = mod_row - factor*A[column,:]
            #put the factor in the correct place in the modified row
            mod_row[column] = factor
            #only take the part of the modified row we need
            mod_row = mod_row[column:N]
            A[row,column:N] = mod_row

    return row_order

def LU_solve(A,b,row_order):
    """Take a LU factorized matrix and solve it for RHS b

    Args:
        A: N by N array that has been LU factored with
        assumed 1's on the diagonal of the L matrix
        b: N by 1 array of righthand side
        row_order: list giving the re-ordered equations
        from the the LU factorization with pivoting
    Returns:
        x: N by 1 array of solutions
    """
    [Nrow, Ncol] = A.shape
    assert Nrow == Ncol
    assert b.size == Ncol
    assert row_order.max() == Ncol-1
    N = Nrow
    #reorder the equations
    tmp = b.copy()
    for row in range(N):
        b[row_order[row]] = tmp[row]

    x = np.zeros(N)
    #temporary vector for L^-1 b
    y = np.zeros(N)
    #forward solve
    for row in range(N):
        RHS = b[row]
        for column in range(0,row):
            RHS -= y[column]*A[row,column]
        y[row] = RHS
    #back solve
    for row in range(N-1,-1,-1):
        RHS = y[row]
        for column in range(row+1,N):
            RHS -= x[column]*A[row,column]
        x[row] = RHS/A[row,row]
    return x


#%%
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation
import matplotlib.gridspec as gridspec
from scipy.stats import multivariate_normal as mn
# %%
plt.rcParams['grid.color'] = '#A8BDB7'
plt.rcParams['grid.linestyle'] = '--'
plt.rcParams['text.usetex'] = True
plt.rcParams['text.latex.preamble'] = r'\usepackage{{amsmath}}'

class UpdateFigure:
    def __init__(self, ax1, ax2,):

        self.colors = dict(
            blue        = '#375492',
            green       = '#88E685',
            dark_green  = '#00683B',
            red         = '#93391E',
            pink        = '#E374B7',
            purple      = '#A268B4',
            black       = '#000000',
        )
        self.cm = plt.cm.turbo
        # ====================
        # config data
        # ====================
        self.mean = np.zeros(2)
        self.cov = np.eye(2)
        self.xx, self.yy = np.meshgrid(np.linspace(-4,4,401), np.linspace(-4,4,401))
        xysurf = mn.pdf(np.dstack((self.xx,self.yy)), self.mean, self.cov)*2
        xysurf[self.xx*self.yy>0] = 0
        self.vmin, self.vmax = 0, np.max(xysurf)

        # ====================
        # plot 3d surface
        # ====================
        self.surf = ax1.plot_surface(self.xx, self.yy, xysurf, cmap=self.cm,
                        rstride=1, cstride=1, vmin=self.vmin, vmax=self.vmax)
        ax1.view_init(10, None)

        ax1.set_xlabel(r'$x$', fontsize=20)
        ax1.set_ylabel(r'$y$', fontsize=20)
        ax1.zaxis.set_rotate_label(False)
        # ax1.set_zlabel(r'$f(y)$', rotation=0, fontsize=20)
        xticks=[-4,-2,0,2,4]
        yticks=[-4,-2,0,2,4]
        zticks=[0,.2,.4]
        ax1.set_xticks(xticks)
        ax1.set_yticks(yticks)
        ax1.set_zticks(zticks)
        ax1.set_xlim(xticks[0], xticks[-1])
        ax1.set_ylim(yticks[0], yticks[-1])
        ax1.set_zlim(zticks[0], zticks[-1])
        # ax1.invert_xaxis()
        ax1.xaxis.pane.fill = False
        ax1.yaxis.pane.fill = False
        ax1.zaxis.pane.fill = False
        ax1.xaxis.pane.set_edgecolor('w')
        ax1.yaxis.pane.set_edgecolor('w')
        ax1.zaxis.pane.set_edgecolor('w')
        self.ax1 = ax1

        proj_xz = xysurf.max(0)
        proj_yz = xysurf.max(1)
        xlims=[-4,4]
        ylims=[-4,4]
        ax1.plot(self.xx[0,:], np.ones(self.yy.shape[0])*ylims[1], proj_xz, color='#E0A419', lw=2.5)
        ax1.plot(np.ones(self.xx.shape[0])*xlims[0], self.yy[:,0], proj_yz, color='#88AD80', lw=2.5)

        # ====================
        # draw 2d pcolor
        # ====================
        self.mesh = ax2.pcolormesh(self.xx, self.yy, xysurf, cmap=self.cm)
        ax2.axis('scaled')
        ax2.set_xticks([-4,-2,0,2,4])
        ax2.set_yticks([-4,-2,0,2,4])
        ax2.set_xlabel(r'$x$', fontsize=20)
        ax2.set_ylabel(r'$y$', fontsize=20)

        self.dangle = 1

    def __call__(self, i):
        self.ax1.view_init(None, -60+i*self.dangle)
        return [self.surf,]

fig = plt.figure(figsize=(6,3),dpi=400,)
spec = gridspec.GridSpec(1, 2,
    left=0.05, right=0.95, top=1, bottom=0.10,
    hspace=0.4, wspace=0.4,
    width_ratios=[6,4], figure=fig)
ax1 = fig.add_subplot(spec[0], projection='3d')
ax2 = fig.add_subplot(spec[1])

# create a figure updater
nframes=360
ud = UpdateFigure(ax1, ax2)
# user FuncAnimation to generate frames of animation
anim = FuncAnimation(fig, ud, frames=nframes+1, blit=True)
# save animation as *.mp4
anim.save('2d_gaussian_rotate.mp4', fps=24, dpi=400, codec='libx264', bitrate=-1, extra_args=['-pix_fmt', 'yuv420p'])
# %%


import numpy as np

class Quaternion(np.ndarray):
    """Only works with unit quaternions (as needed to describe rotations)"""

    def __new__(cls, input_array=[0.,1.,0.,0.]):
        """By default works with 0 rotation about x axis"""

        n = np.linalg.norm(input_array)
        if n == 0.:
            raise ValueError("Quaternion axis must not be 0 (the norm of the [1,2,3] components of the quaternion cannot be 0.)")
        input_array = np.array(input_array)/np.linalg.norm(input_array)

        obj = np.asarray(input_array).view(cls)

        return obj

    def __array__finalize(self, obj):
        pass

    @classmethod
    def from_angle(cls, theta, axis):
        """A rotation of angle `theta` about the axis `ax, ay, az`. Allows the axis to be not normalized"""

        axis = np.array(axis)/np.linalg.norm(axis)

        s = np.cos(theta / 2)
        v = np.sin(theta / 2) * axis

        return cls([s, *v])

    def normalize(self):
        """Convert quaternion to unit vector"""

        self/np.linalg.norm(self)

    def rot_matrix(self):
        """Generate a rotation matrix"""

        R = np.array([[1 - 2 * self[2]**2 - 2 * self[3]**2,
                       2 * self[1] * self[2] - 2 * self[0] * self[3],
                       2 * self[1] * self[3] + 2 * self[0] * self[2]],
                [2 * self[1] * self[2] + 2 * self[0] * self[3],
                 1 - 2 * self[1]**2 - 2 * self[3]**2,
                 2 * self[2] * self[3] - 2 * self[0] * self[1]],
                [2 * self[1] * self[3] - 2 * self[0] * self[2],
                 2 * self[2] * self[3] + 2 * self[0] * self[1],
                 1 - 2 * self[1]**2 - 2 * self[2]**2]])

        return R

    def axis_angle(self):
        """Returns the axis of rotation and the angle of rotation (radians) as a tuple"""

        ax = np.array(self[1:])
        ax = ax/np.linalg.norm(ax)

        angle = 2*np.arccos(self[0])

        return ax, angle

    def rate_of_change(self, omega):
        """Return the rate of change of the quaternion based on an angular velocity"""

        # follows the formulation in Box
        # note that there are other methods, which may be more accurate

        omega = np.array(omega)

        s = self[0]
        v = np.array(self[1:])

        sdot = - 0.5 * (omega @ v) # mistake in Box eqn 7 (based on https://www.sciencedirect.com/topics/computer-science/quaternion-multiplication and http://web.cs.iastate.edu/~cs577/handouts/quaternion.pdf)
        vdot = 0.5 * (s * omega + np.cross(omega, v))

        # returns a numpy object, not a quaternion
        return np.array([sdot, *vdot])

    def __mul__(self, other):
        raise NotImplementedError


from __future__ import division
import pycuda.autoinit
from pycuda import gpuarray
import numpy as np
from skcuda import fft
from skcuda import linalg
from matplotlib import pyplot as plt

def cufft_conv(x , y):

    x = x.astype(np.complex64)
    y = y.astype(np.complex64)

    if (x.shape != y.shape):
        return -1

    plan = fft.Plan(x.shape, np.complex64, np.complex64)
    inverse_plan = fft.Plan(x.shape, np.complex64, np.complex64)

    x_gpu = gpuarray.to_gpu(x)
    y_gpu = gpuarray.to_gpu(y)

    x_fft = gpuarray.empty_like(x_gpu, dtype=np.complex64)
    y_fft = gpuarray.empty_like(y_gpu, dtype=np.complex64)
    out_gpu = gpuarray.empty_like(x_gpu, dtype=np.complex64)

    fft.fft(x_gpu, x_fft, plan)
    fft.fft(y_gpu, y_fft, plan)

    linalg.multiply(x_fft, y_fft, overwrite=True)

    fft.ifft(y_fft, out_gpu, inverse_plan, scale=True)

    conv_out = out_gpu.get()

    return conv_out

def conv_2d(ker, img):

    padded_ker = np.zeros( (img.shape[0] + 2*ker.shape[0],  img.shape[1] + 2*ker.shape[1] )).astype(np.float32)

    padded_ker[:ker.shape[0], :ker.shape[1]] = ker

    padded_ker = np.roll(padded_ker, shift=-ker.shape[0]//2, axis=0)
    padded_ker = np.roll(padded_ker, shift=-ker.shape[1]//2, axis=1)

    padded_img = np.zeros_like(padded_ker).astype(np.float32)

    padded_img[ker.shape[0]:-ker.shape[0], ker.shape[1]:-ker.shape[1]] = img

    out_ = cufft_conv(padded_ker, padded_img)

    output = out_[ker.shape[0]:-ker.shape[0], ker.shape[1]:-ker.shape[1]]

    return output

gaussian_filter = lambda x, y, sigma : (1 / np.sqrt(2*np.pi*(sigma**2)) )*np.exp( -(x**2 + y**2) / (2 * (sigma**2) ))

def gaussian_ker(sigma):
    ker_ = np.zeros((2*sigma+1, 2*sigma+1))

    for i in range(2*sigma + 1):
        for j in range(2*sigma + 1):
            ker_[i,j] = gaussian_filter(i - sigma, j - sigma, sigma)

    total_ = np.sum(ker_.ravel())

    ker_ = ker_ / total_

    return ker_

if __name__ == '__main__':

    akshada = np.float32(plt.imread('akshada.jpg')) / 255
    akshada_blurred = np.zeros_like(akshada)
    ker = gaussian_ker(15)

    for k in range(3):
        akshada_blurred[:,:,k] = conv_2d(ker, akshada[:,:,k])

    fig, (ax0, ax1) = plt.subplots(1,2)
    fig.suptitle('Gaussian Filtering', fontsize=20)
    ax0.set_title('Before')
    ax0.axis('off')
    ax0.imshow(akshada)
    ax1.set_title('After')
    ax1.axis('off')
    ax1.imshow(akshada_blurred)
    plt.tight_layout()
    plt.subplots_adjust(top=.85)
    plt.show()


############################################################
#
# Statistical routines for FillRegression
#
############################################################

import numpy
import scipy
import math

def regression_coef (x,y):
    std_x = numpy.std(x,ddof=1)
    std_y = numpy.std(y,ddof=1)
    r = numpy.corrcoef(x,y)[0,1]
    b = r * (std_y/std_x)
    a = numpy.mean(y) - (b * numpy.mean(x))
    return a, b, r, pow(r,2)

def overlapping (x,y, missingValue):
    x_filtered=[]
    x_n2=[]
    y_filtered=[]
    n2=0
    for i in range(0,len(x)):
        if x[i] != missingValue:
            if y[i] == missingValue:
                n2 = n2 + 1
                if x[i] != missingValue:
                    x_n2.append(x[i])
            if y[i] != missingValue:
                x_filtered.append(x[i])
                y_filtered.append(y[i])
    n1 = len(x_filtered)
    return x_filtered, x_n2, y_filtered, n1, n2

def statistics (x):
    Mean = numpy.mean(x)
    StDev = numpy.std(x,ddof=1)
    return Mean, StDev

def dependent_statistics (y, missingValue):
    y_filtered=[]
    for i in range(0,len(y)):
        if y[i] != missingValue:
            y_filtered.append(y[i])
    MeanY = numpy.mean(y_filtered)
    SY = numpy.std(y_filtered,ddof=1)
    NY = len(y_filtered)
    return MeanY, NY, SY

def calculate_Yest (x, a, b):
    Yest=[]
    for i in range(0,len(x)):
        value=a+(b*x[i])
        Yest.append(value)
    meanY1Est = numpy.mean(Yest)
    SY1Est = numpy.std(Yest,ddof=1)
    return Yest, meanY1Est, SY1Est

def RMSE (y, yest, n):
    sumOfSquare = 0
    for i in range(0,len(y)):
        sumOfSquare = sumOfSquare + pow((y[i]-yest[i]),2)
    RMSE = math.sqrt(sumOfSquare/n)
    return RMSE

def SEE (y, yest, n):
    sumOfSquare = 0
    for i in range(0,len(y)):
        sumOfSquare = sumOfSquare + pow((y[i]-yest[i]),2)
    SEE = math.sqrt(sumOfSquare/(n-2))
    return SEE

def SEP (x, SEE, n):
    SEP=[]
    meanX = numpy.mean(x)
    sumOfSquare = 0
    for i in range(0,len(x)):
        sumOfSquare = sumOfSquare + pow((x[i]-meanX),2)
    for i in range(0,len(x)):
        value = pow((x[i]-meanX),2)/sumOfSquare
        value = math.sqrt(1 + (1.0/n) + value) * SEE
        SEP.append(value)
    return SEP

def SESlope (x, y, yest, n):
    meanX = numpy.mean(x)
    sumOfSquareX = 0
    sumOfSquareY = 0
    for i in range(0,len(x)):
        sumOfSquareX = sumOfSquareX + pow((x[i]-meanX),2)
        sumOfSquareY = sumOfSquareY + pow((y[i]-yest[i]),2)
    SESlope = math.sqrt(sumOfSquareY/(n-2))/math.sqrt(sumOfSquareX)
    return SESlope

def TestScore (b, SE):
    if SE != 0:
        TestScore = b/SE
    else:
        TestScore = "inf"
    return TestScore

def LinearRegression (a, b, x):
    y=[]
    for i in range(0,len(x)):
        temp = a + (b*x[i])
        y.append(temp)
    return y

def Skew (x):
    # bias=False means DO correct for bias (use n - 1? as sample standard deviation, etc.?)
    skew = scipy.stats.skew(x, bias=False)
    return skew


# Robust Bayesian Binary logistic regression in 1d for iris flowers

# Code is based on
# https://github.com/aloctavodia/BAP/blob/master/code/Chp4/04_Generalizing_linear_models.ipynb

import superimport

import pymc3 as pm
import numpy as np
import pandas as pd
import theano.tensor as tt
#import seaborn as sns
import scipy.stats as stats
from scipy.special import expit as logistic
import matplotlib.pyplot as plt
import arviz as az
from sklearn.datasets import load_iris
import pyprobml_utils as pml

iris = load_iris()
X = iris.data
y = iris.target

# Convert to pandas dataframe
df_iris = pd.DataFrame(data=iris.data,
                    columns=['sepal_length', 'sepal_width',
                             'petal_length', 'petal_width'])
df_iris['species'] = pd.Series(iris.target_names[y], dtype='category')

df = df_iris.query("species == ('setosa', 'versicolor')")
y_0 = pd.Categorical(df['species']).codes
x_n = 'sepal_length'
x_0 = df[x_n].values

# Create outliers
x_outliers = np.array([4.2, 4.5, 4.0, 4.3, 4.2, 4.4])
y_outliers = np.ones_like(x_outliers, dtype=int)

Ninliers = len(x_0)
Noutliers = len(x_outliers)
N = Ninliers + Noutliers
inlier_ndx = np.arange(0, Ninliers)
outlier_ndx = np.arange(Ninliers, N)

y_0 = np.concatenate((y_0, y_outliers))
x_0 = np.concatenate((x_0, x_outliers))

xmean = np.mean(x_0)
x_c = x_0 - xmean

def plot_training_data():
    plt.figure()
    for c in [0,1]:
        ndx_c = np.where(y_0==c)[0]
        color = f'C{c}'
        sigma = 0.02 # for vertical jittering
        inliers = np.intersect1d(ndx_c, inlier_ndx)
        plt.scatter(x_c[inliers],
                    np.random.normal(y_0[inliers], sigma),
                    marker='o', color=color)
        outliers = np.intersect1d(ndx_c, outlier_ndx)
        plt.scatter(x_c[outliers],
                    np.random.normal(y_0[outliers], sigma),
                    marker='x', color=color)

    plt.xlabel(x_n)
    plt.ylabel('p(y=1)', rotation=0)
    # use original scale for xticks
    locs, _ = plt.xticks()
    plt.xticks(locs, np.round(locs + xmean, 1))
    plt.tight_layout()

def infer_nonrobust_model():
    with pm.Model() as model_0:
        α = pm.Normal('α', mu=0, sd=10)
        β = pm.Normal('β', mu=0, sd=10)

        μ = α + pm.math.dot(x_c, β)
        θ = pm.Deterministic('θ', pm.math.sigmoid(μ))
        bd = pm.Deterministic('bd', -α/β) # decision boundary

        yl = pm.Bernoulli('yl', p=θ, observed=y_0)

        trace = pm.sample(1000, cores=1, chains=2)

    varnames = ['α', 'β', 'bd']
    az.summary(trace, varnames)
    return trace

def infer_robust_model():
    with pm.Model() as model_0:
        α = pm.Normal('α', mu=0, sd=10)
        β = pm.Normal('β', mu=0, sd=10)

        μ = α + pm.math.dot(x_c, β)
        θ = pm.Deterministic('θ', pm.math.sigmoid(μ))
        bd = pm.Deterministic('bd', -α/β) # decision boundary

        #yl = pm.Bernoulli('yl', p=θ, observed=y_0)
        π = pm.Beta('π', 1., 1.) # probability of contamination
        p = π * 0.5 + (1 - π) * θ # true prob or 0.5
        yl = pm.Bernoulli('yl', p=p, observed=y_0)

        trace = pm.sample(1000, cores=1, chains=2)

    varnames = ['α', 'β', 'bd', 'π']
    az.summary(trace, varnames)
    return trace

def make_plot(trace):
    plot_training_data()
    # plot logistic curve
    theta = trace['θ'].mean(axis=0)
    idx = np.argsort(x_c)
    plt.plot(x_c[idx], theta[idx], color='C2', lw=3)
    az.plot_hdi(x_c, trace['θ'], color='C2')

    # plot decision boundary
    plt.vlines(trace['bd'].mean(), 0, 1, color='k')
    bd_hpd = az.hdi(trace['bd'])
    plt.fill_betweenx([0, 1], bd_hpd[0], bd_hpd[1], color='k', alpha=0.5)

trace =  infer_robust_model()
make_plot(trace)
pml.savefig('logreg_iris_bayes_robust_1d.pdf', dpi=300)

trace =  infer_nonrobust_model()
make_plot(trace)
pml.savefig('logreg_iris_bayes_nonrobust_1d.pdf', dpi=300)

plt.show()

import numpy,import scipy,from scipy,import theano,import pymc3


import numpy as np
import pandas as pd
import math

# this function calculates the distance between two vectors
def distance(v1, v2):
    d = 0

    for i in range(len(v1)):
        d += (v1[i]-v2[i])**2

    return math.sqrt(d)

class Kmeans:

    # once a Kmeans object is created with "k" clusters, an array of zeros is defined
    # in which we will store the coordinates of the centroids and an empty list in which
    # we will insert the points belonging to each of the "k" clusters
    def __init__(self, k, n_features):
        self.C = np.zeros((k, n_features))
        self.K = k
        self.n = n_features
        self.Cluster = {i+1:[] for i in range(k)}

    # to initialize centroids randomly, as explained in main.ipynb
    def random_centers(self, data):

        for f in range(self.n):
            self.C[:,f] = np.random.uniform(data[:,f].min(), data[:,f].max(), size=(self.K,))

    # given a point (or vector) it returns the nearest centroid, or the cluster to which such point belongs
    def closest_center(self, vector):
        distances = np.zeros(self.K)
        center = 1

        # evaluate the distance from each centroid
        for k in range(self.K):
            distances[k] = distance(self.C[k], vector)

            if k > 0 and distances[k] < distances[k-1]:
                # when the evaluated distance becomes the shortest, instead the center is "k+1"
                # because the clusters names start from 1 instead of 0 as the index "k"
                center = k+1

        return center

    # it scrolls all the data and adds each of them to a cluster in the dictionary created initially
    def clustering(self, data):

        for i in range(len(data)):
            c = self.closest_center(data[i])
            self.Cluster[c].append(data[i])

    # updates the coordinates of the centroids with the amidpoint of the data belonging
    # to the cluster defined by that centroid
    def update_centers(self):

        for k in range(self.K):
            midpoint = np.mean(self.Cluster[k+1], axis=0)

            # if a cluster has no element, its centroid is not updated because this gives
            # problem in the main code
            if len(self.Cluster[k+1]) != 0:
                self.C[k] = midpoint
            else:
                continue

        # it is necessary to empty the dictionary to fill it with new clusters
        self.Cluster = {i+1:[] for i in range(self.K)}

    # this function evaluates the cost function of k-means algo
    def cost_function(self):
        cost = 0

        for k in range(self.K):
            for point in self.Cluster[k+1]:
                cost += distance(point, self.C[k])**2

        return cost


import newton as nt
import numpy as np

def centering_step(c, A, b, x0, t):
    """
    c, A, b: parameters in LP

            min   c^T x
            s.t.  Ax = b
                  x >= 0

    x0: feasible initial point for constrained Newton's method
    t:  penalty parameter in barrier method

    This function returns the central point x^*(t)
    """

    def f(x):
        return np.dot(c, x) - np.sum(np.log(x)) / t

    def fp(x):
        return c - 1 / x / t

    def fpp(x):
        return np.diag(x ** -2) / t

    return nt.newton_eq_2(f, fp, fpp, x0, A, b)[-1]

def barrier(c, A, b, x0, tol=1e-8, t0=1, rho=10):
    """
    c, A, b: parameters in LP

            min   c^T x
            s.t.  Ax = b
                  x >= 0

    x0:  feasible initial point for the barrier method
    tol: tolerance parameter for the suboptimality gap. The algorithm stops when

             f(x) - f^* <= tol

    t0:  initial penalty parameter in barrier method
    rho: factor by which the penalty parameter is increased in each centering step

    This function should return a list of the iterates
    """
    t = t0
    x = np.array(x0)
    x_traces = [np.array(x0)]

    m = len(c)
    while True:
        x = centering_step(c, A, b, x, t)
        x_traces.append(x.copy())
        if m < t * tol:
            break
        t *= rho

    return x_traces


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

# # Problema Resuelto 3
# ## Problemas de mezcla de productos
#
# Una compañía fabrica tres productos: crema corporal, crema facial y crema para bebés. Los tres productos
# comparten ingredientes en su elaboración: mezcla base, aceite de almendras, vitamina E y manteca
# de karité. En la siguiente tabla se presenta información acerca de los porcentajes de composición de cada
# uno de los tres productos:
#
# | . | Mezcla base | Aceite de Almendras | Vitamina E | Manteca de karité |
# |---|-------------|---------------------|------------|-------------------|
# |Crema Corporal| 90% | 4% | 1% | 5% |
# |Crema facial  | 85% | 8% | 2.5% | 4.5% |
# |Crema para bebé | 80% | 10% | - | 10% |

# Cada día, la compañía cuenta con 500 litros de la mezcla base, 50 litros de aceite de almendras, 5 litros
# de vitamina E y 30 litros de manteca de karité. Adicionalmente, se tiene la siguiente información sobre
# costos y precios de venta.

# | Ingrediente | Costo por litro |
# |-------------|-----------------|
# |Mezcla base  | \$20 |
# |Aceite de almedras | \$500 |
# |Vitamina E | \$1500 |
# |Manteca de karité | \$200 |

# | Producto | Precio de venta (\$/L) |
# |----------|------------------------|
# |Crema corporal | \$80 |
# |Crema facial | \$120 |
# |Crema para bebé | \$100 |

# La demanda diaria de la crema corporal es de 200 litros; de la crema facial, 150 litros; y de la crema para
# bebé, de 250 litros. Por políticas de la empresa, se deben fabricar al menos 50 litros de crema facial.
# ¿Cuánto de cada producto deberá producir la compañía para maximizar su utilidad?

"""
Índice: producto
0: crema corporal
1: crema facial
2: crema para bebé
"""

from ortools.linear_solver import pywraplp
from sympy import Matrix

solver = pywraplp.Solver.CreateSolver("GLOP")

insumos = ["mezcla", "aceite", "vitamina", "manteca"]
productos = ["corporal", "facial", "bebe"]

"""
Restricciones para productos
° Restricción para la demanda de crema corporal: 0<= x0 <= 200.
° Restricción para la demanda de crema facial: 50<= x1 <= 150.
° Restricción para la demanda de crema para bebé: 0 <= x2 <= 250.
"""

restricciones_productos = dict(
    zip(productos, [[0,200], [50,150], [0,250]])
)

inf = solver.infinity()
x = dict()
for producto, restriccion in restricciones_productos.items():
    a, b = restriccion
    x[producto] = solver.NumVar(a,b,producto)

"""
Beneficio = Ingreso - Costo
"""
pct_values = [
    [0.90, 0.04, 0.01, 0.05],
    [0.85, 0.08, 0.025, 0.045],
    [0.80, 0.10, 0.00, 0.10]
]

pct = dict(
    [(producto,
        dict(
            zip(insumos, renglon)
            )
      )
     for producto, renglon in zip(productos, pct_values)]
)

print(pct)

"""
Cm: costo marginal por litro
"""
Cm = Matrix([
    20, 500, 1500, 200
])

# """
# costo productos = T(costo ingredientes)
# """
C = Matrix(pct_values)*Cm
print("C <-",C)

"""
I : ingresos
"""
I = Matrix([
    80, 120, 100
])
print("I <-",I)

"""
R: Ingresos
"""
R = dict(
    zip(productos, (I-C))
)
print("R <- ", R)

solver.Maximize(
    sum(R[tag]*x[tag] for tag in productos)
)

"""
Restricciones para insumos
° Restricción para la mezcla base: 0.9x0 + 0.85x1 + 0.8x2 <= 500.
° Restricción para el aceite de almendras: 0.04x0 + 0.08x1 + 0.1x2 <= 50.
° Restricción para la vitamina E: 0.01x0 + 0.025x1 <= 5.
° Restricción para la manteca de karité: 0.05x0 + 0.045x1 + 0.1x2 <= 30.
"""
restricciones_insumos = dict(zip(insumos, [500, 50, 5, 30]))

for insumo in insumos:
    #insumo_total = 0
    #for producto in productos:
        #insumo_total = insumo_total + pct[producto][insumo]*x[producto]
    solver.Add(
        sum(
            pct[producto][insumo]*x[producto] for producto in productos
        )<=restricciones_insumos[insumo]
    )

# Resolvemos el sistema.
status = solver.Solve()

if status == pywraplp.Solver.OPTIMAL:
    print('Solución:')
    print('Valor objetivo =', solver.Objective().Value())
    for k, v in x.items():
        print("Variable: {}, valor:{} ".format(k, v.solution_value()))

else:
    print('EL problema no tiene solución óptima.')

print('\nUso avanzado:')
print('Problema resuelto en %f milisegundos' % solver.wall_time())
print('Problema resuelto en %d iteraciones' % solver.iterations())


import numpy

def mkGaussian(*args):
    ''' IM = mkGaussian(SIZE, COVARIANCE, MEAN, AMPLITUDE)

        Compute a matrix with dimensions SIZE (a [Y X] 2-vector, or a
        scalar) containing a Gaussian function, centered at pixel position
        specified by MEAN (default = (size+1)/2), with given COVARIANCE (can
        be a scalar, 2-vector, or 2x2 matrix.  Default = (min(size)/6)^2),
        and AMPLITUDE.  AMPLITUDE='norm' (default) will produce a
        probability-normalized function.  All but the first argument are
        optional.

        Eero Simoncelli, 6/96. Python port by Rob Young, 7/15.  '''

    if len(args) == 0:
        print "mkRamp(SIZE, COVARIANCE, MEAN, AMPLITUDE)"
        print "first argument is required"
        exit(1)
    else:
        sz = args[0]
        if isinstance(sz, (int)):
            sz = (sz, sz)
        elif not isinstance(sz, (tuple)):
            print "first argument must be a two element tuple or an integer"
            exit(1)

    # OPTIONAL args:

    if len(args) > 1:
        cov = args[1]
    else:
        cov = (min([sz[0], sz[1]]) / 6.0) ** 2

    if len(args) > 2:
        mn = args[2]
        if isinstance(mn, int):
            mn = [mn, mn]
    else:
        mn = ( (sz[0]+1.0)/2.0, (sz[1]+1.0)/2.0 )

    if len(args) > 3:
        ampl = args[3]
    else:
        ampl = 'norm'

    #---------------------------------------------------------------

    (xramp, yramp) = numpy.meshgrid(numpy.array(range(1,sz[1]+1))-mn[1],
                                    numpy.array(range(1,sz[0]+1))-mn[0])

    if isinstance(cov, (int, long, float)):
        if 'norm' == ampl:
            ampl = 1.0 / (2.0 * numpy.pi * cov)
        e = ( (xramp**2) + (yramp**2) ) / ( -2.0 * cov )
    elif len(cov) == 2 and isinstance(cov[0], (int, long, float)):
        if 'norm' == ampl:
            if cov[0]*cov[1] < 0:
                ampl = 1.0 / (2.0 * numpy.pi *
                              numpy.sqrt(complex(cov[0] * cov[1])))
            else:
                ampl = 1.0 / (2.0 * numpy.pi * numpy.sqrt(cov[0] * cov[1]))
        e = ( (xramp**2) / (-2 * cov[1]) ) + ( (yramp**2) / (-2 * cov[0]) )
    else:
        if 'norm' == ampl:
            detCov = numpy.linalg.det(cov)
            if (detCov < 0).any():
                detCovComplex = numpy.empty(detCov.shape, dtype=complex)
                detCovComplex.real = detCov
                detCovComplex.imag = numpy.zeros(detCov.shape)
                ampl = 1.0 / ( 2.0 * numpy.pi * numpy.sqrt( detCovComplex ) )
            else:
                ampl = 1.0 / (2.0 * numpy.pi * numpy.sqrt( numpy.linalg.det(cov) ) )
        cov = - numpy.linalg.inv(cov) / 2.0
        e = (cov[1,1] * xramp**2) + (
            (cov[0,1]+cov[1,0])*(xramp*yramp) ) + ( cov[0,0] * yramp**2)

    res = ampl * numpy.exp(e)

    return res


import numpy as np
from PyAstronomy.pyaC import pyaErrors as PE

def quadExtreme(x, y, mode="max", dp=(1,1), exInd=None, fullOutput=False, fullPoint=False):
  """
    Find the extreme (minimum or maximum) by means of a parabolic fit.

    This function searches for the maximum (or minimum) in the given
    ordinate values, fits a second-order polynomial to the
    surroundings of that point, and, finally, determines the
    thus approximated abscissa value of the extreme point.

    Parameters
    ----------
    x : array
        Abscissa values.
    y : array
        Ordinate values.
    mode : string, {min, max}, optional
        Determines whether a minimum or a maximum is searched
        for. The default is a maximum.
    dp : tuple with two integers, optional
        Determines the width around the extreme point, which will
        be used in the fitting. The default is one point to the
        right and left, i.e., dp = (1,1).
    extInd : integer, optional
        If given, the function will assume that this is the
        index of the extreme point and not search for it.
    fullOutput : boolean, optional
        If True, the output will also cover the points used in
        the fitting and the resulting polynomial.
    fullPoint : boolean, optional
        If True, the return value `epos` will be a tuple holding
        the abscissa and ordinate values of the extreme point.
        The default is False.

    Returns
    -------
    epos : float or tuple
        Position of the extreme point. If `fullPoint` is True,
        a tuple with the abscissa and ordinate values of the
        extreme point.
    mi : int
        The index of the extreme point (maximum or minimum).
    xb : array, optional
        Only returned if `fullOutput` is True. The abscissa
        values used in the polynomial fit. Note the the
        x-value of the extreme point has been subtracted.
    yb : array, optional
        Only returned if `fullOutput` is True. The ordinate
        values used in the polynomial fit.
    p : numpy polynomial, optional
        Only returned if `fullOutput` is True. The best-fit
        polynomial. Note that the fit refers to the `xb` axis,
        where the x-value of the extreme point has been
        subtracted.
  """
  if exInd is None:
    if mode == "max":
      mi = np.argmax(y)
    elif mode == "min":
      mi = np.argmin(y)
    else:
      raise(PE.PyAValError("Unknown mode '" + str(mode) + "'.", \
                           solution="Choose 'min' or 'max'"))
  else:
    mi = exInd

  if ((mi-dp[0]) < 0) or ((mi+dp[0]+1 > len(x))):
    raise(PE.PyAValError("The requested fitting range around the extreme is not covered by the data.", \
                         solution=["Adapt 'dp'.", "Do you need to change the mode?"]))

  # Cut out the relevant range
  xb = x[mi-dp[0]:mi+dp[1]+1].copy()
  yb = y[mi-dp[0]:mi+dp[1]+1].copy()

  # Subtract the x-value of the extreme
  xm = x[mi]
  xb -= xm

  # Fit second-order polynomial
  p = np.polyfit(xb, yb, 2)

  # Check whether curvature agrees with mode
  signs = {"max":-1, "min":+1}
  if np.sign(p[0]) != signs[mode]:
    PE.warn(PE.PyAValError("The curvature of the fit does not match the requested mode of '" + str(mode) + "'.", \
                           solution="Did you choose the right mode (min/max)?"))

  # Location of the extreme
  epos = -p[1]/(2.0 * p[0])
  # Shift back the location
  epos += xm

  if fullPoint:
    epos = (epos, -p[1]**2/(4.0*p[0]) + p[2])

  if not fullOutput:
    return epos, mi
  else:
    return epos, mi, xb, yb, p


import numpy as np
from scipy.special import expit
from scipy.sparse import diags

def grad(y, tx, w):
    """ returns logistic regression gradient """
    (N, D) = tx.shape
    return (tx.T @ (expit(tx @ w) - y)).reshape(D, 1)

def loss(y, tx, w):
    """ returns logistic regression loss """
    return -np.sum(np.multiply((tx @ w).flatten(), y.flatten())) + np.sum(np.log1p(np.exp(tx @ w)))

def reg_grad(y, tx, w, lambda_):
    """ returns regularized logistic regression gradient """
    (N, D) = tx.shape
    return grad(y, tx, w) + (lambda_ * w).reshape(D, 1)

def reg_loss(y, tx, w, lambda_):
    """ returns regularized logistic regression loss """
    return loss(y, tx, w) + 0.5 * lambda_ * (w.T @ w)

def newton_grad(y, x, w, lambda_ = 0):
    """ returns newton gradient """
    N, D = x.shape
    sigma = expit(x @ w).flatten()
    S = diags(np.multiply(sigma, 1 - sigma))
    H = x.T @ S @ x
    assert H.shape == (D, D), "H shape"
    return np.linalg.pinv(H + np.eye(D) * lambda_) @ reg_grad(y, x, w, lambda_)

def newton_reg_grad(y, x, w, lambda_):
    """ returns regularized newton gradient """
    return newton_grad(y, x, w, lambda_)

### SECONDARY IMPLEMENTATION

def sigmoid(t):
    numerator = np.exp(t)
    denominator = np.add(1,numerator)
    return (numerator/denominator)

def calculate_loss(y, tx, w):
    return -(np.mean(y.T @ tx @ w) - np.mean(np.sum(np.log(1 + np.exp(tx @ w)))))

def calculate_gradient(y, tx, w):
    """compute the gradient of loss."""
    step1 = sigmoid(tx.dot(w)) - y
    step2 = tx.T
    gradient = np.dot(step2,step1)
    return gradient

### TESTS SECTION

def get_random_data(N = 10000, D = 100, seed = 1):
#    N = 250000
#    D = 73
    N = 100
    D = 5
    np.random.seed(seed)
    x = np.random.randn(N, D)
    y = np.random.rand(N, 1) > 0.5
    w = np.random.randn(D, 1)
    return x, y, w

def test_loss_1(seed = 1):
    x, y, w = get_random_data(seed = seed)
    [x_, y_, w_] = [np.copy(z) for z in (x, y, w)]
    assert np.allclose(calculate_loss(y, x, w), loss(y, x, w)), "log_reg_loss"
    for z, z_ in zip((x, y, w), (x_, y_, w_)):
        assert np.allclose(z, z_), "copy"

def test_grad_1(seed = 1):
    x, y, w = get_random_data(seed = seed)
    [x_, y_, w_] = [np.copy(z) for z in (x, y, w)]
    assert np.allclose(calculate_gradient(y, x, w), grad(y, x, w)), "log_reg_grad"
    for z, z_ in zip((x, y, w), (x_, y_, w_)):
        assert np.allclose(z, z_), "copy"

def test_all():
    seeds = [9,2,3,42,56]
    for s in seeds:
        test_grad_1(seed = s)
    for s in seeds:
        test_loss_1(seed = s)
    return 1

if __name__ == "__main__":
    if test_all():
        print("TESTS PASSED")


'''
Example code showing some properties of error bars and the chi2 goodness-of-fit test
'''

import numpy as np
from matplotlib import pyplot as plt

plt.ion()
plt.clf()

'''
Let's suppose that the following hypothesis were true and led to the existence of some data.
The data is "fake_exp_y_values".
'''

true_mean = 5
true_amp = 100
true_wid = 2
true_offset =10
x_values =  np.array([i for i in np.arange(0,10,0.4)])
exact_y_values = np.array([true_offset + true_amp*
                               np.exp(-((i-true_mean)/true_wid)**2)
                               for i in x_values])
# if this is False, we will look at "fresh" fake data every time we
# run the code. if this is True, there is an immutable fixed dataset
# which you can look at over and over.
if False:
    # when the random seed is hardcoded we get the same Poisson
    # draws every time.  Run np.random.seed() later to reset!
    np.random.seed(123456789)

fake_exp_y_values = np.array([np.random.poisson(y)
                                  for y in exact_y_values])

f, (ax1, ax2,ax3,ax4) = plt.subplots(4,1,figsize=(6,8))
ax1.set_title("Hand-comparing (not fitting) a hypothesis to data")

'''
Normally when you run a fit routine, like scipy.curve_fit, something
*inside* the fitter is picking parameter values and doing some sort
of plausibility check.  "Is this parameter set a good match to the
hypothesis?"   We can do those hypothesis checks ourselves.  This is
not fitting, just testing.

Let's test the plausibility of the following alternate hypothesis.
Try adjusting these values yourself!
The hypothesis makes prediuctions which go into "try_these_y_values".
'''
try_mean = 4.8   # Notice the data is generated centered at 5;
# by comparing it to a 4.8 we expect disagreement.
try_amp = 100
try_wid = 2
try_offset =10
try_these_y_values = np.array([try_offset + try_amp*
                                   np.exp(-((i-try_mean)/try_wid)**2)
                                   for i in x_values])

'''
In experiments like this the appropriate y-error is often the square
root of the number of counts.  However, there are some subtleties.
Here, for example, I will enforce that the error never be zero.
'''
y_err = np.array([max(np.sqrt(value),1) for value in fake_exp_y_values])

'''
Always plot your raw data.   Here, we plot data and residuals on the
same panel.
'''

ax1.plot(x_values,try_these_y_values,label="hypothesis")
ax1.errorbar(x_values,fake_exp_y_values,yerr=y_err,linestyle='None',
                 label="data",marker='.')

residuals = fake_exp_y_values-try_these_y_values
ax1.plot(x_values,[0]*x_values)  # draw a line at zero
ax1.errorbar(x_values,residuals,yerr=y_err,linestyle='None',
                 label="residual",marker='.')

ax1.legend()

'''
Plotting the "pulls" shows where your data and hypothesis disagree,
normalizing the error bars.  The y-axis of the pulls plot is "number
of error bars".  You can see  which points fluctuted 1-error-bar-up
or two-error-bars-down or whatever, even when the data has wildly
varying errors.
'''

pulls = residuals/y_err
ax2.plot(x_values,[0]*x_values) # draw a line at zero
ax2.errorbar(x_values,pulls,yerr=[1]*len(x_values),linestyle='None',
                 label="pulls",marker='o')
ax2.legend()

'''
we can calculate what the value of chi2 is; if you have already
calculated the pulls, just take the sum of squares of pulls.
'''

chi2 = sum(pulls*pulls)
chi2ndof = chi2/len(x_values)
ax2.text(0.5,0.1,"chi2/ndof = %f"%(chi2ndof),transform=ax2.transAxes)

'''
The nice thing about knowing the value of chi2 is that there is a
 standard statistical test. "If the fluctuations are just randomness,
 how often does chi2 fluctuate lower than this?".  This probability
is P(chi2).  If the value of P(chi2) is close to 1.0, then this might
 be telling you something is wrong.

a) if your model is correct, your data must have come from a rare bad
 fluctuation (which happens sometimes)
b) your model is wrong and your data disagrees with it for reasons
other than random fluctuations.

If the value of P(chi2) is very close to 0.0, then your data is
*unusually close* to your model ("it's rare to be lower"), so either:
a) if your model is correct, your data must have come from a rare
good fluctuation (which happens sometimes)
b) maybe you assigned error bars which are too big

However, although we think of P(chi2)=0.9999 as a problem,
and P(chi2)=0.0001 as a different problem, it is not correct to say
"the best P(chi2) is 0.5".  If your fit worked, and your data has
plausible- and ordinary-looking fluctuations away from the predictions,
 P(chi2) winds up being more or less evenly distributed between 0 and
1. (Try running a bunch of "good" evaluations and you'll see this.)

Notice that I get P(chi2) by calling scipy.stats.chi2.cdf.  That's a
two-argument function.  The second argument is the "number of degrees
of freedom".  This means: the number of data points MINUS the number
of parameters your fit-procedure has optimized.  In this code we have
not optimized anything so we are in the unusual ndof=npoints case.
'''
import scipy.stats
prob_chi2_lower_than_this = scipy.stats.chi2.cdf(chi2,len(x_values))
ax2.text(0.5,0.2,"P(chi2) = %f"%(prob_chi2_lower_than_this),
             transform=ax2.transAxes)

ax2.legend()

'''
If your hypothesis agrees with your data as well as possible, the pulls
 should follow a Gaussian distribution centered at zero and with width
1.  Sometimrs a histogram is a good way to look at this. In particular, the pulls distribution lets you see whether a "bad chi2" is dominated
by random-looking behavior (just a wide distribution of pulls, maybe
telling you your datapoints' error bars were assigmed too small) or
dominated by outliers (a nice Gaussian distribution of width 1 _plus_
a couple of points far away) or something else.

In the default settings of this code, you will probably see a bimodal
histogram---the hypothesized curve (and parameter choices) yield many
datapoints overshooting or undershooting the theory (mostly by small
amounts) and "surprisingly few" landing right in the middle.
'''
counts, bins,patches = ax3.hist(pulls,label="histogram of pulls",
                                    histtype='step')
ax3.text(0.1,0.9,"mean=%f"%(np.mean(pulls)),transform=ax4.transAxes)
ax3.text(0.1,0.7,"stdev=%f"%(np.std(pulls)),transform=ax4.transAxes)

'''
we'll draw a Gaussian on here to show what the pulls distribution
should have looked like if the data/curve disagreement had been only
random fluctuations.
'''

norm = sum(counts)
binned_gaussian = np.array([0.0]*len(counts))
bin_centers = np.array([0.0]*len(counts))
for i in range(len(counts)):
    bin_centers[i] = (bins[i+1]+bins[i])/2.0
    bin_width =  (bins[i+1]-bins[i])
    binned_gaussian[i] = (1/np.sqrt(2*np.pi)*np.exp(-0.5*bin_centers[i]**2))*bin_width*norm
ax3.errorbar(bin_centers,binned_gaussian,yerr=np.sqrt(binned_gaussian),label="expectation")

ax3.legend(loc="upper right")

'''
Finally, let's do a classic Monte Carlo exercise.   Imagine we had
done the experiment 50 times.  50 times we'll test it against the
"bad hypothesis" and 50 times we'll do a fit to find the "best
hypothesis" (the actual truth---not even the fit result)

Notice that SOMETIMES the "bad hypothesis" gives a totally-reasonable-
looking P(chi2)---right in the middle of the range that the "best fit"
would give.  That's why we have error bars!  Sometimes the data's
fluctuations genuinely make it "fit well" to something that happens to
be wrong, and there is no diagnostic in the world that will change
that.

Notice that SOMETIMES the "best hypothesis" gives a P(chi2) near the
top of the range---in the region we associate more with bad fits!
That's why we can't be too strict about what sort of fit we call
reasonable and what kind we don't.  P(chi2)=0.9 is often reasonable.
P(chi2)=0.99 is rarely, but not never, obtained from reasonable
fluctuations.  P(chi2)=0.999 is pushing it further but it DOES HAPPEN.
'''

bad_hypothesis_Pchi2 = np.array([0.]*50)
best_hypothesis_Pchi2 = np.array([0.]*50)
for irun in range(50):
    # generate a new fake experiment
    fake_exp_y_values = np.array([np.random.poisson(y)
                                      for y in exact_y_values])
    # try this "bad" hypothesis and see how it matches
    try_these_y_values = np.array([try_offset +
                    try_amp*np.exp(-((i-try_mean)/try_wid)**2)
                    for i in x_values])
    y_err = np.array([max(np.sqrt(value),1)
                        for value in fake_exp_y_values])

    residuals = fake_exp_y_values-try_these_y_values
    pulls = residuals/y_err
    chi2 = sum(pulls**2)
    bad_hypothesis_Pchi2[irun] = scipy.stats.chi2.cdf(chi2,len(x_values))

    # try the "best" hypothesis and see how it matches
    residuals = fake_exp_y_values-exact_y_values
    pulls = residuals/y_err
    chi2 = sum(pulls**2)
    best_hypothesis_Pchi2[irun] = scipy.stats.chi2.cdf(chi2,
                                len(x_values))

ax4.hist(bad_hypothesis_Pchi2,label="bad hypothesis",histtype='step')
ax4.hist(best_hypothesis_Pchi2,label="best hypothesis",histtype='step')
ax4.legend()


try:
    import numpy as _np
    from math import sqrt as _sqrt, pi as _pi, exp as _exp
    from typing import Union, Dict, List
    from univariate._base import Infinite
except Exception as e:
    print(f"some modules are missing {e}")

class Logistic(Infinite):
    """
    This class contains methods concerning Logistic Distirbution [#]_.

    .. math::
        \\text{Logistic}(x;\\mu,s) = \\frac{\\exp{(-(x-\\mu)/s)}} {s(1+\\exp(-(x-\\mu)/s)^2)}

    Args:

        location(float): location parameter (:math:`\\mu`)
        scale(float): scale parameter (:math:`s`) x > 0
        x(float): random variable

    Reference:
        .. [#] Wikipedia contributors. (2020, December 12). Logistic distribution. https://en.wikipedia.org/w/index.php?title=Logistic_distribution&oldid=993793195
    """

    def __init__(self, location: float, scale: float):
        if scale < 0:
            raise ValueError('scale should be greater than 0.')

        self.scale = scale
        self.location = location

    def pdf(self, x: Union[List[float], _np.ndarray, float]) -> Union[float, _np.ndarray]:
        """
        Args:
            x (Union[List[float], numpy.ndarray, float]): random variable(s)

        Returns:
            Union[float, numpy.ndarray]: evaluation of pdf at x
        """
        mu = self.location
        s = self.scale

        if isinstance(x, (_np.ndarray, List)):
            if not type(x) is _np.ndarray:
                x = _np.array(x)
            return _np.exp(-(x - mu) / s) / (s * (1 + _np.exp(-(x - mu) / s))**2)
        return _exp(-(x - mu) / s) / (s * (1 + _exp(-(x - mu) / s))**2)

    def cdf(self, x: Union[List[float], _np.ndarray, float]) -> Union[float, _np.ndarray]:
        """
        Args:
            x (Union[List[float], numpy.ndarray, float]): data point(s) of interest

        Returns:
            Union[float, numpy.ndarray]: evaluation of cdf at x
        """
        mu = self.location
        s = self.scale

        if isinstance(x, (_np.ndarray, List)):
            x = _np.array(x)
            return 1 / (1 + _np.exp(-(x - mu) / s))
        return 1 / (1 + _exp(-(x - mu) / s))

    def mean(self) -> float:
        """
        Returns: Mean of the Logistic distribution.
        """
        return self.location

    def median(self) -> float:
        """
        Returns: Median of the Logistic distribution.
        """
        return self.location

    def mode(self) -> float:
        """
        Returns: Mode of the Logistic distribution.
        """
        return self.location

    def var(self) -> float:
        """
        Returns: Variance of the Logistic distribution.
        """
        return pow(self.scale, 2) * pow(_pi, 2)/3

    def std(self) -> float:
        """
        Returns: Standard deviation of the Logistic distribution.
        """
        return _sqrt(self.var())

    def skewness(self) -> float:
        """
        Returns: Skewness of the Logistic distribution.
        """
        return 0.0

    def kurtosis(self) -> float:
        """
        Returns: Kurtosis of the Logistic distribution.
        """
        return 6 / 5

    def entropy(self) -> float:
        """
        Returns: differential entropy of the Logistic distribution.

        Reference: Park, S.Y. & Bera, A.K.(2009). Maximum entropy autoregressive conditional heteroskedasticity model. Elsivier.
        link: http://wise.xmu.edu.cn/uploadfiles/paper-masterdownload/2009519932327055475115776.pdf
        """
        return 2.0

    def summary(self) -> Dict[str, Union[float, str]]:
        """
        Returns:
            Dictionary of Logistic distirbution moments. This includes standard deviation.
        """
        return {
            'mean': self.mean(), 'median': self.median(), 'mode': self.mode(),
            'var': self.var(), 'std': self.std(), 'skewness': self.skewness(), 'kurtosis': self.kurtosis()
        }


__author__ = 'Prateek'
from quadraticResidues import quadraticResidues
from sympy import isprime

def quadraticNonResidues(n):
    nonResidues = set()
    residues = quadraticResidues(n);
    for elt in range(1, n):
        if elt not in residues:
            nonResidues.add(elt)
    return nonResidues

if __author__ == 'Prateek':
    n = 13
    if isprime(n):
        print quadraticNonResidues(n)
    else:
        print "N is not prime"


"""Signal-to-noise ratio (SNR)

SNR = P_signal / P_noise     # power of signal
    = (A_signal / A_noise)^2 # amplitude
    = sum(x^2) / sum(n^2)
SNR_dB = 10 * log(SNR)
=> SNR = 10 ^ (SNR_dB / 10)
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.fftpack import fft, ifft
from scipy.io import wavfile as wav

# Add white Gaussian noise to signal
def awgn(x, snr):
    # SNR = 10 ^ (SNR_dB / 10)
    snr = 10 ** (snr / 10.0)
    # A_sig ^ 2 = (x1^2 + x2^2 + ... + xn^2) / n
    xpower = np.sum(np.abs(x ** 2)) / len(x)
    npower = np.abs(xpower / snr)

    # np.random.randn():
    #   Return a sample (or samples) from the “standard normal” distribution
    n = np.random.randn(len(x)) * np.sqrt(npower)
    return x+n, n

if __name__ == "__main__":
    # sampling and signal information
    fs, data = wav.read("audio/test.wav")
    Ts = 1.0 / fs
    N = len(data)
    t = N / fs
    amplitude = np.max(data)

    # avoid overflow -> normalization (-1, 1)
    y = data / amplitude

    # channel data
    yl = y[:, 0]
    yr = y[:, 1]

    # use left channel
    x = np.arange(0, t, Ts)
    yn, n = awgn(yl, 6)

    """ Parseval's theorem

    np.sum(np.abs(sig)**2) == np.sum(np.abs(np.fft.fft(sig))**2)/sig.size

    power = np.sum(np.abs(sig)**2)/sig.size
          = np.sum(np.abs(np.fft.fft(sig))**2)/(sig.size**2)
    """
    # By Parseval's theorem,
    # we can calculate the power of a signal through X[k] that is the FFT of x[n].
    yfft = fft(yl)
    ynfft = fft(yn)
    xpower = np.sum(np.abs(yfft ** 2)) / (len(yfft) ** 2)

    noise = yn - yl
    nfft = fft(noise)
    npower = np.sum(np.abs(nfft ** 2)) / (len(nfft) ** 2)

    # SNR
    snr = xpower / npower
    snr_dB = 10 * np.log10(snr)
    print("SNR: {}dB".format(round(snr_dB, 2)))

    plt.figure()
    plt.plot(x[:441], yl[:441], 'k')
    plt.plot(x[:441], yn[:441], 'r')
    plt.title("The original sinal with Gauss White Noise")

    plt.tight_layout()
    plt.show()


# Enneper Surface, brought to you by PharaohCola13

import mpl_toolkits.mplot3d.axes3d as p3
import matplotlib.pyplot as plt
from matplotlib import *
from numpy import *
from mpl_toolkits.mplot3d.art3d import *
from matplotlib.animation import *

name = "Enneper's Surface"

def shape(fig, alpha, color, edge_c, edge_w, grid, sides, edges, multi_pi, figcolor, rotation, rotmagt, rotmagp, save):
	plt.clf()

	def x_(u,v):
		x = u - (u**3/3) + (u * v**2)
		return x

	def y_(u,v):
		y = v - (v**3/3) + (v * u**2)
		return y

	def z_(u,v):
		z = u**2 - v**2
		return z

	u = linspace(-pi,pi, 100)
	v = linspace(-pi,pi, 100)

	u, v = meshgrid(u, v)

	x = x_(u,v)
	y = y_(u,v)
	z = z_(u,v)

	# Figure Properties
	ax = p3.Axes3D(fig)
	ax.set_facecolor(figcolor)

	plt.axis(grid)
	plt.axis('equal')

	# Surface Plot
	interest = ax.plot_surface(x, y, z)

	interest.set_alpha(alpha)
	interest.set_edgecolor(edge_c)
	interest.set_linewidth(edge_w)
	interest.set_facecolor(color)

	def rot_on():
		def animate(i):
			ax.view_init(azim=rotmagt * i, elev=rotmagp * i)

		if save == "MP4":
			# Animate
			ani = FuncAnimation(fig, animate, frames=500,
								interval=100, save_count=50)  # frames=100)#, repeat=True)

			Writer = writers['ffmpeg']
			writer = Writer(fps=30, bitrate=1800)
			ani.save('{}.mp4'.format(name), writer=writer)
		else:
			#save = None
			# Animate
			ani = FuncAnimation(fig, animate,
								interval=1, save_count=50)  # frames=100)#, repeat=True)
			pass

		plt.ion()
		plt.show()
		time.sleep(0)
		plt.close()

	if rotation == "On":
		rot_on()
	elif rotation == "Off":
		pass


import numpy as np

def constant(angle):
    """
    Return a constant emission profile for angles between 0 and pi
    Normalisation factor K = 1
    # The profile is normalized so that $ 2\pi \int_0^\pi x * I(x) dx = 1$

    Parameters
    ----------
    angle: `numpy.ndarray` or float

    Returns
    -------
    `numpy.ndarray` of the same shape as angle or float
    """
    # K = 1. / np.pi ** 3
    K = 1
    if isinstance(angle, np.ndarray):
        assert (angle <= np.pi).all()
        return K * np.ones(len(angle))
    else:
        assert angle <= np.pi
        return K

def heaviside(angle, limit=0.1):
    """
    Return a normalized emission profile I(angle) following a heaviside profile
    (constant between 0 and limit and null above)
    Normalisation factor K = 1
    # The profile is normalized so that $ 2\pi \int_0^\pi x * I(x) dx = 1$

    Parameters
    ----------
    angle: `numpy.ndarray` or float
    limit: float

    Returns
    -------
    `numpy.ndarray` of the same shape as angle or float
    """
    assert limit > 0
    # K = 1. / (np.pi * min(np.pi, limit) ** 2)
    K = 1.
    if isinstance(angle, np.ndarray):
        assert (angle <= np.pi).all()
        return K * (angle <= limit)
    else:
        if angle <= limit:
            return K
        else:
            return 0

def lgdt06(angle, eta=0.001):
    """
    Return a normalized emission profile I(angle) following the law proposed by [Lemoine et al 06]
    but with a normalisation factor K = 1
    # The profile is normalized so that $ 2\pi \int_0^\pi x * I(x) dx = 1$
    # The normalization has been updated to work for any eta < pi.

    Parameters
    ----------
    angle: `numpy.ndarray` or float
    eta: float

    Returns
    -------
    `numpy.ndarray` of the same shape as angle or float
    """

    assert eta > 0

    # K = 1./(9 * np.pi * eta**2) #approx by [Lemoine et al 06], valid if eta<<pi

    # Normalisation for any eta:
    # K = 1. / (np.pi * eta ** 2 * (1 + 8 * (1 - np.exp(0.25 * (1 - np.pi / eta)))))

    K = 1.

    if isinstance(angle, np.ndarray):
        assert (angle <= np.pi).all()
        y = K * np.ones(len(angle))
        y[angle > eta] = K * eta / angle[angle > eta] * np.exp(- (angle[angle > eta] - eta) / (4 * eta))
        return y

    else:
        assert angle <= np.pi
        if angle <= eta:
            return K
        else:
            return K * eta / angle * np.exp(- (angle - eta) / (4 * eta))

def verify_normalisation_solid_angle(profile, **kwargs):
    """
    Quick verification if a profile function I is normalised
    so that $ 2\pi \int_0^\pi x * I(x) dx = 1$
    Disclaimer: the precision might not be enough for certain profiles.

    Parameters
    ----------
    profile: profile function
    **kwargs: profile function arguments

    Returns
    -------
    bool
    """
    import scipy.integrate as integrate
    integral = integrate.quad(lambda x: 2*np.pi*x*profile(x, **kwargs), 0, np.pi)
    return np.isclose(integral[0], 1)

def exp_peak(angle, eta=0.1, alpha=1):
    """
    Return a normalized emission profile I(angle) with an exponential increase from 0 to eta
    followed by an exponential decrease from eta to 2*eta.
    The function is null between 2*eta and pi.
    Normalisation factor K = 1
    # The profile is normalized so that $ 2\pi \int_0^\pi x * I(x) dx = 1$

    Parameters
    ----------
    angle: `numpy.ndarray` or float
    eta: float, peak angle
    alpha: float, exponential rate

    Returns
    -------
    `numpy.ndarray` of the same shape as angle or float
    """

    assert eta > 0
    assert alpha > 0

    # K = alpha / (4 * np.pi * (np.exp(alpha * eta) - alpha * eta - 1))

    K = 1/(np.exp(alpha*eta)-1)  # verify max emission (at angle=eta) = 1

    if isinstance(type(angle), np.ndarray):
        assert (angle <= np.pi).all()
        return K  * (np.exp(alpha * angle) - 1) * (angle < eta) \
               + K * (np.exp(alpha * (2 * eta - angle)) - 1) * (angle >= eta) * (angle <= 2 * eta)

    else:
        assert angle <= np.pi
        if angle <= eta:
            return K * (np.exp(alpha * angle) - 1)
        elif angle < 2 * eta:
            return K * (np.exp(alpha * (2 * eta - angle)) - 1)
        else:
            return 0


"""
Analysis functions
"""

import numpy as np
from PIL import Image
import os
import seaborn as sns

def get_p1_value(x):
    """
    Returns p1 value for the image.
    """
    return x.mean()

def radial_profile(image, center, img_size):
    x, y = np.indices(image.shape)
    rad = np.sqrt((x - center[0])**2 + (y - center[1])**2)
    ind = np.argsort(rad.flat)
    rad_sorted = rad.flat[ind]
    image_sorted = image.flat[ind]
    rad_round = rad_sorted.astype(int)
    deltar = rad_round[1:] - rad_round[:-1]
    nonzero_deltar = np.where(deltar > 0.0)[0]
    nind = nonzero_deltar[1:] - nonzero_deltar[:-1]
    yvalues = np.cumsum(image_sorted, dtype = np.float64)
    yvalues = yvalues[nonzero_deltar[1:]] - yvalues[nonzero_deltar[:-1]]
    radial_var = yvalues/nind
    radial_dis = rad_round[nonzero_deltar]/(min(image.shape))
    return radial_var, radial_dis[:-1]

def get_p2_vector(img):
    """
    Returns a p2 vector.
    We calculate the p2 vector by taking the radial mean of
    the autocorrelation of the input image.
    """
    radvars = []
    dimX = img.shape[0]
    dimY = img.shape[1]
    fftimage = np.fft.fft2(img)
    final_image = np.fft.ifft2(fftimage*np.conj(fftimage))
    finImg = np.abs(final_image)/(dimX*dimY)
    centrdImg = np.fft.fftshift(finImg)
    center = [int(dimX/2), int(dimY/2)]
    radvar, _ = radial_profile(centrdImg, center, (dimX, dimY))
    radvars.append(radvar)
    p2_vec = np.array(radvars)
    return p2_vec[0]

def get_dataset_p1_distribution(ds, bins=100):
    """
    Return the p1 distribution over the dataset or a segment of the same
    Params:
        ds: AIRADataset instance or AIRADataset iterator
        bins: # of bins for the histogram
    Returns:
        Histogram of distribution
    """
    p1_s = []
    for i in ds:
        p1_s.append(get_p1_value(i))
    ax = sns.distplot(p1_s)
    return ax

def get_segmented_ds(ds, p1_range):
    p1_s = np.asarray([get_p1_value(i) for i in ds])
    p1_max = p1_s.max()
    p1_min = p1_s.min()
    if p1_range[0] < p1_min  or p1_range[1] > p1_max:
        raise Exception('P1 values out range. Dataset has samples within p1_ranges: [{}, {}]'.format(p1_min, p1_max))
    indices = [i for i in range(len(p1_s)) if ((p1_s[i] >= p1_range[0]) and (p1_s[i]<=p1_range[1]))]
    ds_filtered = np.asarray([ds[i] for i in indices])
    return ds_filtered


import numpy as np

def compute_brightness_temperature(
    image: np.ndarray, M: float, A: float, k1: float, k2: float, mask: np.ndarray = None
) -> np.ndarray:

    """Converts image raw digital numbers to brightness temperature
        Refer to USGS page https://www.usgs.gov/core-science-systems/nli/landsat/using-usgs-landsat-level-1-data-product
            for more details.

    Args:
        image (np.ndarray): Level 1 quantized and calibrated scaled Digital Numbers
                                (DN) TIR band data (e.g Band 10 landsat 8 data)
        M (float): Band-specific multiplicative rescaling factor from the image
                    folder metadata (RADIANCE_MULT_BAND_x, where x is the band number).
        A (float): Band-specific additive rescaling factor from the image
                    folder metadata (RADIANCE_ADD_BAND_x, where x is the band number).
        k1 (float): Band-specific thermal conversion constant from the image
                    folder metadata (K1_CONSTANT_BAND_x, where x is the thermal band number)
        k2 (float): Band-specific thermal conversion constant from the image
                    folder metadata (K2_CONSTANT_BAND_x, where x is the thermal band number.
        mask (bool): True if you want to mask NaN, O or irregular values from the computation

    Returns:
        np.ndarray: Brightness temperature corrected landsat image
    """
    toa_radiance = (M * image) + A
    brightness_temp = k2 / (np.log((k1 / toa_radiance) + 1))

    if mask is not None:
        brightness_temp[mask] = np.nan
    return brightness_temp


"""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)