### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 1968a1b0-9090-11eb-3424-ad2150889fd1
using PlutoUI

# ╔═╡ c09c0e20-d8be-11ea-11be-0966527e6f2d
using Plots

# ╔═╡ b711bf40-c912-11eb-120d-1bb53ca25d17
using ControlSystems

# ╔═╡ cfbeb720-90a8-11eb-1ac2-3bda60b2285c
using OrdinaryDiffEq

# ╔═╡ 889652b2-dbe8-11ea-0835-e966a39af471
md"""# Problema de estabilidad absoluta de Lur`e
## Partimos de una planta con función de transferencia
"""

# ╔═╡ bf0a5c82-d8bc-11ea-3b1c-b397e3b526fa
H1(s)=(s+3)/(s^2+7s+10)

# ╔═╡ 6b0b4680-9095-11eb-3afe-9b792355080b
md"lo primero que miramos es qué polos tiene, factorizando el denominador:

$H_1=\frac{s+3}{(s+5)(s+2)}$

No hay polos ni ceros en el semiplano derecho ni tampoco en el origen.
"

# ╔═╡ 9999f2c0-dbdd-11ea-04a5-dd56cf1eae5c
md"""# Conjeturas A-K
## Se trata de comprobar el rango para el cuál la realimentación negativa con ganancia K es estable:
+ Se calcula el número de polos $n$ en el semiplano derecho (inestables)
+ El criterio de Nyquist dice que la planta es estable si la respuesta en frecuecia rodea al punto $-1/K+0*i$ $n$ veces en sentido antihorario.
"""

# ╔═╡ 4628f590-dbe3-11ea-2a7a-83cabda49ee7
function conjeturas(H,K)
	#respuesta en frecuencia
	wp=10.0.^(-50:0.01:50) #rango de 10^-50 a 10ymin=^50
	wn=-10.0.^(50:-0.01:-50)
	w=vcat(wn,0,wp)#todas las frecuencias incluidas las negativas
	Hs=H.(w*im)

	#si hay polos en el eje imaginario habrá valores infinitos o NaN para pintar bien esos casos los "recorto", en este caso no se cumplen las condiciones de las conjeturas pero puede igualemnte ser interesante ver el rango de estabilidad.

	recortar(x)=min(max(x,-10^3),10^3)

	#pinta la flecha en la frecuencia w0
if isfinite(maximum(norm.(Hs)))
	#pinto normalmente
	plot(real.(Hs),imag.(Hs), label=false)
	#pinta la flecha
	quiver!([real(H(0))],[imag(H(0))],
			quiver=([real(H(0.01*im)-H(0))],
				[imag(H(0.01*im)-H(0))]),label=false)
else
	#en este caso es cuando hay que recortar
	plot(recortar.(real.(Hs)),recortar.(imag.(Hs)), label=false)
	ylims!((-10,10))
	xlims!((-10,10))
end

#para que el movimiento sea intuitivo el slider ha de ir desde la derecha

#pinta el punto
	scatter!([-1/K],[0],label=false)
	title!("""La respuesta en frecuencia ha de rodear al punto
	en sentido antihorario tantas veces como polos inestables
	hay, K=$(K)""")
end

# ╔═╡ 490c8852-909f-11eb-2522-ad23744ce9d8
	#CSS que hace los Sliders más grandes
    html"""<style type="text/css">
    input[type=range] {
    width: 75%;
    } </style>
    """

# ╔═╡ 203ea980-9090-11eb-25f8-8ff76ab5262a
begin
	K_max_conjeturas=10
	@bind sliderAK Slider((-1:0.001:1)*K_max_conjeturas,default=1.0)
end

# ╔═╡ 93bfc0d0-cb03-11eb-283b-6d01e65bed3c
conjeturas(H1,sliderAK)

# ╔═╡ 29e11c20-90ae-11eb-1348-0f5a981d3361
md"Esto sugiere que es estable en el sector $[-\frac{10}{3}, \infty)$
Podemos ver analíticamente el -10/3 observando que el corte es el $w_c=0$ y por lo tanto $H(w_c)=\frac{0+3}{(0+5)(0+2)}=\frac{3}{10}
 \to H(0)k=-1 \to k=-10/3$

Veamos qué dicen los criterios *de verdad* y no las conjeturas
"

# ╔═╡ 66f0b770-d8c5-11ea-0ee3-7ba9b6dfb585
md"""# Criterio del círculo
## Antes de empezar comprueba que:
+ Como es una funcion de transferencia ya es contolable/observable (OK)
+ No hay integradores ni realimentación directa $$d=0$$
¿Es estable? Usa los casos A y B. Si no lo es usa los casos C y D.

**Cada caso** de los que se puede usar da **un rango** posible de ganancias
"""

# ╔═╡ 00b900b0-d8c6-11ea-0e42-0f4bfaf17fd2
md"## Planta estable"

# ╔═╡ ed754f20-db56-11ea-02fa-3581c244dcbd
md"""### Caso A
$$0=k_1<k_2$$
La recta ha de quedar a la izquierda
"""

# ╔═╡ 9975e6e0-db57-11ea-30e1-3ffbd104839c
K_maxA=100

# ╔═╡ adee0ba0-d988-11ea-19c7-e324f1f86031
@bind k2a Slider((0:0.001:1)*K_maxA,default=1.0)

# ╔═╡ 839b4c72-9092-11eb-2237-1dc284e3041e
k2a

# ╔═╡ 2b1196f0-90dc-11eb-2e20-89a61fcb1498
md"Este caso ha sido más conservador que las conjeturas $[0, \infty)$"

# ╔═╡ b47374ee-db55-11ea-2b8f-5b38e3a2c5a0
md"""### Caso B
$$k_1<0<k_2$$
El círculo ha de rodear la respuesta en frecuencia
"""

# ╔═╡ bd8ef6f0-dbb8-11ea-3fe0-3b17ccdb1fcd
begin
	K_maxB1=10
	K_maxB2=100
end

# ╔═╡ fe58cf80-dbb8-11ea-1c6c-93fef94ed79c
@bind k1b Slider((-1:0.001:0)*K_maxB1,default=-1.0)

# ╔═╡ d2de7080-dbb8-11ea-1d7e-11b40cd2ce8d
@bind k2b Slider((0:0.001:1)*K_maxB2,default=1.0)

# ╔═╡ d3077b80-9092-11eb-0f8b-85c6edde81f5
(k1=k1b,k2=k2b)

# ╔═╡ 65071560-90dc-11eb-24f3-3d03b8461507
md"Este caso es igual que las conjeturas AK es decir $[-\frac{10}{3}, \infty)$ "

# ╔═╡ 1ae51110-d98c-11ea-3b02-1ddccacadb56
md"## Planta inestable
En este caso **la planta de antes no nos sirve** ya que era estable, usamos estra otra

$H_2=\frac{s+3}{(s+5)(s-2)}$

Que tiene **un** polo inestable

"

# ╔═╡ 8054e0c0-90a2-11eb-20fb-ad26bded4db5
H2(s)=(s+3)/(s^2+3s-10)

# ╔═╡ 9ca66610-d985-11ea-0426-13ef5f785091
function circulo(H,K1,K2,caso)
	"""hay cuatro casos se pueden determinar por los signos de k1 y k2
	------ Estables
	A) 0=k1<k2 recta a la izquierda
	B) k1<0<k2 dentro del circulo
	------ Inestables
	C) 0<k1<=k2 La respueta rodea al circulo en sentido antihorario n veces donde n es el número de polos inestables
	D) k1<k2<0 La respuesta de -H rodea en el sentido antihorario n veces el circulo D(-k1,-k2) es como C pero con todo negado"""

	#compruebo que las entradas son correctas
	if !isfinite(H(0))
		return "No puede haber polos en el origen"
	end
	if caso=='A'
		if (K1==0)&(0<K2)
			titulo="Recta a la izquierda"
		else
			return "Error: en el caso A 0=k1<k2"
		end
	elseif caso=='B'
		if (K1<0)&(0<K2)
			titulo="El circulo rodea la respuesta en frecuencia"
		else
			return "Error: en el caso B k1<0<k2"
		end
	elseif caso=='C'
		if (0<K1)&(K1<=K2)
			titulo="La respueta rodea al circulo en sentido antihorario n veces\n donde n es el número de polos inestables"
		else
			return "Error: en el caso C 0<k1<=k2 "
		end
	elseif caso=='D'
		if (K1<=K2)&(K2<0)
			H2(x)=-H(x)
            return circulo( H2,-K2,-K1,'C')
		else
			return "Error: en el caso D k1<k2<0"
		end
	else
		return "No es un caso valido, ha de ser 'A','B','C' o 'D'"
	end
	#por fin los cálculos...
	#respuesta en frecuencia
	    wp=10.0.^(-50:0.01:50) #rango de 10^-50 a 10ymin=^50
	    wn=-10.0.^(50:-0.01:-50)
		w=vcat(wn,wp)#todas las frecuencias incluidas las negativas

	    Hs=H.(w*im)
	    plot(Hs,line=:arrow,label=false)
	    Hs=collect(Hs)
	    plot(real.(Hs),imag.(Hs), line=:arrow,label=false)
	    #pinta la flecha
	    quiver!([real(H(0))],[imag(H(0))],
		        quiver=([real(H(0.01*im)-H(0))],[imag(H(0.01*im)-H(0))])
		        ,label=false)

	#pinta el "circulo"
	if caso=='A'
		#el círculo es una recta 1/K1=inf
		ymin=minimum(imag.(Hs))
	    ymax=maximum(imag.(Hs))
	    y=[ymin,ymax]
		x=[-1/K2,-1/K2]
	    plot!(x,y,label=false)
	    title!(titulo)
	else
		#el círuclo es un círculo de vedad
		p1=-1/K1
		p2=-1/K2 #puntos de corte
	    centro=(p1+p2)/2
	    radio=abs(p2-p1)/2
	    theta=collect(0:0.01:2*pi)
	    circunferencia=centro.+radio*exp.(im*theta)
		scatter!([-1/K1 -1/K2],[0 0],label=false)
	    plot!(circunferencia,label="circulo")
		title!(titulo)
	end
end

# ╔═╡ b642eb40-d988-11ea-26a4-374a74d5dc08
circulo(H1,0,k2a,'A')

# ╔═╡ 9d833cf0-db03-11ea-3b28-33343306f4dc
circulo(H1,k1b,k2b,'B')

# ╔═╡ 7697fac0-90a4-11eb-1687-5ff2a5cf415d
md"Primero veamos que dicen las conjeturas"

# ╔═╡ ab676b50-90a4-11eb-1632-f1f60fd49f07
begin
	K_max_conjeturas2=10
	@bind sliderAK2 Slider((-1:0.001:1)*K_max_conjeturas2,default=1.0)
end

# ╔═╡ 81144e40-90a4-11eb-2941-f53c46a9300c
conjeturas(H2,sliderAK2)

# ╔═╡ c2d18d10-90dc-11eb-1792-3779513d19cb
md"Las conjeturas dicen $[\frac{10}{3}, \infty)$"

# ╔═╡ f9938e60-d98b-11ea-1e6d-2fc1399f7bef
md"""### Caso C
$$0<k_1 \le k_2$$

Como es inestable habrá $n$ polos en el semiplano derecho.
El círculo ha de **estar rodeado** por la respuesta en frecuencia $n$ veces en sentido antihorario.
"""

# ╔═╡ 1be6ec20-dbba-11ea-2479-ed82409ebdd7
begin
	K_maxC1=10
	K_maxC2=100
end

# ╔═╡ 76a9f030-dbba-11ea-3326-95356d82a2c5
@bind k1c Slider((0:0.001:1)*K_maxC1,default=1.0)

# ╔═╡ 783969d0-dbba-11ea-10d0-b1cea6c3b727
@bind k2c Slider((0:0.001:1)*K_maxC2,default=2.0)

# ╔═╡ 5d374e70-9093-11eb-3126-bfa8080d6bc2
(k1=k1c,k2=k2c)

# ╔═╡ e3f469a0-dbb9-11ea-30bb-813e5688c466
circulo(H2,k1c,k2c,'C')

# ╔═╡ e748e0ce-90dc-11eb-14cf-5f0c6c4ae075
md"Lo mismo que las conjeturas   $[\frac{10}{3}, \infty)$"

# ╔═╡ c45e9020-d989-11ea-028f-5f36703d75da
md"""### Caso D
$$k_1<k_2<0$$

Es como el caso C pero con todos los signos cambiados y H cambiado por -H:
El círculo ha de **estar rodeado** por la respuesta en frecuencia (de -H) $n$ veces en sentido antihorario.
"""

# ╔═╡ 245ff590-dbba-11ea-1deb-1976fc8d165d
begin
	K_maxD1=100
	K_maxD2=100
end

# ╔═╡ 8d0b4eee-dbba-11ea-2fad-5d7ed6f8035b
@bind k1d Slider((-1:0.001:0)*K_maxD1,default=-2.0 )

# ╔═╡ 8f55db30-dbba-11ea-0370-ddadc3b11dfe
@bind k2d Slider((-1:0.001:0)*K_maxD1,default=-1.0)

# ╔═╡ 3f2dfc70-9094-11eb-340c-592edc764cd1
(k1d=k1d,k2d=k2d)

# ╔═╡ 11a30cf0-9094-11eb-2af0-6f424889df66
circulo(H2,k1d,k2d,'D')

# ╔═╡ 105992d0-90dd-11eb-1a3d-653111528c53
md"No es posible meter el círculo dentro de la respuesta en frecuencia, este caso **no nos dice nada**"

# ╔═╡ 8e2736c0-d8bb-11ea-2fde-43a428dd9281
md"""# Criterio de popov
## Antes de empezar comprueba que:
+ Como es una funcion de transferencia ya es contolable/observable
+ Hay a lo sumo un integrador y si lo hay la planta se puede poner como $$d/s+ H_{estable}$$ con d>0 es decir
    + La parte del ingregrador es positiva ($$H(s) \to \infty$$  si  $$s\to 0^+$$)
    + Quitando el integrador la planta es estable"""

# ╔═╡ 7c27fd1e-c90f-11eb-0edf-29b2a8b2867b
md"Como ejemplo podemos usar la planta H1 añadiendo un polo en el origen:

$H_3=\frac{s+3}{s(s+5)(s+2)}$"

# ╔═╡ b7044bc0-9880-11eb-31b6-b1578086897a
H3(s)=H1(s)/s

# ╔═╡ bc4b38a0-c913-11eb-2eb9-25b5e6468a61
md"""### Primero a ver que nos dicen las cojneturas AK
Siendo rigurosos no nos dicen **nada** ¿Por qué?

El problema es que no **son aplicables** si hay un polo con parte real igual a cero como es el caso...

Pero al fin y al cabo son conjeturas y lo que queremos es una idea *optimista* de dónde buscar y las conjeturas nos la dan ya que si una función cualquiera del sector es estable ha de serlo también el caso particular $f(x)=kx$ que es lo que calcularemos aquí.
"""

