import numpy

from bokeh.plotting import figure
from bokeh.models import Range1d, BasicTickFormatter
from bokeh.colors import RGB
from bokeh.io.export import export_png

def Shannon_Wiener_Index(clone_counts):
	"""Calculates the Shannon-Wiener index of diversity given an iterable of clone frequencies.

	Parameters
	----------
	clone_counts: iterable of floats
		Frequencies of all clones/members of a population

	Returns
	----------
	sw_index: float
		The Shannon-Wiener index value
	"""

	total_clone_counts = float(sum(clone_counts))
	sw_index = 0.0

	for clone in clone_counts:
		clone_freq = clone / total_clone_counts
		sw_index -= (clone_freq * numpy.log(clone_freq))

	return sw_index

def Hill_Diversity_Index(clone_counts, N = (0.0, 10.0), step = 0.1):
	"""Calculates the Hill Diversity index/indices; will return the indices from orders 0 to 10 by default.

	Parameters
	----------
	clone_counts: iterable of floats
		Frequencies of all clones/members of a population
	N: tuple of (float, float)
		Start and stop point for orders to calculate the diversity index from (both inclusive); default is (0.0, 10.0)
	step: float
		Step value to create the range from the start and stop point in N; default is 0.1

	Returns
	----------
	hill_index/hill_indices: float or list of floats
		The Hill Diversity index/indices for the given repertoire
	"""

	if hasattr(N, "__iter__"):
		if len(N) != 2 or N[1] <= N[0]:
			raise ValueError("If N is an iterable it must be of length 2 for the start/end orders.")
		chunks = numpy.floor(N[1] / step) + 1
		orders = numpy.linspace(start = N[0], stop = N[1], num = chunks)
	else:
		orders = [N]

	total_clone_counts = float(sum(clone_counts))
	hill_indices = []

	for order in orders:
		hill_index = 0.0
		if order == 0.0:
			#Hill index at zero is simply the species richness (total number of clones)
			hill_index = len(clone_counts)
		elif order == 1.0:
			#Hill index at one is the exponential of the Shannon-Wiener index
			sw_index = Shannon_Wiener_Index(clone_counts)
			hill_index = numpy.exp(sw_index)
		else:
			for clone in clone_counts:
				clone_freq = clone / total_clone_counts
				hill_index += (clone_freq ** order)
			order_exponent = 1.0 / (1.0 - order)
			hill_index = hill_index ** order_exponent

		order_index = (order, hill_index)
		hill_indices.append(order_index)

	if len(hill_indices) == 1:
		return hill_indices[0]
	else:
		return hill_indices

def Diversity_Plot(clone_df, png = None, title = "", count_col = "Clustered", split_col = None, line_width = 3,
				   add_control_diversities = True, figsize = (1000, 700)):
	"""Creates a plot comparing clonal repertoire diversity rates, using the Hill Diversity metric.

	Parameters
	----------
	clone_df: pandas DataFrame
		DataFrame of the repertoire(s) to plot
	png: str
		Title of the output PNG filename or None if none should be made; default is None
	title: str
		Title of the output graph; default is ""
	count_col: str
		Column name in clone_df of the clone counts/frequencies; default is "Clustered"
	split_col: str
		Column separating various repertoire subsets in clone_df or None if single repertoire; default is None
	line_width: int
		Width for the plot lines
	add_control_diversities: bool
		Whether to add lines for control diversities of artificial polarity; default is True
	figsize: tuple of (int, int)
		The width and height of the output plot; default is (100, 700)

	Returns
	----------
	plot: bokeh figure
		The figure object for the diversity plot
	"""

	figure_params = {
		"plot_width": figsize[0],
		"plot_height": figsize[1],
		"x_range": Range1d(0, 10),
		"y_axis_type": "log",
		"title": title,
		"tools": "save, help",
		"toolbar_location": "right"
	}

	plot = figure(**figure_params)
	plot.xgrid.grid_line_alpha = 0.0
	plot.xaxis.axis_label = "Order (N)"
	plot.yaxis.axis_label = "Hill Diversity Constant"
	plot.yaxis.formatter = BasicTickFormatter()

	#If comparing multiple samples, add the sample column to split on to the DataFrame
	if split_col is not None:
		diversity_df = clone_df[[count_col, split_col]]

		samples = []
		diversity_dfs = []
		for sample, df in diversity_df.groupby([split_col]):
			samples.append(sample)
			diversity_dfs.append(df)

	else:
		samples = ["Repertoire"]
		diversity_dfs = [clone_df[[count_col]]]

	sample_colors = (RGB(30, 160, 120), RGB(220, 90, 0), RGB(120, 110, 180), RGB(230, 40, 140))
	for sample, df, line_color in zip(samples, diversity_dfs, sample_colors[:len(samples)]):
		hill_indices = Hill_Diversity_Index(df[count_col])
		n_orders = [i[0] for i in hill_indices]
		order_diversities = [i[1] for i in hill_indices]

		#ADD MORE LINE STYLES (dotted, etc.)
		plot.line(x = n_orders, y = order_diversities, color = line_color, line_width = line_width, legend = sample)

	if add_control_diversities:
		total_clones = max([len(i) for i in diversity_dfs])
		total_counts = max([df[count_col].sum() for df in diversity_dfs])

		#Very highly polarized data creates a sample in which the top 20 clones are 20% of the total by prevalence
		top20_20_data = [total_counts * 0.2 / 20] * 20
		top20_20_data += [total_counts * 0.8 / (total_clones - 20) for _ in range(total_clones - 20)]
		#Highly polarized data has the top 20 clones at 15% of the total
		top20_15_data = [total_counts * 0.15 / 20] * 20
		top20_15_data += [total_counts * 0.85 / (total_clones - 20) for _ in range(total_clones - 20)]
		#Moderately polarized data has the top 20 clones at 10% of the total
		top20_10_data = [total_counts * 0.1 / 20] * 20
		top20_10_data += [total_counts * 0.9 / (total_clones - 20) for _ in range(total_clones - 20)]
		#Lowly polarized data has the top 20 clones at 5% of the total
		top20_5_data = [total_counts * 0.05 / 20] * 20
		top20_5_data += [total_counts * 0.95 / (total_clones - 20) for _ in range(total_clones - 20)]

		top20_20_diversities = [i[1] for i in Hill_Diversity_Index(top20_20_data)]
		top20_15_diversities = [i[1] for i in Hill_Diversity_Index(top20_15_data)]
		top20_10_diversities = [i[1] for i in Hill_Diversity_Index(top20_10_data)]
		top20_5_diversities = [i[1] for i in Hill_Diversity_Index(top20_5_data)]
		plot.line(x = n_orders, y = top20_20_diversities, color = RGB(160, 200, 230), alpha = 0.8, line_dash = (12,),
				  line_width = line_width, legend = "Very Highly Polarized (Top 20 Clones 20%)")
		plot.line(x = n_orders, y = top20_15_diversities, color = RGB(30, 120, 180), alpha = 0.8, line_dash = (12,),
				  line_width = line_width, legend = "Highly Polarized (Top 20 Clones 15%)")
		plot.line(x = n_orders, y = top20_10_diversities, color = RGB(180, 220, 140), alpha = 0.8, line_dash = (12,),
				  line_width = line_width, legend = "Moderately Polarized (Top 20 Clones 10%)")
		plot.line(x = n_orders, y = top20_5_diversities, color = RGB(50, 160, 40), alpha = 0.8, line_dash = (12,),
				  line_width = line_width, legend = "Lowly Polarized (Top 20 Clones 5%)")

	if png is not None:
		export_png(plot, png)

	return plot


from fenics import *
import matplotlib.pyplot as plt
import numpy as np

class nonlinearPDE1:
    def __init__(self, N):
        """
        Sets up the 2D PDE: -grad((1+u)^2.grad(u)) = 0

        Inputs:
        N: number of finite elements in mesh
        """
        # set up function space
        self.mesh = UnitSquareMesh(N,N)
        self.V = FunctionSpace(self.mesh, 'CR', 1)
        v = TestFunction(self.V)
        # set up BCs on left and right
        # lambda functions ensure the boundary methods take two variables
        self.B1 = DirichletBC(self.V, Constant(0.0), lambda x, on_boundary : self.left_boundary(x, on_boundary)) # u(0) = 0
        self.B2 = DirichletBC(self.V, Constant(1.0), lambda x, on_boundary : self.right_boundary(x, on_boundary)) # u(1) = 1

        # construct exact solution in C format
        self.uExpr = Expression('pow((pow(2,m+1) - 1)*x[0] + 1,(1/(m+1))) - 1', m = 2, degree=4)
        self.f = Constant(0.0)

    def left_boundary(self, x, on_boundary):
            return on_boundary and abs(x[0]) < 1E-14
    def right_boundary(self, x, on_boundary):
            return on_boundary and abs(x[0]-1) < 1E-14
    def q(self, u):
            return (1+u)**2
    def dqdu(self,u):
            return 2*(1+u)

    def solvePicard(self, iterTol = 1.0e-5, maxIter = 25, dispOutput = False,
                    writeData = True, filePath = 'solution-data/PDE1Picard'):
        """
        Solves the PDE using Picard iterations

        Inputs:
        iterTol: Iterations stop when |u_(k) - u_(k-1)| < iterTol. Default: 1e-5
        maxIter: Maximum number of iterations
        dispOutput(bool): display iteration differences and exact errors at each iteration
        writeData(True/False): write solution and convergence data to files
        ernErrors(True/False): calculate Ern error estimated at each iteration
        filePath: Path AND name of files WITHOUT file extension

        Outputs:
        u: solution to PDE
        iterDiffArray: Differences between iterative solutions (in H1 norm) at each iteration
        exactErrArray: Exact errors (in H1 norm) at each iteration

        Saved data:
        FEniCS solution saved to <filePath>.pvd
        Convergence data saved to <filePath>.csv:
            column 0: iterate differences
            column 1: exact errors
        """

        V = self.V
        bcs = [self.B1, self.B2]
        u = TrialFunction(V)
        v = TestFunction(V)
        u_k = TrialFunction(V)
        a0 = inner(grad(u_k), grad(v))*dx
        L0 = self.f*v*dx
        u_k = Function(V)
        solve(a0 == L0, u_k, bcs)
        a = inner(self.q(u_k)*grad(u), grad(v))*dx
        L = self.f*v*dx

        u = Function(V)     # new unknown function
        itErr = 1.0           # error measure ||u-u_k||
        iterDiffArray = [errornorm(u_k, u, 'H1')]
        exactErrArray = [errornorm(self.uExpr, u, 'H1')]
        R = assemble(self.q(u)*inner(grad(u),grad(v))*dx)
        self.B1.apply(R, u.vector())
        self.B2.apply(R, u.vector())
        residualArr = [R.norm('l2')]
        iter = 0
        # Begin Picard iterations
        while itErr > iterTol and iter < maxIter:
            iter += 1

            solve(a == L, u, bcs)

            # calculate iterate difference and exact error in L2 norm
            itErr = errornorm(u_k, u, 'H1')
            exErr = errornorm(self.uExpr, u, 'H1')
            iterDiffArray.append(itErr) # fill arrays with error data
            exactErrArray.append(exErr)
            R = assemble(self.q(u)*inner(grad(u),grad(v))*dx)
            self.B1.apply(R, u.vector())
            self.B2.apply(R, u.vector())
            residualArr.append(R.norm('l2'))

            if dispOutput:
                print('k = ' + str(iter) + ' | u-diff =  ' + str(itErr) + ', exact error = ' + str(exErr))
            u_k.assign(u) # update for next iteration

        if writeData:
            # save solution
            solution = File(filePath + '.pvd')
            solution << u_k
            # save convergence data
            convergenceData = [iterDiffArray, exactErrArray]
            np.savetxt(filePath + '.csv', convergenceData)

        # save data to object
        self.picardSol = u_k
        self.picardIterDiff = iterDiffArray
        self.picardExactErr = exactErrArray
        self.picardResArr = residualArr
        return [u_k, iterDiffArray, exactErrArray]

    def solveNewton(self, iterTol = 1.0e-5, maxIter = 25, dispOutput = False, writeData = True, filePath = 'solution-data/PDE1Newton'):
        """
        Solves the PDE using Newton iterations

        Inputs:
        iterTol: Iterations stop when |u_(k) - u_(k-1)| < iterTol. Default: 1e-5
        maxIter: Maximum number of iterations
        dispOutput(True/False): display iteration differences and exact errors at each iteration
        writeData(True/False): write solution and convergence data to files
        filePath: Path AND name of files WITHOUT file extension

        Outputs:
        u: solution to PDE
        iterDiffArray: Differences between iterative solutions (in H1 norm) at each iteration
        exactErrArray: Exact errors (in H1 norm) at each iteration

        Saved data:
        FEniCS solution saved to <filePath>.pvd
        Convergence data saved to <filePath>.csv:
            column 0: iterate differences
            column 1: exact errors
        """

        V = self.V
        v = TestFunction(V)

        # construct initial guess (solution to PDE with q(u) = 1)
        u_k = TrialFunction(V)
        a0 = inner(grad(u_k), grad(v))*dx
        f = Constant(0.0)
        L0 = self.f*v*dx

        u_k = Function(V)
        bcs = [self.B1, self.B2]
        solve(a0 == L0, u_k, bcs)

        # construct problem in du
        du = TrialFunction(V) # newton step
        a = (inner(self.q(u_k)*grad(du),grad(v)) + inner(self.dqdu(u_k)*du*grad(u_k),grad(v)))*dx
        L = -inner(self.q(u_k)*grad(u_k),grad(v))*dx
        # du = 0 on boundaries
        B2_du = DirichletBC(self.V, Constant(0.0), lambda x, on_boundary : self.right_boundary(x, on_boundary))
        bcs = [self.B1, B2_du]

        du = Function(V)
        u = Function(V)
        itErr = 1.0
        iterDiffArray = [errornorm(u_k, u, 'H1')]
        exactErrArray = [errornorm(self.uExpr, u, 'H1')]
        R = assemble(self.q(u)*inner(grad(u),grad(v))*dx)
        self.B1.apply(R, u.vector())
        self.B2.apply(R, u.vector())
        residualArr = [R.norm('l2')]
        iter = 0
        while itErr > iterTol and iter < maxIter:
            iter += 1

            solve(a == L, du, bcs)
            u.vector()[:] = u_k.vector() + du.vector()

            # calculate iterate difference and exact error in L2 norm
            itErr = errornorm(u_k, u, 'H1')
            exErr = errornorm(self.uExpr, u, 'H1')
            iterDiffArray.append(itErr) # fill arrays with error data
            exactErrArray.append(exErr)
            R = assemble(self.q(u)*inner(grad(u),grad(v))*dx)
            self.B1.apply(R, u.vector())
            self.B2.apply(R, u.vector())
            residualArr.append(R.norm('l2'))

            if dispOutput:
                print('k = ' + str(iter) + ' | u-diff =  ' + str(itErr) + ', exact error = ' + str(exErr))
            u_k.assign(u)

        if writeData:
            # save solution
            solution = File(filePath + '.pvd')
            solution << u_k
            # save convergence data
            convergenceData = [iterDiffArray, exactErrArray]
            np.savetxt(filePath + '.csv', convergenceData)

        # save data to object
        self.newtonSol = u_k
        self.newtonIterDiff = iterDiffArray
        self.newtonExactErr = exactErrArray
        self.newtonResArr = residualArr
        return [u_k, iterDiffArray, exactErrArray]

    def plotConvergence(self):
        """
        Plots the convergence data (exact errors and iterate differences) for
        Newton and/or Picard soltutions of the given PDE
        """
        from matplotlib import rc
        import matplotlib.pylab as plt

        rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})
        rc('text', usetex=True)
        # check which methods have been used to solve PDE
        if hasattr(self, 'newtonSol'):
            plt.figure(1, figsize = (12,10))
            plt.semilogy(self.newtonExactErr, 'r^-', linewidth=2, markersize=10)
            plt.semilogy(self.newtonIterDiff, 'b^-', linewidth=2, markersize=10)
            plt.semilogy(self.newtonResArr, 'g^-', linewidth=2, markersize=10)
            plt.ylabel('Error', fontsize=40)
            plt.xlabel('Iteration', fontsize=40)
            #plt.legend(['Exact error', 'Iterate difference', 'Residual'], loc=3, fontsize=30)
            plt.tick_params(labelsize=25)
        else:
            print('No Newton solution calculated, run solveNewton method first')
        if hasattr(self, 'picardSol'):
            plt.figure(2, figsize = (12,10))
            plt.semilogy(self.picardExactErr, 'r^-', linewidth=2, markersize=10)
            plt.semilogy(self.picardIterDiff, 'b^-', linewidth=2, markersize=10)
            plt.semilogy(self.picardResArr, 'g^-', linewidth=2, markersize=10)
            plt.ylabel('Error', fontsize=40)
            plt.xlabel('$k$', fontsize=40)
            plt.legend(['Exact error', 'Iterate difference', 'Residual'], loc=3, fontsize=30)
            plt.tick_params(labelsize=25)
        else:
            print('No Picard solution calculated, run solvePicard method first')

myPDE = nonlinearPDE1(30)
myPDE.solvePicard(iterTol = 0, maxIter = 10, dispOutput = True)
myPDE.solveNewton(iterTol = 0, maxIter = 10, dispOutput = True)
myPDE.plotConvergence()


"""Generate networks with two communities."""
# %%
import sys

import numpy as np
import pandas as pd
from scipy import sparse, stats

cave, cdiff = (
    100,
    80,
)  # average degree and difference betw. inter-/intra-community edges
n = 500
output_file = "edges.csv"

#
# Load
#
def sampling_num_edges(n, p):
    try:
        return stats.binom.rvs(n=n, p=p, size=1)[0]
    except ValueError:
        return np.sum(np.random.rand(n) < p)

def generate_dcSBM(cin, cout, N):
    n = int(np.floor(N / 2))
    pin, pout = cin / N, cout / N
    gids = np.concatenate([np.zeros(n), np.ones(n)])
    # d = (cin + cout) / 2

    within_edges = set([])
    target_num = sampling_num_edges(n=int(n * (n - 1) / 2 * 2), p=pin)
    esize = 0
    while esize < target_num:
        r = np.random.choice(2 * n, target_num - esize)
        c = np.random.choice(2 * n, target_num - esize)
        s = (gids[r] == gids[c]) * (r != c)
        r, c = r[s], c[s]
        r, c = np.maximum(r, c), np.minimum(r, c)
        eids = set(r + c * N)
        within_edges = within_edges.union(eids)
        esize = len(within_edges)

    between_edges = set([])
    target_num = sampling_num_edges(n=n * n, p=pout)
    esize = 0
    while esize < target_num:
        r = np.random.choice(2 * n, target_num - esize)
        c = np.random.choice(2 * n, target_num - esize)
        s = (gids[r] != gids[c]) * (r != c)
        r, c = r[s], c[s]
        r, c = np.maximum(r, c), np.minimum(r, c)
        eids = set(r + c * N)
        between_edges = between_edges.union(eids)
        esize = len(between_edges)

    edges = np.array(list(between_edges) + list(within_edges))
    r, c = divmod(edges, N)
    A = sparse.csr_matrix((np.ones_like(r), (r, c)), shape=(N, N))
    A = A + A.T
    A.data = np.ones_like(A.data)

    return A

#
# Preprocess
#
nc = n / 2
cin = (n - nc) / n * cdiff + cave
cout = cave - nc / n * cdiff

A = generate_dcSBM(cin, cout, n)

# %%
# Save
#
r, c, v = sparse.find(A)
pd.DataFrame({"src": r, "trg": c}).to_csv(output_file, index=False)


import numpy as np
from hypothises import hypothises

def cost_function(x, y, theta):
    m = len(y)
    h = hypothises(x, theta)
    hy = h - y
    hy = np.square(hy)
    cost = 1/(2*m) * np.sum(hy)
    return cost


## NUMERICS 3 - ALIASING ERROR

import numpy as np
import matplotlib.pyplot as plt

## grid

dx = 1
x = np.arange(0,10*dx,dx)
xa = np.linspace(0,max(x),1000)

kmax = np.pi/dx

k1 = 1
k2 = 2*kmax - k1

y1 = np.cos(k1*x)
y2 = np.cos(k2*x)

y1a = np.cos(k1*xa)
y2a = np.cos(k2*xa)

## plotting

plt.plot(xa,y1a,label=r'$k = 1, dx \to 0$')
plt.plot(xa,y2a,label=r'$k = 2\pi-1, dx \to 0$')
plt.plot(x,y1,'ko-',lw=2,label=r'$k = 1, dx = 1$')
plt.plot(x,y2,'rd--',lw=1,label=r'$k = 2\pi - 1, dx = 1$')
plt.title('Aliasing error')
plt.xlabel('x')
plt.ylabel(r'$\cos(kx)$')
plt.legend(loc='best')
plt.tight_layout()
plt.show()


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

import numpy as np

def MonteCarlo_double(f, g, x0, x1, y0, y1, n):
    """
    Monte Carlo integration of f over a domain g>=0, embedded
    in a rectangle [x0,x1]x[y0,y1]. n^2 is the number of
    random points.
    """
    # Draw n**2 random points in the rectangle
    x = np.random.uniform(x0, x1, n)
    y = np.random.uniform(y0, y1, n)
    # Compute sum of f values inside the integration domain
    f_mean = 0
    num_inside = 0   # number of x,y points inside domain (g>=0)
    for i in range(len(x)):
        for j in range(len(y)):
            if g(x[i], y[j]) >= 0:
                num_inside += 1
                f_mean += f(x[i], y[j])
    f_mean = f_mean/float(num_inside)
    area = num_inside/float(n**2)*(x1 - x0)*(y1 - y0)
    return area*f_mean


import random
import numpy as np
import math

#==========================================
#构造对称矩阵，包括两个城市之间的距离信息
def distance_p2p_mat(location,num_city):
    dis_mat=[]
    for i in range(num_city):
        dis_mat_each=[]
        for j in range(num_city):
            dis=math.sqrt(pow(location[i][0]-location[j][0],2)+pow(location[i][1]-location[j][1],2))
            dis_mat_each.append(round(dis,2))	#四舍五入保留两位小数
        dis_mat.append(dis_mat_each)
   # print(dis_mat)
    return dis_mat

#计算所有路径对应的距离
def cal_newpath(dis_mat,path,num_city):
    dis=0
    for j in range(num_city-1):
        dis=dis_mat[path[j]][path[j+1]]+dis
    dis=dis_mat[path[num_city-1]][path[0]]+dis#回家
    return dis

#==========================================

def simulatedAnneal(M,dis_mat,path,dis,num_city,t_current):
	while (t_current>lowest_t):#外循环，改变温度
	    count_m=0#M的计数
	    count_iter=0#迭代次数计数
	    while (count_m<M and count_iter<iteration):#内循环，连续多次不接受新的状态或者是迭代多次,跳出内循环
	        i=0
	        j=0
	        while(i==j):#防止随机了同一城市
	            i=random.randint(1,num_city-1)
	            j=random.randint(1,num_city-1)
	        path_new=path.copy()
	        path_new[i],path_new[j]=path_new[j],path_new[i]#任意交换两个城市的位置,产生新解
	        #计算新解的距离
	        dis_new=cal_newpath(dis_mat,path_new,num_city)
	        #求差
	        dis_delta=dis_new-dis
	        #取0-1浮点随机数
	        rand=random.random()
	        #计算指数函数的值
	        exp_d=math.exp(-dis_delta/t_current)
	        #选择
	        if dis_delta<0:
	            path=path_new
	            dis=dis_new
	        elif exp_d>rand:
	            path=path_new
	            dis=dis_new
	        else:
	            count_m=count_m+1
	        count_iter=count_iter+1
	    t_current=0.99*t_current#改变温度
	#外循环结束
	minDistanceDict = {'最短距离':dis, '最短路径':path}
	print(minDistanceDict)
	return dis,path

def getTestResult(location,num_city):
	initial_t = 100 * num_city    #初始温度
	global lowest_t,iteration
	lowest_t = 0.001  #最低温度
	iteration = 500  #设置迭代次数
	M = 50 * num_city  #当连续多次都不接受新的状态，开始改变温度

	#点对点距离矩阵
	dis_mat = distance_p2p_mat(location,num_city)
	#初始路径
	path = list(range(num_city))
	#初始距离
	dis = cal_newpath(dis_mat,path,num_city)
	#初始温度

	simulatedAnneal(M,dis_mat,path,dis,num_city,initial_t)

if __name__ == '__main__':
	location1 = np.loadtxt('city30_location.txt')	#test1
	#print(location1)
	location2 = np.loadtxt('city31_location.txt')	#test2
	num_city1 = 30	#城市总数
	num_city2 = 31	#城市总数

	getTestResult(location1,num_city1)
	getTestResult(location2,num_city2)


"""
bezier.py
by Ted Morin

a simple program to play with bezier curves
"""
import numpy as np
import sympy as sy
from scipy.special import comb as choose

def bezier1d(points):
#    d = len(points[0])
    N = len(points) - 1
    bez_x = lambda t: sum([
            choose(N,k)*np.power(t,k)*np.power(1-t,N-k)*points[k][0]
                                for k in range(N+1) ])
    bez_y = lambda t: sum([
            choose(N,k)*np.power(t,k)*np.power(1-t,N-k)*points[k][1]
                                for k in range(N+1) ])
    return bez_x, bez_y

if __name__ == "__main__":
    # prepare to plot some data
    import matplotlib.pyplot as plt
    plt.plot([0,1,1,0,0], [0,0,1,1,0], 'k-')
    n = 5
    plt.title("Click on {} points!".format(n))
    points = np.array(plt.ginput(n))

    x = points[:,0]
    y = points[:,1]
    #x = np.array([point[0] for point in points])
    #y = np.array([point[1] for point in points])
    bez_x, bez_y = bezier1d(points)

    # make the bezier curve data
    fit_points_per_input_point = 10
    t = np.linspace(0,1, len(x)*fit_points_per_input_point)
    x_bez = bez_x(t)
    y_bez = bez_y(t)

    plt.close()
    plt.plot([0,1,1,0,0], [0,0,1,1,0], 'k-')
    plt.plot(x_bez, y_bez, 'b--')
    plt.plot(x,y, 'r+')
    plt.title("Bezier!")
    plt.show()


# Program 14b: Bifurcation diagram of the logistic map.
# See Figures 14.15 and 14.16.

import numpy as np
import matplotlib.pyplot as plt

def f(x, r):
    return r * x * (1 - x)

ys = []
rs = np.linspace(0, 4, 2000)
#rs = np.linspace(3.5, 4, 2000) # For Figure 14.16.
for r in rs:
    x = 0.1
    for i in range(500):
        x = f(x, r)
    for i in range(50):
        x = f(x, r)
        ys.append([r, x])

ys = np.array(ys)

plt.plot(ys[:, 0], ys[:, 1], 'r.', markersize=0.05)
plt.xlabel('x', fontsize=15)
plt.ylabel('$\mu$', fontsize=15)
plt.tick_params(labelsize=15)
plt.show()


import numpy as np
import matplotlib.pyplot as plt
import cv2
import math
import pandas as pd
from numpy import linalg as linear

# Importing data from csv file using pandas

data = pd.read_csv('sumedh.csv')

# storing data in indivdual lists
ages = data['age'].to_numpy(int)
insurance_cost = data['charges'].to_numpy(np.float64)