# ╔═╡ 31d44760-c90f-11eb-3a7a-ffb4d6c20b53
begin
	K_max_conjeturas3=10
	@bind sliderAK3 Slider((-1:0.001:1)*K_max_conjeturas3,default=1.0)
end

# ╔═╡ 439f2e10-c90f-11eb-1b76-eb23238ba418
begin
	conjeturas(H3,sliderAK3)
	ylims!(-1,1) #hay que hacer zoom para ver algo
	xlims!(()) #esto directamente quita el límite de x y lo pone automático
end

# ╔═╡ 710513f0-c910-11eb-2abb-3df94293cdb4
md"**OJO** Ahora no está claro cuando la respuesta *rodea* al punto ya que es abierta..

Lo que si es importante es ver cuando **cambiamos de lado** (la estabilidad cambia al pasar de un lado de la curva al otro lado).

En este caso se cambia de lado cuando $k=0$ y nos aproximamos a la curva cuando $k=\pm \infty$

Luego basta con tomar un ejemplo concreto de ganancia ($k=1$ por ejemplo) y ver si es estable para ver en qué lado es estable y en cual inestable)

Realimentando con $k=1$ tenemos

$H_{LazoCerrado}=\frac{H_3}{1+H_3}=\frac{\frac{s+3}{s(s+5)(s+2)}}{1+\frac{s+3}{s(s+5)(s+2)}}=\frac{s+3}{s(s+5)(s+2)+(s+3)}$

$H_{LazoCerrado}=\frac{s+3}{s^3+7s^2+11s+3}$

Que es estable por Routh–Hurwitz ya que

$7>0$
$3>0$
$7 \cdot 11>3$

Luego el lado que contine $k=1$ es estable y el otro inestable por tanto el intervalo estable es $[0, \infty)$.

Si no queremos hacerlo a meno podemos usar la toolbox de control:"

# ╔═╡ d8394bbe-c912-11eb-250f-1b7385c0a536
S=tf([1.0, 0.0],[1.0])

# ╔═╡ 76c049a0-90ac-11eb-1892-a31cb390383d
md"Evaluando H(s) con S nos da la función de transferencia (esto funciona sólo por que H es racional)"

# ╔═╡ 33a205b0-c913-11eb-22a9-d51437c6252c
H3s=H3(S)

# ╔═╡ 420138b0-c913-11eb-2509-558e30b8048b
#el lazo cerrado con ganancia 1
H_LC=minreal(H3s/(1+H3s)) #minreal es para cancelar los ceros y polos sobrantes

# ╔═╡ 5c92fab0-c913-11eb-0542-2bd8ac7560de
#y por fin los ceros y los polos
zpk(H_LC)

# ╔═╡ 2cfce6fe-c915-11eb-1f34-5be3aed8045a
md"## Ahora a ver qué dice Popov"

# ╔═╡ 497f9250-d8bc-11ea-194d-77979b94814d
function popov(H,K,alpha)
	if K<0
		return "Error, en Popov K>=0"
	end
	W(w)=real(H(w*im))+w*im*imag(H(w*im)) #funcion asociada
	wp=10.0.^(-50:0.01:50) #rango de 10^-50 a 10ymin=^50
	wn=-10.0.^(-50:0.01:50)
	w=hcat(wn,wp) #todas las frecuencias incluidas las negativas
	Ws=W.(w)
	p=plot(Ws, line = :arrow,label=false)
	#recta
	xmin=minimum(real.(Ws))
	xmin=min(xmin,-1/K) #me aseguro de pintar el corte en el eje
	xmax=maximum(real.(Ws))
	x=[xmin,xmax]
	y=(x.+1/K)./alpha
	scatter!([-1/K],[0],label=false)
	plot!(x.+y.*im,label=false)
	title!("La recta ha de estar por encima de la función asociada")
end

# ╔═╡ fee20250-d8bb-11ea-299c-6d093b0edd20
K_max=1000

# ╔═╡ cb075200-d8bb-11ea-234b-6d4b6fca9083
@bind K Slider(0.0:0.01:K_max,default=10.0)

# ╔═╡ f2567110-dbe0-11ea-24a0-4d1a72151516
@bind pendiente Slider(0.0:0.1:90.0,default=45.0)

# ╔═╡ f8ad6aa0-d8bb-11ea-0dea-85d5051066fe
md"K=$(K)"

# ╔═╡ e443a3e0-d8c0-11ea-362f-75314f8b2cbb
alpha=1/tan(pendiente*pi/180)

# ╔═╡ dbcceee0-d985-11ea-2a3e-f9c3ec08cb75
popov(H3,K,alpha)

# ╔═╡ 367c9ac0-90dd-11eb-011b-0f7c4fca2c89
md"Es decir  $[0, \infty)$"

# ╔═╡ 3b303230-90a5-11eb-148b-d9c93a8c3772
md"# Simulando el sistema
Hasta ahora hemos usado funciones de tranferencia, pero ¿cómo lo simulamos?
Lo más fácil es pasar al espacio de estados

$\dot x =Ax+Bu$

$y=Cx$

Asumiremos siempre en este tema que $D=0$
"

# ╔═╡ e7264d30-90a6-11eb-07fc-d1f2120385cd
H=H1(S)

# ╔═╡ 98a5d8b2-90dd-11eb-3061-1713c5934c8c
md"Obtenido esto podemos pasar al espacio de estados"

# ╔═╡ f2cdb330-90a6-11eb-3169-8fb43ab0c79c
estados=ss(H)

# ╔═╡ a4378a10-90ac-11eb-0041-a9145a07928c
md"Definimos ahora el sistema"

# ╔═╡ 17f33990-90ad-11eb-2fa5-b762b4f1138b
md"# Sector
Una función f(x) pertenece al sector $[k_1 k_2]$ si está entre las rectas $y=k_1x$ y $y=k_2x$, algebráicamente:

$k_1x^2 \le xf(x) \le k_2x^2$

u obviamdo la división por cero (suponiendo que $\frac{f(x)}{x}$ sea continua en $0$ definiendola como el $lim_{x\to 0}\frac{f(x)}{x}$)

$k_1\le \frac{f(x)}{x} \le k_2 \  \forall x\ne0$
"

# ╔═╡ b1a75e90-90ad-11eb-0f7d-1d2b98b9fce4
f(x)=x*abs(x)-10/3*sin(x)

# ╔═╡ a9dd2000-90df-11eb-0edb-e364fefafe32
@bind k1 Slider(-10:0.01:10,default=-1.0)

# ╔═╡ 3387bd60-90e0-11eb-1963-7dfc892de25c
@bind k2 Slider(0:0.01:100,default=1.0)

# ╔═╡ a134fad0-90e0-11eb-1145-df3b4d6da246
@bind maximo Slider(1:0.01:100,default=1.0)

# ╔═╡ 041d31e0-90b3-11eb-1d07-c3b6bb15e4ae
begin
	x=-maximo:0.01:maximo
	fx=f.(x)
	f_max=maximum(fx)
	f_min=minimum(fx)
	fig1=plot(x,fx)
	plot!(x,k1*x,label="$(k1)*x",color=:red)
	plot!(x,k2*x,label="$(k2)*x",color=:red)
	sector1 = Shape([(0.0, 0.0), (maximo, k1*maximo),
			   (maximo, k2*maximo), (0.0, 0.0)])
	sector2 = Shape([(0.0, 0.0), (-maximo, -k1*maximo),
			   (-maximo, -k2*maximo), (0.0, 0.0)])
    plot!(sector1, fillcolor = plot_color(:yellow, 0.3),label=false)
	plot!(sector2, fillcolor = plot_color(:yellow, 0.3),label=false)
	plot!(ylims = (f_min,f_max))
	title!("La función ha de estar dentro de la zona amarilla")

	fig2=plot(x,fx./x,label="f(x)/x")
	sector3 = Shape([(-maximo, k1), (maximo, k1),
			   (maximo, k2), (-maximo, k2),(-maximo, k1)])
	plot!(sector3, fillcolor = plot_color(:yellow, 0.3),label=false)

	l = @layout [a; b]
	plot(fig1,fig2, layout=l)
end

# ╔═╡ 0eae14a0-c923-11eb-115a-616a0e34d031
max(1,NaN)

# ╔═╡ d6fd4a4e-90e0-11eb-34bf-ab3bb29cc905
md"Jugando un poco es fácil ver que es  $[-\frac{10}{3}, \infty)$
Para verlo analíticamente habría que hacer

$\frac{f(x)}{x}=|x|-\frac{10}{3}\frac{sin(x)}{x}$

Y ver que toma cualquier valor entre $-\frac{10}{3}$ (cuando $x \to 0$) e $\infty$ (cuando $x \to \infty$)
"

# ╔═╡ af179aae-90a7-11eb-2a7b-1b96a28c8ebc
function control(y)
	r=0
    u=r-f(y)
end

# ╔═╡ 21c1cfa0-90a7-11eb-3c54-23fa9b677080
function derivadas(x,p,t)
	A,B,C=p #extraemos los parámetros
	y=C*x
	y=y[1] #no quiero un vector sino un escalar [2.5] no es lo mismo que 2.5
	u=control(y) #calculamos el control
	dx=A*x+B*u
end

# ╔═╡ 1fc9b3b0-90a8-11eb-3b34-77d759a763d5
let
	x0 = zeros(size(estados.B))
	x0[1]=1
	parametros=(estados.A, estados.B, estados.C)
	tspan = (0.0, 30.0)
	prob = ODEProblem(derivadas,x0,tspan,parametros);
	sol = solve(prob,Tsit5());

	# Salida y referencia
	#la salida y el control no se guardan así que los reconstruyo
	ts=tspan[1]:0.01:tspan[end]
	X=sol.(ts)
	y=[(estados.C*x)[1] for x in X] #explicar paso a paso
	u=control.(y)

	plot(ts,y,xaxis="t",yaxis="y(t)",label="y(t)")
	#referencia
	fig1=plot!(ts,0.0*ts, linestyle=:dot, linewidth=2, label="referencia")

	#Señal de control, u no se guarda, hay que recalcularlo
	fig2=plot(ts,u,xaxis="t",yaxis="u(t)",legend=false)
	l = @layout [a ; b]
	plot(fig1,fig2, layout=l)
end

# ╔═╡ Cell order:
# ╟─889652b2-dbe8-11ea-0835-e966a39af471
# ╠═bf0a5c82-d8bc-11ea-3b1c-b397e3b526fa
# ╟─6b0b4680-9095-11eb-3afe-9b792355080b
# ╟─9999f2c0-dbdd-11ea-04a5-dd56cf1eae5c
# ╠═4628f590-dbe3-11ea-2a7a-83cabda49ee7
# ╠═1968a1b0-9090-11eb-3424-ad2150889fd1
# ╠═490c8852-909f-11eb-2522-ad23744ce9d8
# ╠═203ea980-9090-11eb-25f8-8ff76ab5262a
# ╠═c09c0e20-d8be-11ea-11be-0966527e6f2d
# ╠═93bfc0d0-cb03-11eb-283b-6d01e65bed3c
# ╟─29e11c20-90ae-11eb-1348-0f5a981d3361
# ╟─66f0b770-d8c5-11ea-0ee3-7ba9b6dfb585
# ╠═9ca66610-d985-11ea-0426-13ef5f785091
# ╟─00b900b0-d8c6-11ea-0e42-0f4bfaf17fd2
# ╟─ed754f20-db56-11ea-02fa-3581c244dcbd
# ╠═9975e6e0-db57-11ea-30e1-3ffbd104839c
# ╠═adee0ba0-d988-11ea-19c7-e324f1f86031
# ╠═839b4c72-9092-11eb-2237-1dc284e3041e
# ╠═b642eb40-d988-11ea-26a4-374a74d5dc08
# ╟─2b1196f0-90dc-11eb-2e20-89a61fcb1498
# ╟─b47374ee-db55-11ea-2b8f-5b38e3a2c5a0
# ╠═bd8ef6f0-dbb8-11ea-3fe0-3b17ccdb1fcd
# ╠═fe58cf80-dbb8-11ea-1c6c-93fef94ed79c
# ╠═d2de7080-dbb8-11ea-1d7e-11b40cd2ce8d
# ╟─d3077b80-9092-11eb-0f8b-85c6edde81f5
# ╠═9d833cf0-db03-11ea-3b28-33343306f4dc
# ╟─65071560-90dc-11eb-24f3-3d03b8461507
# ╟─1ae51110-d98c-11ea-3b02-1ddccacadb56
# ╠═8054e0c0-90a2-11eb-20fb-ad26bded4db5
# ╟─7697fac0-90a4-11eb-1687-5ff2a5cf415d
# ╟─ab676b50-90a4-11eb-1632-f1f60fd49f07
# ╠═81144e40-90a4-11eb-2941-f53c46a9300c
# ╟─c2d18d10-90dc-11eb-1792-3779513d19cb
# ╟─f9938e60-d98b-11ea-1e6d-2fc1399f7bef
# ╠═1be6ec20-dbba-11ea-2479-ed82409ebdd7
# ╠═76a9f030-dbba-11ea-3326-95356d82a2c5
# ╠═783969d0-dbba-11ea-10d0-b1cea6c3b727
# ╟─5d374e70-9093-11eb-3126-bfa8080d6bc2
# ╠═e3f469a0-dbb9-11ea-30bb-813e5688c466
# ╟─e748e0ce-90dc-11eb-14cf-5f0c6c4ae075
# ╟─c45e9020-d989-11ea-028f-5f36703d75da
# ╠═245ff590-dbba-11ea-1deb-1976fc8d165d
# ╠═8d0b4eee-dbba-11ea-2fad-5d7ed6f8035b
# ╠═8f55db30-dbba-11ea-0370-ddadc3b11dfe
# ╟─3f2dfc70-9094-11eb-340c-592edc764cd1
# ╠═11a30cf0-9094-11eb-2af0-6f424889df66
# ╟─105992d0-90dd-11eb-1a3d-653111528c53
# ╟─8e2736c0-d8bb-11ea-2fde-43a428dd9281
# ╟─7c27fd1e-c90f-11eb-0edf-29b2a8b2867b
# ╠═b7044bc0-9880-11eb-31b6-b1578086897a
# ╟─bc4b38a0-c913-11eb-2eb9-25b5e6468a61
# ╟─31d44760-c90f-11eb-3a7a-ffb4d6c20b53
# ╠═439f2e10-c90f-11eb-1b76-eb23238ba418
# ╟─710513f0-c910-11eb-2abb-3df94293cdb4
# ╠═b711bf40-c912-11eb-120d-1bb53ca25d17
# ╠═d8394bbe-c912-11eb-250f-1b7385c0a536
# ╟─76c049a0-90ac-11eb-1892-a31cb390383d
# ╠═33a205b0-c913-11eb-22a9-d51437c6252c
# ╠═420138b0-c913-11eb-2509-558e30b8048b
# ╠═5c92fab0-c913-11eb-0542-2bd8ac7560de
# ╟─2cfce6fe-c915-11eb-1f34-5be3aed8045a
# ╠═497f9250-d8bc-11ea-194d-77979b94814d
# ╠═fee20250-d8bb-11ea-299c-6d093b0edd20
# ╠═cb075200-d8bb-11ea-234b-6d4b6fca9083
# ╠═f2567110-dbe0-11ea-24a0-4d1a72151516
# ╟─f8ad6aa0-d8bb-11ea-0dea-85d5051066fe
# ╟─e443a3e0-d8c0-11ea-362f-75314f8b2cbb
# ╠═dbcceee0-d985-11ea-2a3e-f9c3ec08cb75
# ╟─367c9ac0-90dd-11eb-011b-0f7c4fca2c89
# ╟─3b303230-90a5-11eb-148b-d9c93a8c3772
# ╠═e7264d30-90a6-11eb-07fc-d1f2120385cd
# ╟─98a5d8b2-90dd-11eb-3061-1713c5934c8c
# ╠═f2cdb330-90a6-11eb-3169-8fb43ab0c79c
# ╠═a4378a10-90ac-11eb-0041-a9145a07928c
# ╠═21c1cfa0-90a7-11eb-3c54-23fa9b677080
# ╟─17f33990-90ad-11eb-2fa5-b762b4f1138b
# ╠═b1a75e90-90ad-11eb-0f7d-1d2b98b9fce4
# ╠═a9dd2000-90df-11eb-0edb-e364fefafe32
# ╠═3387bd60-90e0-11eb-1963-7dfc892de25c
# ╠═a134fad0-90e0-11eb-1145-df3b4d6da246
# ╟─041d31e0-90b3-11eb-1d07-c3b6bb15e4ae
# ╠═0eae14a0-c923-11eb-115a-616a0e34d031
# ╟─d6fd4a4e-90e0-11eb-34bf-ab3bb29cc905
# ╠═af179aae-90a7-11eb-2a7b-1b96a28c8ebc
# ╠═cfbeb720-90a8-11eb-1ac2-3bda60b2285c
# ╠═1fc9b3b0-90a8-11eb-3b34-77d759a763d5