# Creating Matrix with age and cost dataw
data = np.vstack((ages, insurance_cost)).T

# Formulating Covariance
def covariance(x, y):
    x_i, y_i = x.mean(), y.mean()
    return np.sum((x - x_i)*(y - y_i))/(len(x) - 1)

# Covariance matrix
def covariance_mat(matrix):
    return np.array([[covariance(matrix[0], matrix[0]), covariance(matrix[0], matrix[1])],
                     [covariance(matrix[1], matrix[0]), covariance(matrix[1], matrix[1])]])

def standardLS(points):

    x_axis=points[:,0]
    y_axis=points[:,1]

    # Implementing line equation
    M = np.stack((x_axis, np.ones((len(x_axis)), dtype = int)), axis = 1)

    M_transpose = M.transpose()
    MTM = M_transpose.dot(M)
    MTY = M_transpose.dot(M)
    ls_estimate = (np.linalg.inv(MTM)).dot(MTY)
    ls_value= M.dot(ls_estimate)

    return ls_value

def leastSquares(E,Y):

    E_transpose = E.transpose()
    ETE = E_transpose.dot(E)
    ETY = E_transpose.dot(Y)
    ls_est = (np.linalg.inv(ETE)).dot(ETY)
    return ls_est

def totalLeastSquares (x_ls,y_ls):

    #Constructing U matrix
    n = len(x_ls)

    x_mean=np.mean(x_ls)
    y_mean=np.mean(y_ls)

    U=np.vstack(((x_ls-x_mean),(y_ls-y_mean))).T

    # Evaluating UtU
    UTU=np.dot(U.transpose(),U)

    # Calculating coefficents of line equation y=ax+b
    beta=np.dot(UTU.transpose(),UTU)

    # Calculating eigen values of beta
    val,vect=linear.eig(beta)

    # Sorting the index of smallest eigen value
    ind =np.argmin(val)

    # eigen vector for smallest egien value
    coef =vect[:,ind]

    a,b= coef
    Ym =a*x_mean+b

    total_least_square=[]
    for i in range(0,n):
        temp=(Ym-(a*x_ls[i]))/b
        total_least_square.append(temp)

    return total_least_square

def ransacGenerator(ransac_array):

    x=ransac_array[:,0]
    y=ransac_array[:,1]

    # generating a random line array
    a =np.stack((x,np.ones((len(x)),dtype=int)),axis=1)

    #Setting threshold for inliers and outliers
    threshold=np.std(y)/2

    #Fiting a  line to the selected 2 points
    ransac_voting =ransacVoter(a,y,2, threshold)

    ransac_=a.dot(ransac_voting)

    return ransac_

def ransacVoter(line_array, y, size, threshold):
    max_i=math.inf
    iteration=0
    desired_prob=0.95
    best_fit=None
    max=0
    outlier=0

    ransac_data=np.column_stack((line_array,y))
    dsize=len(ransac_data)

    # finding points close to line
    while max_i >iteration:

        np.random.shuffle(ransac_data)
        data_set=ransac_data[:size,:]

        #creating line using least square
        x_t=data_set[:,:-1]
        y_t=data_set[:,-1:]
        voting_scheme=leastSquares(x_t,y_t)

        #counting inliers with defined range of threshold
        inliers=line_array.dot(voting_scheme)
        error =np.abs(y-inliers.T)
        count=np.count_nonzero(error<threshold)
        print("number of Inliers: ", count, "for iteration ", iteration)

        #best fit
        if count >max:
            max=count
            best_fit=voting_scheme

        #Outliers
        outlier=1-count/dsize
        max_i =math.log(1-desired_prob)/math.log(1-(1-outlier)**size)

        iteration+=1

    return best_fit

#Covariance for given data
cov =covariance_mat(data.T)

c_value, c_vector = np.linalg.eig(cov)

# Plotting data scatter plot
plt.scatter(data[:, 0], data[:, 1])

origin=[np.mean(data[:,0]),np.mean(data[:,1])]
eigen_1=c_vector[:,0]
eigen_2=c_vector[:,1]

#Calculating Least Square
ls = standardLS(data)
#Calculating total Least Square
total_ls = totalLeastSquares(ages, insurance_cost)
#Calculating RANSAC
rans = ransacGenerator(data)

#Plotting data and curve fitting
#Plotting Covariance vector and egien vectors
plt.figure(1)
plt.quiver(*origin, *eigen_1, color=['r'], scale=15)
plt.quiver(*origin, *eigen_2, color=['b'], scale=15)
plt.title('Covariance Matrix')
plt.xlabel("age")
plt.ylabel('Insurance cost')

#Curve fitting using Least Square
plt.figure(2)
plt.xlabel('Age')
plt.ylabel('Insurance cost')
plt.title('curve fitting using Least Square ')
plt.scatter(ages,insurance_cost,c='g',label='given data')
plt.plot(ages,ls, c='red', label='Linear Least Squares')
plt.legend()

#Curve fitting using Total Least Square
plt.figure(3)
plt.xlabel('Age')
plt.ylabel('Insurance cost')
plt.title('curve fitting using Total Least Square ')
plt.scatter(ages,insurance_cost,c='y',label='given data')
plt.plot(ages,total_ls, c='red', label='Linear Least Squares')
plt.legend()

plt.figure(4)
plt.xlabel('Age')
plt.ylabel('Insurance cost')
plt.title('curve fitting using RANSAC ')
plt.scatter(ages,insurance_cost,c='y',label='given data')
plt.plot(ages,rans, c='red', label='RANSAC')
plt.legend()

plt.show()


import cv2
import numpy as np
from scipy import ndimage
from sys import argv

if len(argv) ==  2:
    file_name = argv[1]
else:
    file_name = "../images/statue_small.jpg"

kernel_3x3 = np.array([[-1, -1, -1],
                   [-1,  8, -1],
                   [-1, -1, -1]])

kernel_5x5 = np.array([[-1, -1, -1, -1, -1],
                       [-1,  1,  2,  1, -1],
                       [-1,  2,  4,  2, -1],
                       [-1,  1,  2,  1, -1],
                       [-1, -1, -1, -1, -1]])

# 转换为灰度图像，numpy只接收一维数组
img = cv2.imread(file_name, 0)

# 3x3 卷积 实现高通滤波器
k3 = ndimage.convolve(img, kernel_3x3)
# 5x5 卷积 实现高通滤波器
k5 = ndimage.convolve(img, kernel_5x5)

# 通过对图像应用 低通滤波器 后，与原图像计算差值
blurred = cv2.GaussianBlur(img, (17,17), 0)
g_hpf = img - blurred

cv2.imwrite("../samples/hpf_3x3.jpg",k3)
cv2.imwrite("../samples/hpf_5x5.jpg",k5)
cv2.imwrite("../samples/hpf_g_hpf.jpg",g_hpf)
# cv2.imshow("3x3", k3)
# cv2.imshow("5x5", k5)
# cv2.imshow("g_hpf", g_hpf)
cv2.waitKey()
cv2.destroyAllWindows()


import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
#i线性回归
plt.rcParams['figure.figsize']=(14,8)

n_observations=100
xs = np.linspace(-3,3,n_observations)
ys = np.sin(xs)+np.random.uniform(-0.5,0.5,n_observations)
plt.scatter(xs,ys)
plt.show()

X = tf.placeholder(tf.float32,name='X')
Y = tf.placeholder(tf.float32,name='Y')

W=tf.Variable(tf.random_normal([1]),name='weight')
b=tf.Variable(tf.random_normal([1]),name='bias')\

Y_pred=tf.add(tf.multiply(X,W),b)

loss=tf.square(Y-Y_pred,name='loss')

learning_rate=0.01
optimizer=tf.train.GradientDescentOptimizer(learning_rate).minimize(loss)

n_sample=xs.shape[0]
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    writer=tf.summary.FileWriter('./graphs/linear_reg',sess.graph)

    for i in range(50):
        total_loss=0
        for x,y in zip(xs,ys):
            _,los=sess.run([optimizer,loss],feed_dict={X:x,Y:y})
            total_loss+=los

        if i%5==0:
            print('Epoch {0}: {1}'.format(i,total_loss/n_sample))

    writer.close()

    W,b=sess.run([W,b])

print(W,b)
print('W:'+str(W[0]))
print('b:'+str(b[0]))

plt.plot(xs,ys,'bo',label='Real data')
plt.plot(xs,xs*W+b,'r',label='Predicted data')
plt.legend()
plt.show()


"""
  Multiobjective Branin and Currin Exponential Functions.
  -- kandasamy@cs.cmu.edu
"""

# You should have one of the following two in this file:
#   1. objectives: A list of functions, each of which corresponds to an objective.
#   2. A function called compute_objectives which returns the value of all objectives when
#      called as a list, and an integer called num_objectives which indicates the number
#      of objectives.

# pylint: disable=invalid-name

import numpy as np

def branin_with_params(x, a, b, c, r, s, t):
  """ Computes the Branin function. """
  x1 = x[0]
  x2 = x[1]
  neg_ret = float(a * (x2 - b*x1**2 + c*x1 - r)**2 + s*(1-t)*np.cos(x1) + s)
  return - float(neg_ret)

def branin(x):
  """ Branin function."""
  a = 1
  b = 5.1/(4*np.pi**2)
  c = 5/np.pi
  r = 6
  s = 10
  t = 1/(8*np.pi)
  return branin_with_params(x, a, b, c, r, s, t)

def currin_exp_01(x):
  """ Currin exponential function. """
  x1 = x[0]
  x2 = x[1]
  val_1 = 1 - np.exp(-1/(2 * x2))
  val_2 = (2300*x1**3 + 1900*x1**2 + 2092*x1 + 60) / (100*x1**3 + 500*x1**2 + 4*x1 + 20)
  return float(val_1 * val_2)

def currin_exp(x):
  """ Currint exponential in branin bounds. """
  return currin_exp_01([x[0] * 15 - 5, x[1] * 15])

# # Option 1: Define objectives to be a list of the functions
# objectives = [branin, currin_exp]

# Option 2: Define num_objectives to be the number of objectives and compute_objectives to
# return a list of values of each function.
num_objectives = 2
def compute_objectives(x):
  """ Computes the objectives. """
  return [branin(x), currin_exp_01(x)]


from __future__ import division

import numpy as np
import numpy.random as rn
import scipy.special as sp
import scipy.stats as st

########################################################################################################################

def betaln(x, a=1, b=1):
    """Returns the log-density of the beta distribution at x"""

    ## return
    return sp.gammaln(a+b) - sp.gammaln(a) - sp.gammaln(b) + (a-1) * np.log(x) + (b-1) * np.log(1-x)

########################################################################################################################

def normalln(x, mean=0, var=1):
    """Returns the log-density of the normal distribution at x"""

    ## return
    return - 0.5 * (x - mean)**2 / var - 0.5 * np.log(2 * np.pi * var)

########################################################################################################################

def lognormalln(x, mean=0, var=1):
    """Returns the log-density of the lognormal distribution at x"""

    ## return
    return -np.log(x) - 0.5 * (np.log(x) - mean)**2 / var - 0.5 * np.log(2 * np.pi * var)

########################################################################################################################

def poissonln(x, rate=1):
    """Returns the log-density of the Poisson distribution at x"""

    ## return
    return x * np.log(rate) - sp.gammaln(x + 1) - rate

########################################################################################################################

def binomln(x, n=1, p=0.5):
    """Returns the log-density of the binomial distribution at x"""

    ## return
    return sp.gammaln(n + 1) - sp.gammaln(x + 1) - sp.gammaln(n - x + 1) + x * np.log(p) + (n - x) * np.log(1 - p)

########################################################################################################################

def nbinomln(x, alpha=1, p=0.5):
    """Returns the log-density of the negative binomial distribution at x"""

    ## return
    return sp.gammaln(x + alpha) - sp.gammaln(alpha) - sp.gammaln(x + 1) + alpha * np.log(p) + x * np.log1p(-p)

########################################################################################################################

def bbinomln(x, n=1, alpha=0.5, beta=0.5):
    """Returns the log-density of the beta binomial distribution at x"""

    ## compute intermediate quantities
    c1 = sp.gammaln(n + 1) - sp.gammaln(x + 1) - sp.gammaln(n - x + 1)
    c2 = sp.gammaln(x + alpha) + sp.gammaln(n - x + beta) - sp.gammaln(n + alpha + beta)
    c3 = sp.gammaln(alpha + beta) - sp.gammaln(alpha) - sp.gammaln(beta)

    ## return
    return c1 + c2 + c3

########################################################################################################################

def exponentialln(x, scale=1):
    """Returns the log-density of the exponential distribution at x"""

    ##
    return -np.log(scale) - x / scale

########################################################################################################################

def sample_normal_mean(s1, ndata, prec, m0=0, t0=0):
    """Samples the mean of a normal distribution"""

    ##
    t = t0 + prec * ndata
    m = (t0 * m0 + prec * s1) / t

    ## return
    return rn.normal(m, 1 / np.sqrt(t))

########################################################################################################################

def sample_normal_prec(s1, s2, ndata, mean, a0=0, b0=0):
    """Samples the precision of a normal distribution"""

    ##
    dot = s2 - 2 * mean * s1 + ndata * mean**2

    a = a0 + ndata * 0.5
    b = b0 + 0.5 * dot

    ## return
    return rn.gamma(a, 1 / b)

########################################################################################################################

def sample_normal_mean_prec(s1, s2, ndata, m0=0, l0=0, a0=0, b0=0):
    """Samples the precision of a normal distribution"""

    ##
    avg = s1 / ndata
    dot = s2 - 2 * avg * s1 + ndata * avg**2

    l = l0 + ndata
    m = (l0 * m0 + s1) / l
    a = a0 + 0.5 * ndata
    b = b0 + 0.5 * dot + 0.5 * l0 * ndata * (avg - m0)**2 / l

    ##
    prec = rn.gamma(a, 1 / b)
    mean = rn.normal(m, 1 / np.sqrt(l * prec))

    ## return
    return mean, prec