using Test
using MotifSequenceGenerator

@testset "possiblesums" begin
    a = [1,1,3,4]
    mina = minimum(a); maxa = maximum(a)
    for n in 1:3
        allsum = all_possible_sums(a, n)
        m = length(a)
        for i in 1:binomial(m + n - 1, n)
            @test mina*n ≤ allsum[i][1] ≤ maxa*n
        end
    end
end

include("integer_length_tests.jl")
include("float_length_tests.jl")


using Graphs
using CustomVertexGraphs

g = set_vertices(grid((4,)), ["A", "B", "C", "D"])
@show g
@show g[["A", "B"]]
@show has_edge(g, "A" => "B")
@show add_edge!(g, "A" => "C")
@show has_edge(g, "A" => "C")


using Base.Meta
using BenchmarkTools

struct Variable
    index::Int
end

# Represents sum(coefficients[i]*variables[i] for i in ...) + offset
mutable struct AffineExpression
    variables::Vector{Variable}
    coefficients::Vector{Float64}
    offset::Float64
end

Base.zero(::Type{AffineExpression}) = AffineExpression(Variable[],Float64[],0.0)
Base.:*(c::Float64,v::Variable) = AffineExpression([v],[c],0.0)
function Base.:+(aff1::AffineExpression,aff2::AffineExpression)
    return AffineExpression([aff1.variables;aff2.variables],[aff1.coefficients;aff2.coefficients],aff1.offset+aff2.offset)
end

function build_generator(n)
    return sum(i^1.5*Variable(i) for i in 1:n)
end

function build_loop(n)
    aff = zero(AffineExpression)
    for i in 1:n
        aff += i^1.5*Variable(i)
    end
    return aff
end

function build_manual(n)
    aff = zero(AffineExpression)
    sizehint!(aff.variables, n)
    sizehint!(aff.coefficients, n)
    for i in 1:n
        push!(aff.variables, Variable(i))
        push!(aff.coefficients, i^1.5)
    end
    return aff
end

@enum NodeClass CONST VAR MUL ADD

struct OperationNode
    nc::NodeClass
    children::Vector{OperationNode}
    metadata::Float64
end

OperationNode(v::Variable) = OperationNode(VAR, [], v.index)
OperationNode(n::Float64) = OperationNode(CONST, [], n)

Base.:*(o1::OperationNode,o2::OperationNode) = OperationNode(MUL, [o1,o2], 0.0)
Base.:+(o1::OperationNode,o2::OperationNode) = OperationNode(ADD, [o1,o2], 0.0)

# special structure:
#          ADD
#         /   \
#       MUL    ADD
#      /   \     \
#  CONST   VAR   ...

function graph_to_aff(g::OperationNode, aff::AffineExpression = zero(AffineExpression))
    if g.nc == ADD
        for c in g.children
            graph_to_aff(c, aff)
        end
    elseif g.nc == MUL
        # we know it is CONST*VAR
        coef = 0.0
        varidx = 0
        @assert length(g.children) == 2
        for c in g.children
            if c.nc == CONST
                coef = c.metadata
            elseif c.nc == VAR
                varidx = Int(c.metadata)
            end
        end
        push!(aff.variables, Variable(varidx))
        push!(aff.coefficients, coef)
    else
        error()
    end
    return aff
end

function build_graph(n)

    g = sum(OperationNode(i^1.5)*OperationNode(Variable(i)) for i in 1:n)

    return graph_to_aff(g)
end

@benchmark build_generator(100)
@benchmark build_loop(100)
@benchmark build_manual(100)
@benchmark build_graph(100)

@benchmark build_generator(1000)
@benchmark build_loop(1000)
@benchmark build_manual(1000)
@benchmark build_graph(1000)

ex = :(sum(i^1.5*Variable(i) for i in 1:n))
dump(ex)
ex.args[2]
ex.args[2].args

function rewrite_generator(ex)
    @assert isexpr(ex,:call)
    @assert ex.args[1] == :sum
    @assert isexpr(ex.args[2], :generator)
    return quote
        aff = zero(AffineExpression)
        len = length($(ex.args[2].args[2].args[2]))
        sizehint!(aff.variables, len)
        sizehint!(aff.coefficients, len)
        for $(ex.args[2].args[2].args[1]) in $(ex.args[2].args[2].args[2])
            push!(aff.variables, $(ex.args[2].args[1].args[3]))
            push!(aff.coefficients, $(ex.args[2].args[1].args[2]))
        end
        aff
    end
end

rewrite_generator(:(sum(i^1.5*Variable(i) for i in 1:n)))

macro build_expression(ex)
    return rewrite_generator(ex)
end

function build_macro(n)
    return @build_expression sum(i^1.5*Variable(i) for i in 1:n)
end

@benchmark build_macro(100)
@benchmark build_macro(1000)


@testset "tuples" begin
    for n = 1:4
        for t in unique(combinations(repeat(1:n, n), n))
            @test all(Bennu.tuplesort(tuple(t...)) .== sort(t))
            @test all(t[collect(Bennu.tuplesortperm(tuple(t...)))] .== sort(t))
        end

        for t in permutations(1:n)
            @test all(Bennu.tuplesortperm(tuple(t...)) .== sortperm(t))
        end

        @test Bennu.permutationtuples(n) == Tuple(Tuple.(collect(permutations(1:n))))
    end

    @test (1, 2, false) == Bennu.minmaxflip(1, 2)
    @test (1, 1, false) == Bennu.minmaxflip(1, 1)
    @test (1, 2, true) == Bennu.minmaxflip(2, 1)
end


"""
Struct for deterministic state feedback policies of type ``a=\\pi(s)``
"""
struct FeedbackPolicy <: AbstractPolicy
    π
end

"""
Implement interface required for RLCore agent
"""
(p::FeedbackPolicy)(env) = p.π(get_state(env))

# policy stays constant
RLBase.update!(p::FeedbackPolicy, experience) = nothing
RLBase.update!(p::FeedbackPolicy, args...) = nothing

"""
Linear state feedback policy ``a=\\pi(s)=Ks``
"""
LinearFeedbackPolicy(K) = FeedbackPolicy(s -> K * reshape(s, :))
LinearFeedbackPolicy(K, a_lb, a_ub) = FeedbackPolicy(s -> clamp.(K * reshape(s, :), a_lb, a_ub))

"""
Discrete LQR state feedback policy.

Discretizes the continuous sytem ``\\dot{s}=As+Ba`` and
calculates the optimal state-feedback law ``a_k = -K s_k`` that
minimizes the cost function:
``J = \\sum_{k=0}^\\infty(s_k'Qs_k + a_k'Ra_k)``
For the discretized time model ``s_{k+1} = \\bar{A}s_k + \\bar{B}a_k``.
"""
function LQRPolicy(A, B, Q, R, dt)
    K = _dlqr_from_cont_ss(A, B, Q, R, dt)
    return LinearFeedbackPolicy(-K)
end

function LQRPolicy(A, B, Q, R, dt, a_lb, a_ub)
    K = _dlqr_from_cont_ss(A, B, Q, R, dt)
    return LinearFeedbackPolicy(-K, a_lb, a_ub)
end

function _dlqr_from_cont_ss(A, B, Q, R, dt)
    # Build discrete state-space model using place-holder matrices C, D
    nx = size(A)[1]
    na = size(B[:,:])[2]

    C = zeros(1, nx)
    D = zeros(1, na)
    sysd = ss(A, B, C, D, dt)

    # Compute discrete LQR gain matrix
    K = lqr(sysd, Q, R)
end


import TriangleMesh, UnicodePlots, Plots
using Printf, LoopVectorization, LinearAlgebra, SparseArrays

include("CreateMeshFCFV.jl")
include("VisuFCFV.jl")
include("DiscretisationFCFV.jl")

function SetUpProblem!( mesh, T, Tdir, Tneu, se, a, b, c, d, alp, bet )
    # Evaluate T analytic on cell faces
    @avx for in=1:mesh.nf
        x        = mesh.xf[in]
        y        = mesh.yf[in]
        Tdir[in] = exp(alp*sin(a*x + c*y) + bet*cos(b*x + d*y))
        dTdx     = Tdir[in] * (a*alp*cos(a*x + c*y) - b*bet*sin(b*x + d*y))
        dTdy     = Tdir[in] * (alp*c*cos(a*x + c*y) - bet*d*sin(b*x + d*y))
        Tneu[in] = -dTdy # nx*dTdx + nt*dTdy on SOUTH face
    end
    # Evaluate T analytic on barycentres
    @avx for iel=1:mesh.nel
        x       = mesh.xc[iel]
        y       = mesh.yc[iel]
        T       = exp(alp*sin(a*x + c*y) + bet*cos(b*x + d*y))
        T[iel]  = T
        se[iel] = T*(-a*alp*cos(a*x + c*y) + b*bet*sin(b*x + d*y))*(a*alp*cos(a*x + c*y) - b*bet*sin(b*x + d*y)) + T*(a^2*alp*sin(a*x + c*y) + b^2*bet*cos(b*x + d*y)) + T*(-alp*c*cos(a*x + c*y) + bet*d*sin(b*x + d*y))*(alp*c*cos(a*x + c*y) - bet*d*sin(b*x + d*y)) + T*(alp*c^2*sin(a*x + c*y) + bet*d^2*cos(b*x + d*y))
    end
    return
end

function ComputeError( mesh, Te, qx, qy, a, b, c, d, alp, bet )
    eT  = zeros(mesh.nel)
    eqx = zeros(mesh.nel)
    eqy = zeros(mesh.nel)
    Ta  = zeros(mesh.nel)
    qxa = zeros(mesh.nel)
    qya = zeros(mesh.nel)
    @avx for iel=1:mesh.nel
        x        = mesh.xc[iel]
        y        = mesh.yc[iel]
        Ta[iel]       = exp(alp*sin(a*x + c*y) + bet*cos(b*x + d*y))
        qxa[iel]      = -Ta[iel] * (a*alp*cos(a*x + c*y) - b*bet*sin(b*x + d*y))
        qya[iel]      = -Ta[iel] * (alp*c*cos(a*x + c*y) - bet*d*sin(b*x + d*y))
        eT[iel]  = Te[iel] - Ta[iel]
        eqx[iel] = qx[iel] - qxa[iel]
        eqy[iel] = qy[iel] - qya[iel]
    end
    errT  = norm(eT)/norm(Ta)
    errqx = norm(eqx)/norm(qxa)
    errqy = norm(eqy)/norm(qya)
    return errT, errqx, errqy
end

function StabParam(tau, dA, Vol, mesh_type)
    if mesh_type=="Quadrangles";        taui = tau;    end
    # if mesh_type=="UnstructTriangles";  taui = tau*dA; end
    if mesh_type=="UnstructTriangles";  taui = tau end
    return taui
end