########################################################################################################################

def sample_normal_mean_jeffreys(s1, ndata, prec):
    """Samples the mean of a normal distribution"""

    ##
    return rn.normal(s1 / ndata, 1 / np.sqrt(prec * ndata))

########################################################################################################################

def sample_normal_prec_jeffreys(s1, s2, ndata):
    """Samples the precision of a normal distribution"""

    ##
    avg = s1 / ndata
    dot = s2 - 2 * avg * s1 + ndata * avg**2

    ##
    return rn.gamma(ndata * 0.5, 2 / dot)

########################################################################################################################

def sample_normal_mean_prec_jeffreys(s1, s2, ndata):
    """Samples the precision of a normal distribution"""

    ##
    avg = s1 / ndata
    dot = s2 - 2 * avg * s1 + ndata * avg**2

    mean = st.t.rvs(ndata+1, avg, dot / (ndata * ndata + ndata))
    prec = rn.gamma((ndata+1) * 0.5, 2 / dot)

    ##
    return mean, prec

########################################################################################################################

def sample_gamma_rate(s, ndata, shape, a0=0, b0=0):
    """Samples the scale of the gamma distribution from its posterior, when shape is known"""

    ## return
    return rn.gamma(a0 + ndata * shape, 1 / (b0 + s))

########################################################################################################################

def sample_gamma_shape(sl, ndata, shape, rate, la0=0, b0=0, c0=0):
    """Samples the shape of the gamma distribution from its posterior, when scale is known"""

    ## compute updated params
    la = la0 + sl
    b = b0 + ndata
    c = c0 + ndata

    ## make proposal
    shape_ = shape * np.exp(0.01 * rn.randn())

    ## compute logpost and logpost_
    logpost = (shape - 1) * la + shape * c * np.log(rate) - b * sp.gammaln(shape)
    logpost_ = (shape_ - 1) * la + shape_ * c * np.log(rate) - b * sp.gammaln(shape_)

    ## do Metropolis step
    if logpost_ > logpost or rn.rand() < np.exp(logpost_ - logpost):
        shape = shape_

    ## return
    return shape

########################################################################################################################

def sample_dirichlet(a):
    """Sample from multiple Dirichlet distributions given the matrix of concentration parameter columns a"""

    x = rn.gamma(a, 1)
    w = x / np.sum(x, 0)

    ##
    return w

########################################################################################################################

def sample_categorical(w, nsamples=1):
    """Samples from the categorical distribution with matrix of weight columns w"""

    _, ncols = w.shape

    ws = w.cumsum(0)
    ws[-1] = 1                                      # sum of ws along rows should be equal to 1

    idxs = np.sum(ws[:, :, np.newaxis] < rn.rand(ncols, nsamples), 0)

    ## return
    return idxs.T

########################################################################################################################

def sample_stick(cluster_occupancies, eta):
    """Samples random stick lengths (in logarithmic scale) given a vector of cluster occupancies"""

    ## compute the cumulative sum of the count vector
    cs = cluster_occupancies.cumsum()

    ## generate beta variates
    v = rn.beta(1 + cluster_occupancies, eta + cs[-1] - cs)
    v[-1] = 1                  # this ensures that sum(w) = 1
    v = np.clip(v, 1e-12, 1 - 1e-12)

    ## compute weights
    lv = np.log(v)
    lcp = np.log(1-v).cumsum()

    lw = np.r_[lv[0], lv[1:] + lcp[:-1]]

    ## return
    return lw, lv

########################################################################################################################

def sample_eta_ishwaran(lw, eta, a=0, b=0):
    """Samples the concentration parameter eta given a vector of mixture log-weights"""

    eta = rn.gamma(lw.size + a - 1, 1 / (b - lw[-1])) if np.isfinite(lw[-1]) else eta

    ##
    return eta

########################################################################################################################

def sample_eta_west(eta, nact, n0, a=1, b=0):
    """Samples the concentration parameter eta"""

    ## compute x, r and p
    x = rn.beta(eta + 1, n0)
    lx = np.log(x)
    r = (a + nact - 1) / (n0 * (b - lx))
    p = r / (r + 1)

    ## return
    return rn.gamma(a + nact, 1 / (b - lx)) if rn.rand() < p else rn.gamma(a + nact - 1, 1 / (b - lx))

########################################################################################################################

def sample_eta(eta, nact, n0, a=0, b=0):
    """Samples the concentration parameter eta"""

    ## proposal
    eta_ = eta * np.exp(0.01 * rn.randn())

    ## posterior densities
    lp = sp.gammaln(eta) - sp.gammaln(eta + n0) + (nact + a - 1) * np.log(eta) - eta * b
    lp_ = sp.gammaln(eta_) - sp.gammaln(eta_ + n0) + (nact + a - 1) * np.log(eta_) - eta_ * b

    ## return
    return eta_ if lp_ > lp or rn.rand() < np.exp(lp_ - lp) else eta

########################################################################################################################


import numpy as np
import matplotlib.pyplot as plt

# Random values make for crappy fitting
#input_data = 5 * np.random.rand(10,2)

# Make a perfect linear relationship y = 2x
input_data = np.array([
			[1,1],[2,4],[3,6],[4,8],
			[5,10],[6,12],[7,14],[8,16]
			])

# Add in random error
error = np.zeros((8,2))
error[:,1] = 5 * np.random.rand(8).T
input_data = input_data + error

m = len(input_data)
# X is composed of the transpose of a row of 1s and the x values as such:
# [ 1 1 1 1 ]    [ 1 1 ]
# [ 1 2 3 4 ] -> [ 2 1 ]
#                [ 3 1 ]
#                [ 4 1 ]
X = np.array([np.ones(m), input_data[:,0]]).T

# y is composed of the second column of the input data, reshaped to a single
# column.  Since we gave -1 as the first argument to reshape, it will infer
# the amount of rows based on the size of the input.  In this case, that will
# become 4, so we get the 4x1 column matrix with values 6,5,7,and 10
y = np.array(input_data[:,1]).reshape(-1,1)

# Numpy has a built-in solver for scalar linear equations of the form Ax=B.
# numpy.linalg.solve(a,b) expects a square matrix a of coefficients, and an
# ordinate vector b sharing a dimension with a.  The resultant matrix x is
# the exact solution

# If the coefficients happen to violate either condition (not square or
# singular, use numpy.linalg.lstsq() instead to find the "best" solution x

# Here, we use it to solve as such:  [X^t][X][params] = [X^t][y]
params = np.linalg.solve(X.T.dot(X), X.T.dot(y))

# Build a plot
plt.figure(1)
xvals = np.linspace(0,10,2)
yvals = np.array(params[0] + params[1] * xvals)
plt.plot(xvals, yvals.T, color='b')
plt.scatter(input_data[:,0], input_data[:,1], color='r')

# Create labels with the exact and best fit equation
plt.text(0.5, 22, "Exact: y = 2*x", fontsize=10)
plt.text(0.5, 20, "Best Fit: y = {:.2f}*x + {:.2f}".format(float(params[1]), float(params[0])), fontsize=10)

plt.show()


# ------------------------------------------------------------------------------------------------------------------------------------
# Title             : Multivariate Regression using Gradient Descent
# Description       : This is an implementation of gradient descent algorithm to minimize the least squares error for Multivariate Linear regression.
# Author            : Subburam Rajaram
# IMAT Number       : 03658122
# ------------------------------------------------------------------------------------------------------------------------------------

import numpy as np

# Initialize constants
stepsize = 0.01
iterations = 1000

# Read dataset from disk and structure it
Data = np.loadtxt("airfoil_self_noise.dat")
rows = Data.shape[0]
columns = Data.shape[1]
Y = Data[:,columns - 1 ]        # Labels
Y = Y.reshape((rows,1))
F = Data[:,0:(columns - 1)]     # Features

# Feature scaling and Mean normalization
F_normed = np.empty_like(F)
F_mean = F.mean(0)
F_std = F.std(0)
for i in range(0, rows):
    F_normed[i,:] = np.divide(( F[i,:] - F_mean ), F_std)

# stack column of ones to X for the constant term
ones = np.ones(rows)
X = np.column_stack((ones,F_normed))

# Weight Attribute Vector where W[0] would hold the value of constant term
Weight = np.zeros((columns, 1))
Cost = np.zeros((iterations, 1))

for iter in range(0, iterations):
    print "\nIteration:", iter
    Cost[iter] = (0.5/rows) * np.dot((np.dot(X, Weight) - Y).T, (np.dot(X, Weight) - Y))
    print "Cost:", Cost[iter]
    gradient = np.dot(X.T, (np.dot(X, Weight) - Y))/rows
    Weight = Weight - stepsize * gradient
    print "Intercept b: ", Weight[0]
    print "Weight W: ", Weight[1:columns].T


# =========================================================================
#
#    Graduate Course: Finite Element Analysis - Spring 2022 @JBNU
#
#  ========================================================================
#
#   Example 8.2: 2D Truss Problem
#
#   Last updated: 24/05/2022 by Minh-Chien Trinh (mctrinh@jbnu.ac.kr)
#
#   Copyright 2022 Minh-Chien Trinh. All rights reserved.
#
# =========================================================================

import numpy as np

def Stiffness2DTruss(gDof, nElem, eNode, nNode, nCoord, xCoord, yCoord, EA):

    # Initialization
    Kmat = np.zeros((gDof,gDof))

    # Computation of the system stiffness matrix
    for i in range(nElem):
        # Element degree of freedom
        indice = eNode[i,:]

        # Element DOFs
        eDof = np.array([indice[0]*2-1, indice[0]*2, indice[1]*2-1, indice[1]*2])
        nDof = np.size(eDof)

        # Element length
        xa = xCoord[indice[1]-1] - xCoord[indice[0]-1]
        ya = yCoord[indice[1]-1] - xCoord[indice[0]-1]
        eLeng = np.sqrt(xa**2 + ya**2)

        # The contribution of the element to the stiffness matrix
        C = xa / eLeng
        S = ya / eLeng
        k1 = EA / eLeng * np.array([[C*C, C*S, -C*C, -C*S], [C*S, S*S, -C*S, -S*S],
                                [-C*C, -C*S, C*C, C*S], [-C*S, -S*S, C*S, S*S]])

        # Assemble the contribution to the global stiffness matrix
        # Row index
        rIndex = np.zeros((nDof,nDof), dtype=int)
        rIndex[0:4,0] = eDof - 1
        rIndex[0:4,1] = eDof - 1
        rIndex[0:4,2] = eDof - 1
        rIndex[0:4,3] = eDof - 1
        # Column index
        cIndex = np.zeros((nDof,nDof), dtype=int)
        cIndex[0,0:4] = eDof - 1
        cIndex[1,0:4] = eDof - 1
        cIndex[2,0:4] = eDof - 1
        cIndex[3,0:4] = eDof - 1
        # Assemble
        Kmat[rIndex,cIndex] = Kmat[rIndex,cIndex] + k1

    return Kmat

def Stresses2DTruss(nElem, eNode, xCoord, yCoord, Umat, E):

    # Initialization
    sigma = np.zeros((nElem, 1))

    for i in range(nElem):
        # Element degree of freedom
        indice = eNode[i,:]

        # Element DOFs
        eDof = np.array([indice[0]*2-1, indice[0]*2, indice[1]*2-1, indice[1]*2])
        nDof = np.size(eDof)

        # Element length
        xa = xCoord[indice[1]-1] - xCoord[indice[0]-1]
        ya = yCoord[indice[1]-1] - xCoord[indice[0]-1]
        eLeng = np.sqrt(xa**2 + ya**2)

        C = xa / eLeng
        S = ya / eLeng

        # Element stress
        sigma[i,0] = E / eLeng * np.matmul(np.array([-C, -S, C, S]),
                        np.array([Umat[eDof[0]-1,0], Umat[eDof[1]-1,0], Umat[eDof[2]-1,0], Umat[eDof[3]-1,0]]))

    return sigma

# Given paramerers
E = 30e6; A = 2; EA = E*A

# Number of elements
nElem = 3

# Number of nodes
nNode = 4

# Element nodes
eNode = np.array([[1,2], [1,3], [1,4]], dtype=int)

# Node coordinates
nCoord = np.array([[0,0], [0,120], [120,120], [120,0]])
xCoord = nCoord[:,0]
yCoord = nCoord[:,1]

# Total number of degree of freedom
gDof  = 2 * nNode

# Initialization
Umat = np.zeros((gDof,1))              # Displacement vector
Fmat = np.zeros((gDof,1))              # Force vector
Kmat = np.zeros((gDof,gDof))           # Stiffness matrix

# Applied load at node 1, 2nd dof
Fmat[1,0] = -10000

# Computation of the system stiffness matrix
Kmat = Stiffness2DTruss(gDof, nElem, eNode, nNode, nCoord, xCoord, yCoord, EA)
print("---------------------------------------------------------------------")
print("Stiffness matrix (Kmat): \n\n", Kmat)
print("---------------------------------------------------------------------")

# Apply boundary condition
# Fix/prescribed degree of freedom
fixDof = np.array([[2],[3],[4],[5],[6],[7]])
# Free/active degree of freedom
activeDof = np.setdiff1d(np.arange(0,gDof,1), fixDof)       # 02 active Dofs

# Solution
# Indexing of activeDof
iIndex = np.zeros((np.size(activeDof),np.size(activeDof)), dtype=int)
iIndex[0:2,0] = activeDof
iIndex[0:2,1] = activeDof
jIndex = np.zeros((np.size(activeDof),np.size(activeDof)), dtype=int)
jIndex[0,0:2] = activeDof
jIndex[1,0:2] = activeDof
# Evaluate displacement at activeDof
Umat[activeDof,0] = np.dot(np.linalg.inv(Kmat[iIndex,jIndex]),Fmat[activeDof,0])
print("Displacement vector (Umat): \n\n", Umat)
print("-------------------------------------------------")

# Force vector
Fmat = np.dot(Kmat,Umat)
print("Force vector (Fmat): \n\n", Fmat)
print("-------------------------------------------------")

# Reactions
Reaction = np.zeros(np.size(fixDof))
Reaction = Fmat[fixDof,0]
print("Reactions: \n\n", Reaction)
print("-------------------------------------------------")

# Stresses of elements
sigma = np.zeros((nElem, 1))
sigma = Stresses2DTruss(nElem, eNode, xCoord, yCoord, Umat, E)
print("Stresses: \n\n", sigma)
print("-------------------------------------------------")

# -------------------------------------------------------------------------
#
# For students:
#
# 1. Plot the intact and deformed shapes of the system
#
# -------------------------------------------------------------------------

# Add your code from there ...


from scipy.optimize import linear_sum_assignment

import numpy as np

def abs_dist_cost_matrix(points_1, points_2):
    """Generates the cost matrix suitable for scipy.optimize.linear_sum_assignment between two sets of points

    Args:
        points_1 (np.ndarray): A set of points with dimensions (point index, number of spatial dimensions)
        points_2 (np.ndarray): A set of points with dimensions (point index, number of spatial dimensions)

    Returns:
        (np.ndarray of shape (points_1.shape[0], points_2.shape[0])): The Euclidean distances between
        pairs of points in points_1 and points_2 such that the distance from pair points_1[i] points_2[j] occurs at
        position (j, i) in output

    """
    vect_dists = points_1 - points_2.reshape([points_2.shape[0], 1, points_2.shape[1]])
    return np.linalg.norm(vect_dists, axis=-1)

def match_indices(points_0, points_1):
    """Gets the zeroth-dimension indices in points_1 that match zeroth-dimension elements of points_2 such that the linear
    sum assignment of Euclidean distance costs are minimised

    Args:
        points_0 (np.ndarray): np.ndarray
        points_1 (np.ndarray): np.ndarray

    Returns:
        (tuple of two lists of ints): Each list will contain the integer indicies such that matching, ordered point
        sets are recovered via points_0[return_val[1]] closest matches are points_1[return_val[0]]

    """
    cost_mat = abs_dist_cost_matrix(points_0, points_1)
    return linear_sum_assignment(cost_mat)


from prepare_data import stock_data, benchmark_data
import matplotlib.pyplot as plt
import numpy as np

"""
The inputs for the Sharpe Ratio: Starting with Daily Stock Returns

The Sharpe Ratio uses the difference in returns between the two investment opportunities under consideration.

However, our data show the historical value of each investment, not the return.
To calculate the return, we need to calculate the percentage change in value from one day to the next.

"""

# calculate daily stock_data returns
stock_returns = stock_data.pct_change()

# The pandas method .pct_change() calculates the relative change
# from one value to the next in a DataFrame.
stock_returns.plot(subplots = True)
plt.show()

# summarize the daily returns
print(stock_returns.describe())

# calculate daily benchmark_data returns
sp_returns = benchmark_data['S&P 500'].pct_change()

# plot the daily returns
sp_returns.plot();
plt.show()

# summarize the daily returns
print(sp_returns.describe())

"""
Calculating Excess Returns for Amazon and Facebook vs. S&P 500

Next, we need to calculate the relative performance of stocks vs. the S&P 500 benchmark.
This is calculated as the difference in returns between stock_returns and sp_returns for each day.

"""

# calculate the difference in daily returns
excess_returns = stock_returns.sub(sp_returns, axis = 0)

# plot the excess_returns
excess_returns.plot(subplots = True)

# summarize the excess_returns
print(excess_returns.describe())

"""

The Sharpe Ratio, Step 1: The Average Difference in Daily Returns Stocks vs S&P 500

Now we can finally start computing the Sharpe Ratio.
First we need to calculate the average of the excess_returns.
This tells us how much more or less the investment yields per day compared to the benchmark.

"""

# calculate the mean of excess_returns
avg_excess_return = excess_returns.mean()

# plot avg_excess_returns
avg_excess_return.plot.bar(title = "Mean of the Return Difference")
plt.show()

"""

The Sharpe Ratio, Step 2: Standard Deviation of the Return Difference

t looks like there was quite a bit of a difference between average daily returns for Amazon and Facebook.

Next, we calculate the standard deviation of the excess_returns.
This shows us the amount of risk an investment in the stocks implies as compared to an investment in the S&P 500.

"""

# calculate the standard deviations
sd_excess_return = excess_returns.std()

# plot the standard deviations
sd_excess_return.plot.bar(title = "Standard Deviation of the Return Difference");
plt.show()

"""

Putting it all together

Now we just need to compute the ratio of avg_excess_returns and sd_excess_returns.
The result is now finally the Sharpe ratio and indicates how much more (or less)
return the investment opportunity under consideration yields per unit of risk.

"""

# calculate the daily sharpe ratio
daily_sharpe_ratio = avg_excess_return.div(sd_excess_return)

# annualize the sharpe ratio
annual_factor = np.sqrt(252)
annual_sharpe_ratio = daily_sharpe_ratio.mul(annual_factor)

# plot the annualized sharpe ratio
annual_sharpe_ratio.plot.bar(title = "Annualized Sharpe Ratio: Stocks vs S&P 500");
plt.show()


import numpy as np

def img2col(x, fh, fw, stride=1, pad=0, pad_value=0):
    x = np.pad(x, pad_width=((0, 0), (0, 0), (pad, pad), (pad, pad)), mode='constant', constant_values=pad_value)
    N, C, H, W = x.shape
    OH = (H - fh) / stride + 1
    OW = (W - fw) / stride + 1

    result = []
    for n in range(N):
        for h in range(int(OH)):
            for w in range(int(OW)):
               result.append(x[n, :, h * stride : h * stride + fh, w * stride : w * stride + fw].flatten())

    return np.array(result)

if __name__ == '__main__':
    x = np.arange(64).reshape(2,2,4,4)
    print(x)
    print(np.array(img2col(x, 3, 3, 1, 0)))


#The OLG Model
#%%
# 1. Setup

#importing the necessary packages
import numpy as np
import sympy as sm
import matplotlib.pyplot as plt
import ipywidgets as widgets

# 2. Symbolic solution of the Household's (i) and Firm's (ii) problems
# (i) Household's problem

#%%
# Activating pretty printing
sm.init_printing(use_unicode=True)

# Defining symbolically the necessary variables
c1 = sm.symbols('c_1_t')
c2 = sm.symbols('c_2_t+1')
beta = sm.symbols('beta')
z = sm.symbols('z_t+1')
w = sm.symbols('w_t') # wage
r1 = sm.symbols('r_t+1')
u = sm.symbols('U_t') #Utility
L = sm.symbols('L') #Langrange
omega = sm.symbols('omega') #The Langrangian multiplier

# Checking that our variables have the proper format
c1,c2, beta, w, z, r1, L, omega

#%%
# Defining the utility function
def utility(c1, c2, beta):
    return (sm.log(c1) + beta*sm.log(c2))

#%%
#Defining the two budget constraints symbolically
budget_constraint_1 = sm.Eq(z, w - c1)
budget_constraint_2 = sm.Eq(c2, z*(1+r1))
budget_constraint_1, budget_constraint_2

#%%
#Solve the two budget constraints to obtain the Lifetime Budget Constraint in three steps:
# 1.Solve  for 𝑧𝑡+1  from the budget constraint 1
# 2.Substitute in budget constraint 2
# 3.Solve for  𝑤𝑡

# Step 1
z_from_constraint = sm.solve(budget_constraint_1, z)
z_from_constraint

#%%
# Step 2
z_subs_constraint = budget_constraint_2.subs(z, z_from_constraint[0])
z_subs_constraint

#%%
# Step 3
lifetime_budget_constraint = sm.solve(z_subs_constraint, w)
lifetime_budget_constraint

#%%
# Defining the Lifetime Budget Constraint symbolically
budget_constraint = w - c1 - c2/(1+r1)
budget_constraint

#%%
#Constructing the Langrangian
Langrangian = utility(c1, c2, beta) + omega * budget_constraint
Langrangian

#%%
#Solve in six steps:
# 1.Calculate the two FOCs wrt.  𝑐1,𝑡  and  𝑐2,𝑡+1
# 2.Isolate the first FOC for  𝜔
# 3.Substitute in the second FOC
# 4.Isolate the second FOC for  𝑐1,𝑡
# 5.Substitute into the Lifetime Budget Constraint
# 6.Solve for  𝑐2,𝑡+1  to find the consumption in period 2
# 7.Substitute  𝑐2,𝑡+1  in Lifetime Budget Constraint
# 8.Solve for  𝑐1,𝑡  to find the consumption in period 1
# 9.Substitute in the first budget constraint ( 𝑧𝑡+1=𝑤𝑡−𝑐1,𝑡 )
# 10.Solve for  𝑧𝑡+1  to find the assets held by the household at period 2

# Step 1
foc1 = sm.diff(Langrangian, c1)
foc2 = sm.diff(Langrangian, c2)
foc1, foc2

#%%
# Step 2
foc_1 = sm.solve(foc1, omega)
foc_1

#%%
#Step 3
foc_2 = foc2.subs(omega, foc_1[0])
foc_2

#%%
# Step 4
AB = sm.solve(foc_2, c1)
AB

#%%
# Step 5
ABC = budget_constraint.subs(c1, AB[0])
ABC

#%%
# Step 6
consumption_2 = sm.solve(ABC,c2)
consumption_2

#%%
# Step 7
ABCD = budget_constraint.subs(c2, consumption_2[0])
ABCD

#%%
# Step 8
consumption_1 = sm.solve(ABCD, c1)
consumption_1

#%%
# Step 9
ABCDE = budget_constraint_1.subs(c1, consumption_1[0])
ABCDE

#%%
# Step 10
Assets = sm.solve(ABCDE, z)
Assets

# (ii) Firm's Problem

#%%
# Defining symbolically the necessary variables
alpha = sm.symbols('alpha') # share of physical capital
k = sm.symbols ('k_t')  # capital per effective worker
r = sm.symbols('r_t') # interest rate
delta = sm.symbols('delta') # depreciation rate of capital per period
A = sm.symbols('A_t') # technological index at period t
A1 = sm.symbols('A_{t+1}')
g = sm.symbols ('g') #rate of technological growth
L0 = sm.symbols('L_{t}')
L1 = sm.symbols('L_{t+1}')
n = sm.symbols('n') #growth of labor force

# Checking that our variables have the proper format
alpha, k, r, delta, A, A1, g, L0, L1, n

#%%
tech = sm.Eq((1+g)*A, A1)
labor = sm.Eq((1+n)*L0, L1)
tech, labor

#%%
tech0 = sm.solve(tech, A)
labor0 = sm.solve(labor, L0)
tech0, labor0

#%%
# Defining the two FOC's and setting them equal with their rental rates
MPK = sm.Eq((alpha*k**(alpha-1)), r+delta)
MPL = sm.Eq (((1-alpha)*A*k**alpha), w)
MPK, MPL

#%%
Wage = sm.solve(MPL, w)
Wage

# 3. Symbolic (iii) and numerical (iv) solution for the aggregate level of the model
# (iii) Symbolic solution for the aggregate level of the model

#%%
# Defining symbolically the necessary variables
Ks = sm.symbols('K_{t+1}^{s}') #aggregate supply of physical capital at period t+1
Kd = sm.symbols('K_{t+1}^{s}')#aggregate demand of physical capital at period t+1
K = sm.symbols('K_{t+1}')
g1 = sm.symbols('g^{*}')
k1 = sm.symbols('k_{t+1}')

# Checking that our variables have the proper format
Ks, Kd, K, g1, k1

#%%
#Solve in seven steps:
# 1.Define the capital supply  𝐾𝑠𝑡+1
# 2.Substitute for  𝑧𝑡+1  as calculated in the Household's problem
# 3.Substitute for  𝑤𝑡  as calculated in the Firm's problem
# 4.Substitute  𝐴𝑡  with  𝐴𝑡+1
# 5.Substitute  𝐿𝑡  with  𝐿𝑡+1
# 6.Substitute  (1+𝑛)(1+𝑔)  with  𝑔∗
# 7.Equalize  𝐾𝑠𝑡+1=𝐾𝑑𝑡+1=𝐾𝑡+1

# Step 1
capital_supply = sm.Eq(L0*z ,Ks)
capital_supply

#%%
# Step 2
H = capital_supply.subs(z, Assets[0])
H

#%%
# Step 3
I = H.subs(w, Wage[0])
I

#%%
# Step 4
J = I.subs(A, tech0[0])
J

#%%
# Step 5
M = J.subs(L0, labor0[0])
M

#%%
# Step 6
N = M.subs(((1+g)*(1+n)), (g1+1))
N

#%%
# Step 7
P = N.subs(Ks, K)
P

#%%
sm.init_printing(use_unicode=False)

#(iv) Numerical solution for the aggregate level of the model

#%%
# Converting the calculated equations to python function
g1 = lambda n, g: (1+g)*(1+n)
solution = lambda k, alpha = 0.4 , g = 0.01, n = 0.06, beta = 0.8: (1-alpha)*k**alpha / ((1+g1(n,g))*(1+1/beta))
y = lambda x: x # this will be used to plot a 45o degree line at the first plot

#%%
# Creating values for the python functions
k_values = np.linspace(0, 0.04, 10000)
x = np.linspace(0, 0.04, 10000)
solution(k_values)

#%%
fig = plt.figure()
ax = fig.add_subplot(1,1,1)

ax.plot(k_values, solution(k_values))
ax.plot(x, y(x))

plt.xlabel('$k_{t}$')
plt.ylabel('$k_{t+1}$')
plt.title('Key Transition Equation')