@views function main()

    println("\n******** FCFV POISSON ********")

    # Create sides of mesh
    xmin, xmax = 0, 1
    ymin, ymax = 0, 1
    nx, ny     = 16, 16
    R          = 0.5
    inclusion  = 0
    mesh_type  = "Quadrangles"
    # mesh_type  = "UnstructTriangles"

    # Generate mesh
    if mesh_type=="Quadrangles"
        tau  = 1
        mesh = MakeQuadMesh( nx, ny, xmin, xmax, ymin, ymax, inclusion, R )
    elseif mesh_type=="UnstructTriangles"
        tau  = 100
        mesh = MakeTriangleMesh( nx, ny, xmin, xmax, ymin, ymax, inclusion, R )
    end
    println("Number of elements: ", mesh.nel)

    # Source term and BCs etc...
    Tanal  = zeros(mesh.nel)
    se     = zeros(mesh.nel)
    Tdir   = zeros(mesh.nf)
    Tneu   = zeros(mesh.nf)
    alp = 0.1; bet = 0.3; a = 5.1; b = 4.3; c = -6.2; d = 3.4;
    println("Model configuration :")
    @time SetUpProblem!(mesh , Tanal, Tdir, Tneu, se, a, b, c, d, alp, bet)

    # Compute some mesh vectors
    println("Compute FCFV vectors:")
    @time ae, be, ze = ComputeFCFV(mesh, se, Tdir, tau)

    # Assemble element matrices and RHS
    println("Compute element matrices:")
    @time Kv, fv = ElementAssemblyLoop(mesh, ae, be, ze, Tdir, Tneu, tau)

    # Assemble triplets and sparse
    println("Assemble triplets and sparse:")
    @time K, f = CreateTripletsSparse(mesh, Kv, fv)
    # display(UnicodePlots.spy(K))

    # Solve for hybrid variable
    println("Direct solve:")
    # @time Th   = K\f
    PC  = 0.5.*(K.+K')
    PCc = cholesky(PC)
    Th  = zeros(mesh.nf)
    dTh = zeros(mesh.nf,1)
    r   = zeros(mesh.nf,1)
    r  .= f - K*Th
    # @time Th   = K\f
    for rit=1:5
        r    .= f - K*Th
        println("It. ", rit, " - Norm of residual: ", norm(r)/length(r))
        if norm(r)/length(r) < 1e-10
            break
        end
        dTh  .= PCc\r
        Th  .+= dTh[:]
    end

    # Reconstruct element values
    println("Compute element values:")
    @time Te, qx, qy = ComputeElementValues(mesh, Th, ae, be, ze, Tdir, tau)
    # @time Te, qx, qy = ComputeElementValuesFaces(mesh, Th, ae, be, ze, Tdir, tau)

    # Compute discretisation errors
    err_T, err_qx, err_qy = ComputeError( mesh, Te, qx, qy, a, b, c, d, alp, bet )
    println("Error in T:  ", err_T )
    println("Error in qx: ", err_qx)
    println("Error in qy: ", err_qy)

    # Visualise
    println("Visualisation:")
    @time PlotMakie( mesh, Te, xmin, xmax, ymin, ymax, :viridis )

end

main()


module Complex #hide
# # Complex numbers
# Trivial example to demonstrate the capability of handling complex numbers
using PetscWrap
using SlepcWrap

# Start by checking that you're using a PETSc/SLEPc build with complex number
(PetscScalar <: Real) && error("You must configure PETSc/SLEPc with complex numbers to run this example")

# Initialize SLEPc
SlepcInitialize()

# Create matrix
A = create_matrix(4, 4; auto_setup = true)

# Get rows handled by the local processor
A_rstart, A_rend = get_range(A)

# Create diag matrix with complex numbers
J = I = A_rstart:A_rend
V = im .* I
set_values!(A, I, J, V)
assemble!(A)

# Now we set up the eigenvalue solver
eps = create_eps(A; auto_setup = true)

# Then we solve
solve!(eps)

# Show converged eivenvalues
@show get_eigenvalues(eps)

# Finally, let's free the memory
destroy!(A)
destroy!(eps)

# And call finalize when you're done
SlepcFinalize()

end #hide


using Rasch
using Optim
using Test
using DelimitedFiles
using LinearAlgebra

@testset "Rasch.jl" begin
    X = readdlm(joinpath(dirname(pathof(Rasch)), "../test/lsat.txt"));
    n, m = size(X)
    r = vec(mapslices(sum, X, dims=2))
    s = Float64.(vec(mapslices(sum, X, dims=1)))
    f = [sum(r.==j) for j in 1:m]
    t = Float64.(vec(mapslices(sum, X .* r, dims=1)))
    q = Float64.(vec(mapslices(sum, X .* r .^2, dims=1)))

    res0, delta_hat0, πs0 =  haberman(m,reshape(s,(m,1)),f,zeros(m-1),Optim.Newton(), Optim.Options(f_tol=1e-6, g_tol=0))
    res1, delta_hat1, πs1 =  haberman(m,hcat(s,t),f,zeros(2m-2),Optim.Newton(), Optim.Options(f_tol=1e-6, g_tol=0))
    res2, delta_hat2, πs2 =  haberman(m,hcat(s,t,q),f,zeros(3m-3),Optim.Newton(), Optim.Options(f_tol=1e-6, g_tol=0)).
end


using gcSolver
using Plots
using OrdinaryDiffEq
using DiffEqDevTools

rxntfR = ones(gcSolver.Nparams) * 0.2

prob = gcSolver.runCkineSetup([500.0], rxntfR)

setups = [
    Dict(:alg => AutoTsit5(Rodas5())),
    Dict(:alg => AutoTsit5(Kvaerno5())),
    Dict(:alg => AutoVern9(Rodas5())),
    Dict(:alg => AutoVern9(KenCarp5())),
    Dict(:alg => AutoVern9(Kvaerno5())),
]

test_sol = solve(prob, AutoTsit5(Rodas5()), reltol = 1.0e-16, abstol = 1.0e-16)

tols = 1.0 ./ 10.0 .^ LinRange(4, 12, 30)

wp = WorkPrecisionSet(prob, tols, tols, setups; save_everystep = false, appxsol = test_sol, maxiters = Int(1e6))

plot(wp)


export pca_session

function pca_session(spikes::Array{SpikeTrain,1},binsize::Float64,kern::Float64)

    myrate=rate_KD(spikes,binsize,kern)

    pcamat=zeros(Float64,length(spikes[1].trials[1].time:binsize:spikes[1].trials[end].time)-1,length(spikes))

for i=1:length(spikes)
    pcamat[:,i]=rate_session(myrate,i)
end

    mypca=pca(pcamat)

    (mypca.rotation, mypca.proportion_of_variance)

end

function project_trials(spikes::Array{SpikeTrain,1},inds::Array{Int64,1},timeseq::FloatRange{Float64},components::Array{Float64,2},mypc::Int64,kern::Float64)

    binsize=timeseq.step/timeseq.divisor

    myrate=rate_KD(spikes,binsize,kern)

    pca_vec=components[:,mypc]

    projections=zeros(Float64,length(inds),length(timeseq)-1)

    trialrates=zeros(Float64,length(timeseq)-1,length(spikes))

    for i=1:length(inds)

        trialrates=rate_window_pop(myrate,inds[i],timeseq)

        projections[i,:]=trialrates*pca_vec

    end

    projections

end


export stnKD, PointNet

stnKD(K::Int) = Chain(
    Conv((1,), K=>64, relu),
    BatchNorm(64),
    Conv((1,), 64=>128, relu),
    BatchNorm(128),
    Conv((1,), 128=>1024, relu),
    BatchNorm(1024),
    x -> maximum(x, dims=1),
    x -> reshape(x, :, size(x, 3)),
    Dense(1024, 512, relu),
#     BatchNorm(512),
    Dense(512, 256, relu),
    BatchNorm(256),
    Dense(256,K*K),
    x -> reshape(x, K, K, size(x, 2)),
    x -> x .+ reshape(Matrix{Float64}(I, K, K), K, K, 1),
    x -> PermutedDimsArray(x, (2,1,3)),
)

# PointNet(k::Int) = Chain(
#     x -> batched_mul(x, stn(x)),
#     Conv((1,), 3=>64),
#     BatchNorm(64),
#     x -> batched_mul(x, fstn(x)),
#     Conv((1,), 64=>128, relu),
#     BatchNorm(128),
#     Conv((1,), 128=>1024),
#     BatchNorm(1024),
#     x -> maximum(x, dims=1),
#     x -> reshape(x, 1024, :),
#     Dense(1024, 512, relu),
#     BatchNorm(512),
#     Dense(512, 256, relu),
#     Dropout(0.4),
#     BatchNorm(256),
#     Dense(256, k, relu),
#     x -> softmax(x, dims=1),
# )

struct PointNet
    stn
    fstn
    conv_block1
    feat
    cls
end

function PointNet(num_classes::Int=10, K::Int=64)
    stn = stnKD(3)
    fstn = stnKD(K)
    conv_block1 = conv_bn_block(3, 64, (1,))
    feat = Chain(
        Conv((1,), 64=>128, relu),
        BatchNorm(128),
        Conv((1,), 128=>1024),
        BatchNorm(1024),
        x -> maximum(x, dims=1),
        x -> reshape(x, 1024, :),
        Dense(1024, 512, relu),
        BatchNorm(512),
        Dense(512, 256, relu),
        Dropout(0.4),
        BatchNorm(256),
        )
    cls = Dense(256, num_classes, relu)
    PointNet(stn,fstn, conv_block1, feat, cls)
end

function (m::PointNet)(X)

        # X: [N, 3, B]

        X = batched_mul(X, m.stn(X))
        # X: [3, 3, B]

        X = m.conv_block1(X)
        # X: [3, 64, B]

        X = batched_mul(X, m.fstn(X))
        # X: [3, 64, B]

        X = m.feat(X)
        # X: [256, B]

        X = m.cls(X)
        # X: [num_classes, B]

        return softmax(X, dims=1)
end

@functor PointNet


module InverseOptimalControl

import ..DynamicsModelInterface
import ..JuMPUtils
import ..CostUtils
import Ipopt
import JuMP

using ..ForwardOptimalControl: linear_dynamics_constraints, forward_quadratic_objective
using JuMP: @variable, @constraint, @objective
using UnPack: @unpack

export solve_inverse_optimal_control

#============================================= LQ case =============================================#

"The lagrangian of the forward LQR problem."
function lqr_lagrangian(x, u, λ; Q, R, A, B)
    c_dyn = linear_dynamics_constraints(x, u; A, B)
    @views forward_quadratic_objective(x, u; Q, R) +
           sum(λ[:, t]' * c_dyn[:, t] for t in axes(x)[2][1:(end - 1)])
end

"The hand-written gradient of the forward LQR problem in x."
function lqr_lagrangian_grad_x(x, u, λ; Q, R, A, B)
    T = size(x, 2)
    hcat(
        2 * Q * x[:, 1] - A[1]' * λ[:, 1], # special handling of first time step
        reduce(hcat, (2 * Q * x[:, t] + λ[:, t - 1] - A[t]' * λ[:, t] for t in 2:(T - 1))),
        2 * Q * x[:, T] + λ[:, T - 1],
    )
end

"The hand-written gradient of the forward LQR problem in u."
function lqr_lagrangian_grad_u(x, u, λ; Q, R, A, B)
    T = size(x, 2)
    hcat(reduce(hcat, (2 * R * u[:, t] - B[t]' * λ[:, t] for t in 1:(T - 1))), 2 * R * u[:, T])
end

"""
Solves aninverse LQR problem using JuMP.

`x̂` is the (observed) state trajectory for which we wish to find the cost parametrization.

`Q̃` and `R̃` are iterables of quadatic state and control cost matrices for which the weight
vectors`q` and `r` are to be estimated.
"""
function solve_inverse_lqr(
    x̂,
    Q̃,
    R̃;
    A,
    B,
    r_sqr_min = 1e-5,
    solver = Ipopt.Optimizer,
    solver_attributes = (; print_level = 3),
)
    T = size(x̂)[2]
    n_states, n_controls = size(only(unique(B)))
    opt_model = JuMP.Model(solver)
    JuMPUtils.set_solver_attributes!(opt_model; solver_attributes...)

    # decision variable
    q = @variable(opt_model, [1:length(Q̃)], lower_bound = 0)
    r = @variable(opt_model, [1:length(R̃)], lower_bound = 0)
    x = @variable(opt_model, [1:n_states, 1:T])
    u = @variable(opt_model, [1:n_controls, 1:T])
    λ = @variable(opt_model, [1:n_states, 1:(T - 1)]) # multipliers of the forward LQR condition

    # initial condition
    @constraint(opt_model, x[:, 1] .== x̂[:, 1])
    # dynamics
    @constraint(opt_model, linear_dynamics_constraints(x, u; A, B) .== 0)
    # Optimality conditions (KKT) of forward LQR show up as a constraints
    Q = sum(q .* Q̃)
    R = sum(r .* R̃)
    ∇ₓL = @constraint(opt_model, lqr_lagrangian_grad_x(x, u, λ; Q, R, A, B)[:, 2:end] .== 0)
    ∇ᵤL = @constraint(opt_model, lqr_lagrangian_grad_u(x, u, λ; Q, R, A, B) .== 0)
    # regularization
    @constraint(opt_model, r' * r >= r_sqr_min)
    @constraint(opt_model, r' * r + q' * q == 1)

    @objective(opt_model, Min, sum((x .- x̂) .^ 2))
    @time JuMP.optimize!(opt_model)

    JuMPUtils.isconverged(opt_model),
    JuMPUtils.get_values(; q, r, x, u, λ, ∇ₓL, ∇ᵤL),
    opt_model,
    JuMP.value.(Q),
    JuMP.value.(R)
end

#=========================================== Non-LQ-Case ===========================================#

function solve_inverse_optimal_control(
    y;
    control_system,
    cost_model,
    observation_model,
    fixed_inputs = (),
    init = (),
    solver = Ipopt.Optimizer,
    solver_attributes = (; print_level = 3),
    cmin = 1e-5,
    max_observation_error = nothing,
    verbose = false,
    init_with_observation = true,
)
    T = size(y)[2]
    @unpack n_states, n_controls = control_system

    opt_model = JuMP.Model(solver)
    JuMPUtils.set_solver_attributes!(opt_model; solver_attributes...)

    # decision variable
    weights = @variable(opt_model, [keys(cost_model.weights)],)
    x = @variable(opt_model, [1:n_states, 1:T])
    u = @variable(opt_model, [1:n_controls, 1:T])
    λ = @variable(opt_model, [1:n_states, 1:(T - 1)]) # multipliers of the forward optimality condition

    # initialization
    if hasproperty(init, :weights) && !isnothing(init.weights)
        verbose && @info "Using weight guess: $(init.weights)"
        for k in keys(init.weights)
            JuMP.set_start_value(weights[k], init.weights[k])
        end
    else
        verbose && @info "Using weight default"
        JuMP.set_start_value.(weights, 1 / length(weights))
    end

    # Initialization
    if init_with_observation
        # Note: This initialization is only works if `observation_model.expected_observation`
        # effectively creates an array view into x (extracting sub-components of the state).
        JuMP.set_start_value.(observation_model.expected_observation(x), y)
    end
    JuMPUtils.init_if_hasproperty!(u, init, :u)

    # constraints
    if iszero(observation_model.σ)
        @constraint(opt_model, observation_model.expected_observation(x[:, 1]) .== y[:, 1])
    end
    DynamicsModelInterface.add_dynamics_constraints!(control_system, opt_model, x, u)
    # figure out which inputs we control and fix all others
    controlled_inputs = filter(i -> i ∉ fixed_inputs, 1:n_controls)
    for i in fixed_inputs
        JuMP.fix.(u[i, :], init.u[i, :])
    end
    # Require forward-optimality for all *controlled* inputs.
    df = DynamicsModelInterface.add_dynamics_jacobians!(control_system, opt_model, x, u)
    dJ = cost_model.add_objective_gradients!(opt_model, x, u; weights)
    @constraint(
        opt_model,
        [t = 2:(T - 1)],
        dJ.dx[:, t] + λ[:, t - 1] - (λ[:, t]' * df.dx[:, :, t])' .== 0
    )
    @constraint(opt_model, dJ.dx[:, T] + λ[:, T - 1] .== 0)
    @constraint(
        opt_model,
        [t = 1:(T - 1)],
        dJ.du[controlled_inputs, t] - (λ[:, t]' * df.du[:, controlled_inputs, t])' .== 0
    )
    @constraint(opt_model, dJ.du[controlled_inputs, T] .== 0)
    # regularization
    @constraint(opt_model, weights .>= cmin)
    @constraint(opt_model, sum(weights) == 1)

    y_expected = observation_model.expected_observation(x)
    # Sometimes useful for debugging: Only search in a local neighborhood of the demonstration if we
    # have an error-bound on the noise.
    if !isnothing(max_observation_error)
        @constraint(opt_model, (y_expected - y) .^ 2 .<= max_observation_error^2)
    end

    # The inverse objective: match the observed demonstration
    @objective(opt_model, Min, sum(el -> el^2, y_expected .- y))

    time = @elapsed JuMP.optimize!(opt_model)
    verbose && @info time

    solution = merge(
        CostUtils.namedtuple(JuMPUtils.get_values(; weights)),
        JuMPUtils.get_values(; x, u, λ),
    )

    JuMPUtils.isconverged(opt_model), solution, opt_model
end

end


using LinearAlgebra, Rotations

function nullspace_method(jac_joints, delta; regularization_strength=0.0)
    hess_approx = transpose(jac_joints)*jac_joints
    joint_delta = transpose(jac_joints)*delta

    if regularization_strength > 0.0
        hess_approx += I(size(hess_approx)[1])*regularization_strength
        return hess_approx\joint_delta
    else
        return hess_approx\joint_delta
    end
end

function inverse_kinematics!(sim::MJSim, site_name::String,
    target_pos::Union{Vector{Float64},Nothing},
    target_quat::Union{Vector{Float64},Nothing}=nothing,
    joint_names=nothing,
    tol::Float64=1E-14,
    rot_weight::Float64=1.0,
    regularization_threshold::Float64=0.1,
    regularization_strength::Float64=3E-2,
    max_update_norm::Float64=2.0,
    progress_thresh::Float64=20.0,
    max_steps::Int64=100)

    if isnothing(target_quat)
        jac = Array{Float64}(undef, 3, sim.m.nv)#zeros((3, sim.m.nv))
        err = Array{Float64}(undef, 3)#zeros(3)
        jac_pos, jac_rot = jac, nothing
        err_pos, err_rot = err, nothing

    else
        jac = Array{Float64}(undef, 6, sim.m.nv)#zeros((6, sim.m.nv))
        err = Array{Float64}(undef, 6)#zeros(6)
        jac_pos, jac_rot = jac[1:4], jac[4:end]
        err_pos, err_rot = err[1:4], err[4:end]
    end

    update_nv = zeros(sim.m.nv)

    if !isnothing(target_quat)
        site_xquat = Array{Float64}(undef, 4)#zeros(4)
        neg_site_xquat = Array{Float64}(undef, 4)#zeros(4)
        err_rot_quat = Array{Float64}(undef, 4)
    end

    mj_fwdPosition(sim.m, sim.d)

    site_id = jl_name2id(sim.m, MJCore.mjOBJ_SITE, site_name)

    site_xpos = sim.d.site_xpos[:,site_id]
    site_xmat = sim.d.site_xmat[:,site_id]
    println("site pos: ",site_xpos)

    if isnothing(joint_names)
        dof_indices = sim.m.dof_jntid
    else
        dof_indices = []
        for joint_name in joint_names
            jid = jl_name2id(sim.m, MJCore.mjOBJ_JOINT, joint_name)
            push!(dof_indices, jid)
        end
        println("dof inds: ",dof_indices)
    end

    steps = 0
    success = false

    for steps = 1:max_steps

        err_norm = 0.0

        if !isnothing(target_pos)
            err_pos = target_pos - site_xpos
            err_norm += norm(err_pos)
        end

        if !isnothing(target_quat)
            # q = UnitQuaternion(SMatrix{3,3}(transpose(reshape(site_xmat, (3,3)))))
            MJCore.mju_mat2Quat(site_xquat, site_xmat)
            MJCore.mju_negQuat(neg_site_xquat, site_xquat)
            MJCore.mju_mulQuat(err_rot_quat, target_quat, neg_site_xquat)
            println("error rot: ",err_rot)
            MJCore.mju_quat2Vel(err_rot, err_rot_quat, 1.)
            err_norm += norm(err_rot)*rot_weight
        end

        if err_norm < tol
            success = true
            break
        else
            mj_jacSite(sim.m, sim.d, jac_pos, jac_rot, site_id)
            jac_joints = Array{Float64}(undef, 6, length(dof_indices))
            for i = 1:length(dof_indices)
                jac_joints[:,i] = jac[:,dof_indices[i]]
            end

            reg_strength = err_norm > regularization_threshold ?
                            regularization_strength : 0.0

            update_joints = nullspace_method(jac_joints, err,
                                regularization_strength=reg_strength)
            update_norm = norm(update_joints)
            println(update_joints)
            progress_criterion = err_norm/update_norm
            if progress_criterion > progress_thresh
                break
            end
            if update_norm > max_update_norm
                update_joints *= max_update_norm/update_norm
            end

            for i in dof_indices
                update_nv[i] = update_joints[i]
            end

            mj_integratePos(sim.m, sim.d.qpos, update_nv, 1.)
            mj_fwdPosition(sim.m, sim.d)
        end
    end

    if !success && steps == max_steps
        println("Failed to converge")
    end

    qpos = sim.d.qpos
    return qpos
end


# ProximalOperators.jl - library of nonsmooth functions and associated proximal mappings

__precompile__()

module ProximalOperators

typealias RealOrComplex Union{Real, Complex}
typealias HermOrSym{T,S} Union{Hermitian{T,S}, Symmetric{T,S}}

export prox, prox!

export ProximableFunction

export IndAffine, IndHalfspace,
       IndBallLinf, IndBallL0, IndBallL1, IndBallL2, IndBallRank,
       IndBox, IndNonnegative, IndNonpositive,
       IndExpPrimal, IndExpDual, IndPSD, IndSOC, IndRotatedSOC,
       IndFree,
       IndPoint, IndZero,
       IndSimplex,
       IndSphereL2, IndBinary,
       SumPositive, HingeLoss, HuberLoss,
       LogBarrier,
       LeastSquares,
       Maximum,
       NormL0, NormL1, NormL2, NormL21, NormLinf, NuclearNorm, SqrNormL2, ElasticNet,
       FirmThreshold,
       DistL2, SqrDistL2,
       Zero

export Conjugate,
       Postcompose,
       Precompose,
       SlicedSeparableSum,
       SeparableSum,
       Tilt,
       Regularize

abstract ProximableFunction
abstract ProximableConvex <: ProximableFunction
abstract IndicatorConvex <: ProximableConvex
abstract IndicatorConvexCone <: IndicatorConvex
abstract IndicatorNonconvex <: ProximableFunction

include("utilities/symmetricpacked.jl")

include("calculus/conjugate.jl")
include("calculus/postcompose.jl")
include("calculus/precompose.jl")
include("calculus/separableSum.jl")
include("calculus/slicedSeparableSum.jl")
include("calculus/tilt.jl")
include("calculus/regularize.jl")

include("functions/distL2.jl")
include("functions/elasticNet.jl")
include("functions/firmThreshold.jl")
include("functions/logBarrier.jl")
include("functions/normL2.jl")
include("functions/normL1.jl")
include("functions/normL21.jl")
include("functions/normL0.jl")
include("functions/nuclearNorm.jl")
include("functions/hingeLoss.jl")
include("functions/huberLoss.jl")
include("functions/indAffine.jl")
include("functions/indBallL0.jl")
include("functions/indBallL1.jl")
include("functions/indBallL2.jl")
include("functions/indBallRank.jl")
include("functions/indBinary.jl")
include("functions/indBox.jl")
include("functions/indNonnegative.jl")
include("functions/indNonpositive.jl")
include("functions/indZero.jl")
include("functions/indExp.jl")
include("functions/indFree.jl")
include("functions/indPoint.jl")
include("functions/indPSD.jl")
include("functions/indSimplex.jl")
include("functions/indSOC.jl")
include("functions/indSphereL2.jl")
include("functions/indHalfspace.jl")
include("functions/sqrDistL2.jl")
include("functions/sqrNormL2.jl")
include("functions/sumLargest.jl")
include("functions/sumPositive.jl")
include("functions/maximum.jl")
include("functions/normLinf.jl")
include("functions/leastSquares.jl")

function Base.show(io::IO, f::ProximableFunction)
  println(io, "description : ", fun_name(f))
  println(io, "domain      : ", fun_dom(f))
  println(io, "expression  : ", fun_expr(f))
  print(  io, "parameters  : ", fun_params(f))
end

fun_name(  f) = "n/a"
fun_dom(   f) = "n/a"
fun_expr(  f) = "n/a"
fun_params(f) = "n/a"

is_prox_accurate(f::ProximableFunction) = true

"""
  prox(f::ProximableFunction, x::AbstractArray, γ::Real=1.0)

Computes the proximal point of `x` with respect to function `f`
and parameter `γ > 0`, that is

  y = argmin_z { f(z) + 1/(2γ)||z-x||² }

and returns `y` and `f(y)`.
"""

function prox(f::ProximableFunction, x::AbstractArray, gamma::Real=1.0)
  y = similar(x)
  fy = prox!(f, x, y, gamma)
  return y, fy
end

"""
  prox!(f::ProximableFunction, x::AbstractArray, y::AbstractArray, γ::Real=1.0)

Computes the proximal point of `x` with respect to function `f`
and parameter `γ > 0`, and writes the result in `y`. Returns `f(y)`.
"""

function prox!(f::ProximableFunction, x::AbstractArray, y::AbstractArray, gamma::Real=1.0)
  throw(MethodException(
    "prox! is not implemented for f::", typeof(f),
    ", x::", typeof(x), ", y::", typeof(y), ", gamma::", typeof(gamma)
  ))
end

"""
  prox!(f::ProximableFunction, x::AbstractArray, γ::Real=1.0)

Computes the proximal point of `x` with respect to function `f`
and parameter `γ > 0` *in place*, that is

  x ← argmin_z { f(z) + 1/(2γ)||z-x||² }

and returns `f(x)`.
"""

prox!(f::ProximableFunction, x::AbstractArray, gamma::Real=1.0) = prox!(f, x, x, gamma)

end


using OctreeBH
import OctreeBH: AbstractData, assign_additional_node_data!

using StaticArrays

mutable struct Data2{N,T<:Real} <: AbstractData{N,T}
	pos::SVector{N,T}
	idx::Int64
	hsml::T
	mass::T
	#define additional members below
	massH2::T
end

Data2{N,T}() where {N,T} = Data2{N,T}(zero(SVector{N,T}),0,0,0,0)

function assign_additional_node_data!(n::Data2, old::Data2, new::Data2)
	n.massH2 = old.massH2 + new.massH2
end

const BOXSIZE = 1.0

const N = 3 #spatial dimension
const T = Float64

const Npart = 100000 #number of particles

#desired number of neighbor particles (ngbs) within a search radius (hsml)
const Nngb0 = 32

#searching radius that should contain roughly Nngb0 neighboring particles
const hsml0 = BOXSIZE * (Nngb0/(4*pi/3*Npart))^(1/N)

#randomly distributing particles
X = [@SVector rand(N) for _ in 1:Npart]

topnode_length = @SVector(ones(N)) * BOXSIZE  #actual length of tree

center = topnode_length .* 0.5

#individual smoothing length (only relevant for scatter ngbs)
δ = 0.0 #add a small perturbation
hsml = ones(Npart) .* hsml0 .* ( 1 .+ δ .* (rand(Npart) .- 0.5) )

mass = ones(Npart) #particle mass

part = [Data2{N,T}(SVector(X[i]), i, hsml[i], mass[i], 0.6*mass[i]) for i in eachindex(X)]

#build the tree
tree = buildtree(part, center, topnode_length);

#search center
x = @SVector rand(N)

boxsizes = @SVector(ones(N)) * BOXSIZE  #for periodic B.C.

#get indices of the gather neighbors
idx_ngbs_g = get_gather_ngb_tree(x, hsml0, tree, boxsizes)

#get indices of the scatter neighbors
idx_ngbs_s = get_scatter_ngb_tree(x, tree, boxsizes)

#for constant hsml, the two ngbs are identical
@show length(idx_ngbs_g), length(idx_ngbs_s)


using LinearAlgebra

using Oceananigans.Architectures: array_type

function can_solve_single_tridiagonal_system(arch, N)
    ArrayType = array_type(arch)

    a = rand(N-1)
    b = 3 .+ rand(N) # +3 to ensure diagonal dominance.
    c = rand(N-1)
    f = rand(N)

    # Solve the system with backslash on the CPU to avoid scalar operations on the GPU.
    M = Tridiagonal(a, b, c)
    ϕ_correct = M \ f

    # Convert to CuArray if needed.
    a, b, c, f = ArrayType.([a, b, c, f])

    ϕ = reshape(zeros(N), (1, 1, N)) |> ArrayType

    grid = RectilinearGrid(arch, size=(1, 1, N), extent=(1, 1, 1))

    btsolver = BatchedTridiagonalSolver(arch, grid;
                                        lower_diagonal = a,
                                        diagonal = b,
                                        upper_diagonal = c)

    solve!(ϕ, btsolver, f)

    return Array(ϕ[:]) ≈ ϕ_correct
end

function can_solve_single_tridiagonal_system_with_functions(arch, N)
    ArrayType = array_type(arch)

    grid = RectilinearGrid(arch, size=(1, 1, N), extent=(1, 1, 1))

    a = rand(N-1)
    c = rand(N-1)

    @inline b(i, j, k, grid) = 3 .+ cos(2π * grid.zᵃᵃᶜ[k])  # +3 to ensure diagonal dominance.
    @inline f(i, j, k, grid) = sin(2π * grid.zᵃᵃᶜ[k])

    bₐ = [b(1, 1, k, grid) for k in 1:N]
    fₐ = [f(1, 1, k, grid) for k in 1:N]

    # Solve the system with backslash on the CPU to avoid scalar operations on the GPU.
    M = Tridiagonal(a, bₐ, c)
    ϕ_correct = M \ fₐ

    # Convert to CuArray if needed.
    a, c = ArrayType.((a, c))

    ϕ = reshape(zeros(N), (1, 1, N)) |> ArrayType

    btsolver = BatchedTridiagonalSolver(arch, grid;
                                        lower_diagonal = a,
                                        diagonal = b,
                                        upper_diagonal = c)

    solve!(ϕ, btsolver, f)

    return Array(ϕ[:]) ≈ ϕ_correct
end

function can_solve_batched_tridiagonal_system_with_3D_RHS(arch, Nx, Ny, Nz)
    ArrayType = array_type(arch)

    a = rand(Nz-1)
    b = 3 .+ rand(Nz) # +3 to ensure diagonal dominance.
    c = rand(Nz-1)
    f = rand(Nx, Ny, Nz)

    M = Tridiagonal(a, b, c)
    ϕ_correct = zeros(Nx, Ny, Nz)

    # Solve the systems with backslash on the CPU to avoid scalar operations on the GPU.
    for i = 1:Nx, j = 1:Ny
        ϕ_correct[i, j, :] .= M \ f[i, j, :]
    end

    # Convert to CuArray if needed.
    a, b, c, f = ArrayType.([a, b, c, f])

    grid = RectilinearGrid(arch, size=(Nx, Ny, Nz), extent=(1, 1, 1))
    btsolver = BatchedTridiagonalSolver(arch, grid;
                                        lower_diagonal = a,
                                        diagonal = b,
                                        upper_diagonal = c)

    ϕ = zeros(Nx, Ny, Nz) |> ArrayType

    solve!(ϕ, btsolver, f)

    return Array(ϕ) ≈ ϕ_correct
end

function can_solve_batched_tridiagonal_system_with_3D_functions(arch, Nx, Ny, Nz)
    ArrayType = array_type(arch)

    grid = RectilinearGrid(arch, size=(Nx, Ny, Nz), extent=(1, 1, 1))

    a = rand(Nz-1)
    c = rand(Nz-1)

    @inline b(i, j, k, grid) = 3 + grid.xᶜᵃᵃ[i] * grid.yᵃᶜᵃ[j] * cos(2π * grid.zᵃᵃᶜ[k])
    @inline f(i, j, k, grid) = (grid.xᶜᵃᵃ[i] + grid.yᵃᶜᵃ[j]) * sin(2π * grid.zᵃᵃᶜ[k])

    ϕ_correct = zeros(Nx, Ny, Nz)

    # Solve the system with backslash on the CPU to avoid scalar operations on the GPU.
    for i = 1:Nx, j = 1:Ny
        bₐ = [b(i, j, k, grid) for k in 1:Nz]
        M = Tridiagonal(a, bₐ, c)

        fₐ = [f(i, j, k, grid) for k in 1:Nz]
        ϕ_correct[i, j, :] .= M \ fₐ
    end

    # Convert to CuArray if needed.
    a, c = ArrayType.([a, c])

    btsolver = BatchedTridiagonalSolver(arch, grid;
                                        lower_diagonal = a,
                                        diagonal = b,
                                        upper_diagonal = c)

    ϕ = zeros(Nx, Ny, Nz) |> ArrayType
    solve!(ϕ, btsolver, f)

    return Array(ϕ) ≈ ϕ_correct
end

@testset "Batched tridiagonal solvers" begin
    @info "Testing BatchedTridiagonalSolver..."

    for arch in archs
        @testset "Batched tridiagonal solver [$arch]" begin
            for Nz in [8, 11, 18]
                @test can_solve_single_tridiagonal_system(arch, Nz)
                @test can_solve_single_tridiagonal_system_with_functions(arch, Nz)
            end

            for Nx in [3, 8], Ny in [5, 16], Nz in [8, 11]
                @test can_solve_batched_tridiagonal_system_with_3D_RHS(arch, Nx, Ny, Nz)
                @test can_solve_batched_tridiagonal_system_with_3D_functions(arch, Nx, Ny, Nz)
            end
        end
    end
end


module Triangle
export basic_triangulation
export basic_triangulation_vertices
export constrained_triangulation
export constrained_triangulation_vertices

if !isfile(joinpath(dirname(@__FILE__), "..", "deps", "deps.jl"))
  error("Triangle is not properly installed. Please run\nPkg.build(\"Triangle\")")
end

include("common.jl")

"""
    basic_triangulation_vertices(vertices::Array{Float64,2})

Compute a Delaunay triangulation for a list of `vertices` in the form of `[x1 y1; x2 y2; ... ; xn yn]`

The function will return an array of array of 3-vertices lists (triangles with the correct vertices order) using the vertex coordinates in each triangle definition.

# Example
```jldoctest
julia> using Triangle

julia> points = Array{Float64,2}([0. 0.; 1. 0.; 0. 1.])
3×2 Array{Float64,2}:
 0.0  0.0
 1.0  0.0
 0.0  1.0

julia> Triangle.basic_triangulation_vertices(points)
1-element Array{Array{Float64,2},1}:
 [0.0 0.0; 1.0 0.0; 0.0 1.0]
```
"""
function basic_triangulation_vertices(vertices::Array{Float64,2})
  basic_triangulation_vertices(vertices, Array{Int64,1}(collect(1:1:size(vertices)[1])))
end

"""
    basic_triangulation(vertices::Array{Float64,2},vertices_map::Array{Int64,1})

Compute a Delaunay triangulation for a list of `vertices` in the form of `[x1 y1; x2 y2; ... ; xn yn]`

A list of indexes is provided in `vertices_map` so that each vertex can have a custom integer identifier.

The function will return an array of array of 3-vertices lists (triangles with the correct vertices order) using the vertex identifiers in each triangle definition.

# Example
```jldoctest
julia> using Triangle

julia> points = Array{Float64,2}([0. 0.; 1. 0.; 0. 1.])
3×2 Array{Float64,2}:
 0.0  0.0
 1.0  0.0
 0.0  1.0

julia> points_map = [1, 2, 3]
3-element Array{Int64,1}:
 1
 2
 3

julia> Triangle.basic_triangulation(points,points_map)
1-element Array{Array{Int64,1},1}:
 [1, 2, 3]
```
"""
function basic_triangulation(vertices::Array{Float64,2}, vertices_map::Array{Int64,1})
  Triangulate.basic_triangulation(vertices, vertices_map)
end

"""
    basic_triangulation_vertices(vertices::Array{Float64,2},vertices_map::Array{Int64,1})

Compute a Delaunay triangulation for a list of `vertices` in the form of `[x1 y1; x2 y2; ... ; xn yn]`

A list of indexes is provided in `vertices_map` so that each vertex can have a custom integer identifier.

The function will return an array of array of 3-vertices lists (triangles with the correct vertices order) using the vertex coordinates in each triangle definition.

# Example
```jldoctest
julia> using Triangle

julia> points = Array{Float64,2}([0. 0.; 1. 0.; 0. 1.])
3×2 Array{Float64,2}:
 0.0  0.0
 1.0  0.0
 0.0  1.0

julia> points_map = [1, 2, 3]
3-element Array{Int64,1}:
 1
 2
 3

julia> Triangle.basic_triangulation_vertices(points,points_map)
1-element Array{Array{Float64,2},1}:
 [0.0 0.0; 1.0 0.0; 0.0 1.0]
```
"""
function basic_triangulation_vertices(vertices::Array{Float64,2}, vertices_map::Array{Int64,1})
  map_to_vertices( vertices, vertices_map, basic_triangulation(vertices, vertices_map) )
end

"""
    constrained_triangulation(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2})

Compute a Constrained Delaunay triangulation for a list of `vertices` in the form of `[x1 y1; x2 y2; ... ; xn yn]` and a list of edges that will be kept.

A list of indexes is provided in `vertices_map` so that each vertex can have a custom integer identifier.

A list of edges (to be included in the final triangulation) is passed in `edges_list` in the form of `[ vertex-identifier-1 vertex-identifier-2; vertex-identifier-1 vertex-identifier-3; ... ; vertex-identifier-N vertex-identifier-M ]`

The function will return an array of array of 3-vertices lists (triangles with the correct vertices order) using the vertex identifiers in each triangle definition.

# Example
```jldoctest
julia> using Triangle

julia> points = [0. 0.; 0. 3.; 1. 3.; 1. 1.; 2. 1.; 2. 0.]
6×2 Array{Float64,2}:
 0.0  0.0
 0.0  3.0
 1.0  3.0
 1.0  1.0
 2.0  1.0
 2.0  0.0

julia> points_map = Array{Int64,1}(collect(1:1:size(points)[1]))
6-element Array{Int64,1}:
 1
 2
 3
 4
 5
 6

julia> edges_list = Array{Int64,2}([1 2; 2 3; 3 4; 4 5; 5 6; 6 1])
6×2 Array{Int64,2}:
 1  2
 2  3
 3  4
 4  5
 5  6
 6  1

julia> Triangle.constrained_triangulation(points,points_map,edges_list)
4-element Array{Array{Int64,1},1}:
 [1, 4, 2]
 [4, 1, 6]
 [2, 4, 3]
 [5, 4, 6]
```
"""
function constrained_triangulation(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2})
  Triangulate.constrained_triangulation(vertices, vertices_map, edges_list, fill(true, size(edges_list)[1])  )
end

"""
    constrained_triangulation_vertices(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2})

Compute a Constrained Delaunay triangulation for a list of `vertices` in the form of `[x1 y1; x2 y2; ... ; xn yn]` and a list of edges that will be kept.

A list of indexes is provided in `vertices_map` so that each vertex can have a custom integer identifier.

A list of edges (to be included in the final triangulation) is passed in `edges_list` in the form of `[ vertex-identifier-1 vertex-identifier-2; vertex-identifier-1 vertex-identifier-3; ... ; vertex-identifier-N vertex-identifier-M ]`

The function will return an array of array of 3-vertices lists (triangles with the correct vertices order) using the vertex coordinates in each triangle definition.

# Example
```jldoctest
julia> using Triangle

julia> points = [0. 0.; 0. 3.; 1. 3.; 1. 1.; 2. 1.; 2. 0.]
6×2 Array{Float64,2}:
 0.0  0.0
 0.0  3.0
 1.0  3.0
 1.0  1.0
 2.0  1.0
 2.0  0.0

julia> points_map = Array{Int64,1}(collect(1:1:size(points)[1]))
6-element Array{Int64,1}:
 1
 2
 3
 4
 5
 6

julia> edges_list = Array{Int64,2}([1 2; 2 3; 3 4; 4 5; 5 6; 6 1])
6×2 Array{Int64,2}:
 1  2
 2  3
 3  4
 4  5
 5  6
 6  1

julia> Triangle.constrained_triangulation_vertices(points,points_map,edges_list)
4-element Array{Array{Float64,2},1}:
 [0.0 0.0; 1.0 1.0; 0.0 3.0]
 [1.0 1.0; 0.0 0.0; 2.0 0.0]
 [0.0 3.0; 1.0 1.0; 1.0 3.0]
 [2.0 1.0; 1.0 1.0; 2.0 0.0]
```
"""
function constrained_triangulation_vertices(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2})
  map_to_vertices( vertices, vertices_map, constrained_triangulation(vertices, vertices_map, edges_list, fill(true, size(edges_list)[1])) )
end

"""
    constrained_triangulation(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2}, edges_boundary::Array{Bool,1})

Compute a Constrained Delaunay triangulation for a list of `vertices` in the form of `[x1 y1; x2 y2; ... ; xn yn]` and a list of edges that will be kept. Some of those edges can be marked as the boundary of the mesh.

A list of indexes is provided in `vertices_map` so that each vertex can have a custom integer identifier.

A list of edge (to be included in the final triangulation) is passed in `edges_list` in the form of `[ vertex-identifier-1 vertex-identifier-2; vertex-identifier-1 vertex-identifier-3; ... ; vertex-identifier-N vertex-identifier-M ]`

A list of boundary markers passed in `edges_boundary` in the form of booleans that tell the triangulator if the edge is on the boundary or not (the indexing is the same of `edges_list`).

The function will return an array of array of 3-vertices lists (triangles with the correct vertices order) using the vertex identifiers in each triangle definition.

# Example
```jldoctest
julia> using Triangle

julia> points = [0. 0.; 0. 3.; 1. 3.; 1. 1.; 2. 1.; 2. 0.]
6×2 Array{Float64,2}:
 0.0  0.0
 0.0  3.0
 1.0  3.0
 1.0  1.0
 2.0  1.0
 2.0  0.0

julia> points_map = Array{Int64,1}(collect(1:1:size(points)[1]))
6-element Array{Int64,1}:
 1
 2
 3
 4
 5
 6

julia> edges_list = Array{Int64,2}([1 2; 2 3; 3 4; 4 5; 5 6; 6 1])
6×2 Array{Int64,2}:
 1  2
 2  3
 3  4
 4  5
 5  6
 6  1

julia> edge_boundary = [false, false, true, true, false, false]
6-element Array{Bool,1}:
 false
 false
  true
  true
 false
 false

julia> Triangle.constrained_triangulation(points,points_map,edges_list,edge_boundary)
4-element Array{Array{Int64,1},1}:
 [1, 4, 2]
 [4, 1, 6]
 [2, 4, 3]
 [5, 4, 6]
```
"""
function constrained_triangulation(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2}, edges_boundary::Array{Bool,1})
  Triangulate.constrained_triangulation(vertices, vertices_map, edges_list, edges_boundary)
end

"""
    constrained_triangulation_vertices(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2}, edges_boundary::Array{Bool,1})

Compute a Constrained Delaunay triangulation for a list of `vertices` in the form of `[x1 y1; x2 y2; ... ; xn yn]` and a list of edges that will be kept. Some of those edges can be marked as the boundary of the mesh.

A list of indexes is provided in `vertices_map` so that each vertex can have a custom integer identifier.

A list of edge (to be included in the final triangulation) is passed in `edges_list` in the form of `[ vertex-identifier-1 vertex-identifier-2; vertex-identifier-1 vertex-identifier-3; ... ; vertex-identifier-N vertex-identifier-M ]`

A list of boundary markers passed in `edges_boundary` in the form of booleans that tell the triangulator if the edge is on the boundary or not (the indexing is the same of `edges_list`).

The function will return an array of array of 3-vertices lists (triangles with the correct vertices order) using the vertex coordinates in each triangle definition.

# Example
```jldoctest
julia> using Triangle

julia> points = [0. 0.; 0. 3.; 1. 3.; 1. 1.; 2. 1.; 2. 0.]
6×2 Array{Float64,2}:
 0.0  0.0
 0.0  3.0
 1.0  3.0
 1.0  1.0
 2.0  1.0
 2.0  0.0

julia> points_map = Array{Int64,1}(collect(1:1:size(points)[1]))
6-element Array{Int64,1}:
 1
 2
 3
 4
 5
 6

julia> edges_list = Array{Int64,2}([1 2; 2 3; 3 4; 4 5; 5 6; 6 1])
6×2 Array{Int64,2}:
 1  2
 2  3
 3  4
 4  5
 5  6
 6  1

julia> edge_boundary = [false, false, true, true, false, false]
6-element Array{Bool,1}:
 false
 false
  true
  true
 false
 false

julia> Triangle.constrained_triangulation_vertices(points,points_map,edges_list,edge_boundary)
4-element Array{Array{Float64,2},1}:
 [0.0 0.0; 1.0 1.0; 0.0 3.0]
 [1.0 1.0; 0.0 0.0; 2.0 0.0]
 [0.0 3.0; 1.0 1.0; 1.0 3.0]
 [2.0 1.0; 1.0 1.0; 2.0 0.0]
```
"""
function constrained_triangulation_vertices(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2}, edges_boundary::Array{Bool,1})
  map_to_vertices( vertices, vertices_map, constrained_triangulation(vertices, vertices_map, edges_list, edges_boundary) )
end

"""
    constrained_triangulation(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2}, edges_boundary::Array{Bool,1}, holes::Array{Float64,2})

Compute a Constrained Delaunay triangulation for a list of `vertices` in the form of `[x1 y1; x2 y2; ... ; xn yn]` and a list of edges that will be kept. Some of those edges can be marked as the boundary of the mesh.

A list of `holes` in the form of _[x1 y1; x2 y2; ... ; xn yn]_ can be passed to avoid triangulation in that part of the mesh.

A list of indexes is provided in `vertices_map` so that each vertex can have a custom integer identifier.

A list of edge (to be included in the final triangulation) is passed in `edges_list` in the form of `[ vertex-identifier-1 vertex-identifier-2; vertex-identifier-1 vertex-identifier-3; ... ; vertex-identifier-N vertex-identifier-M ]`

A list of boundary markers passed in `edges_boundary` in the form of booleans that tell the triangulator if the edge is on the boundary or not (the indexing is the same of `edges_list`).

The function will return an array of array of 3-vertices lists (triangles with the correct vertices order) using the vertex identifiers in each triangle definition.

# Example
```jldoctest
julia> using Triangle

julia> points = Array{Float64,2}([0. 0.; 4. 0.; 2. 3.; 8. 0.; 6. 3.; 4. 6.])
6×2 Array{Float64,2}:
 0.0  0.0
 4.0  0.0
 2.0  3.0
 8.0  0.0
 6.0  3.0
 4.0  6.0

julia> points_map = [1, 2, 4, 3, 5, 6]
6-element Array{Int64,1}:
 1
 2
 4
 3
 5
 6

julia> edges_list = Array{Int64,2}([1 2; 2 3; 3 1; 2 4; 4 5; 5 2; 3 5; 5 6; 6 3])
9×2 Array{Int64,2}:
 1  2
 2  3
 3  1
 2  4
 4  5
 5  2
 3  5
 5  6
 6  3

julia> edge_boundary = [false,true,false,false,false,true,true,false,false]
9-element Array{Bool,1}:
 false
  true
 false
 false
 false
  true
  true
 false
 false

julia> holes = [4. 2.]
1×2 Array{Float64,2}:
 4.0  2.0

julia> Triangle.constrained_triangulation(points, points_map, edges_list, edge_boundary, holes)
3-element Array{Array{Int64,1},1}:
 [1,2,3]
 [5,6,3]
 [4,5,2]
```
"""
function constrained_triangulation(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2}, edges_boundary::Array{Bool,1}, holes::Array{Float64,2})
  Triangulate.constrained_triangulation(vertices, vertices_map, edges_list, edges_boundary, holes)
end

"""
    constrained_triangulation_vertices(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2}, edges_boundary::Array{Bool,1}, holes::Array{Float64,2})

Compute a Constrained Delaunay triangulation for a list of `vertices` in the form of `[x1 y1; x2 y2; ... ; xn yn]` and a list of edges that will be kept. Some of those edges can be marked as the boundary of the mesh.

A list of `holes` in the form of _[x1 y1; x2 y2; ... ; xn yn]_ can be passed to avoid triangulation in that part of the mesh.

A list of indexes is provided in `vertices_map` so that each vertex can have a custom integer identifier.

A list of edge (to be included in the final triangulation) is passed in `edges_list` in the form of `[ vertex-identifier-1 vertex-identifier-2; vertex-identifier-1 vertex-identifier-3; ... ; vertex-identifier-N vertex-identifier-M ]`

A list of boundary markers passed in `edges_boundary` in the form of booleans that tell the triangulator if the edge is on the boundary or not (the indexing is the same of `edges_list`).

The function will return an array of array of 3-vertices lists (triangles with the correct vertices order) using the vertex coordinates in each triangle definition.

# Example
```jldoctest
julia> using Triangle

julia> points = Array{Float64,2}([0. 0.; 4. 0.; 2. 3.; 8. 0.; 6. 3.; 4. 6.])
6×2 Array{Float64,2}:
 0.0  0.0
 4.0  0.0
 2.0  3.0
 8.0  0.0
 6.0  3.0
 4.0  6.0

julia> points_map = [1, 2, 4, 3, 5, 6]
6-element Array{Int64,1}:
 1
 2
 4
 3
 5
 6

julia> edges_list = Array{Int64,2}([1 2; 2 3; 3 1; 2 4; 4 5; 5 2; 3 5; 5 6; 6 3])
9×2 Array{Int64,2}:
 1  2
 2  3
 3  1
 2  4
 4  5
 5  2
 3  5
 5  6
 6  3

julia> edge_boundary = [false,true,false,false,false,true,true,false,false]
9-element Array{Bool,1}:
 false
  true
 false
 false
 false
  true
  true
 false
 false

julia> holes = [4. 2.]
1×2 Array{Float64,2}:
 4.0  2.0

julia> Triangle.constrained_triangulation_vertices(points, points_map, edges_list, edge_boundary, holes)
3-element Array{Array{Float64,2},1}:
 [0.0 0.0; 4.0 0.0; 2.0 3.0]
 [6.0 3.0; 4.0 6.0; 2.0 3.0]
 [8.0 0.0; 6.0 3.0; 4.0 0.0]
```
"""
function constrained_triangulation_vertices(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, edges_list::Array{Int64,2}, edges_boundary::Array{Bool,1}, holes::Array{Float64,2})
  map_to_vertices( vertices, vertices_map, constrained_triangulation(vertices, vertices_map, edges_list, edges_boundary, holes) )
end

function map_to_vertices(vertices::Array{Float64,2}, vertices_map::Array{Int64,1}, triangle_list::Array{Array{Int64, 1},1})
  vert_size = size(vertices)

  map_positions = Array{Int64,1}(undef, length(vertices_map))
  for i in 1:1:length(vertices_map)
    map_positions[vertices_map[i]] = i
  end

  triangle_list_vertices = Array{Array{Float64,2},1}()
  for i in 1:1:length(triangle_list)
        triangle = Array{Float64}(undef, 3,2)
        triangle[1] = vertices[map_positions[triangle_list[i][1]]]
        triangle[4] = vertices[map_positions[triangle_list[i][1]]+vert_size[1]]
        triangle[2] = vertices[map_positions[triangle_list[i][2]]]
        triangle[5] = vertices[map_positions[triangle_list[i][2]]+vert_size[1]]
        triangle[3] = vertices[map_positions[triangle_list[i][3]]]
        triangle[6] = vertices[map_positions[triangle_list[i][3]]+vert_size[1]]
        push!(triangle_list_vertices, triangle)
  end

  return triangle_list_vertices
end

end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
using LinearAlgebra
using LoopVectorization
using LinearMaps
using Arpack: Arpack
using ArnoldiMethod

function prepareDiag(L)
    diag = zeros(2^L)
    for state = 1:2^L
        for i = 1:L
            j = i==L ? 1 : i+1
            @inbounds diag[state] -= (((state >> (i-1))&1) == ((state >> (j-1))&1)) ? 1 : -1
        end
    end
    diag
end

function Hfunc!(C, B, diag, L)
    N = length(diag)
    @tturbo for state = 1:N
        C[state] = diag[state] * B[state]
    end
    for i = 1:L
        @tturbo for state = 1:N
            newstate = (state&(~(2^L))) ⊻ (1<<(i-1))
            c = newstate == 0
            newstate = !c*newstate + c*N # remove if statement
            C[newstate] -= B[state]
        end
    end
end
prepareHfunc!(diag, L) = (C, B) -> Hfunc!(C, B, diag, L)

struct MyLinMap{T, F, I<:Integer}
    f!::F; N::I; issymmetric::Bool; isposdef::Bool
end
MyLinMap(f!, N; issymmetric=false, isposdef=false) =
    MyLinMap{Float64, typeof(f!), typeof(N)}(f!, N, issymmetric, isposdef)
LinearAlgebra.mul!(y, A::MyLinMap, x) = A.f!(y, x)
Base.size(A::MyLinMap) = (A.N, A.N)
Base.size(A::MyLinMap, i::Integer) = size(A::MyLinMap)[i]
Base.eltype(A::MyLinMap{T}) where T = T
LinearAlgebra.issymmetric(A::MyLinMap) = A.issymmetric
LinearAlgebra.isposdef(A::MyLinMap) = A.isposdef

L = 20
diag_ = prepareDiag(L)
H! = prepareHfunc!(diag_, L)

H_LinearMap = LinearMap(H!, 2^L, ismutating=true, issymmetric=true, isposdef=false)
H_MyLinMap = MyLinMap(H!, 2^L, issymmetric = true, isposdef = false)
sol_exact = -2sum(abs(sin((n-1/2) * pi/L)) for n in 1:L)
@show sol_exact;

# %%
@time e1, v1 = Arpack.eigs(H_LinearMap, nev=1, which=:SR)
@time e1, v1 = Arpack.eigs(H_LinearMap, nev=1, which=:SR)
e1[1]

# %%
@time e2, v2 = Arpack.eigs(H_MyLinMap, nev=1, which=:SR)
@time e2, v2 = Arpack.eigs(H_MyLinMap, nev=1, which=:SR)
e2[1]

# %%
@time d1, h1 = partialschur(H_LinearMap, nev=1, which=SR())
@time d1, h1 = partialschur(H_LinearMap, nev=1, which=SR())
d1.eigenvalues[1]

# %%
@time d2, h2 = partialschur(H_MyLinMap, nev=1, which=SR())
@time d2, h2 = partialschur(H_MyLinMap, nev=1, which=SR())
d2.eigenvalues[1]

# %%


blas_set_num_threads(1) # because small matrices
using DataFrames, Distributions#, Gadfly

const dat = readdlm("../data/dat.txt")
n,k1 = size(dat)
const k = k1-1

const y = dat[:,1]
const X = dat[:,2:end]

const XXi = inv(X'X)
const s = 100
const mle = XXi*X'y

a = 1
b = 1
s2 = 10
csb = 4XXi
S = chol(csb)'
css = 1

function ll(be::Array{Float64,1},sig2::Float64)
  c = y-X*be
  c'c/(-2sig2)-n*log(sig2)/2
end

lpb(be::Array{Float64,1}) = -be'XXi*be/2s2
lps(sig2::Float64) = (a-1)*log(s2)-s2/b
mvrnorm(M::Array{Float64,1}) = M+S*randn(k)

function mh(B=100000)
  accb = 0; accs = 0
  bb = Array(Float64,(k,B))
  ss = Array(Float64,B)
  bcur = bb[:,1]
  scur = 1.0

  println("Starting Metropolis...")
  @time for i in 1:B

    # Update β̂:
    candb = mvrnorm(bcur)
    q = ll(candb,scur)+lpb(candb) -ll(bcur,scur)-lpb(bcur)
    if q[1]>log(rand())
      bcur = candb
      accb += 1
    end

    # Update ŝ²:
    cands = rand(Normal(scur,sqrt(css)))
    if cands>0
      q = ll(bcur,cands)+lps(cands) -ll(bcur,scur)-lps(scur)
      if q[1]>log(rand())
        scur = cands
        accs += 1
      end
    end

    @inbounds bb[:,i] = bcur
    @inbounds ss[i]  = scur

    if i%(B/10)==0 print("\r",round(100*i/B),"%") end
  end # End of loops
  println("End of Metropolis.\n\n")

  return (bb,ss,accb,accs,B)
end # End of mh function

bb,ss,accb,accs,B = @time mh();

println("β̂: \n", mean(bb[:,Int(B*.9):end],2),"\n")
println("ŝ²: ",  mean(ss[Int(B*.9):end]),"\n")

println("Acceptance rate for β̂: ",accb/B)
println("Acceptance rate for ŝ²:",accs/B)


using Gadfly, DataArrays, RDatasets

plot(dataset("ggplot2", "diamonds"), x="Price", color="Cut", Geom.density)


using Test, OMEinsum
using OMEinsum: match_rule, PairWise, Sum, Tr, DefaultRule,
                Permutedims, PTrace, Hadamard, MatMul, nopermute,
                Identity, BatchedContract

@testset "match rule" begin
    ixs = ((1,2), (2,3))
    iy = (1,3)
    @test match_rule(ixs, iy) == MatMul()
    @test match_rule(EinCode(ixs, iy)) == MatMul()

    ixs = ((1,2), (2,3), (3,4))
    iy = (1,4)
    @test match_rule(ixs, iy) == PairWise()

    ixs = ((1,1),)
    iy = ()
    @test match_rule(ixs, iy) == Tr()

    ixs = ((1,1,2),)
    iy = (2,)
    @test match_rule(ixs, iy) == PTrace()

    ixs = ((1,2),)
    iy = (1,)
    @test match_rule(ixs, iy) == Sum()

    ixs = ((1,2,3),)
    iy = (2,1,3)
    @test match_rule(ixs, iy) != Sum()

    ixs = ((1,2),(1,2),(1,2))
    iy = (1,2)
    @test match_rule(ixs, iy) == Hadamard()

    ixs = ((1,2,3),)
    iy = (1,3,2)
    @test match_rule(ixs, iy) == Permutedims()

    ixs = ((1,2),)
    iy = (1,1)
    @test match_rule(ixs, iy) == DefaultRule()

    ix = (1,2,3)
    iy = (1,2)
    @test nopermute(ix,iy)

    ix = (1,2,3)
    iy = (2,1)
    @test !nopermute(ix,iy)

    # 3-argument match_rule
    @test match_rule(Tr, ((1,1),), ())
    @test !match_rule(Tr, ((1,2),), ())
    @test !match_rule(Tr, ((1,1),), (1,))

    @test match_rule(PairWise, ((1,2),(2,3),(3,4),(4,5)), (1,5))
    @test match_rule(PairWise, ((1,2),(2,3),(5,4),(3,4)), (1,5))
    @test !match_rule(PairWise, ((1,2),(2,3),(3,4),(3,5)), (1,5))

    @test match_rule(Sum, ((1,2,3,4),), (1,2))
    @test match_rule(Sum, ((1,2,3,4),), (1,2))
    @test match_rule(Sum, ((1,2,3,4),), (1,2,3,4))
    @test match_rule(Sum, ((1,2,3,4),), (2,1))
    @test !match_rule(Sum, ((1,2,3,4),), (2,1,1))

    @test match_rule(Permutedims, ((1,2),), (2,1))
    @test !match_rule(Permutedims, ((1,2,3),), (2,1))
    @test !match_rule(Permutedims, ((3,3),), (1,3))

    @test match_rule(Hadamard, ((1,2),(1,2)), (1,2))
    @test !match_rule(Hadamard, ((1,2),(2,1)), (1,2))
    @test !match_rule(Hadamard, ((1,2),(1,2)), (1,2,1))

    @test match_rule(PTrace, ((1,1,2),), (2,))
    @test match_rule(PTrace, ((1,1,2,3),), (2,3))
    @test match_rule(PTrace, ((1,1,2,3),), (3,2))
    @test !match_rule(PTrace, ((1,1,2),), (1,2,))

    @test match_rule(MatMul, ((1,2),(2,3)), (1,3))
    @test match_rule(MatMul, ((1,2),(2,3)), (3,1))
    @test match_rule(MatMul, ((2,1),(2,3)), (3,1))
    @test match_rule(MatMul, ((2,1),(3,2)), (3,1))
    @test !match_rule(MatMul, ((1,2,3),(2,3,4)), (1,4))

    @test match_rule(Identity, ((1,2,3),), (1,2,3))
    @test !match_rule(Identity, ((1,2,3),), (1,3,2))

end

@testset "isbatchmul" begin
    @test match_rule(BatchedContract,((1,2), (2,3)), (1,3)) # matmul
    @test match_rule(BatchedContract,((1,2,3), (2,3)), (1,3)) # matvec-batched
    @test match_rule(BatchedContract,((7,1,2,3), (2,4,3,7)), (1,4,3)) # matmat-batched
    @test match_rule(BatchedContract,((3,), (3,)), (3,))  # pure batch
    @test match_rule(BatchedContract,((3,1), (3,)), (3,1))  # batched vector-vector
    @test !match_rule(BatchedContract,((2,2), (2,3)), (1,3))
    @test !match_rule(BatchedContract,((2,1), (2,3)), (1,1))
    @test !match_rule(BatchedContract,((1,2,3), (2,4,3)), (1,3))
end

@testset "match_rule eye candies" begin
    @test match_rule(ein"ij,jk,kl->il") == PairWise()
    @test match_rule(ein"(ij,jk),kl->il") == (OMEinsum.MatMul(), ((OMEinsum.MatMul(), (1, 2)), 3))
end


function myBubbleSort(x)
if length(x) < 1
return "Your original is not an array!"
else
println(" original: $x")
s=x
index = [1:length(s)...]
rank = 0 * index
for i in 2:length(s), j in 1:length(s)-1
if s[j] > s[j+1]
s[j], s[j+1] = s[j+1], s[j]
index[j], index[j+1] = index[j+1], index[j]
end
end
for i in 1:length(index)
rank[index[i]] = i
end
end
println(" sorted:   $s\n index:    $index\n rank:     $rank")
end


include("day5_solution.jl")
using Test

@test parse_rows("FBFBBFF") == 44
@test parse_columns("RLR") == 5

@test get_seat_id("BFFFBBFRRR") == 567
@test get_seat_id("FFFBBBFRRR") == 119
@test get_seat_id("BBFFBBFRLL") == 820

@test id_to_row_column(357) == (44, 5)


using SimpleTraits

@testset "basic functionality test" begin
    x = BooleanSemiringElement(true)
    y = BooleanSemiringElement(false)
    z = BooleanSemiringElement{Bool}(true)

    @test x == z
    @test x != y != z
    @test x+y isa BooleanSemiringElement
    @test x*y isa BooleanSemiringElement
    @test x+y == y+x == x
    @test x+z == z+x == z == x

    @test zero(BooleanSemiringElement) == y
    @test one(BooleanSemiringElement) == x

    @test iszero(zero(BooleanSemiringElement))
    @test isone(one(BooleanSemiringElement))

    @test x⊕y == x+y
    @test x⊙y == x*y
end

@testset "conversions between elements of the same semiring with different types" begin
    x = RealSemiringElement{Int64}(3)
    y = RealSemiringElement{Float64}(11.1)
    z = RealSemiringElement{Int16}(1)

    xc = convert(RealSemiringElement{Float64}, x)
    yc = convert(RealSemiringElement{Float64}, y)
    zc = convert(RealSemiringElement{Float64}, z)

    @test x == xc
    @test y == yc
    @test z == zc

    @test x+x isa RealSemiringElement{Int64}
    @test y+y isa RealSemiringElement{Float64}
    @test z+z isa RealSemiringElement{Int16}
    @test x+y isa RealSemiringElement{Float64}
    @test x+z isa RealSemiringElement{Int64}
    @test y+z isa RealSemiringElement{Float64}

    @test xc isa RealSemiringElement{Float64}
    @test yc isa RealSemiringElement{Float64}
    @test zc isa RealSemiringElement{Float64}
end

@testset "operations between different semirings" begin
    x = MaxPlusSemiringElement(3.0)
    y = MinPlusSemiringElement(5.0)
    z = RealSemiringElement(2)

    # TODO: make these domain errors
    @test_throws ErrorException x+y
    @test_throws ErrorException x+z
    @test_throws ErrorException y+z
end

@testset "semiring type methods" begin
    x = RealSemiringElement(3)
    y = RealSemiringElement(5.0)
    z = MaxPlusSemiringElement(11.2)

    @test issamesemiring(x, y)
    @test !issamesemiring(x, z)
    @test !issamesemiring(y, z)

    @test semiring(x) <: RealSemiringElement
    @test eltype(semiring(x)) == Int64
    @test eltype(semiring(z)) == Float64
end

@testset "vectors and matrices" begin
    A = MaxPlusSemiringElement{Int64}[3 5; 1 2]
    B = MaxPlusSemiringElement{Int64}[10 6; -1 -2]

    @test A+B == MaxPlusSemiringElement{Int64}[10 6; 1 2]
    @test A*B == MaxPlusSemiringElement{Int64}[13 9; 11 7]

    C = MaxPlusSemiringElement{Float64}[1.0 2.0; 3.0 4.0]
    D = MaxPlusSemiringElement[1.0 2.0; 3.0 4.0]
    E = MaxPlusSemiringElement{Float64}[1 2; 3 4]
    @test C == D == E

    x = rand(3,3)
    y = rand(3)
    F = convert(Matrix{RealSemiringElement{Float64}}, x)
    G = convert(Matrix{RealSemiringElement{Float64}}, x+x)
    H = convert(Matrix{RealSemiringElement{Float64}}, x*x)
    I = convert(Vector{RealSemiringElement{Float64}}, y)
    J = convert(Vector{RealSemiringElement{Float64}}, x*y)

    @test F isa Matrix{RealSemiringElement{Float64}}
    @test G isa Matrix{RealSemiringElement{Float64}}
    @test H isa Matrix{RealSemiringElement{Float64}}
    @test I isa Vector{RealSemiringElement{Float64}}
    @test J isa Vector{RealSemiringElement{Float64}}

    @test F+F == G
    @test F*F == H
    @test F*I == J

    @test A+C isa Matrix{MaxPlusSemiringElement{Float64}}
end

@testset "promotion" begin
    x = MinPlusSemiringElement(3)
    y = MinPlusSemiringElement(4.0)
    z = MinPlusSemiringElement(5//3)
    a = MaxPlusSemiringElement(1)

    @test promote_type(typeof(x), typeof(y), typeof(z)) == MinPlusSemiringElement{Float64}
    xp, yp, zp = promote(x, y, z)
    @test xp isa promote_type(typeof(x), typeof(y), typeof(z))
    @test yp isa promote_type(typeof(x), typeof(y), typeof(z))
    @test zp isa promote_type(typeof(x), typeof(y), typeof(z))
    @test_throws DomainError promote(x, y, z, a)
end

@testset "test each implementation" begin
    for S in (RealSemiringElement, MaxPlusSemiringElement, MinPlusSemiringElement)
        for elty in (Int32, Int64, Float32, Float64, Rational)
            x = S(elty(5))
            y = S(elty(6))
            z = S(elty(2))
            @test x+zero(x) == zero(x) + x == x
            @test x*one(x) == one(x)*x == x
            @test zero(x)*(x+y) == (x+y)*zero(x) == zero(x)
            @test (x+y)*z == x*z + y*z
            @test z*(x+y) == z*x + z*y
        end
    end

    x = BooleanSemiringElement(true)
    y = BooleanSemiringElement(false)
    z = BooleanSemiringElement(true)

    @test x+zero(x) == zero(x) + x == x
    @test x*one(x) == one(x)*x == x
    @test zero(x)*(x+y) == (x+y)*zero(x) == zero(x)
    @test (x+y)*z == x*z + y*z
    @test z*(x+y) == z*x + z*y

    # optional methods defined for RealSemiringElements

    x = RealSemiringElement(3.0)
    y = RealSemiringElement(0.5)

    @test x-x == zero(x)
    @test x/x == one(x)
    @test x/y == RealSemiringElement(3.0/0.5)
    @test x*inv(x) == one(x)
    @test_throws DomainError star(x)
    @test star(y) == RealSemiringElement(1/(1-.5))

    @test y+addinv(y) == addinv(y)+y == zero(y)
    @test y*mulinv(y) == mulinv(y)*y == one(y)
end

@testset "traits" begin
    @traitfn commutativetest(x::X) where {X; IsCommutative{X}} = true
    @traitfn commutativetest(x::X) where {X; !IsCommutative{X}} = false

    @traitfn idempotenttest(x::X) where {X; IsIdempotent{X}} = true
    @traitfn idempotenttest(x::X) where {X; !IsIdempotent{X}} = false

    @traitfn startest(x::X) where {X; IsStar{X}} = true
    @traitfn startest(x::X) where {X; !IsStar{X}} = false

    w = BooleanSemiringElement(true)
    x = MaxPlusSemiringElement(3.0)
    y = MinPlusSemiringElement(5.0)
    z = RealSemiringElement(2)

    @test commutativetest(w)
    @test commutativetest(x)
    @test commutativetest(y)
    @test commutativetest(z)

    @test idempotenttest(w)
    @test idempotenttest(x)
    @test idempotenttest(y)
    @test !idempotenttest(z)

    @test startest(w)
    @test startest(x)
    @test startest(y)
    @test startest(z)
end


module MCMC

using Distributions
using LinearAlgebra
include("PropCov.jl")

function _metropolis(x::Float64, logprob::Function, propsd::Float64)
  cand = curr + randn() * propsd
  met_ratio = logprob(cand) - logprob(curr)
  accept = met_ratio > log(rand())
  draw = accept ? cand : curr
  return (draw, accept)
end

function metropolis(x::Float64, logprob::Function, propsd::Float64)
  return _metropolis(x, logprob, propsd)[1]
end

function metropolis(curr::Vector{Float64}, logprob::Function,
                    propcov::Matrix{Float64})
  cand = rand(MvNormal(curr, propcov))
  met_ratio = logprob(cand) - logprob(curr)
  out = met_ratio > log(rand()) ? cand : curr
  return out
end

function mcmc(init::Vector{T}, ll::Function, lp::Function;
              nmcmc::Int, nburn::Int, batchsize::Int=50,
              thin::Int=1,
              propcov=nothing,
              propcov_init=eye(length(init)),
              temper_power=1,
              momentum=0.1, min_temper=1.0, max_temper=1.0,
              verbose=0) where T <: AbstractFloat
  @assert min_temper <= max_temper

  # Initialize MCMC
  state = deepcopy(init)

  # Number of parameters
  nparam = length(state)

  # Preallocate memory
  out = zeros(nmcmc, nparam)

  # Log density of target distribution, for a given temperature
  logprob(s::Vector{Float64}, temper::Float64) = ll(s) / temper + lp(s)

  # Total number of iterations
  total_iters = nmcmc + nburn

  # proposal covariance
  if propcov == nothing
    propcov = PropCov(propcov_init, momentum=momentum, batchsize=batchsize)
  end

  # Update function
  function update(s, iter, propcov)
    # Compute temperature
    p = ((nburn - iter)^temper_power + 1) / (nburn^temper_power - 1)
    temper = iter < nburn ? max_temper^p : min_temper

    # Log density of target distribution
    _logprob(s::Vector{Float64}) = logprob(s, temper)

    # Update state
    s .= metropolis(s, _logprob, propcov)

    # Build a message
    msg = "temper: $(round(temper, digits=3))"

    return msg
  end

  # Run MCMC
  msg = ""
  for i in 1:total_iters
    print("\r$(i)/$(total_iters)")
    # Update state
    for t in 1:thin
      msg = update(state, i, propcov.value)
    end

    # Update proposal covariance
    update!(propcov, state)

    # Sanity check
    if verbose > 0 && mod(i, batchsize) == 0
      println(" | proposal covariance: $(round.(propcov.value[1], digits=3)) | $(msg)")
    end

    # Keep samples
    if i > nburn
      out[i - nburn, :] .= state
    end
  end

  println()
  return out, propcov
end

end # module MCMC


export OneDGrid

struct OneDGrid

    dev :: AbstractDevice
    len  :: Int
    start :: Real
    stop  :: Real
    points :: LinRange
    step :: Real

    function OneDGrid( dev, len, start, stop )

        points = LinRange( start, stop, len+1 )[1:end-1]
        step = ( stop - start ) / len
        new( dev, len, start, stop, points, step)

    end

end


# Memoisation for Approximate Computation
#
# #######################################

function ApproximateHashingMemoise(fn, hashmap, hashfn, inputs...)
    hashid = hashfn(fn, inputs...)
    if( haskey(hashmap, hashid))
       return hashmap[hashid]
    else
        res = fn(inputs...)
        hashmap[hashid] = res
        return res
    end
end

function TrendingMemoisation(fn, hashfn, hasharray, limit, inputs...)
    index = hashfn(fn, inputs...)

    # Training is complete
    if(hasharray[index][1] > limit)
        return hasharray[index][3]
    end

    # Continue Training
    res = fn(inputs...)
    hasharray[index][1] += 1
    hasharray[index][2] += res

    if(hasharray[index][1] > limit)
        hasharray[index][3] = hasharray[index][2]/Float64.(hasharray[index][1])
        return hasharray[index][3]
    end

    return res
end