plt.show()

# 4.Implementation of the Bisection method at the steady state level of the model

#%%
# Creating the Key Transition Equation for the steady state
f = lambda k, alpha=0.4, n=0.01, g=0.06, beta=0.8: k - ((1-alpha)/((1+g1(n,g))*(1+1/beta)))** (1/(1-alpha))

#%%
# Implementing the Bisection method
def bisection(f,a,b,max_iter=500,tol=1e-6,full_info=False):
    """ bisection

    Args:

        f (function): function
        a (float): left bound
        b (float): right bound
        tol (float): tolerance on solution

    Returns:

        m (float): root

    """

    # test inputs
    if f(a)*f(b) >= 0:
        print("bisection method fails.")
        return None

    # step 1: initialize
    _a = a
    _b = b
    a = np.zeros(max_iter)
    b = np.zeros(max_iter)
    m = np.zeros(max_iter)
    fm = np.zeros(max_iter)
    a[0] = _a
    b[0] = _b

    # step 2-4: main
    i = 0
    while i < max_iter:

        # step 2: midpoint and associated value
        m[i] = (a[i]+b[i])/2
        fm[i] = f(m[i])

        # step 3: determine sub-interval
        if abs(fm[i]) < tol:
            break
        elif f(a[i])*fm[i] < 0:
            a[i+1] = a[i]
            b[i+1] = m[i]
        elif f(b[i])*fm[i] < 0:
            a[i+1] = m[i]
            b[i+1] = b[i]
        else:
            print("bisection method fails.")
            return None

        i += 1

    if full_info:
        return m,i,a,b,fm
    else:
        return m[i],i

#%%
m,i = bisection(f,-10,10)
print(i,m,f(m))

#%%
def plot_bisection(f, a, b, xmin=-10, xmax=10, xn=100):

    # a. find root and return all information
    m, max_iter, a, b, fm = bisection(f, a, b, full_info=True)

    # b. compute function on grid
    xvec = np.linspace(xmin, xmax, xn)
    fxvec = f(xvec)

    # c. figure
    def _figure(i):
        fig = plt.figure(dpi=100)
        ax = fig.add_subplot(1,1,1)

        ax.plot(xvec,fxvec) # on grid
        ax.plot(m[i],fm[i],'o',color='black',label='current') # mid
        ax.plot([a[i],b[i]],[fm[i],fm[i]],'--',color='black',label='range') # range
        ax.axvline(a[i],ls='--',color='black')
        ax.axvline(b[i],ls='--',color='black')

        ax.legend(loc='lower right')
        ax.grid(True)
        ax.set_ylim([fxvec[0],fxvec[-1]])

    widgets.interact(_figure,
        i=widgets.IntSlider(description="iterations", min=0, max=max_iter, step=1, value=0))

plot_bisection(f, -10, 10)


# %%
import numpy as np
import numpy.linalg as la
import time
import scipy.optimize as sciopt

# %%
# Soft-thresholding
def sft(x, t):
    return x * (1.0 - np.minimum(t / np.maximum(np.abs(x), 1E-32), 1.0))

# Projection onto the l1-ball
def project_l1_ball(x, t):
    xL1Nrm  = la.norm(x.ravel(), 1)
    if xL1Nrm > t:
        # If the input is outside the l1-ball we project.
        nx = x.size;
        # Sort magnitudes in decreasing order
        mx = np.flip(np.sort(np.absolute(x.ravel()), kind='mergesort'))
        # Obtain unique values (this is needed for some borderline cases)
        umx = np.flip(np.unique(mx))
        if t <= umx[0] - umx[1]:
            # If the radius of the l1-ball is small and the input is far away
            # and the shrinkage is large.
            vp = umx[0]
            sroot = vp - t
            return sft(x, sroot)
        else:
            # This is the average case
            cmx = np.cumsum(mx)
            smx = cmx - np.array(range(1, nx+1)) * mx
        if smx[nx - 1] < t:
            # This condition handles some cases when the input is close to
            # the boundary and the shrinkage is small.
            sroot = (xL1Nrm - t) / nx
            return sft(x, sroot)
        idp, = np.where( smx > t )
        idp = idp[0];
        sroot = (cmx[idp - 1] - t) / idp
        return sft(x, sroot)
    # If the input is within the l1-ball we do nothing.
    return x

# Minimize Sum-of-Squares
#   min_x || Ax - y ||_2^2
def MinimizeSumOfSquares(y, A, xinit=None,
                            maxItns=1E4,
                            dxAbsTol=1E-6, dxRelTol=1E-9,
                            gradAbsTol=1E-5, gradRelTol=1E-8,
                            disp=False, printEvery=100):
    # Find Lipschitz constant of the gradient
    L = 2.05 * A.norm() ** 2
    # Initialize variables
    if xinit is None:
        x = np.zeros(A.wavShape, dtype=np.complex)
    else:
        x = xinit
    xNrm = la.norm(x.ravel(), 2)
    dxNrm = np.inf
    itn = 0
    fx = la.norm((y - A.eval(x)).ravel(), 2) ** 2
    # Optimization loop
    stop = False
    if disp:
        print('[MinimizeSumOfSquares]')
        print(' Initial objective:   {:1.6e}'.format(fx))
    while not stop:
        itn = itn + 1
        xp = x - (2.0/L) * A.adjoint(A.eval(x) - y)
        fxp = la.norm((y - A.eval(xp)).ravel(), 2) ** 2
        dxNrm = la.norm((xp - x).ravel(), 2)
        gradNrm = L * dxNrm
        dxNrmTol = np.max([ dxAbsTol,
                            dxRelTol * xNrm,
                            L * gradAbsTol,
                            L * gradRelTol * xNrm ])
        if dxNrm < dxNrmTol or itn > maxItns:
            stop = True
        x = xp
        fx = fxp
        xNrm = la.norm(x.ravel(), 2)
        if disp and (itn == 1 or np.mod(itn, printEvery) == 0 or stop):
            print(' {:04d} | obj: {:1.3E} | gradNrm: {:1.3E}'.format(itn, fx, gradNrm))
    return x, fx

# Minimize Sum-of-Squares subject to l1-norm
#   min_x || Ax - y ||_2^2 s.t. ||x||_1 <= t
def MinimizeSumOfSquaresL1Ball(t, y, A, xinit=None, L=None,
                                maxItns=1E4,
                                dxAbsTol=1E-4, dxRelTol=1E-6,
                                dpAbsTol=1E-5, dpRelTol=1E-8,
                                disp=False, printEvery=10,
                                restart=True):
    # Project initial iterate on the l1-ball
    if xinit is None:
        x = np.zeros(A.inShape, dtype=np.complex)
    else:
        x = project_l1_ball(xinit, t)
    if L is None:
        L = 2.05 * A.norm() ** 2
    rx = A.eval(x) - y
    fx = la.norm(rx.ravel()) ** 2
    gx = 2.0 * A.adjoint(rx)
    rpx = x - project_l1_ball(x - gx/L, t)
    rpxNrm = la.norm(rpx.ravel())
    # Initialize variables
    xNrm = la.norm(x.ravel())
    dxNrm = np.inf
    dfx = -np.inf
    itn = 0
    s = 1
    z = x
    # Optimization loop
    stop = False
    if disp:
        print('[PGD-SoS-L1Ball]')
        print(' Initial objective:   {:1.6e}'.format(fx))
        print(' Initial residual:    {:1.6e}'.format(rpxNrm))
    while not stop:
        itn = itn + 1
        gz = 2.0 * A.adjoint(A.eval(z, 1) - y)
        xp = project_l1_ball(z - gz/L, t)
        sp = 0.5 * (1.0 + np.sqrt(1.0 + 4.0 * s ** 2))
        zp = xp + ((s - 1.0)/sp) * (xp - x)

        # Update
        rxp = A.eval(xp) - y
        fxp = la.norm(rxp.ravel()) ** 2

        dfx = fxp - fx
        dxNrm = la.norm((xp - x).ravel())

        if restart and dfx > 1E-12:
            z = x
            s = 1.0
            dfx = np.inf
            dxNrm = np.inf
        else:
            x = xp
            gx = 2.0 * A.adjoint(rxp)
            rpx = x - project_l1_ball(x - gx/L, t)
            rpxNrm = la.norm(rpx.ravel())

            z = zp
            fx = fxp
            s = sp

        xNrm = la.norm(x.ravel())
        dxTol = np.max([ dxAbsTol, dxRelTol * xNrm ])
        rpxTol = np.max([ dpAbsTol, dpRelTol * xNrm ])
        if (dxNrm < dxTol and rpxNrm < rpxTol) or itn > maxItns:
            stop = True
        if disp and (itn == 1 or np.mod(itn, printEvery) == 0 or stop):
            print(' {:04d} | obj: {:1.3E} | dfx: {:+1.3E} | dxNrm: {:1.3E} | rpxNrm: {:5.3E}'.format(itn, fx, dfx, dxNrm, rpxNrm))
    return x, fx

# Minimize Sum-of-Squares
#   min_x || Ax - y ||_2^2 + t ||x||_1
def MinimizeBPDN(t, y, A, xinit=None, L=None,
                    maxItns=1E4,
                    dxAbsTol=1E-4, dxRelTol=1E-6,
                    dpAbsTol=1E-5, dpRelTol=1E-8,
                    disp=False, printEvery=10,
                    restart=True):
    # Project initial iterate on the l1-ball
    if xinit is None:
        x = np.zeros(A.inShape, dtype=np.complex)
    else:
        x = project_l1_ball(xinit, t)
    z = x
    if L is None:
        L = 2.05 * A.norm() ** 2
    rx = A.eval(x) - y
    fx = la.norm(rx.ravel()) ** 2 + t * la.norm(x.ravel(), 1)
    gx = 2.0 * A.adjoint(rx)
    rpx = x - sft(x - gx, t)
    rpxNrm = la.norm(rpx.ravel())
    # Initialize variables
    xNrm = la.norm(x.ravel())
    dxNrm = np.inf
    dfx = -np.inf
    itn = 0
    s = 1.0
    # Optimization loop
    stop = False
    if disp:
        print('[PGD-BPDN]')
        print(' Initial objective:   {:1.6e}'.format(fx))
        print(' Initial residual:    {:1.6e}'.format(rpxNrm))
    while not stop:
        itn = itn + 1
        gz = 2.0 * A.adjoint(A.eval(z, 1) - y)
        xp = sft(z - gz/L, t/L)
        sp = 0.5 * (1.0 + np.sqrt(1.0 + 4.0 * s ** 2))
        zp = xp + ((s - 1.0)/sp) * (xp - x)

        # Update
        rxp = A.eval(xp) - y
        fxp = la.norm(rxp.ravel()) ** 2 + t * la.norm(xp.ravel(), 1)

        dfx = fxp - fx
        dxNrm = la.norm((xp - x).ravel())

        if restart and dfx > 1E-12:
            z = x
            s = 1.0
            dfx = np.inf
            dxNrm = np.inf
        else:
            x = xp
            gx = 2.0 * A.adjoint(rxp)
            rpx = x - sft(x - gx, t)
            rpxNrm = la.norm(rpx.ravel())

            z = zp
            fx = fxp
            s = sp

        xNrm = la.norm(x.ravel())
        dxTol = np.max([ dxAbsTol, dxRelTol * xNrm ])
        rpxTol = np.max([ dpAbsTol, dpRelTol * xNrm ])
        if (dxNrm < dxTol and rpxNrm < rpxTol) or itn > maxItns:
            stop = True

        if disp and (itn == 1 or np.mod(itn, printEvery) == 0 or stop):
            print(' {:04d} | obj: {:1.3E} | dfx: {:+1.3E} | dxNrm: {:1.3E} | rpxNrm: {:5.3E}'.format(itn, fx, dfx, dxNrm, rpxNrm))
    return x, fx

class RootSolverL1NormNoisy():
    def __init__(self, eta, y, A,
                    maxItns=1E4,
                    dxAbsTol=1E-4, dxRelTol=1E-6,
                    dpAbsTol=1E-5, dpRelTol=1E-8,
                    disp=False, printEvery=10,
                    restart=True,
                    method='SoS-L1Ball', disp_method=False):
        # Problem parameters
        self.eta = eta
        self.y = y
        self.A = A
        self.disp_method = disp_method

        self.ANrm = A.norm()
        self.L = 2.05 * self.ANrm ** 2

        # For bracketing
        self._itn = 0
        self._t = [ None, None ]
        self._ft = [ None, None ]
        self._xinit = [ None, None ]

        _xopt, _ = MinimizeSumOfSquares(y, A,
                                            maxItns=maxItns,
                                            dxAbsTol=dxAbsTol, dxRelTol=dxRelTol,
                                            gradAbsTol=dpAbsTol, gradRelTol=dpRelTol,
                                            disp=False)
        _ropt = A.eval(_xopt) - y
        if method == 'SoS-L1Ball':
            # Lower bound
            self._t[0] = 0.0
            self._ft[0] = la.norm(y.ravel()) - eta
            self._xinit[0] = np.zeros(A.inShape, dtype=np.complex)
            # Upper bound
            self._t[1] = la.norm(_xopt.ravel(), 1)
            self._ft[1] = la.norm(_ropt.ravel(), 2) - eta
            self._xinit[1] = _xopt
        elif method == 'BPDN':
            # Lower bound
            self._t[0] = 0.0
            self._ft[0] = -eta
            self._xinit[0] = _xopt
            # Upper bound
            _z = A.adjoint(y)
            self._t[1] = la.norm(_z.ravel(), np.inf)
            self._ft[1] = la.norm(y.ravel(), 2) - eta
            self._xinit[1] = np.zeros(A.inShape, dtype=np.complex)
        else:
            raise ValueError('Method can be either SoS-L1Ball or BPDN.')
        self.method = method
        # Solver options
        self.restart = restart
        # Solver parameters
        self.maxItns = maxItns
        self.dxAbsTol = dxAbsTol
        self.dxRelTol = dxRelTol
        self.dpAbsTol = dpAbsTol
        self.dpRelTol = dpRelTol
        # Display options
        self.disp = disp
        self.printEvery = printEvery

    def solve(self, t):
        self._itn = self._itn + 1
        if self.disp:
            print(' {:02d} | t* = {:1.4E} in [{:1.4E}, {:1.4E}] | f: [{:+1.4E}, {:+1.4E}]'.format(self._itn, t, self._t[0], self._t[1], self._ft[0], self._ft[1]))

        if np.abs(t - self._t[0]) < np.abs(t - self._t[1]):
            tNear = 0
        else:
            tNear = 1

        theta = np.abs(t - self._t[0]) / (np.abs(t - self._t[0]) + np.abs(t - self._t[1]))
        xinit = theta * self._xinit[0] + (1.0 - theta) * self._xinit[1]

        if self.method == 'SoS-L1Ball':
            xopt, _ = MinimizeSumOfSquaresL1Ball(t, self.y, self.A, xinit=xinit, L=self.L,
                                                    maxItns=self.maxItns,
                                                    dxAbsTol=self.dxAbsTol, dxRelTol=self.dxRelTol,
                                                    dpAbsTol=self.dpAbsTol, dpRelTol=self.dpRelTol,
                                                    disp=self.disp_method, printEvery=self.printEvery,
                                                    restart=self.restart)
        elif self.method == 'BPDN':
            xopt, _ = MinimizeBPDN(t, self.y, self.A, xinit=xinit, L=self.L,
                                    maxItns=self.maxItns,
                                    dxAbsTol=self.dxAbsTol, dxRelTol=self.dxRelTol,
                                    dpAbsTol=self.dpAbsTol, dpRelTol=self.dpRelTol,
                                    disp=self.disp_method, printEvery=self.printEvery,
                                    restart=self.restart)

        ft = la.norm((self.A.eval(xopt) - self.y).ravel(), 2) - self.eta

        if t > self._t[1]:
            self._t[0] = self._t[1]
            self._ft[0] = self._ft[1]
            self._xinit[0] = self._xinit[1]

            self._t[1] = t
            self._ft[1] = ft
            self._xinit[1] = xopt
        elif t < self._t[0]:
            self._t[1] = self._t[0]
            self._ft[1] = self._ft[0]
            self._xinit[1] = self._xinit[0]

            self._t[0] = t
            self._ft[0] = ft
            self._xinit[0] = xopt
        else:
            self._t[tNear] = t
            self._ft[tNear] = ft
            self._xinit[tNear] = xopt

        return ft

def RecoveryL1NormNoisy(eta, y, A,
                    maxItns=1E4,
                    dxAbsTol=1E-4, dxRelTol=1E-6,
                    dpAbsTol=1E-5, dpRelTol=1E-8,
                    disp=False, printEvery=100,
                    method='BPDN', disp_method=False):

    # Create solver object
    T = RootSolverL1NormNoisy(eta, y, A,
                            maxItns=maxItns,
                            dxAbsTol=dxAbsTol, dxRelTol=dxRelTol,
                            dpAbsTol=dpAbsTol, dpRelTol=dpRelTol,
                            disp=disp, printEvery=printEvery,
                            restart=True, method=method, disp_method=disp_method)
    if disp:
        print('[RECOVERY: L1-NOISY]' )
        print(' Method:          {:s}'.format(method) )
        print(' Bound:           {:1.4E}'.format(eta) )
        print(' Initial values:')
        print('     t           [{:1.4E},   {:1.4E}]'.format(T._t[0], T._t[1]))
        print('     f(t)        [{:+1.4E}, {:+1.4E}]'.format(T._ft[0], T._ft[1]))
    # Root finding via TOMS748
    tStart = time.time()
    tMin = T._t[0]
    tMax = T._t[1]
    topt, rnfo = sciopt.toms748(lambda t: T.solve(t), tMin, tMax, k=1, xtol=1E-4, full_output=disp, disp=disp)
    if np.abs(topt - T._t[0]) < np.abs(topt - T._t[1]):
        xopt = T._xinit[0]
    else:
        xopt = T._xinit[1]
    tStop = time.time() - tStart
    if disp:
        print('  Summary');
        print('      Iterations:          {:d}'.format(rnfo.iterations) )
        print('      Function calls:      {:d}'.format(rnfo.function_calls) )
        print('      Optimal value:       {:5.3E}'.format(topt) )
        print('      Elapsed time:        {:8.4f} seconds'.format(tStop) )
    return xopt, la.norm(xopt.ravel(), 1)

# %%
if __name__ == '__main__':
    # Import
    from maps import OperatorWaveletToFourier, OperatorWaveletToFourierX4

    do_single_image = True
    do_x4_images = False

    print('[TESTING SOLVER ROUTINES]')
    # Set up map
    imShape = (128, 128)
    # *** EXPERIMENTS FOR SINGLE IMAGE AND SUBSAMPLING - L1-NORM
    if do_single_image:
        samplingSet = np.where(np.random.uniform(size=imShape) < 0.75, True, False)
        A = OperatorWaveletToFourier(imShape, samplingSet=samplingSet, waveletName='db4')
        xinit = np.random.normal(size=A.inShape) + 1j * np.random.normal(size=A.inShape)
        y = A.eval(xinit)
        print('[SINGLE IMAGE - RANDOM UNDERSAMPLING - L1-NORM RECOVERY]')
        print('    Input shape:    {:d} x {:d}'.format(A.inShape[0], A.inShape[1]))
        print('    Output shape:   {:d} x {:d}'.format(A.outShape[0], 1))
        print('    Operator norm:  {:1.3e}'.format(A.norm()))
        # Tests
        print('\n[TEST: MinimizeSumOfSquares]')
        xsol, fmin = MinimizeSumOfSquares(y, A, disp=True)

        print('\n[TEST: MinimizeSumOfSquaresL1Ball]')
        print('  ---> with no restart...')
        xsol, fmin = MinimizeSumOfSquaresL1Ball(0.25, y, A, disp=True, restart=False)
        print('  ---> with restart...')
        _xsol, _fmin = MinimizeSumOfSquaresL1Ball(0.25, y, A, disp=True, restart=True)
        print('  ---> Difference:     {:1.5e}'.format(la.norm((xsol - _xsol).ravel())))

        print('\n[TEST: RecoveryL1Noisy - SoS-L1Ball]')
        xsol, fmin = RecoveryL1NormNoisy(10, y, A, maxItns=1E4, disp=True, printEvery=100, method='SoS-L1Ball')

        print('\n[TEST: RecoveryL1Noisy - BPDN]')
        _xsol, _fmin = RecoveryL1NormNoisy(10, y, A, maxItns=1E4, disp=True, printEvery=100, method='BPDN')
        print('  ---> Difference:')
        print('         Objective:    {:1.5e}'.format(np.abs(fmin - _fmin)))
        print('         Solution:     {:1.5e}'.format(la.norm((xsol - _xsol).ravel())))

    # *** EXPERIMENTS FOR MULTIPLE IMAGES AND SUBSAMPLING - L1-NORM
    if do_x4_images:
        samplingSet = np.where(np.random.uniform(size=imShape + (4,)) < 0.75, True, False)
        A = OperatorWaveletToFourierX4(imShape, samplingSet=samplingSet, waveletName='db4')
        xinit = np.random.normal(size=A.inShape) + 1j * np.random.normal(size=A.inShape)
        y = A.eval(xinit)
        print('[4 IMAGES - RANDOM UNDERSAMPLING - L1-NORM RECOVERY]')
        print('    Input shape:    {:d} x {:d} x {:d}'.format(A.inShape[0], A.inShape[1], A.inShape[2]))
        print('    Output shape:   {:d} x {:d}'.format(A.outShape[0], 1))
        print('    Operator norm:  {:1.3e}'.format(A.norm()))
        # Tests
        print('\n[TEST: MinimizeSumOfSquares]')
        xsol, fmin = MinimizeSumOfSquares(y, A, disp=True)

        print('\n[TEST: MinimizeSumOfSquaresL1Ball]')
        print('  ---> with no restart...')
        xsol, fmin = MinimizeSumOfSquaresL1Ball(0.25, y, A, disp=True, restart=False)
        print('  ---> with restart...')
        _xsol, _fmin = MinimizeSumOfSquaresL1Ball(0.25, y, A, disp=True, restart=True)
        print('  ---> Difference:     {:1.5e}'.format(la.norm((xsol - _xsol).ravel())))

        print('\n[TEST: RecoveryL1Noisy - SoS-L1Ball]')
        xsol, fmin = RecoveryL1NormNoisy(10, y, A, maxItns=1E4, disp=True, printEvery=100, method='SoS-L1Ball')

        print('\n[TEST: RecoveryL1Noisy - BPDN]')
        _xsol, _fmin = RecoveryL1NormNoisy(10, y, A, maxItns=1E4, disp=True, printEvery=100, method='BPDN')
        print('  ---> Difference:')
        print('         Objective:    {:1.5e}'.format(np.abs(fmin - _fmin)))
        print('         Solution:     {:1.5e}'.format(la.norm((xsol - _xsol).ravel())))
# %%


#================================================
#Solutions To Value Function Iteration Lab
#================================================
#"Problem 1"
#import scipy as sp
#from matplotlib import pyplot as plt
#from matplotlib import cm
#from mpl_toolkits . mplot3d import Axes3D
#
#
#beta = 0.9;
#T = 10;
#N = 100;
#u = lambda c: sp.sqrt(c);
#W = sp.linspace(0,1,N);
#X, Y = sp.meshgrid(W,W);
#Wdiff = Y-X
#index = Wdiff <0;
#Wdiff[index] = 0;
#util_grid = u(Wdiff);
#util_grid[index] = -10**10;
#V = sp.zeros((N,T+2));
#psi = sp.zeros((N,T+1));
#
#
#for k in xrange(T,-1,-1):
#    val = util_grid + beta*sp.tile(sp.transpose(V[:,k+1]),(N,1));
#    vt = sp.amax(val, axis = 1);
#    psi_ind = sp.argmax(val,axis = 1)
#    V[:,k]    = vt;
#    psi[:,k]    = W[psi_ind];
#
#
#
#x=sp.arange(0,N)
#y=sp.arange(0,T+2)
#X,Y=sp.meshgrid(x,y)
#fig1 = plt.figure()
#ax1= Axes3D(fig1)
#ax1.plot_surface(W[X],Y,sp.transpose(V), cmap=cm.coolwarm)
#plt.show ()
#
#fig2 = plt.figure()
#ax2 = Axes3D(fig2)
#y = sp.arange(0,T+1)
#X,Y=sp.meshgrid(x,y)
#ax2.plot_surface(W[X],Y,sp.transpose(psi), cmap = cm.coolwarm)
#plt.show()

#================================================
"Problem 2"
#import scipy as sp
#from matplotlib import pyplot as plt
#
#beta = 0.9;
#T = 1000;
#N = 100;
#u = lambda c: sp.sqrt(c);
#W = sp.linspace(0,1,N);
#X, Y = sp.meshgrid(W,W);
#Wdiff = sp.transpose(X-Y);
#index = Wdiff <0;
#Wdiff[index] = 0;
#util_grid = u(Wdiff);
#util_grid[index] = -10**10;
#V = sp.zeros((N,T+2));
#psi = sp.zeros((N,T+1));
#
#
#for k in xrange(T,-1,-1):
#    val = util_grid + beta*sp.tile(sp.transpose(V[:,k+1]),(N,1));
#    vt = sp.amax(val, axis = 1);
#    psi_ind = sp.argmax(val,axis = 1)
#    V[:,k]    = vt;
#    psi[:,k]    = W[psi_ind];
#
#
#
#plt.plot(psi[99,:])

#================================================
#"Problem 3"
#import scipy as sp
#from matplotlib import pyplot as plt
#
#beta = 0.99
#N = 1000
#u = lambda c: sp.sqrt(c)
#W = sp.linspace(0,1,N)
#X, Y = sp.meshgrid(W,W)
#Wdiff = sp.transpose(X-Y)
#index = Wdiff <0
#Wdiff[index] = 0
#util_grid = u(Wdiff)
#util_grid[index] = -10**10
#
#Vprime = sp.zeros((N,1))
#psi = sp.zeros((N,1))
#delta = 1.0
#tol = 10**-9
#it = 0
#max_iter = 500
#
#while (delta >= tol) and (it < max_iter):
#    V = Vprime
#    it += 1;
#    print(it)
#    val = util_grid + beta*sp.transpose(V)
#    Vprime = sp.amax(val, axis = 1)
#    Vprime = Vprime.reshape((N,1))
#    psi_ind = sp.argmax(val,axis = 1)
#    psi    = W[psi_ind]
#    delta = sp.dot(sp.transpose(Vprime - V),Vprime-V)

#plt.plot(W,psi)
#plt.show()


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

# In[4]:

import numpy as np
import seaborn as sns
import pandas as pd
import timeit
import matplotlib.pyplot as plt
from numba import jit
from scipy.stats import t
from scipy.stats import norm
from scipy.stats import multivariate_normal
from mpl_toolkits import mplot3d
from scipy.special import gamma
from sklearn.neighbors import KernelDensity
from sklearn.model_selection import GridSearchCV

def ratio(x,mu,sigma,df,c):
    return t.pdf(x,df)/(c*norm.pdf(x,mu,sigma))

def rejection_sampling1(L,c,mu,sigma,df):
    y=np.zeros(L)
    i=0
    j=0
    while i<L:
        j+=1
        x=np.random.normal(mu,sigma,1)
        if np.random.uniform(0,1,1)<=ratio(x,mu,sigma,df,c):
            y[i]=x
            i+=1
    return y,j

def rejection_sampling2(L,c,mu,sigma,df):
    y=np.zeros(L)
    i=j=0
    while i<L:
        j+=1
        x=np.random.normal(mu,sigma,1)
        if np.random.uniform(0,1,1)<=ratio(x,mu,sigma,df,c):
            y[i]=x
            c_star=np.maximum(c,c*ratio(x,mu,sigma,df,c))
            c=c_star
            i+=1
    return y,c,j

def sup_norm(y,z,z_min,z_max):
    grid = GridSearchCV(KernelDensity(),{'bandwidth': np.linspace(0.1, 1.0, 30)},cv=10)
    grid.fit(y[:,None])
    kde = grid.best_estimator_
    z_grid = np.linspace(z_min,z_max,10000)
    pdf_kde = np.exp(kde.score_samples(z_grid[:,None]))
    error = np.max(np.absolute(np.subtract(z,pdf_kde)))
    return error

#Description of target and candidate distribution
mu=0
sigma=2
df=2
z_min=-5
z_max=5
z=t.pdf(np.linspace(z_min,z_max,10000),df)
f=plt.figure(figsize=(4,4))
plt.plot(np.linspace(z_min,z_max,10000),z,c='C1',label=r't (df=%d)' %df )
plt.plot(np.linspace(z_min,z_max,10000),norm.pdf(np.linspace(z_min,z_max,10000),mu,sigma),c='C4',label='N($\mu=%1.1f$,$\sigma$=%d)' %(mu,sigma))
plt.xlim([-4,4])
plt.yticks(np.linspace(0,0.5,3))
plt.xticks(np.linspace(-4,4,3))
plt.ylabel('PDF value')
plt.xlabel('x')
plt.legend(frameon=False)
plt.gca().spines['right'].set_color('none')
plt.gca().spines['top'].set_color('none')
plt.title('Candidate (normal) and target (t) distributions',y=1.05)
plt.show()

L=[10,100,1000]
c=[3,6,9]
M=len(c)   #M is length of vector c
N=len(L)   #N is length of vector L

error = np.zeros([M,N])
f, axes = plt.subplots(M,N,figsize=(8,6))
for i in range(M):   #i is c index
    for j in range(N):  #j is L index
        start = timeit.default_timer()
        y,iteration=rejection_sampling1(L[j],c[i],mu,sigma,df)
        stop = timeit.default_timer()
        error[i,j]=sup_norm(y,z,z_min,z_max)
        sns.kdeplot(y,label='estimate',shade=True,ax=axes[i,j])
        axes[i,j].plot(np.linspace(z_min,z_max,10000),z,label='true density')
        axes[i,j].text(2.5,0.2,f'run time=\n%1.2f sec.' %(stop-start),fontsize=9)
        axes[i,j].text(2.5,0.45,f'accept rate=\n%1.2f' %(L[j]/iteration),fontsize=9)
        axes[i,j].set_xlim([z_min,z_max])
        axes[i,j].set_ylim([0,0.6])
        axes[i,j].set_xticks(np.linspace(z_min,z_max,3))
        axes[i,j].set_yticks(np.linspace(0,0.6,3))
        axes[i,j].set_xticklabels([])
        axes[i,j].set_yticklabels([])
        axes[i,j].set_xlabel('x')
        axes[i,j].legend([],frameon=False)
        axes[i,j].spines['right'].set_color('none')
        axes[i,j].spines['top'].set_color('none')
        if i==0:
            axes[i,j].set_title('L = %d' %(L[j]))
        if j==0:
            axes[i,j].set_ylabel('c = %d' %(c[i]))
            axes[i,j].set_yticklabels(np.linspace(0,0.6,3))
        if i==M-1:
            axes[i,j].set_xticklabels(np.linspace(z_min,z_max,3).astype(int))
plt.tight_layout(h_pad=0.1,w_pad=1)
axes[0,0].legend(loc='upper left',fontsize=9)
f.suptitle('Density estimates using rejection sampling 1',y=1.05)
plt.show()
#f.savefig('algorithm1.png', format='png', dpi=500,bbox_inches="tight")

f = plt.figure(figsize=(5,2))
col=['C0','C1','C2']
labels=['c=3','c=6','c=9']
for j in range(N):
    for i in range(M):
        plt.scatter([j+1],error[i,j],c=col[i])
        if j==N-1:
            plt.scatter([j+1],error[i,j],c=col[i],label=labels[i])
plt.xlim([0,4])
plt.xticks([1,2,3],L)
plt.yticks(np.linspace(0,np.max(error),3))
plt.legend()
plt.ylabel(r'$\|| f-g \||_\infty$')
plt.xlabel(r'$L$')
plt.gca().spines['right'].set_color('none')
plt.gca().spines['top'].set_color('none')
plt.title('Distance between the true and estimated densities',y=1.05)
plt.show()

mu=0
sigma=[2,10,20]
df=2
z=t.pdf(np.linspace(-5,5,10000),2)
c=6
L=5000
f, axes = plt.subplots(1,N,figsize=(8,3))
for i in range(3):
    start = timeit.default_timer()
    y,iteration=rejection_sampling1(L,c,mu,sigma[i],df)
    stop = timeit.default_timer()
    sns.kdeplot(y,shade=True,ax=axes[i])
    axes[i].text(2,0.2,f'run time=\n%1.2f sec.' %(stop-start),fontsize=9)
    axes[i].text(2,0.35,f'accept rate=\n%1.2f' %(L/iteration),fontsize=9)
    axes[i].plot(np.linspace(z_min,z_max,10000),z,label='true density')
    axes[i].set_xlim([z_min,z_max])
    axes[i].set_ylim([0,0.5])
    axes[i].set_xticks(np.linspace(z_min,z_max,3))
    axes[i].set_yticks(np.linspace(0,0.5,3))
    axes[i].set_xlabel('x')
    axes[i].spines['right'].set_color('none')
    axes[i].spines['top'].set_color('none')
    axes[i].set_title(r'$N(0,%1.0f)$' %(sigma[i]))
    if i>0:
        axes[i].set_yticklabels([])
plt.suptitle('Effect of candidate distribution',y=1.1)
plt.tight_layout(h_pad=0.1,w_pad=0.5)
plt.show()

mu=0
sigma=2
df=2
c=6
L=[10,100,1000]
f, axes = plt.subplots(1,N,figsize=(8,2.5))
for j in range(N):  #j is L index
    start = timeit.default_timer()
    y,c_star,iteration=rejection_sampling2(L[j],c,mu,sigma,df)
    stop = timeit.default_timer()
    sns.kdeplot(y,label='estimate',shade=True,ax=axes[j])
    axes[j].plot(np.linspace(z_min,z_max,10000),z,label='true density')
    axes[j].text(2,0.45,f'accept rate=\n%1.2f' %(L[j]/iteration),fontsize=9)
    axes[j].text(2,0.3,f'run time =\n%1.2f sec.' %(stop-start),fontsize=9)
    axes[j].text(2,0.2,r'$\hat{c}$ = %1.2f' %(c_star),fontsize=9)
    axes[j].set_xlim([z_min,z_max])
    axes[j].set_ylim([0,0.6])
    axes[j].set_xticks(np.linspace(z_min,z_max,3))
    axes[j].set_yticks(np.linspace(0,0.6,3))
    axes[j].set_xticklabels(np.linspace(z_min,z_max,5).astype(int))
    axes[j].set_yticklabels([])
    axes[j].legend([],frameon=False)
    axes[j].spines['right'].set_color('none')
    axes[j].spines['top'].set_color('none')
    axes[j].set_title('L = %d' %(L[j]))
    axes[j].set_label('x')
axes[0].set_ylabel('PDF value')
axes[0].set_yticklabels(np.linspace(0,0.6,3))
plt.tight_layout(h_pad=0.1,w_pad=1)
axes[0].legend(loc='upper left',fontsize=9)
f.suptitle('Density estimates using rejection sampling 2',y=1.05)
plt.show()
#f.savefig('algorithm2.png', format='png', dpi=500,bbox_inches="tight")

#### Multivariate Version
def plot_density(X,Y,z):
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    ax.plot_wireframe(X, Y, z, color='green')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    ax.set_xticks([-3,0,3])
    ax.set_yticks([-3,0,3])
    ax.set_zticks(np.linspace(0,0.14,3))
    ax.plot_surface(X, Y, z, rstride=1, cstride=1,cmap='winter_r', edgecolor='none')
    ax.set_title('True density')

def plot_samples(sample):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    hist, xedges, yedges = np.histogram2d(sample[0,:], sample[1,:],range=[[-3, 3], [-3, 3]],density=True)
    xpos, ypos = np.meshgrid(xedges[:-1] + 0.25, yedges[:-1] + 0.25, indexing="ij")
    xpos = xpos.ravel()
    ypos = ypos.ravel()
    zpos = 0
    dx = dy = 0.5 * np.ones_like(zpos)
    dz = hist.ravel()
    ax.bar3d(xpos, ypos, zpos, dx, dy, dz, zsort='average',color='C0')
    ax.set_xticks([-3,3])
    ax.set_yticks([-3,3])
    ax.set_zticks(np.linspace(0,0.14,3))
    ax.set_title('Density estimate')

def plot_contour(X,Y,z,sample):
    sns.kdeplot(sample[0,:],sample[1,:],cmap='winter_r',shade=True,ax=axes[1])
    axes[0].contourf(X,Y,z,cmap='winter_r')
    for i in range(2):
        axes[i].set_xlim([-3,3])
        axes[i].set_ylim([-3,3])
        axes[i].set_xticks([-3,0,3])
        axes[i].set_yticks([-3,0,3])
        axes[i].spines['right'].set_color('none')
        axes[i].spines['top'].set_color('none')
    axes[1].set_title('Density contour estimate')
    axes[0].set_title('True density contour')
    axes[1].text(1,2,f'run time =\n%1.2f sec.' %(stop-start),fontsize=9)
    axes[i].text(1,1,f'accept rate=\n%1.2f' %(L/iteration),fontsize=9)

def t_pdf(x,mu,df,sigma):
    p=np.shape(sigma)[1]
    ex=p/2
    A=gamma((df+p)/2)
    B=gamma(df/2)*(df**ex)*(np.pi**ex)*np.sqrt(np.linalg.det(sigma))
    C=(1+((1/df)*np.dot(np.dot((x-mu),np.linalg.inv(sigma)),np.transpose(x-mu))))**(-(df+p)/2)
    return A*C/B

def twod_ratio(x,df,mu_t,mu_normal,cov_t,cov_normal,c):
    return t_pdf(x,mu_t,df,cov_t)/(c*multivariate_normal.pdf(x,mu_normal,cov_normal))

def twod_rejection_sampling_1(df,mu_t,mu_normal,cov_t,cov_normal,c,L):
    y=np.zeros([len(mu_t),L])
    i=j=0
    while i<L:
        j+=1
        x=np.random.multivariate_normal(mu_normal,cov_normal,1)
        if np.random.uniform(0,1,1)<=twod_ratio(x,df,mu_t,mu_normal,cov_t,cov_normal,c):
            y[:,i]=x
            i+=1
    return y,j

def twod_rejection_sampling_2(df,mu_t,mu_normal,cov_t,cov_normal,c,L):
    y=np.zeros([len(mu_t),L+1])
    i=j=0
    rho=cov_normal[0,1]/np.sqrt(cov_normal[0,0]*cov_normal[1,1])
    while i<L:
        j+=1
        cond_mu = mu_normal[i%2]+np.sqrt(cov_normal[i%2,i%2])*rho*(y[(i+1)%2,i]-mu_normal[(i+1)%2])/np.sqrt(cov_normal[(i+1)%2,(i+1)%2])
        x=np.random.normal(cond_mu,np.sqrt(cov_normal[i%2,i%2]*(1-rho**2)),1)
        if i%2==0:
            prop=np.array([y[0,i],x])
        else:
            prop=np.array([x,y[1,i]])
        if np.random.uniform(0,1,1)<=twod_ratio(prop,df,mu_t,mu_normal,cov_t,cov_normal,c):
                y[:,i+1]=prop
                i+=1
    return y,j

mu_t = np.array([0, 0])
mu_normal = np.array([0, 0])
cov_normal = np.array([[2, 0], [0, 2]])
cov_t = np.array([[1, 0], [0, 1]])
df=2
c=3
L=2000
start = timeit.default_timer()
sample,iteration=twod_rejection_sampling_1(df,mu_t,mu_normal,cov_t,cov_normal,c,L)
stop = timeit.default_timer()

f, axes = plt.subplots(1,2,figsize=(8,2.5))
x = np.linspace(-3, 3, 500)
y = np.linspace(-3, 3, 500)
X, Y = np.meshgrid(x, y)
z=np.zeros([500,500])
for i in range(500):
        for j in range(500):
            z[i,j]= t_pdf([X[i,j],Y[i,j]],mu_t,df,cov_t)
plot_contour(X,Y,z,sample)
plot_density(X,Y,z)
plot_samples(sample)

mu_t = np.array([0, 0])
mu_normal = np.array([0, 0])
cov_normal = np.array([[2, 0], [0, 2]])
cov_t = np.array([[1, 0], [0, 1]])
df=2
c=3
L=2000
start = timeit.default_timer()
sample,iteration=twod_rejection_sampling_2(df,mu_t,mu_normal,cov_t,cov_normal,c,L)
stop = timeit.default_timer()

f, axes = plt.subplots(1,2,figsize=(8,2.5))
x = np.linspace(-3, 3, 500)
y = np.linspace(-3, 3, 500)
X, Y = np.meshgrid(x, y)
z=np.zeros([500,500])
for i in range(500):
        for j in range(500):
            z[i,j]= t_pdf([X[i,j],Y[i,j]],mu_t,df,cov_t)
plot_contour(X,Y,z,sample)
plot_density(X,Y,z)
plot_samples(sample)