# Simulations parameters
include("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/1-Inputs/src/decision_making.jl")
include("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/1-Inputs/src/reaction_time_analysis_functions.jl")
include("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/1-Inputs/src/distributions.jl")

param = Dict("T_rest"=>10000,
"mu0"=>30,
"noise_amp"=>0.02,
"I0E1"=>0.3255,
"I0E2"=>0.3255,
"JN11"=>0.2609,
"JN22"=>0.2609,
"JN12"=>0.0497,
"JN21"=>0.0497,
"JAext" => 0.00052,
"Tnmda" => 100,
"Tampa" => 2,
"gamma" => 0.641,
"Tstim"=>5000,
"coh"=>5,
"false_time"=>300,
"Threshold"=>20
)

fixed_param=Dict("a"=>270,"b"=>108,"d"=>0.1540)
#Parameters of the simulation for EUler methods
simu_param=Dict("dt"=>0.5,"time_wind"=>4,"slide_wind"=>4) #4=2/dt

using DataFrames
using CSV

#a=2
Tr=500
println(Tr)
# de 0.02 à 0.025 déjàa sauvegardé (voir 0.02,0.3 50)

Irlist=[0.035,0.047]
Cohlist=[2,10,20]

using Plots
gr()

############################
## Fig neuronale PES
############################
a=300
PEsameL=zeros(251+a)#
PEdiffL=zeros(251+a)
PEsameW=zeros(251+a)
PEdiffW=zeros(251+a)
PCsameL=zeros(251+a)#
PCdiffL=zeros(251+a)
PCsameW=zeros(251+a)
PCdiffW=zeros(251+a)
numPEsame=0
numPEdiff=0
numPEsame=0
numPEdiff=0
numPCsame=0
numPCdiff=0
numPCsame=0
numPCdiff=0
Ir=Irlist[2]
c=Cohlist[3]

numr=0
numr2=1

for i=1:25
    r=CSV.read("./datatempAll/Coh$c-Final200Exp$i-T$Tr-Ir$Ir.csv")
    r2=CSV.read("./datatempAll/NondynCoh$c-Final200Exp$i-T$Tr-Ir$Ir.csv")

temp = 250
for i=1:2500
temp= temp + Int(r2[:RTs][i]/2)
if r2[:NotError][i] == 1 && r2[:Distri][i] >0 && r2[:Distri][i+1] >0
PCsameL=PCsameL + r[:R2][temp:(temp+250+a)]
PCsameW=PCsameW + r[:R1][temp:(temp+250+a)]

numPCsame=numPCsame+1
#rlist2=rlist2 + r[:R2][temp:(temp+250)]

elseif r2[:NotError][i] == -1 && r2[:Distri][i] <0 && r2[:Distri][i+1] >0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    PEsameL=PEsameL + r[:R2][temp:(temp+250+a)]
    PEsameW=PEsameW + r[:R1][temp:(temp+250+a)]

    numPEsame=numPEsame+1

elseif r2[:NotError][i] == -1 && r2[:Distri][i] >0 && r2[:Distri][i+1] <0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    PEsameL=PEsameL + r[:R1][temp:(temp+250+a)]
    PEsameW=PEsameW + r[:R2][temp:(temp+250+a)]

    numPEsame=numPEsame+1
elseif r2[:NotError][i] == 1 && r2[:Distri][i] <0 && r2[:Distri][i+1] <0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    PCsameL=PCsameL + r[:R1][temp:(temp+250+a)]
    PCsameW=PCsameW + r[:R2][temp:(temp+250+a)]

    numPCsame=numPCsame+1

elseif r2[:NotError][i] == 1 && r2[:Distri][i] >0 && r2[:Distri][i+1] <0 && r2[:NotError][i+1] == 1
    PCdiffL=PCdiffL + r[:R2][temp:(temp+250+a)]
    PCdiffW=PCdiffW + r[:R1][temp:(temp+250+a)]

    numPCdiff=numPCdiff+1
    #rlist2=rlist2 + r[:R2][temp:(temp+250)]

elseif r2[:NotError][i] == -1 && r2[:Distri][i] <0 && r2[:Distri][i+1] <0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        PEdiffL=PEdiffL + r[:R2][temp:(temp+250+a)]
        PEdiffW=PEdiffW + r[:R1][temp:(temp+250+a)]

        numPEdiff=numPEdiff+1

    elseif r2[:NotError][i] == -1 && r2[:Distri][i] >0 && r2[:Distri][i+1] >0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        PEdiffL=PEdiffL + r[:R1][temp:(temp+250+a)]
        PEdiffW=PEdiffW + r[:R2][temp:(temp+250+a)]

        numPEdiff=numPEdiff+1
    elseif r2[:NotError][i] == 1 && r2[:Distri][i] <0 && r2[:Distri][i+1] >0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        PCdiffL=PCdiffL + r[:R1][temp:(temp+250+a)]
        PCdiffW=PCdiffW + r[:R2][temp:(temp+250+a)]

        numPCdiff=numPCdiff+1

end
temp = temp+250
end

end

PEsameL=PEsameL./numPEsame
PEdiffL=PEdiffL./numPEdiff
PEsameW=PEsameW./numPEsame
PEdiffW=PEdiffW./numPEdiff
PCsameL=PCsameL./numPCsame
PCdiffL=PCdiffL./numPCdiff
PCsameW=PCsameW./numPCsame
PCdiffW=PCdiffW./numPCdiff

plot(PEsameL,linewidth=3)
plot!(PCsameL,linewidth=3)

plot(PEdiffL[5:250],linewidth=4,color=:darkorange,line=:dash)
plot!(PCdiffL[5:250],linewidth=4,color=:blue,line=:dash)
plot!(PEdiffW[5:250],linewidth=4,color=:darkorange)
plot!(PCdiffW[5:250],linewidth=4,color=:blue)
plot!(legend=:bottomleft)

savefig("./DiffPEQTrue.svg")

mean(r2[:RTs])

plot!(PEsameW,linewidth=3)
plot!(PCsameW,linewidth=3)

plot(PEsameL[5:250],linewidth=4,color=:darkorange,line=:dash)
plot!(PCsameL[5:250],linewidth=4,color=:blue,line=:dash)
plot!(PEsameW[5:250],linewidth=4,color=:darkorange)
plot!(PCsameW[5:250],linewidth=4,color=:blue)
plot!(legend=:topleft)

savefig("./SamePEQ.svg")

plot(PEdiffL[30:320]-PCdiffL[30:320],linewidth=3,color=:lightgray)
plot!(-PEsameW[30:320]+PCsameW[30:320],linewidth=3,color=:black)
vline!([220],linewidth=4,color=:black,line=:dash)

savefig("./DifferencePES.svg")

############################
## Fig neuronale PES: time from sacccade
############################
a=80
PEsameL=zeros(a+1)#
PEdiffL=zeros(a+1)
PEsameW=zeros(a+1)
PEdiffW=zeros(a+1)
PCsameL=zeros(a+1)#
PCdiffL=zeros(a+1)
PCsameW=zeros(a+1)
PCdiffW=zeros(a+1)
numPEsame=0
numPEdiff=0
numPEsame=0
numPEdiff=0
numPCsame=0
numPCdiff=0
numPCsame=0
numPCdiff=0
Ir=Irlist[2]
c=Cohlist[1]

numr=0
numr2=1

for i=1:25
    r=CSV.read("./datatempAll/Coh$c-Final200Exp$i-T$Tr-Ir$Ir.csv")
    r2=CSV.read("./datatempAll/NondynCoh$c-Final200Exp$i-T$Tr-Ir$Ir.csv")

temp = 250
for i=1:2500
temp= temp + Int(r2[:RTs][i]/2) + 250
if r2[:NotError][i] == 1 && r2[:Distri][i] >0 && r2[:Distri][i+1] >0
PCsameL=PCsameL + r[:R2][temp:temp+a]
PCsameW=PCsameW + r[:R1][temp:temp+a]

numPCsame=numPCsame+1
#rlist2=rlist2 + r[:R2][temp:(temp+250)]

elseif r2[:NotError][i] == -1 && r2[:Distri][i] <0 && r2[:Distri][i+1] >0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    PEsameL=PEsameL + r[:R2][temp:temp+a]
    PEsameW=PEsameW + r[:R1][temp:temp+a]

    numPEsame=numPEsame+1

elseif r2[:NotError][i] == -1 && r2[:Distri][i] >0 && r2[:Distri][i+1] <0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    PEsameL=PEsameL + r[:R1][temp:temp+a]
    PEsameW=PEsameW + r[:R2][temp:temp+a]

    numPEsame=numPEsame+1
elseif r2[:NotError][i] == 1 && r2[:Distri][i] <0 && r2[:Distri][i+1] <0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    PCsameL=PCsameL + r[:R1][temp:temp+a]
    PCsameW=PCsameW + r[:R2][temp:temp+a]

    numPCsame=numPCsame+1

elseif r2[:NotError][i] == 1 && r2[:Distri][i] >0 && r2[:Distri][i+1] <0 && r2[:NotError][i+1] == 1
    PCdiffL=PCdiffL + r[:R2][temp:temp+a]
    PCdiffW=PCdiffW + r[:R1][temp:temp+a]

    numPCdiff=numPCdiff+1
    #rlist2=rlist2 + r[:R2][temp:(temp+250)]

elseif r2[:NotError][i] == -1 && r2[:Distri][i] <0 && r2[:Distri][i+1] <0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        PEdiffL=PEdiffL + r[:R2][temp:temp+a]
        PEdiffW=PEdiffW + r[:R1][temp:temp+a]

        numPEdiff=numPEdiff+1

    elseif r2[:NotError][i] == -1 && r2[:Distri][i] >0 && r2[:Distri][i+1] >0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        PEdiffL=PEdiffL + r[:R1][temp:temp+a]
        PEdiffW=PEdiffW + r[:R2][temp:temp+a]

        numPEdiff=numPEdiff+1
    elseif r2[:NotError][i] == 1 && r2[:Distri][i] <0 && r2[:Distri][i+1] >0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        PCdiffL=PCdiffL + r[:R1][temp:temp+a]
        PCdiffW=PCdiffW + r[:R2][temp:temp+a]

        numPCdiff=numPCdiff+1

end

end

end

PEsameL=PEsameL./numPEsame
PEdiffL=PEdiffL./numPEdiff
PEsameW=PEsameW./numPEsame
PEdiffW=PEdiffW./numPEdiff
PCsameL=PCsameL./numPCsame
PCdiffL=PCdiffL./numPCdiff
PCsameW=PCsameW./numPCsame
PCdiffW=PCdiffW./numPCdiff

using Plots
gr()
plot(PEsameL,linewidth=3,color=:darkorange,line=:dash)
plot!(PCsameL,linewidth=3,color=:blue,line=:dash)

plot!(PEsameW,color=:darkorange,linewidth=3)
plot!(PCsameW,color=:blue,linewidth=3)
plot!(legend=:topleft)

#savefig("./PEQSameMotiononset.svg")
savefig("./NoSameMotiononset.svg")

plot(PEdiffL,linewidth=3,color=:darkorange,line=:dash)
plot!(PCdiffL,linewidth=3,color=:blue,line=:dash)
plot!(PEdiffW,color=:darkorange,linewidth=3)
plot!(PCdiffW,color=:blue,linewidth=3)
plot!(legend=:topleft)
savefig("./NoDiffTrueMotiononset.svg")

c_min=-30
 c_max = 30
 number = 40

   distri = make_distributions(c_min,c_max,number,20000)

 param["T_rest"] = 1500
 Ir=0.035
 If=0

  result,r2= sim_mf_decisionmaking_pattern(distri,Ir,If,
    param,fixed_param,simu_param)

Accrept=[]
Accaltt=[]

for a in linspace(-30,30,40)
    Accrep=Float64[]
    Accalt=Float64[]
for i=2:1000
    if r2[:Distri][i]==a
        if r2[:NotError][i-1]*sign(r2[:Distri][i-1])==1

            push!(Accrep,r2[:NotError][i]*sign(r2[:Distri][i]))
        else
            push!(Accalt,r2[:NotError][i]*sign(r2[:Distri][i]))
        end
    end
end

push!(Accrept,mean(Accrep))
push!(Accaltt,mean(Accalt))

end
r2

scatter(linspace(-30,30,40),0.5+0.5.*Accrept,color=:blue)
scatter!(linspace(-30,30,40),0.5+0.5.*Accaltt,color=:darkorange)

using LsqFit

model(x, p) = 1./(1+exp.(-(p[1]+p[2].*x)))

xdata = linspace(-30,30,40)
ydata = 0.5+0.5.*Accrept
p0 = [0.5, 0.5]
fit = curve_fit(model, xdata, ydata, p0)

fit.param

plot!(linspace(-30,30,40),1./(1+exp.(-(fit.param[1]+fit.param[2].*linspace(-30,30,40)))),color=:blue,label="Repetition",linewidth=3)

fit = curve_fit(model, xdata, 0.5+0.5.*Accaltt, p0)

fit.param

plot!(linspace(-30,30,40),1./(1+exp.(-(fit.param[1]+fit.param[2].*linspace(-30,30,40)))),color=:darkorange,label="Alternation",linewidth=3)
ylabel!("Percentage Rightward choice")
plot!(legend=:topleft)
savefig("./accuracy.pdf")

c_min=-30
 c_max = 30
 number = 20

   distri = make_distributions(c_min,c_max,number,10000)

 param["T_rest"] = 500
 Ir=0.04
 If=0

  result= sim_mf_decisionmaking_pattern_fast(distri,Ir,If,
    param,fixed_param,simu_param)

using LsqFit
r2=0

function list_elements(l)
      a = Float64[]
        i=1
        while (length(a) <20)
           if any(a.==l[i])

            else
                push!(a,l[i])
            end
            i=i+1
        end
        return sort(a)
    end

colorsT=[:red,:blue,:darkorange,:violet,:green]
o=1

param0=zeros(10)
paramR0=zeros(10)

paramR1=zeros(10)
param1=zeros(10)
fig=plot(legend=false)
for Tr in [200,500,1500,2500,5000]
    fig=plot()
    o=2
    Accrept=zeros(20,10)
    Accaltt=zeros(20,10)
for u=1:10
    r2 = CSV.read("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/2-Data/AccuracySequential/Accuracyc51Tau200Exp$u-T$Tr-Ir0.04.csv")
j=0
for a in list_elements(r2[:Distri])
    j=j+1
    Accrep=Float64[]
    Accalt=Float64[]
for i=2:10000
    if r2[:Distri][i]==a
        if r2[:NotError][i-1]*sign(r2[:Distri][i-1])==1

            push!(Accrep,r2[:NotError][i]*sign(r2[:Distri][i]))
        else
            push!(Accalt,r2[:NotError][i]*sign(r2[:Distri][i]))
        end
    end
end

Accrept[j,u]=mean(Accrep)
Accaltt[j,u]=mean(Accalt)

end
end

scatter!(list_elements(r2[:Distri]),0.5+0.5.*mean(Accrept,2),yerr=sqrt.(var(Accrept,2)),color=colorsT[o])
scatter!(list_elements(r2[:Distri]),0.5+0.5.*mean(Accaltt,2),yerr=sqrt.(var(Accaltt,2)),color=colorsT[o])

model(x, p) = 1./(1+exp.(-(p[1]+p[2].*x)))
for u=1:10
xdata = linspace(-51,51,20)
ydata = 0.5+0.5.*Accrept[:,u]
p0 = [0.5, 0.5]
fit = curve_fit(model, xdata, ydata, p0)

paramR0[u]=fit.param[1]
paramR1[u]=fit.param[2]
plot!(linspace(-51,51,40),1./(1+exp.(-(fit.param[1]+fit.param[2].*linspace(-51,51,40)))),alpha=0.3,label="Repetition",linewidth=3,color=colorsT[o])
ydata = 0.5+0.5.*Accaltt[:,u]
p0 = [0.5, 0.5]
fit = curve_fit(model, xdata, ydata, p0)
plot!(linspace(-51,51,40),1./(1+exp.(-(fit.param[1]+fit.param[2].*linspace(-51,51,40)))),alpha=0.3,label="Repetition",linewidth=3,line=:dash,color=colorsT[o])
param0[u]=fit.param[1]
param1[u]=fit.param[2]
end
plot!(legend=false)
fig
savefig(fig,"./Acc$Tr.svg")
end

#### Left right indecision stuff
# compute param0/param1

histogram(+paramR0./paramR1-param0./param1,bins=10)
-param0./param1

scatter(linspace(-30,30,20),0.5+0.5.*Accrept,color=:blue)
scatter!(linspace(-30,30,20),0.5+0.5.*Accaltt,color=:darkorange)

using LsqFit

model(x, p) = 1./(1+exp.(-(p[1]+p[2].*x)))
for u=1:10
xdata = linspace(-20,20,12)
ydata = 0.5+0.5.*Accrept[:,u]
p0 = [0.5, 0.5]
fit = curve_fit(model, xdata, ydata, p0)

fit.param
plot!(linspace(-20,20,40),1./(1+exp.(-(fit.param[1]+fit.param[2].*linspace(-20,20,40)))),alpha=0.1,color=:blue,label="Repetition",linewidth=3)
end

fig

fit = curve_fit(model, xdata, 0.5+0.5.*Accaltt, p0)

fit.param

plot!(linspace(-30,30,40),1./(1+exp.(-(fit.param[1]+fit.param[2].*linspace(-30,30,40)))),alpha=0.5,color=:darkorange,label="Alternation",linewidth=3)
ylabel!("Percentage Rightward choice")
plot!(legend=:topleft)
savefig("./accuracy.pdf")

#######################################################
##########################################
#Mean stuff for neuronal activity
################################################
#################################################

############################
## Fig neuronale PES
############################
a=300
PEsameL=zeros(251+a,25)#
PEdiffL=zeros(251+a,25)
PEsameW=zeros(251+a,25)
PEdiffW=zeros(251+a,25)
PCsameL=zeros(251+a,25)#
PCdiffL=zeros(251+a,25)
PCsameW=zeros(251+a,25)
PCdiffW=zeros(251+a,25)

j=2
for u=2:25
numPEsame=0
numPEdiff=0
numPEsame=0
numPEdiff=0
numPCsame=0
numPCdiff=0
numPCsame=0
numPCdiff=0
Ir=Irlist[1]
c=Cohlist[2]

numr=0
numr2=1

    r=CSV.read("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/2-Data/FiringRates/datatempAll/V$j-Coh$c-Final200Exp$u-T$Tr-Ir$Ir.csv")
    r2=CSV.read("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/2-Data/FiringRates/datatempAll/V$j-NondynCoh$c-Final200Exp$u-T$Tr-Ir$Ir.csv")

temp = 250
for i=1:2990
temp= temp + Int(r2[:RTs][i]/2)
if r2[:NotError][i] == 1 && r2[:Distri][i] >0 && r2[:Distri][i+1] >0
PCsameL[:,u]=PCsameL[:,u] + r[:R2][temp:(temp+250+a)]
PCsameW[:,u]=PCsameW[:,u] + r[:R1][temp:(temp+250+a)]

numPCsame=numPCsame+1
#rlist2=rlist2 + r[:R2][temp:(temp+250)]

elseif r2[:NotError][i] == -1 && r2[:Distri][i] <0 && r2[:Distri][i+1] >0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    PEsameL[:,u]=PEsameL[:,u] + r[:R2][temp:(temp+250+a)]
    PEsameW[:,u]=PEsameW[:,u] + r[:R1][temp:(temp+250+a)]

    numPEsame=numPEsame+1

elseif r2[:NotError][i] == -1 && r2[:Distri][i] >0 && r2[:Distri][i+1] <0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    PEsameL[:,u]=PEsameL[:,u] + r[:R1][temp:(temp+250+a)]
    PEsameW[:,u]=PEsameW[:,u] + r[:R2][temp:(temp+250+a)]

    numPEsame=numPEsame+1
elseif r2[:NotError][i] == 1 && r2[:Distri][i] <0 && r2[:Distri][i+1] <0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    PCsameL[:,u]=PCsameL[:,u] + r[:R1][temp:(temp+250+a)]
    PCsameW[:,u]=PCsameW[:,u] + r[:R2][temp:(temp+250+a)]

    numPCsame=numPCsame+1

elseif r2[:NotError][i] == 1 && r2[:Distri][i] >0 && r2[:Distri][i+1] <0 && r2[:NotError][i+1] == 1
    PCdiffL[:,u]=PCdiffL[:,u] + r[:R2][temp:(temp+250+a)]
    PCdiffW[:,u]=PCdiffW[:,u] + r[:R1][temp:(temp+250+a)]

    numPCdiff=numPCdiff+1
    #rlist2=rlist2 + r[:R2][temp:(temp+250)]

elseif r2[:NotError][i] == -1 && r2[:Distri][i] <0 && r2[:Distri][i+1] <0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        PEdiffL[:,u]=PEdiffL[:,u] + r[:R2][temp:(temp+250+a)]
        PEdiffW[:,u]=PEdiffW[:,u] + r[:R1][temp:(temp+250+a)]

        numPEdiff=numPEdiff+1

    elseif r2[:NotError][i] == -1 && r2[:Distri][i] >0 && r2[:Distri][i+1] >0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        PEdiffL[:,u]=PEdiffL[:,u] + r[:R1][temp:(temp+250+a)]
        PEdiffW[:,u]=PEdiffW[:,u] + r[:R2][temp:(temp+250+a)]

        numPEdiff=numPEdiff+1
    elseif r2[:NotError][i] == 1 && r2[:Distri][i] <0 && r2[:Distri][i+1] >0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        PCdiffL[:,u]=PCdiffL[:,u] + r[:R1][temp:(temp+250+a)]
        PCdiffW[:,u]=PCdiffW[:,u] + r[:R2][temp:(temp+250+a)]

        numPCdiff=numPCdiff+1

end
temp = temp+250
end

PEsameL[:,u]=PEsameL[:,u]./numPEsame
PEdiffL[:,u]=PEdiffL[:,u]./numPEdiff
PEsameW[:,u]=PEsameW[:,u]./numPEsame
PEdiffW[:,u]=PEdiffW[:,u]./numPEdiff
PCsameL[:,u]=PCsameL[:,u]./numPCsame
PCdiffL[:,u]=PCdiffL[:,u]./numPCdiff
PCsameW[:,u]=PCsameW[:,u]./numPCsame
PCdiffW[:,u]=PCdiffW[:,u]./numPCdiff

end

#############################"
##########" Same with S
#################

a=300
SPEsameL=zeros(251+a,25)#
SPEdiffL=zeros(251+a,25)
SPEsameW=zeros(251+a,25)
SPEdiffW=zeros(251+a,25)
SPCsameL=zeros(251+a,25)#
SPCdiffL=zeros(251+a,25)
SPCsameW=zeros(251+a,25)
SPCdiffW=zeros(251+a,25)

j=2
for u=2:25
numPEsame=0
numPEdiff=0
numPEsame=0
numPEdiff=0
numPCsame=0
numPCdiff=0
numPCsame=0
numPCdiff=0
Ir=Irlist[1]
c=Cohlist[2]

numr=0
numr2=1

    r=CSV.read("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/2-Data/FiringRates/datatempAll/V$j-Coh$c-Final200Exp$u-T$Tr-Ir$Ir.csv")
    r2=CSV.read("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/2-Data/FiringRates/datatempAll/V$j-NondynCoh$c-Final200Exp$u-T$Tr-Ir$Ir.csv")

temp = 250
for i=1:2990
temp= temp + Int(r2[:RTs][i]/2)
if r2[:NotError][i] == 1 && r2[:Distri][i] >0 && r2[:Distri][i+1] >0
SPCsameL[:,u]=SPCsameL[:,u] + r[:S2][temp:(temp+250+a)]
SPCsameW[:,u]=SPCsameW[:,u] + r[:S1][temp:(temp+250+a)]

numPCsame=numPCsame+1
#rlist2=rlist2 + r[:R2][temp:(temp+250)]

elseif r2[:NotError][i] == -1 && r2[:Distri][i] <0 && r2[:Distri][i+1] >0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    SPEsameL[:,u]=SPEsameL[:,u] + r[:S2][temp:(temp+250+a)]
    SPEsameW[:,u]=SPEsameW[:,u] + r[:S1][temp:(temp+250+a)]

    numPEsame=numPEsame+1

elseif r2[:NotError][i] == -1 && r2[:Distri][i] >0 && r2[:Distri][i+1] <0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    SPEsameL[:,u]=SPEsameL[:,u] + r[:S1][temp:(temp+250+a)]
    SPEsameW[:,u]=SPEsameW[:,u] + r[:S2][temp:(temp+250+a)]

    numPEsame=numPEsame+1
elseif r2[:NotError][i] == 1 && r2[:Distri][i] <0 && r2[:Distri][i+1] <0
    #rlist=rlist + r[:R2][temp:(temp+250)]
    SPCsameL[:,u]=SPCsameL[:,u] + r[:S1][temp:(temp+250+a)]
    SPCsameW[:,u]=SPCsameW[:,u] + r[:S2][temp:(temp+250+a)]

    numPCsame=numPCsame+1

elseif r2[:NotError][i] == 1 && r2[:Distri][i] >0 && r2[:Distri][i+1] <0 && r2[:NotError][i+1] == 1
    SPCdiffL[:,u]=SPCdiffL[:,u] + r[:S2][temp:(temp+250+a)]
    SPCdiffW[:,u]=SPCdiffW[:,u] + r[:S1][temp:(temp+250+a)]

    numPCdiff=numPCdiff+1
    #rlist2=rlist2 + r[:R2][temp:(temp+250)]

elseif r2[:NotError][i] == -1 && r2[:Distri][i] <0 && r2[:Distri][i+1] <0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        SPEdiffL[:,u]=SPEdiffL[:,u] + r[:S2][temp:(temp+250+a)]
        SPEdiffW[:,u]=SPEdiffW[:,u] + r[:S1][temp:(temp+250+a)]

        numPEdiff=numPEdiff+1

    elseif r2[:NotError][i] == -1 && r2[:Distri][i] >0 && r2[:Distri][i+1] >0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        SPEdiffL[:,u]=SPEdiffL[:,u] + r[:S1][temp:(temp+250+a)]
        SPEdiffW[:,u]=SPEdiffW[:,u] + r[:S2][temp:(temp+250+a)]

        numPEdiff=numPEdiff+1
    elseif r2[:NotError][i] == 1 && r2[:Distri][i] <0 && r2[:Distri][i+1] >0 && r2[:NotError][i+1] == 1
        #rlist=rlist + r[:R2][temp:(temp+250)]
        SPCdiffL[:,u]=SPCdiffL[:,u] + r[:S1][temp:(temp+250+a)]
        SPCdiffW[:,u]=SPCdiffW[:,u] + r[:S2][temp:(temp+250+a)]

        numPCdiff=numPCdiff+1

end
temp = temp+250
end

SPEsameL[:,u]=SPEsameL[:,u]./numPEsame
SPEdiffL[:,u]=SPEdiffL[:,u]./numPEdiff
SPEsameW[:,u]=SPEsameW[:,u]./numPEsame
SPEdiffW[:,u]=SPEdiffW[:,u]./numPEdiff
SPCsameL[:,u]=SPCsameL[:,u]./numPCsame
SPCdiffL[:,u]=SPCdiffL[:,u]./numPCdiff
SPCsameW[:,u]=SPCsameW[:,u]./numPCsame
SPCdiffW[:,u]=SPCdiffW[:,u]./numPCdiff

end

using Bootstrap
PEL = Float64[]
PELB=Float64[]
PELU = Float64[]

PCL = Float64[]
PCLB=Float64[]
PCLU = Float64[]

PEW = Float64[]
PEWB=Float64[]
PEWU = Float64[]

PCW = Float64[]
PCWB=Float64[]
PCWU = Float64[]

for i=5:300
    n_boot = 2000

      ## basic bootstrap
      #bs1 = bootstrap(PEsameL[i,2:end], mean, BasicSampling(n_boot))
      bs1 = bootstrap(PEdiffL[i,2:end], mean, BasicSampling(n_boot))

      cil = 0.95;

      ## basic CI
      bci1 = ci(bs1, BasicConfInt(cil));
      push!(PEL,bci1[1][1])
      push!(PELB,-bci1[1][3]+bci1[1][1])
      push!(PELU,bci1[1][2]-bci1[1][1])

     # bs1 = bootstrap(PCsameL[i,2:end], mean, BasicSampling(n_boot))
      bs1 = bootstrap(PCdiffL[i,2:end], mean, BasicSampling(n_boot))
      cil = 0.95;

      ## basic CI
      bci1 = ci(bs1, BasicConfInt(cil));
      push!(PCL,bci1[1][1])
      push!(PCLB,-bci1[1][3]+bci1[1][1])
      push!(PCLU,bci1[1][2]-bci1[1][1])

     # bs1 = bootstrap(PEsameW[i,2:end], mean, BasicSampling(n_boot))
     bs1 = bootstrap(PEdiffW[i,2:end], mean, BasicSampling(n_boot))

      cil = 0.95;

      ## basic CI
      bci1 = ci(bs1, BasicConfInt(cil));
      push!(PEW,bci1[1][1])
      push!(PEWB,-bci1[1][3]+bci1[1][1])
      push!(PEWU,bci1[1][2]-bci1[1][1])

    #  bs1 = bootstrap(PCsameW[i,2:end], mean, BasicSampling(n_boot))

      bs1 = bootstrap(PCdiffW[i,2:end], mean, BasicSampling(n_boot))
      cil = 0.95;

      ## basic CI
      bci1 = ci(bs1, BasicConfInt(cil));
      push!(PCW,bci1[1][1])
      push!(PCWB,-bci1[1][3]+bci1[1][1])
      push!(PCWU,bci1[1][2]-bci1[1][1])
end

fig=plot()
plot!(PCW,ribbon=(PCWB,PCWU),linewidth=1,color=:blue)
plot!(PEW,ribbon=(PEWB,PEWU),linewidth=1,color=:darkorange)

plot!(PEL,ribbon=(PELB,PELU),linewidth=1,line=:dash,color=:darkorange)
plot!(PCL,ribbon=(PCLB,PCLU),linewidth=1,line=:dash,color=:blue)

plot!(legend=false)

plot!(mean(PEsameL[5:300,2:end],2),line=:dash,ribbon=sqrt.(var(PEsameL[5:300,2:end],2)),color=:darkorange)
plot!(mean(PCsameL[5:300,2:end],2),line=:dash,ribbon=sqrt.(var(PCsameL[5:300,2:end],2)),color=:blue)

plot(linspace(5,300,296),mean(PEsameW[5:300,2:end],2),ribbon=sqrt.(var(PEsameW[5:300,2:end],2)),color=:darkorange,linewidth=3)
plot!(linspace(5,300,296),mean(PCsameW[5:300,2:end],2),ribbon=sqrt.(var(PCsameW[5:300,2:end],2)),color=:blue,linewidth=3)

plot!(legend=false)
savefig("./PESsame.svg")

plot(mean(log.(PEsameW[200:300,2:end]),2),ribbon=sqrt.(var(log.(PEsameW[200:300,2:end]),2)),color=:darkorange,linewidth=3)
plot!(mean(log.(PCsameW[200:300,2:end]),2),ribbon=sqrt.(var(log.(PCsameW[200:300,2:end]),2)),color=:blue,linewidth=3)

plot!(legend=false)

n_boot=1000
MSPE=Float64[]
MSPC=Float64[]
VUSPE=Float64[]
VDSPE=Float64[]
VUSPC=Float64[]
VDSPC=Float64[]
DMSPE=Float64[]
DVUSPE=Float64[]
DVDSPE=Float64[]

using Bootstrap
for i=10:350
    bs1 = bootstrap(SPEsameW[i,2:end], mean, BasicSampling(n_boot))

    cil = 0.95;

    ## basic CI
    bci1 = ci(bs1, BasicConfInt(cil));
    push!(MSPE,bci1[1][1])
    push!(VUSPE,-bci1[1][3]+bci1[1][1])
    push!(VDSPE,bci1[1][2]-bci1[1][1])

    bs1 = bootstrap(SPCsameW[i,2:end], mean, BasicSampling(n_boot))

    cil = 0.95;

    ## basic CI
    bci1 = ci(bs1, BasicConfInt(cil));
    push!(MSPC,bci1[1][1])
    push!(VUSPC,-bci1[1][3]+bci1[1][1])
    push!(VDSPC,bci1[1][2]-bci1[1][1])

    bs1 = bootstrap(SPEsameW[i,2:end]-SPCsameW[i,2:end], mean, BasicSampling(n_boot))

    cil = 0.95;

    ## basic CI
    bci1 = ci(bs1, BasicConfInt(cil));
    push!(DMSPE,bci1[1][1])
    push!(DVUSPE,-bci1[1][3]+bci1[1][1])
    push!(DVDSPE,bci1[1][2]-bci1[1][1])

end
MSPE
gr()

p1=plot(linspace(10,300,141),MSPE[1:141],ribbon=(VDSPE[1:141],VUSPE[1:141]),color=:darkorange,linewidth=2)
plot!(p1,linspace(10,300,141),MSPC[1:141],ribbon = (VDSPC[1:141],VUSPC[1:141]),color=:blue,linewidth=2)
ps=twinx(p1)
plot!(ps,linspace(600,700,51),MSPE[291:end],ribbon=(VDSPE[291:end],VUSPE[291:end]),color=:darkorange,linewidth=2)
plot!(ps,linspace(600,700,51),MSPC[291:end],ribbon=(VDSPC[291:end],VUSPC[291:end]),color=:blue,linewidth=2)
plot!(legend=false)
vline!([500],line=:dash,color=:black)
plot!(legend=false)

savefig("./PESsameS.svg")

plot(linspace(10,700,341),DMSPE./MSPE.*100,ribbon=(DVDSPE./MSPE.*100,DVUSPE./MSPE.*100),linewidth=3)

plot!(twinx(),linspace(10,700,341),DMSPE./MSPE+sqrt.(linspace(10,700,341)),ribbon=(DVDSPE./MSPE,DVUSPE./MSPE),linewidth=3)

plot!(legend=false)
savefig("./PESsameDiff.svg")

plot(mean(log.(SPEsameW[200:300,2:end]),2),ribbon=sqrt.(var(log.(SPEsameW[200:300,2:end]),2)),color=:darkorange,linewidth=3)
plot!(mean(log.(SPCsameW[200:300,2:end]),2),ribbon=sqrt.(var(log.(SPCsameW[200:300,2:end]),2)),color=:blue,linewidth=3)

plot!(legend=false)

plot(MSPE,ribbon=(VDSPE,VUSPC))
plot!(MSPC,ribbon = (VDSPC,VUSPC))

n_boot=1000
MSPE=Float64[]
MSPC=Float64[]
DMSPE=Float64[]
DVUSPE=Float64[]
DVDSPE=Float64[]

VUSPE=Float64[]
VDSPE=Float64[]
VUSPC=Float64[]
VDSPC=Float64[]
using Bootstrap
for i=10:350
    bs1 = bootstrap(PEsameW[i,2:end], mean, BasicSampling(n_boot))

    cil = 0.95;

    ## basic CI
    bci1 = ci(bs1, BasicConfInt(cil));
    push!(MSPE,bci1[1][1])
    push!(VUSPE,-bci1[1][3]+bci1[1][1])
    push!(VDSPE,bci1[1][2]-bci1[1][1])

    bs1 = bootstrap(PCsameW[i,2:end], mean, BasicSampling(n_boot))

    cil = 0.95;

    ## basic CI
    bci1 = ci(bs1, BasicConfInt(cil));
    push!(MSPC,bci1[1][1])
    push!(VUSPC,-bci1[1][3]+bci1[1][1])
    push!(VDSPC,bci1[1][2]-bci1[1][1])

    bs1 = bootstrap(PEsameW[i,2:end]-PCsameW[i,2:end], mean, BasicSampling(n_boot))

    cil = 0.95;

    ## basic CI
    bci1 = ci(bs1, BasicConfInt(cil));
    push!(DMSPE,bci1[1][1])
    push!(DVUSPE,-bci1[1][3]+bci1[1][1])
    push!(DVDSPE,bci1[1][2]-bci1[1][1])

end

p1=plot(linspace(10,300,141),MSPE[1:141],ribbon=(VDSPE[1:141],VUSPE[1:141]),color=:darkorange,linewidth=2)
plot!(p1,linspace(10,300,141),MSPC[1:141],ribbon = (VDSPC[1:141],VUSPC[1:141]),color=:blue,linewidth=2)
ps=twinx(p1)
plot!(ps,linspace(420,700,141),MSPE[201:end],ribbon=(VDSPE[201:end],VUSPE[201:end]),color=:darkorange,linewidth=2)
plot!(ps,linspace(420,700,141),MSPC[201:end],ribbon=(VDSPC[201:end],VUSPC[201:end]),color=:blue,linewidth=2)
vline!([500],line=:dash,color=:black)
plot!(legend=false)
savefig("./PESBootstrapsame.svg")

plot(linspace(10,700,341),MSPE,ribbon=(VDSPE,VUSPC),linewidth=3,color=:darkorange)
plot!(linspace(10,700,341),MSPC,ribbon = (VDSPC,VUSPC),linewidth=3,color=:blue)
plot!(legend=false)
savefig("./PESBootstrapsame.svg")

plot(linspace(10,700,341),DMSPE./MSPE.*100,ribbon=(DVDSPE./MSPE.*100,DVUSPE./MSPE.*100),linewidth=3)

plot(linspace(10,700,341),DMSPE,ribbon=(DVDSPE,DVUSPE),linewidth=3)
plot!(legend=false)
savefig("./DiffPESsame.svg")

savefig(fig,"./PEQdiffdec.png")

plot(mean(PEdiffL[5:300,2:end],2),ribbon=sqrt.(var(PEdiffL[5:300,2:end],2)))
plot!(mean(PCdiffL[5:300,2:end],2),ribbon=sqrt.(var(PCdiffL[5:300,2:end],2)))

plot!(legend=false)

plot!(mean(PEdiffW[5:300,2:end],2),ribbon=sqrt.(var(PEdiffW[5:300,2:end],2)))
plot!(mean(PCdiffW[5:300,2:end],2),ribbon=sqrt.(var(PCdiffW[5:300,2:end],2)))

plot!(legend=false)

mean(PEsameL[:,2:3],2)
var(PEsameL[:,2:3],2)

#### Indiv stuffs
i=1
r=CSV.read("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/2-Data/FiringRates/datatempAll/Coh$c-Final200Exp$i-T$Tr-Ir$Ir.csv")
r2=CSV.read("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/2-Data/FiringRates/datatempAll/NondynCoh$c-Final200Exp$i-T$Tr-Ir$Ir.csv")

for i=1:2300
if r2[:NotError][i]==-1
    print(i)
end
end

r2[:NotError][35]
r2[:Distri][35]
r2[:Distri][36]

r2[:NotError][270]
r2[:Distri][144]
r2[:Distri][145]
tinit = sum(r2[:RTs][1:357]./2)+250*357
fig=plot()
#plot(r[:R1][Int(tinit)-100:Int(tinit)+450],linewidth=2,color=:darkorange)
plot!(linspace(0,1000,501),r[:R2][Int(tinit)-50:Int(tinit)+450],linewidth=2,color=:darkorange)
r2[:NotError][103]
r2[:Distri][108]
r2[:Distri][119]
tinit = sum(r2[:RTs][1:118]./2)+250*118
plot!(linspace(0,1000,501),r[:R1][Int(tinit)-50:Int(tinit)+450],linewidth=2,color=:blue)
#plot!(r[:R2][Int(tinit)-100:Int(tinit)+450],linewidth=2,color=:blue)
plot!(legend=false)
vline!([100],line=:dash,linewidth=2,color=:black)
vline!([600],line=:dash,linewidth=2,color=:black)

fig
savefig(fig,"./PEQdiffFR.svg")

r2[:NotError][521]
r2[:Distri][521]
r2[:Distri][522]
tinit = sum(r2[:RTs][1:521]./2)+250*521
fig=plot()
#plot(r[:R1][Int(tinit)-100:Int(tinit)+450],linewidth=2,color=:darkorange)
plot!(linspace(0,900,451),r[:R2][Int(tinit)-50:Int(tinit)+400],linewidth=2,color=:darkorange)
r2[:NotError][103]
r2[:Distri][108]
r2[:Distri][111]
tinit = sum(r2[:RTs][1:111]./2)+250*111
plot!(linspace(0,900,451),r[:R2][Int(tinit)-50:Int(tinit)+400],linewidth=2,color=:blue)
#plot!(r[:R2][Int(tinit)-100:Int(tinit)+450],linewidth=2,color=:blue)
plot!(legend=false)
vline!([100],line=:dash,linewidth=2,color=:black)
vline!([600],line=:dash,linewidth=2,color=:black)

fig
savefig(fig,"./neeffFR.svg")

#### Indiv stuffs
i=1
c=Cohlist[2]
Ir=Irlist[1]
r=CSV.read("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/2-Data/FiringRates/datatempAll/Coh$c-Final200Exp$i-T$Tr-Ir$Ir.csv")
r2=CSV.read("/home/kevin/Documents/1 - PhD Projects/1 - Sequential Decision making/3-SequentialEffects/2-Data/FiringRates/datatempAll/NondynCoh$c-Final200Exp$i-T$Tr-Ir$Ir.csv")

for i=1:2300
if r2[:NotError][i]==-1
    print(i)
end
end

r2[:NotError][53]
r2[:Distri][52]
r2[:Distri][53]

r2[:NotError][270]
r2[:Distri][144]
r2[:Distri][145]
fig=plot()
tinit = sum(r2[:RTs][1:499]./2)+250*499
#plot(r[:R1][Int(tinit)-100:Int(tinit)+450],linewidth=2,color=:darkorange)
#linspace(0,1000,501),
plot!(linspace(0,1000,501),r[:R2][Int(tinit)-50:Int(tinit)+450],linewidth=2,color=:darkorange)
r2[:NotError][161]
r2[:Distri][160]
r2[:Distri][161]
tinit = sum(r2[:RTs][1:365]./2)+250*365
plot!(linspace(0,1000,501),r[:R2][Int(tinit)-50:Int(tinit)+450],linewidth=2,color=:blue)
#plot!(r[:R2][Int(tinit)-100:Int(tinit)+450],linewidth=2,color=:blue)
plot!(legend=false)
vline!([100],line=:dash,linewidth=2,color=:black)
vline!([600],line=:dash,linewidth=2,color=:black)
savefig(fig,"./PESsameFR.svg")

for i=1:length(r2[:NotError])-1
    if r2[:NotError][i]==1 && r2[:NotError][i+1]==1 &&
         sign(r2[:Distri][i])==sign(r2[:Distri][i+1])
print(i)
print("\n")
end
end

for i=1:length(r2[:NotError])-1
    if r2[:NotError][i]==-1 && r2[:NotError][i+1]==1 &&
         sign(r2[:Distri][i])!=sign(r2[:Distri][i+1])
print(i)
print("\n")
end
end

fig

savefig(fig,"./PESdiffFR.svg")


export extend, NfToNfMor, automorphisms

mutable struct NfToNfMor <: Map{AnticNumberField, AnticNumberField, HeckeMap, NfToNfMor}
  header::MapHeader{AnticNumberField, AnticNumberField}
  prim_img::nf_elem
  prim_preimg::nf_elem

  function NfToNfMor()
    z = new()
    z.header = MapHeader()
    return r
  end

  function NfToNfMor(K::AnticNumberField, L::AnticNumberField, y::nf_elem, isomorphism::Bool = false)
    z = new()
    z.prim_img = y

    function _image(x::nf_elem)
      g = parent(K.pol)(x)
      return evaluate(g, y)
    end

    if !isomorphism
      z.header = MapHeader(K, L, _image)
      return z
    end

    M = zero_matrix(FlintQQ, degree(L), degree(L))
    b = basis(K)
    for i = 1:degree(L)
      c = _image(b[i])
      for j = 1:degree(L)
        M[j, i] = coeff(c, j - 1)
      end
    end
    t = zero_matrix(FlintQQ, degree(L), 1)
    if degree(L) == 1
      t[1, 1] = coeff(gen(L), 0)
    else
      t[2, 1] = fmpq(1) # coefficient vector of gen(L)
    end

    s = solve(M, t)
    z.prim_preimg = K(parent(K.pol)([ s[i, 1] for i = 1:degree(K) ]))

    function _preimage(x::nf_elem)
      g = parent(L.pol)(x)
      return evaluate(g, z.prim_preimg)
    end

    z.header = MapHeader(K, L, _image, _preimage)
    return z
  end

  function NfToNfMor(K::AnticNumberField, L::AnticNumberField, y::nf_elem, y_inv::nf_elem)
    z = new()
    z.prim_img = y
    z.prim_preimg = y_inv

    function _image(x::nf_elem)
      g = parent(K.pol)(x)
      return evaluate(g, y)
    end

    function _preimage(x::nf_elem)
      g = parent(L.pol)(x)
      return evaluate(g, y_inv)
    end

    z.header = MapHeader(K, L, _image, _preimage)
    return z
  end
end

function inv(f::NfToNfMor)
  if degree(domain(f)) != degree(codomain(f))
    error("The map is not invertible")
  end
  if isdefined(f, :prim_preimg)
    return hom(codomain(f), domain(f), f.prim_preimg, check = false)
  end
  L = codomain(f)
  K = domain(f)
  M = zero_matrix(FlintQQ, degree(L), degree(L))
  a = f.prim_img
  el = one(L)
  M[1, 1] = 1
  for i = 2:degree(L)
    mul!(el, el, a)
    for j = 1:degree(L)
      M[j, i] = coeff(el, j - 1)
    end
  end
  t = zero_matrix(FlintQQ, degree(L), 1)
  if degree(L) == 1
    t[1, 1] = coeff(gen(L), 0)
  else
    t[2, 1] = fmpq(1) # coefficient vector of gen(L)
  end

  s = solve(M, t)
  img = K(parent(K.pol)([ s[i, 1] for i = 1:degree(K) ]))
  return hom(L, K, img, check = false)
end

function _compute_preimg(m::NfToNfMor)
  # build the matrix for the basis change
  K = domain(m)
  L = codomain(m)
  M = zero_matrix(FlintQQ, degree(L), degree(L))
  b = basis(K)
  for i = 1:degree(L)
    c = m(b[i])
    for j = 1:degree(L)
      M[j, i] = coeff(c, j - 1)
    end
  end
  t = zero_matrix(FlintQQ, degree(L), 1)
  t[2, 1] = fmpq(1) # coefficient vector of gen(L)
  s = solve(M, t)
  m.prim_preimg = K(parent(K.pol)([ s[i, 1] for i = 1:degree(K) ]))
  return m.prim_preimg
end

function Base.:(==)(f::NfToNfMor, g::NfToNfMor)
  if (domain(f) != domain(g)) || (codomain(f) != codomain(g))
    return false
  end

  return f.prim_img == g.prim_img
end

#_D = Dict()

function evaluate(f::fmpq_poly, a::nf_elem)
  R = parent(a)
  if iszero(f)
    return zero(R)
  end
  l = length(f) - 1
  s = R(coeff(f, l))
  for i in l-1:-1:0
    #s = s*a + R(coeff(f, i))
    mul!(s, s, a)
    # TODO (easy): Once fmpq_poly_add_fmpq is improved in flint, remove the R(..)
    add!(s, s, R(coeff(f, i)))
  end
  return s
end

function *(f::NfToNfMor, g::NfToNfMor)
#  global _D
#  _s = Base.stacktrace()[2:3]
#  if !(_s in keys(_D))
#    _D[_s] = true
#    println("Called ...")
#    Base.show_backtrace(stdout, Base.stacktrace()[2:3])
#  end
  codomain(f) == domain(g) || throw("Maps not compatible")

  a = gen(domain(f))
  y = g(f(a))

  return hom(domain(f), codomain(g), y, check = false)
end

function ^(f::NfToNfMor, b::Int)
  K = domain(f)
  @assert K == codomain(f)
  d = degree(K)
  b = mod(b, d)
  if b == 0
    return NfToNfMor(K, K, gen(K))
  elseif b == 1
    return f
  else
    bit = ~((~UInt(0)) >> 1)
    while (UInt(bit) & b) == 0
      bit >>= 1
    end
    z = f
    bit >>= 1
    while bit != 0
      z = z * z
      if (UInt(bit) & b) != 0
        z = z * f
      end
      bit >>= 1
    end
    return z
  end
end

Base.copy(f::NfToNfMor) = f

Base.hash(f::NfToNfMor, h::UInt) = Base.hash(f.prim_img, h)

function show(io::IO, h::NfToNfMor)
  if domain(h) == codomain(h)
    println(io, "Automorphism of ", domain(h))
  else
    println(io, "Injection of ", domain(h), " into ", codomain(h))
  end
  println(io, "defined by ", gen(domain(h)), " -> ", h.prim_img)
end

function automorphisms(K::AnticNumberField; copy::Bool = true)
  try
    Aut = _get_automorphisms_nf(K)::Vector{NfToNfMor}
    if copy
      return Base.copy(Aut)
    else
      return Aut
    end
  catch e
    if !isa(e, AccessorNotSetError)
      rethrow(e)
    end
  end
  if degree(K) == 1
    auts = NfToNfMor[hom(K, K, one(K))]
  else
    f = K.pol
    Kt, t = PolynomialRing(K, "t", cached = false)
    f1 = change_ring(f, Kt)
    divpol = Kt(nf_elem[-gen(K), K(1)])
    f1 = divexact(f1, divpol)
    lr = roots(f1, max_roots = div(degree(K), 2))
    Aut1 = Vector{NfToNfMor}(undef, length(lr)+1)
    for i = 1:length(lr)
      Aut1[i] = hom(K, K, lr[i], check = false)
    end
    Aut1[end] = id_hom(K)
    auts = closure(Aut1, degree(K))
  end
  _set_automorphisms_nf(K, auts)
  if copy
    return Base.copy(auts)
  else
    return auts
  end
end

function hom(K::AnticNumberField, L::AnticNumberField, a::nf_elem; check::Bool = true, compute_inverse::Bool = false)
 if check
   if !iszero(evaluate(K.pol, a))
     error("The data does not define a homomorphism")
   end
 end
 return NfToNfMor(K, L, a, compute_inverse)
end

function hom(K::AnticNumberField, L::AnticNumberField, a::nf_elem, a_inv::nf_elem; check::Bool = true)
 if check
   if !iszero(evaluate(K.pol, a))
     error("The data does not define a homomorphism")
   end
   if !iszero(evaluate(L.pol, a_inv))
     error("The data does not define a homomorphism")
   end
 end
 return NfToNfMor(K, L, a, a_inv)
end

id_hom(K::AnticNumberField) = hom(K, K, gen(K), check = false)

morphism_type(::Type{AnticNumberField}) = NfToNfMor

isinjective(m::NfToNfMor) = true

issurjective(m::NfToNfMor) = (degree(domain(m)) == degree(codomain(m)))

isbijective(m::NfToNfMor) = issurjective(m)


is_ch_type(::Val{:Array})  = true
can_be_nullable(::Val{:Array}) = false
result_type(::Val{:Array}, nested::TypeAst) = Vector{result_type(nested)}

"""
    Nested arrays written in flatten form after information about their
    sizes (offsets really).
    One element of array of arrays can be represented as tree:
    (0 depth)          [[3, 4], [5, 6]]
                      |               |
    (1 depth)      [3, 4]           [5, 6]
                   |    |           |    |
    (leaf)        3     4          5     6

    Offsets (sizes) written in breadth-first search order. In example above
    following sequence of offset will be written: 4 -> 2 -> 4
    1) size of whole array: 4
    2) size of array 1 in depth=1: 2
    3) size of array 2 plus size of all array before in depth=1: 2 + 2 = 4

    After sizes info comes flatten data: 3 -> 4 -> 5 -> 6
"""

read_state_prefix(sock::ClickHouseSock, ::Val{:Array}, nested::TypeAst) =
                    read_state_prefix(sock, nested)

write_state_prefix(sock::ClickHouseSock, ::Val{:Array}, nested::TypeAst) =
                write_state_prefix(sock, nested)

function read_offsets!(dest::Vector{Vector{UInt64}}, sock, nest::TypeAst)
    prev_level = dest[end]
    new_level = Vector{UInt64}()
    last_offset = 0
    for offset in prev_level
        for i in 1:(offset - last_offset)
            push!(new_level, chread(sock, UInt64))
        end
        last_offset = offset
    end
    push!(dest, new_level)

    return nest.name == :Array ?
        read_offsets!(dest, sock, nest.args[1]) :
        nest

end

function split_vector(data::T, offsets) where {T <: AbstractVector}
    result = Vector{T}(undef, length(offsets))
    last_offset = 1
    for (i,offset) in enumerate(offsets)
        result[i] = data[last_offset:offset]
        last_offset = offset + 1
    end
    return result
end

function read_col_data(sock::ClickHouseSock, num_rows::VarUInt,
                                            ::Val{:Array}, nest::TypeAst)

    (UInt64(num_rows) == 0) && return
    prev_offset = 0
    cur_deph = 0

    size = UInt64(num_rows)

    offsets = [UInt64[size]]

    data_type = read_offsets!(offsets, sock, nest)

    data_size = offsets[end][end]

    data = read_col_data(sock, VarUInt(data_size), data_type)

    result = data
    #top offset is the size of column, so we don't take it into account
    for off in Iterators.reverse(offsets[2:end])
        result = split_vector(result, off)
    end
    return result
end

const PossibleVectors{T} =
     Union{<:AbstractVector{T}, <:AbstractCategoricalVector{T}}
function get_base_type(
    ::Type{<:AbstractVector{T}},
    nest ::TypeAst,
) where {T}
    nest.name == :Array && return get_base_type(T, nest.args[1])
    return (T, nest)
end
function flatten_array!(offsets, itr, nest)
    new_level = UInt64[]
    last_offset = 0
    for part in itr
        push!(new_level, last_offset + length(part))
        last_offset = new_level[end]
    end
    push!(offsets, new_level)
    (nest.name != :Array) && return (Iterators.flatten(itr), nest)
    return flatten_array!(offsets, Iterators.flatten(itr), nest.args[1])

end
function write_col_data(sock::ClickHouseSock,
                                data::AbstractVector{T},
                                ::Val{:Array}, nest::TypeAst) where {T}

    offsets = [[UInt64(length(data))]]
    flatten_itr, base_ast = flatten_array!(offsets, data, nest)
    for i in 2:length(offsets)
        chwrite(sock, offsets[i])
    end
    write_col_data(sock, collect(flatten_itr), base_ast)
end


using WaveletMatrices

@inline function raw_getindex(seq::BioSequence{A}, i::Integer) where A
    j = BioSequences.bitindex(seq, i)
    @inbounds return convert(UInt8, (seq.data[BioSequences.index(j)] >> BioSequences.offset(j)) & BioSequences.mask(A))
end

## This function is necessary as a result of changes made in Bio.jl as of Julia v0.5
## Since all nucleotides are converted to UInt8 as one hot 4-bit encodings by default
## we can override FMIndexes.sa_range function for those sequences specifically
## to re-encode as DNAAlphabet{2}
@inline function FMIndexes.sa_range(query::BioSequence{A}, index::FMIndex, init_range::UnitRange{Int}) where A <: BioSequences.Alphabet
    sp, ep = init_range.start, init_range.stop
    # backward search
    i = length(query)
    @inbounds while sp ≤ ep && i ≥ 1
        char = raw_getindex(query, i)
        c = index.count[char+1]
        sp = c + WaveletMatrices.rank(char, index.bwt, (sp > index.sentinel ? sp - 1 : sp) - 1) + 1
        ep = c + WaveletMatrices.rank(char, index.bwt, (ep > index.sentinel ? ep - 1 : ep))
        i -= 1
    end
    return sp:ep
end


@noinline not_planar_err() = throw(ArgumentError("not a planar diagram expression"))
@noinline not_planar_err(ex) = throw(ArgumentError("not a planar diagram expression: $ex"))

# Preprocessors used by `@planar` and `@plansor`
function _conj_to_adjoint(ex::Expr)
    if ex.head == :call && ex.args[1] == :conj && TO.istensor(ex.args[2])
        obj, leftind, rightind = TO.decomposetensor(ex.args[2])
        return Expr(:typed_vcat, Expr(TO.prime, obj),
                        Expr(:tuple, rightind...), Expr(:tuple, leftind...))
    else
        return Expr(ex.head, [_conj_to_adjoint(a) for a in ex.args]...)
    end
end
_conj_to_adjoint(ex) = ex

# replacement of TensorOperations functionality:
# adds checks for matching number of domain and codomain indices
# special cases adjoints so that t and t' are considered the same object
# ignore braiding tensors
function _extract_tensormap_objects(ex)
    inputtensors = collect(filter(!=(:τ), _remove_adjoint.(TO.getinputtensorobjects(ex))))
    outputtensors = _remove_adjoint.(TO.getoutputtensorobjects(ex))
    newtensors = TO.getnewtensorobjects(ex)
    (any(==(:τ), newtensors) || any(==(:τ), outputtensors)) &&
        throw(ArgumentError("The name τ is reserved for the braiding, and should not be assigned to."))
    @assert !any(_is_adjoint, newtensors)
    existingtensors = unique!(vcat(inputtensors, outputtensors))
    alltensors = unique!(vcat(existingtensors, newtensors))
    tensordict = Dict{Any,Any}(a => gensym() for a in alltensors)
    pre = Expr(:block, [Expr(:(=), tensordict[a], a) for a in existingtensors]...)
    pre2 = Expr(:block)
    ex = TO.replacetensorobjects(ex) do obj, leftind, rightind
        _is_adj = _is_adjoint(obj)
        if _is_adj
            leftind, rightind = rightind, leftind
            obj = _remove_adjoint(obj)
        end
        newobj = get(tensordict, obj, obj)
        if (obj in existingtensors)
            nl = length(leftind)
            nr = length(rightind)
            nlsym = gensym()
            nrsym = gensym()
            objstr = string(obj)
            errorstr1 = "incorrect number of input-output indices: ($nl, $nr) instead of "
            errorstr2 = " for $objstr."
            checksize = quote
                $nlsym, $nrsym = numout($newobj), numin($newobj)
                ($nlsym == $nl && $nrsym == $nr) ||
                    throw(IndexError($errorstr1 * string(($nlsym, $nrsym)) * $errorstr2))
            end
            push!(pre2.args, checksize)
        end
        return _is_adj ? _add_adjoint(newobj) : newobj
    end
    post = Expr(:block, [Expr(:(=), a, tensordict[a]) for a in newtensors]...)
    pre = Expr(:macrocall, Symbol("@notensor"), LineNumberNode(@__LINE__, Symbol(@__FILE__)), pre)
    pre2 = Expr(:macrocall, Symbol("@notensor"), LineNumberNode(@__LINE__, Symbol(@__FILE__)), pre2)
    post = Expr(:macrocall, Symbol("@notensor"), LineNumberNode(@__LINE__, Symbol(@__FILE__)), post)
    return Expr(:block, pre, pre2, ex, post)
end
_is_adjoint(ex) = isa(ex, Expr) && ex.head == TO.prime
_remove_adjoint(ex) = _is_adjoint(ex) ? ex.args[1] : ex
_add_adjoint(ex) = Expr(TO.prime, ex)

# used by `@planar`: realize explicit braiding tensors
function _construct_braidingtensors(ex::Expr)
    if TO.isdefinition(ex) || TO.isassignment(ex)
        lhs, rhs = TO.getlhs(ex), TO.getrhs(ex)
        if TO.istensorexpr(rhs)
            list = TO.gettensors(_conj_to_adjoint(rhs))
            if TO.isassignment(ex) && TO.istensor(lhs)
                obj, l, r = TO.decomposetensor(lhs)
                lhs_as_rhs = Expr(:typed_vcat, _add_adjoint(obj),
                                    Expr(:tuple, r...), Expr(:tuple, l...))
                push!(list, lhs_as_rhs)
            end
        else
            return ex
        end
    elseif TO.istensorexpr(ex)
        list = TO.gettensors(_conj_to_adjoint(ex))
    else
        return Expr(ex.head, map(_construct_braidingtensors, ex.args)...)
    end

    i = 1
    translatebraidings = Dict{Any,Any}()
    while i <= length(list)
        t = list[i]
        if _remove_adjoint(TO.gettensorobject(t)) == :τ
            translatebraidings[t] = Expr(:call, GlobalRef(TensorKit, :BraidingTensor))
            deleteat!(list, i)
        else
            i += 1
        end
    end

    unresolved = Any[] # list of indices that we couldn't yet figure out
    indexmap = Dict{Any,Any}()
    # indexmap[i] contains the expression to resolve the space for index i
    for (t, construct_expr) in translatebraidings
        obj, leftind, rightind = TO.decomposetensor(t)
        length(leftind) == length(rightind) == 2 ||
            throw(ArgumentError("The name τ is reserved for the braiding, and should have two input and two output indices."))
        if _is_adjoint(obj)
            i1b, i2b, = leftind
            i2a, i1a, = rightind
        else
            i2b, i1b, = leftind
            i1a, i2a, = rightind
        end

        obj_and_pos1a = _findindex(i1a, list)
        obj_and_pos2a = _findindex(i2a, list)
        obj_and_pos1b = _findindex(i1b, list)
        obj_and_pos2b = _findindex(i2b, list)

        if !isnothing(obj_and_pos1a)
            indexmap[i1b] = Expr(:call, :space, obj_and_pos1a...)
            indexmap[i1a] = Expr(:call, :space, obj_and_pos1a...)
        elseif !isnothing(obj_and_pos1b)
            indexmap[i1b] = Expr(TO.prime, Expr(:call, :space, obj_and_pos1b...))
            indexmap[i1a] = Expr(TO.prime, Expr(:call, :space, obj_and_pos1b...))
        else
            push!(unresolved,(i1a,i1b))
        end

        if !isnothing(obj_and_pos2a)
            indexmap[i2b] = Expr(:call, :space, obj_and_pos2a...)
            indexmap[i2a] = Expr(:call, :space, obj_and_pos2a...)
        elseif !isnothing(obj_and_pos2b)
            indexmap[i2b] = Expr(TO.prime, Expr(:call, :space, obj_and_pos2b...))
            indexmap[i2a] = Expr(TO.prime, Expr(:call, :space, obj_and_pos2b...))
        else
            push!(unresolved,(i2a,i2b))
        end
    end
    # simple loop that tries to resolve as many indices as possible
    changed = true
    while changed == true
        changed = false
        i = 1
        while i <= length(unresolved)
            (i1, i2) = unresolved[i]
            if i1 in keys(indexmap)
                changed = true
                indexmap[i2] = indexmap[i1]
                deleteat!(unresolved, i)
            elseif i2 in keys(indexmap)
                changed = true
                indexmap[i1] = indexmap[i2]
                deleteat!(unresolved, i)
            else
                i += 1
            end
        end
    end
    !isempty(unresolved) &&
        throw(ArgumentError("cannot determine the spaces of indices " *
                                string(tuple(unresolved...)) *
                                    "for the braiding tensors in $(ex)"))

    pre = Expr(:block)
    for (t, construct_expr) in translatebraidings
        obj, leftind, rightind = TO.decomposetensor(t)
        if _is_adjoint(obj)
            i1b, i2b, = leftind
            i2a, i1a, = rightind
        else
            i2b, i1b, = leftind
            i1a, i2a, = rightind
        end
        push!(construct_expr.args, indexmap[i1b])
        push!(construct_expr.args, indexmap[i2b])
        s = gensym()
        push!(pre.args, Expr(:(=), s, construct_expr))
        ex = TO.replacetensorobjects(ex) do o, l, r
            if o == obj && l == leftind && r == rightind
                return obj  == :τ ? s : Expr(TO.prime, s)
            else
                return o
            end
        end
    end
    return Expr(:block, pre, ex)
end
_construct_braidingtensors(x) = x

# used by non-planar parser of `@plansor`: remove explicit braiding tensors
function _remove_braidingtensors(ex::Expr)
    outgoing = []

    if TO.isdefinition(ex) || TO.isassignment(ex)
        lhs, rhs = TO.getlhs(ex), TO.getrhs(ex)
        if TO.istensorexpr(rhs)
            list = TO.gettensors(_conj_to_adjoint(rhs))
            if TO.istensor(lhs)
                obj, l, r = TO.decomposetensor(lhs)
                outgoing = [l; r]
            end
        else
            return ex
        end
    elseif TO.istensorexpr(ex)
        list = TO.gettensors(_conj_to_adjoint(ex))
    else
        return Expr(ex.head, map(_remove_braidingtensors, ex.args)...)
    end

    τs = Any[]
    i = 1
    while i <= length(list)
        t = list[i]
        if _remove_adjoint(TO.gettensorobject(t)) == :τ
            push!(τs, t)
            deleteat!(list, i)
        else
            i += 1
        end
    end

    indexmap = Dict{Any,Any}()
    # to remove the braidingtensors, we need to map certain indices to other indices
    for t in τs
        obj, leftind, rightind = TO.decomposetensor(t)
        length(leftind) == length(rightind) == 2 ||
            throw(ArgumentError("The name τ is reserved for the braiding, and should have two input and two output indices."))
        if _is_adjoint(obj)
            i1b, i2b, = leftind
            i2a, i1a, = rightind
        else
            i2b, i1b, = leftind
            i1a, i2a, = rightind
        end

        i1a = get(indexmap, i1a, i1a)
        i1b = get(indexmap, i1b, i1b)
        i2a = get(indexmap, i2a, i2a)
        i2b = get(indexmap, i2b, i2b)

        obj_and_pos1a = _findindex(i1a, list)
        obj_and_pos2a = _findindex(i2a, list)
        obj_and_pos1b = _findindex(i1b, list)
        obj_and_pos2b = _findindex(i2b, list)

        if i1a in outgoing
            indexmap[i1a] = i1a
            indexmap[i1b] = i1a
        elseif i1b in outgoing
            indexmap[i1a] = i1b
            indexmap[i1b] = i1b
        else
            if i1a isa Int && i1b isa Int
                indexmap[i1a] = max(i1a, i1b)
                indexmap[i1b] = max(i1a, i1b)
            else
                indexmap[i1a] = i1a
                indexmap[i1b] = i1a
            end
        end

        if i2a in outgoing
            indexmap[i2a] = i2a
            indexmap[i2b] = i2a
        elseif i2b in outgoing
            indexmap[i2a] = i2b
            indexmap[i2b] = i2b
        else
            if i2a isa Int && i2b isa Int
                indexmap[i2a] = max(i2a, i2b)
                indexmap[i2b] = max(i2a, i2b)
            else
                indexmap[i2a] = i2a
                indexmap[i2b] = i2a
            end
        end
    end

    # simple loop that tries to simplify the indicemaps (a=>b,b=>c -> a=>c,b=>c)
    changed = true
    while changed == true
        changed = false
        i = 1
        for (k,v) in indexmap
            if v in keys(indexmap) && indexmap[v] != v
                changed = true
                indexmap[k] = indexmap[v]
            end
        end
    end

    ex = TO.replaceindices(i -> get(indexmap, i, i), ex)
    return _purge_braidingtensors(ex)
end
_remove_braidingtensors(x) = x

function _purge_braidingtensors(ex::Expr) # actually remove the braidingtensors
    args = collect(filter(ex.args) do a
        if a isa Expr && a.head == :call && a.args[1] == :conj
            a = a.args[2]
        end
        if a isa Expr && TO.istensor(a) && _remove_adjoint(TO.gettensorobject(a)) == :τ
            _, leftind, rightind = TO.decomposetensor(a)
            (leftind[1] == rightind[2] && leftind[2] == rightind[1]) ||
                throw(ArgumentError("unable to remove braiding tensor $a"))
            return false
        end
        return true
    end)

    # multiplication with only a single argument is (rightfully) seen as invalid syntax
    if ex.head == :call && args[1] == :* && length(args) == 2
        return _purge_braidingtensors(args[2])
    else
        return Expr(ex.head, map(_purge_braidingtensors, args)...)
    end
end
_purge_braidingtensors(x) = x

function _check_planarity(ex::Expr)
    if ex.head == :macrocall && ex.args[1] == Symbol("@notensor")
    elseif TO.isassignment(ex) || TO.isdefinition(ex)
        lhs, rhs = TO.getlhs(ex), TO.getrhs(ex)
        if TO.istensorexpr(rhs)
            if TO.istensorexpr(lhs)
                @assert TO.istensor(lhs)
                indlhs = only(get_possible_planar_indices(lhs)) # should have only one element
            else
                indlhs = Any[]
            end
            indsrhs = get_possible_planar_indices(rhs)
            isempty(indsrhs) && not_planar_err(rhs)
            i = findfirst(ind -> iscyclicpermutation(ind, indlhs), indsrhs)
            i === nothing && not_planar_err(ex)
        end
    else
        foreach(ex.args) do a
            _check_planarity(a)
        end
    end
    return ex
end
_check_planarity(ex) = ex

# decompose contraction trees in order to fix index order of temporaries
# to ensure that planarity is guaranteed
_decompose_planar_contractions(ex, temporaries) = ex
function _decompose_planar_contractions(ex::Expr, temporaries)
    if ex.head == :macrocall && ex.args[1] == Symbol("@notensor")
        return ex
    end
    if TO.isassignment(ex) || TO.isdefinition(ex)
        lhs, rhs = TO.getlhs(ex), TO.getrhs(ex)
        if TO.istensorexpr(rhs)
            pre = Vector{Any}()
            rhs = _extract_contraction_pairs(rhs, lhs, pre, temporaries)
            return Expr(:block, pre..., Expr(ex.head, lhs, rhs))
        else
            return ex
        end
    end
    if TO.istensorexpr(ex)
        pre = Vector{Any}()
        rhs = _extract_contraction_pairs(ex, (Any[], Any[]), pre, temporaries)
        return Expr(:block, pre..., rhs)
    end
    if ex.head == :block
        return Expr(ex.head,
                    [_decompose_planar_contractions(a, temporaries) for a in ex.args]...)
    end
    if ex.head == :for || ex.head == :function
        return Expr(ex.head, ex.args[1],
                        _decompose_planar_contractions(ex.args[2], temporaries))
    end
    return ex
end

# decompose a contraction into elementary binary contractions of tensors without inner traces
# if lhs is an expression, it contains the existing lhs and thus the index order
# if lhs is a tuple, the result is a temporary object and the tuple (lind, rind) gives a suggestion for the preferred index order
function _extract_contraction_pairs(rhs, lhs, pre, temporaries)
    if TO.isscalarexpr(rhs)
        return rhs
    elseif TO.isgeneraltensor(rhs)
        if TO.hastraceindices(rhs) && lhs isa Tuple
            s = gensym()
            newlhs = Expr(:typed_vcat, s, Expr(:tuple, lhs[1]...), Expr(:tuple, lhs[2]...))
            push!(temporaries, s)
            push!(pre, Expr(:(:=), newlhs, rhs))
            return newlhs
        else
            return rhs
        end
    elseif rhs.head == :call && rhs.args[1] == :*
        @assert length(rhs.args) == 3

        if lhs isa Expr
            _, leftind, rightind = TO.decomposetensor(lhs)
        else
            leftind, rightind = lhs
        end
        lhs_ind = vcat(leftind, reverse(rightind))

        # find possible planar order
        rhs_inds = Any[]
        for ind1 in get_possible_planar_indices(rhs.args[2])
            for ind2 in get_possible_planar_indices(rhs.args[3])
                for (oind1, oind2, cind1, cind2) in possible_planar_complements(ind1, ind2)
                    if iscyclicpermutation(vcat(oind1, oind2), lhs_ind)
                        push!(rhs_inds, (ind1, ind2, oind1, oind2, cind1, cind2))
                    end
                    isempty(rhs_inds) || break
                end
                isempty(rhs_inds) || break
            end
            isempty(rhs_inds) || break
        end
        ind1, ind2, oind1, oind2, cind1, cind2 = only(rhs_inds) # inds_rhs should hold exactly one match
        if all(in(leftind), oind2) && all(in(rightind), oind1) # reverse order
            a1 = _extract_contraction_pairs(rhs.args[3], (oind2, reverse(cind2)), pre, temporaries)
            a2 = _extract_contraction_pairs(rhs.args[2], (cind1, reverse(oind1)), pre, temporaries)
            oind1, oind2 = oind2, oind1
            cind1, cind2 = cind2, cind1
        else
            a1 = _extract_contraction_pairs(rhs.args[2], (oind1, reverse(cind1)), pre, temporaries)
            a2 = _extract_contraction_pairs(rhs.args[3], (cind2, reverse(oind2)), pre, temporaries)
        end

        if TO.isscalarexpr(a1) || TO.isscalarexpr(a2)
            rhs = Expr(:call, :*, a1, a2)
            s = gensym()
            newlhs = Expr(:typed_vcat, s, Expr(:tuple, oind1...),
                                        Expr(:tuple, reverse(oind2)...))
            push!(temporaries, s)
            push!(pre, Expr(:(:=), newlhs, rhs))
            return newlhs
        end

        # note that index order in _extract... is only a suggestion, now we have actual index order
        _, l1, r1, = TO.decomposegeneraltensor(a1)
        _, l2, r2, = TO.decomposegeneraltensor(a2)
        if all(in(r1), oind1) && all(in(l2), oind2) # reverse order
            a1, a2 = a2, a1
            ind1, ind2 = ind2, ind1
            oind1, oind2 = oind2, oind1
        end
        if lhs isa Tuple
            rhs = Expr(:call, :*, a1, a2)
            s = gensym()
            newlhs = Expr(:typed_vcat, s, Expr(:tuple, oind1...),
                                        Expr(:tuple, reverse(oind2)...))
            push!(temporaries, s)
            push!(pre, Expr(:(:=), newlhs, rhs))
            return newlhs
        else
            if leftind == oind1 && rightind == reverse(oind2)
                rhs = Expr(:call, :*, a1, a2)
                return rhs
            elseif leftind == oind2 && rightind == reverse(oind1) # probably this can not happen anymore
                rhs = Expr(:call, :*, a2, a1)
                return rhs
            else
                rhs = Expr(:call, :*, a1, a2)
                s = gensym()
                newlhs = Expr(:typed_vcat, s, Expr(:tuple, oind1...),
                                            Expr(:tuple, reverse(oind2)...))
                push!(temporaries, s)
                push!(pre, Expr(:(:=), newlhs, rhs))
                return newlhs
            end
        end
    elseif rhs.head == :call && rhs.args[1] ∈ (:+, :-)
        args = [_extract_contraction_pairs(a, lhs, pre, temporaries) for
                    a in rhs.args[2:end]]
        return Expr(rhs.head, rhs.args[1], args...)
    else
        throw(ArgumentError("unknown tensor expression"))
    end
end

function _findindex(i, list) # finds an index i in a list of tensor expressions
    for t in list
        obj, l, r = TO.decomposetensor(t)
        pos = findfirst(==(i), l)
        isnothing(pos) || return (obj, pos)
        pos = findfirst(==(i), r)
        isnothing(pos) || return (obj, pos + length(l))
    end
    return nothing
end


#=
surface_test:
- Julia version: 1.1.0
- Author: Jirik
- Date: 2019-03-19
=#
# using Revise
using Test
using Logging
using SparseArrays
using Distributed
if nprocs() == 1
    addprocs(3)
end

@everywhere using LarSurf
# using Plasm
# using LinearAlgebraicRepresentation
# Lar = LinearAlgebraicRepresentation

@testset "Get Fchar test" begin

    segmentation = LarSurf.data234()
    block_size = [2,2,2]
    ch = RemoteChannel(()->Channel{Int}(32));
    Flin2 = LarSurf.__grid_get_surface_Fchar_per_block_parallel_channel(segmentation, block_size, ch)
    # @test nnz(Flin2) == 6*4
    # println("get fchar test 5")

    # Flin1 = LarSurf.__grid_get_surface_Fchar(segmentation, block_size)
    # println("get fchar test 2")
    # Flin1 = LarSurf.__grid_get_surface_Fchar_per_block(segmentation, block_size)
    # @test nnz(Flin1) == 6*4
    #
    # println("get fchar test 3")
    # Flin2 = LarSurf.__grid_get_surface_Fchar_per_block_parallel_pmap(segmentation, block_size)
    # @test nnz(Flin2) == 6*4
    # println("get fchar test 4")
end


if BACKEND == "GPU"
    using CUDArt
    using CUBLAS
end
function init(material=train_data, num_buf_factor=20, prob_raise=0.75)
    global sampleFactory
    global genBufFac
    global init_model

    info(@sprintf("Using %s-%d", BACKEND, DEVICE))
    idx2wf = jldopen(init_model) do f
         read(f, "idx2wf")
    end

    for i = 1:length(idx2wf)
        @inbounds idx2wf[i] = idx2wf[i]^prob_raise
    end
    # sumx = sum(idx2wf)
    sumx = sum(idx2wf[i] for i = length(idx2wf):-1:1)
    for i = 1:length(idx2wf)
        @inbounds idx2wf[i] /= sumx
    end
    # info(length(idx2wf))

    info("Using $material as training data")
    set!(sampleFactory, SampleFactory(material, idx2wf))
    set!(genBufFac, convert(UInt32,num_buf_factor))
    if BACKEND == "GPU"
        # CUDArt.device(DEVICE)  # already set in config_global.jl
        CUDArt.init(DEVICE)
    end
end

function close_device()
    if BACKEND == "GPU"
        CUBLAS.destroy()
        CUDArt.close(DEVICE)
        CUDArt.device_reset(DEVICE)
    end
end


import Dagger: debug_compute, get_logs!, LocalEventLog

export start_tracking_time, stop_tracking_time, tracktime, fetch_timings!

function time_table(log; profile=false)

    idx = Columns(proc=map(x->x.pid, getfields(log, :timeline)),
                  event_type=getfields(log, :category),
                  event_id=getfields(log, :id))

    if profile
        data = Columns(start=getfields(log, :start),   finish=getfields(log, :finish),
                       gc_diff=getfields(log, :gc_diff), profile=getfields(log, :profiler_samples))
    else
        data = Columns(start=getfields(log, :start),   finish=getfields(log, :finish),
                       gc_diff=getfields(log, :gc_diff))
    end

    NDSparse(idx, data)
end

function add_gc_diff(x,y)
    Base.GC_Diff(
        x.allocd     + y.allocd,
        x.malloc     + y.malloc,
        x.realloc    + y.realloc,
        x.poolalloc  + y.poolalloc,
        x.bigalloc   + y.bigalloc,
        x.freecall   + y.freecall,
        x.total_time + y.total_time,
        x.pause      + y.pause,
        x.full_sweep + y.full_sweep
    )
end

function aggregate_profile(xs)
    treereduce(Dagger.mix_samples, xs)
end

function aggregate_events(xs)
    sort!(xs, by=x->x.start)
    gc_diff = reduce(add_gc_diff, map(x -> x.gc_diff, xs))
    time_spent = sum(map(x -> x.finish - x.start, xs))
    if isdefined(xs[1], :profile)
        time_spent, gc_diff, aggregate_profile(map(x->x.profile, xs))
    end
    time_spent, gc_diff
end

function show_timings(t; maxdepth=5)
    # first aggregate by type of event
    t1 = reducedim_vec(aggregate_events, t, [:proc, :event_id])

    foreach(t1.index, t1.data) do i, x
        time_spent, gc_diff = x
        print(string(i[1]),": ")
        Base.time_print(time_spent, gc_diff.allocd, gc_diff.total_time, Base.gc_alloc_count(gc_diff))
    end

    t2 = reducedim_vec(aggregate_events, t, :event_id)
    println("Breakdown:")
    println(map(x->first(x)/1e9, t2))
    if isdefined(t.data.columns, :profile)
        p = aggregate_profile(t.data.columns.profile)
        if !isempty(p.samples)
            println("\nProfile output:")
            Profile.print(p.samples, p.lineinfo, maxdepth=maxdepth)
        end
    end
end

function getfields(log, fieldname)
    map(x->getfield(x, fieldname), log)
end

function start_tracking_time(;profile=false)
    ctx = get_context()
    dbgctx = Context(procs(ctx), LocalEventLog(), profile)
    compute_context[] = dbgctx
end

function stop_tracking_time()
    compute_context[] = nothing
end

"""
`tracktime(f)`

Track the time spent on different processes in different
categories in running `f`.
"""
function tracktime(f; profile=false, maxdepth=5)
    start_tracking_time(profile=profile)
    res = f()
    ctx = compute_context[]
    stop_tracking_time()
    t = fetch_timings!(ctx, profile=profile)
    show_timings(t, maxdepth=maxdepth)
    t, res
end

function fetch_timings!(ctx=get_context(); profile=true)
    time_table(Dagger.get_logs!(ctx.log_sink), profile=profile)
end


module Turkie

using AbstractPlotting: Scene, Point2f0
using AbstractPlotting: barplot!, lines!, scatter! # Plotting tools
using AbstractPlotting: Observable, Node, lift, on # Observable tools
using AbstractPlotting: recordframe! # Recording tools
using AbstractPlotting.MakieLayout # Layouting tool
using Colors, ColorSchemes # Colors tools
using KernelDensity # To be able to give a KDE
using OnlineStats # Estimators
using DynamicPPL: VarInfo, Model

export TurkieParams, TurkieCallback

export addIO!, record

include("online_stats_plots.jl")

const std_colors = ColorSchemes.seaborn_colorblind

name(s::Symbol) = string(s)
name(s::OnlineStat) = nameof(typeof(s))

"""
    TurkieCallback(model::DynamicPPL.Model, plots::Series/AbstractVector = )

    ## Keyword arguments
    - `showtrace=true` : Show the trace of the variable
    - `window=0` : Use a window for plotting the trace, 0 will not use a window

"""
TurkieCallBack

struct TurkieCallback
    scene::Scene
    data::Dict{Symbol, MovingWindow}
    axis_dict::Dict
    vars::Dict{Symbol, Any}
    params::Dict{Any, Any}
    iter::Observable{Int64}
end

function TurkieCallback(model::Model, plots::Union{Series, AbstractVector} = [:histkde, Mean(), Variance(), AutoCov(20)]; kwargs...)
    variables = VarInfo(model).metadata
    return TurkieCallback(Dict(Pair.(keys(variables), Ref(plots))),
    Dict(kwargs...))
end

function TurkieCallback(varsdict::Dict; kwargs...)
    return TurkieCallback(varsdict, Dict(kwargs...))
end

function TurkieCallback(vars::Dict, params::Dict)
# Create a scene and a layout
    outer_padding = 5
    scene, layout = layoutscene(outer_padding, resolution = (1200, 700))
    display(scene)

    window = get!(params, :window, 1000)

    n_rows = length(keys(vars))
    n_cols = maximum(length.(values(vars)))
    n_plots = n_rows * n_cols
    iter = Node(0)
    data = Dict{Symbol, MovingWindow}(:iter => MovingWindow(window, Int64))
    obs = Dict{Symbol, Any}()
    axis_dict = Dict()
    for (i, (variable, plots)) in enumerate(vars)
        data[variable] = MovingWindow(window, Float32)
        axis_dict[(variable, :varname)] = layout[i, 1, Left()] = Label(scene, string(variable), textsize = 30)
        axis_dict[(variable, :varname)].padding = (0, 50, 0, 0)
        onlineplot!(scene, layout, axis_dict, plots, iter, data, variable, i)
    end
    on(iter) do i
        if i > 10 # To deal with autolimits a certain number of samples are needed
            for (variable, plots) in vars
                for p in plots
                    autolimits!(axis_dict[(variable, p)])
                end
            end
        end
    end
    MakieLayout.trim!(layout)
    TurkieCallback(scene, data, axis_dict, vars, params, iter)
end

function addIO!(cb::TurkieCallback, io)
    cb.params[:io] = io
end

function (cb::TurkieCallback)(rng, model, sampler, transition, iteration)
    fit!(cb.data[:iter], iteration)
    for (vals, ks) in values(transition.θ)
        for (k, val) in zip(ks, vals)
            if haskey(cb.data, Symbol(k))
                fit!(cb.data[Symbol(k)], Float32(val))
            end
        end
    end
    cb.iter[] += 1
    if haskey(cb.params, :io)
        recordframe!(cb.params[:io])
    end
end

end


using Test
using PlayingCards
using TexasHoldem
import TexasHoldem
TH = TexasHoldem

include("tester_bots.jl")

@testset "Game: Play (invalid number of players)" begin
    @test_throws AssertionError("Invalid number of players") Game(ntuple(i->Player(Bot5050(), i), 1))
    @test_throws AssertionError("Invalid number of players") Game(ntuple(i->Player(Bot5050(), i), 11))
end

@testset "Game: Play (BotCheckCall)" begin
    play!(Game(ntuple(i->Player(BotCheckCall(), i), 3)))
end

@testset "Game: Play (BotCheckFold)" begin
    play!(Game(ntuple(i->Player(BotCheckFold(), i), 3)))
end

@testset "Game: Play (BotBetSB) - breaks lowest allowable bet" begin
    game = Game((Player(BotBetSB(),1), Player(BotCheckCall(),2)))
    # Cannot raise small blind on pre-flop! The big blind is the raise min raise, so we must raise at least big-blind
    @test_throws AssertionError("Raise must be between [\$4.0, \$200.0]") play!(game)
end

@testset "Game: Play (BotBetBB)" begin
    play!(Game((Player(BotBetBB(), 1), Player(BotCheckCall(), 2),)))
end

struct NoActionBot <: AbstractAI end

TH.player_option!(game::Game, player::Player{NoActionBot}, ::AbstractGameState, ::PlayerOptions) = nothing

@testset "Game: Play (NoActionBot)" begin
    game = Game((Player(BotCheckCall(), 1), Player(NoActionBot(), 2),))
    @test_throws AssertionError("Must take exactly 1 action.") play!(game)
end

@testset "Non-valid option using BotCheckOnCallRaiseFold" begin
    game = Game((Player(BotCheckCall(), 1), Player(BotCheckOnCallRaiseFold(), 2),))
    @test_throws ErrorException("Cannot check. Available options: CallRaiseFold") play!(game)
end
@testset "Non-valid option using BotCheckOnCallAllInFold" begin
    game = Game((Player(BotCheckOnCallAllInFold(), 1), Player(BotRaiseAlmostAllIn(), 2)))
    @test_throws ErrorException("Cannot check. Available options: CallAllInFold") play!(game)
end
@testset "Non-valid option using BotCheckOnCallFold" begin
    game = Game((Player(BotCheckOnCallFold(), 1), Player(BotRaiseAllIn(), 2)))
    @test_throws ErrorException("Cannot check. Available options: CallFold") play!(game)
end
@testset "Non-valid option using BotCallOnCheckRaiseFold" begin
    game = Game((Player(BotCheckCall(), 1), Player(BotCallOnCheckRaiseFold(), 2),))
    # We catch this incorrect option error before it's completed,
    # so we can't error in `validate_action`
    # @test_throws ErrorException("Cannot call. Available options: CheckRaiseFold") play!(game)
    @test_throws AssertionError("Cannot contribute \$0.0 to the pot!") play!(game)
end
@testset "Non-valid option using BotRaiseOnCallFold" begin
    game = Game((Player(BotRaiseOnCallFold(), 1), Player(BotRaiseAllIn(), 2)))
    @test_throws AssertionError("Cannot raise 0.") play!(game)
end

@testset "Player limps all in" begin
    # Note that the order of the players seems to be important here.
    game = Game((Player(BotCheckCall(), 2; bank_roll = 11), Player(BotLimpAllIn(), 1; bank_roll = 1)))
    play!(game)
end

n_check_actions = [0]
n_call_actions = [0]
@testset "N-actions (3 players)" begin
    play!(Game(ntuple(i->Player(BotNActions(), i), 3)))
    @test n_call_actions[1] == 2 # dealer + small blind
    # 1 check pre-flop (big blind)
    # 3 checks (flop)
    # 3 checks (turn)
    # 3 checks (river)
    @test n_check_actions[1] == 10
end

n_check_actions = [0]
n_call_actions = [0]
@testset "N-actions (4 players)" begin
    play!(Game(ntuple(i->Player(BotNActions(), i), 4)))
    @test n_call_actions[1] == 3 # dealer + small blind + 1 other
    # 1 check pre-flop (big blind)
    # 4 checks (flop)
    # 4 checks (turn)
    # 4 checks (river)
    @test n_check_actions[1] == 13
end

n_check_actions = [0]
n_call_actions = [0]
@testset "N-actions (custom 1)" begin
    # Inspired by:
    # [ Info: Initial bank roll summary: (260.0, 0.0, 37.0, 0.0, 203.0)
    # [ Info: Buttons (dealer, small, big, 1ˢᵗToAct): (1, 3, 5, 1)
    play!(Game((
        Player(BotNActions(), 1),
        Player(BotNActions(), 2; bank_roll = 0),
        Player(BotNActions(), 3),
        Player(BotNActions(), 4; bank_roll = 0),
        Player(BotNActions(), 5),
    ); dealer_id = 1))

    @test n_call_actions[1] == 2 # dealer + small blind
    # 1 check pre-flop (big blind)
    # 3 checks (flop)
    # 3 checks (turn)
    # 3 checks (river)
    @test n_check_actions[1] == 10
end

n_check_actions = [0]
n_call_actions = [0]
@testset "N-actions (custom 2)" begin
    # Inspired by:
    # [ Info: Initial bank roll summary: (195.0, 0.0, 256.0, 4.0, 45.0)
    # [ Info: Buttons (dealer, small, big, 1ˢᵗToAct): (5, 1, 3, 4)
    play!(Game((
        Player(BotNActions(), 1; bank_roll = 195.0),
        Player(BotCheckCall(), 2; bank_roll = 0),
        Player(BotPreFlopRaise(7.0), 3; bank_roll = 256.0),
        Player(BotCheckCall(), 4; bank_roll = 4.0),
        Player(BotCheckFold(), 5; bank_roll = 45.0),
    ); dealer_id = 1))

    @test n_call_actions[1] == 2 # dealer + small blind
    # 1 check pre-flop (big blind)
    # 3 checks (flop)
    # 3 checks (turn)
    # 3 checks (river)
    @test n_check_actions[1] == 3
end


function EigenDecomposition(arg0::RealMatrix)
    return EigenDecomposition((RealMatrix,), arg0)
end

function EigenDecomposition(arg0::Vector{jdouble}, arg1::Vector{jdouble})
    return EigenDecomposition((Vector{jdouble}, Vector{jdouble}), arg0, arg1)
end

function get_d(obj::EigenDecomposition)
    return jcall(obj, "getD", RealMatrix, ())
end

function get_determinant(obj::EigenDecomposition)
    return jcall(obj, "getDeterminant", jdouble, ())
end

function get_eigenvector(obj::EigenDecomposition, arg0::jint)
    return jcall(obj, "getEigenvector", RealVector, (jint,), arg0)
end

function get_imag_eigenvalue(obj::EigenDecomposition, arg0::jint)
    return jcall(obj, "getImagEigenvalue", jdouble, (jint,), arg0)
end

function get_imag_eigenvalues(obj::EigenDecomposition)
    return jcall(obj, "getImagEigenvalues", Vector{jdouble}, ())
end

function get_real_eigenvalue(obj::EigenDecomposition, arg0::jint)
    return jcall(obj, "getRealEigenvalue", jdouble, (jint,), arg0)
end

function get_real_eigenvalues(obj::EigenDecomposition)
    return jcall(obj, "getRealEigenvalues", Vector{jdouble}, ())
end

function get_solver(obj::EigenDecomposition)
    return jcall(obj, "getSolver", DecompositionSolver, ())
end

function get_square_root(obj::EigenDecomposition)
    return jcall(obj, "getSquareRoot", RealMatrix, ())
end

function get_v(obj::EigenDecomposition)
    return jcall(obj, "getV", RealMatrix, ())
end

function get_vt(obj::EigenDecomposition)
    return jcall(obj, "getVT", RealMatrix, ())
end

function has_complex_eigenvalues(obj::EigenDecomposition)
    return jcall(obj, "hasComplexEigenvalues", jboolean, ())
end


using CUDA

# Detect whether there is a display available
display_mode = true
try
   run(pipeline(`xset -q`, `tail -n1`, `grep "Monitor is On"`));
catch e
    global display_mode = false
end
if display_mode
    using ImageView, Images
end

function visualise(image::Array{Float32,3})
    display_mode ? imshow(image) : println("No display available")
end

function visualise(image::Array{Float32,1}, DIMX::Integer, DIMY::Integer, DIMZ::Integer)
    output = zeros(Float32, DIMX, DIMY, DIMZ);
    for i = 1:DIMX
        for j = 1:DIMY
            for k = 1:DIMZ
                index = calculate_index(i, j, k, DIMX, DIMY, DIMZ, 4)
                output[i, j, k] = image[index]
            end
        end
    end
    display_mode ? imshow(output) : println("No display available")
end

function visualise(c_image::CuArray{Float32,3})
    image = Array(c_image)
    display_mode ? imshow(image) : println("No display available")
end

function visualise(c_image::CuArray{Float32,1}, DIMX::Integer, DIMY::Integer, DIMZ::Integer)
    image = Array(c_image);
    visualise(image, DIMX, DIMY, DIMZ)
end

function save_image(image::Array{Float32,1})
    io = open("recon.raw","w")
    for i in eachindex(image)
        write(io, image[i])
    end
    close(io)
end

function save_image(image::Array{Float32,3})
    save_image(image[:])
end

function save_image(image::Array{Float32,1}, DIMX::Integer, DIMY::Integer, DIMZ::Integer)
    output = zeros(Float32, DIMX, DIMY, DIMZ);
    for i = 1:DIMX
        for j = 1:DIMY
            for k = 1:DIMZ
                index = calculate_index(i, j, k, DIMX, DIMY, DIMZ, 4)
                output[i, j, k] = image[index]
            end
        end
    end
    save_image(output)
end

function save_image(c_image::CuArray{Float32,1}, DIMX::Integer, DIMY::Integer, DIMZ::Integer)
    image = Array(c_image);
    save_image(image, DIMX, DIMY, DIMZ)
end


using LightGraphs
import Base.==

input = joinpath(@__DIR__, "input")
raw = readlines(input)

struct Coord
    x::Int
    y::Int
end

==(left::Coord, right::Coord) = left.x == right.x && left.y == right.y

function adjacent(coord::Coord)
    x, y = coord.x, coord.y
    return [
        Coord(x, y + 1),
        Coord(x, y - 1),
        Coord(x + 1, y),
        Coord(x - 1, y)
    ]
end

function map_maze!(cs::Array{String,1}, c2t::Dict{Coord,Char})
    for j = 1:length(cs)
        line = cs[j]
        for i = 1:length(line)
            t = line[i]
            c2t[Coord(i, j)] = t
        end
    end
end

function isalpha(char::Union{Char,Nothing})
    return char !== nothing && char >= 'A' && char <= 'Z'
end

function portals(c2t::Dict{Coord, Char})
    potentials = []
    for ct in c2t
        c, t = ct
        if t >= 'A' && t <= 'Z'
            push!(potentials, c)
        end
    end

    portals = Array{Pair{String, Coord}, 1}(undef, (0))
    for p in potentials
        x, y = p.x, p.y
        i = get(c2t, Coord(x - 1, y), nothing)
        j = c2t[p]
        k = get(c2t, Coord(x + 1, y), nothing)
        if i == '.' && isalpha(j)
            push!(portals, Pair(string(j, k), Coord(x - 1, y)))
        elseif isalpha(i) && k == '.'
            push!(portals, Pair(string(i, j), Coord(x + 1, y)))
        end
    end
    for p in potentials
        x, y = p.x, p.y
        i = get(c2t, Coord(x, y - 1), nothing)
        j = c2t[p]
        k = get(c2t, Coord(x, y + 1), nothing)
        if i == '.' && isalpha(j)
            push!(portals, Pair(string(j, k), Coord(x, y - 1)))
        elseif isalpha(i) && k == '.'
            push!(portals, Pair(string(i, j), Coord(x, y + 1)))
        end
    end
    return portals
end

function warp(coord, portals)
    label, _ = first(filter(kv -> kv.second.x == coord.x && kv.second.y == coord.y, portals))
    for w in filter(kv -> kv.first == label, portals)
        if w.second.x != coord.x && w.second.y != coord.y
            return w.second
        end
    end
end

function seen!(coord::Coord, level::Int, cache::Set{Pair{Coord, Int}})
    if Pair(coord, level) in cache
        return true
    else
        push!(cache, Pair(coord, level))
        return false
    end
end

function explore(start::Coord, finish::Coord, level::Int, c2t::Dict{Coord, Char}, portals::Array{Pair{String, Coord}, 1})
    nodes = [(start, 0, level)]
    seen = Set{Pair{Coord, Int}}()

    warps = [w.second for w in portals]
    xs = map(kv -> kv.second.x, portals)
    ys = map(kv -> kv.second.y, portals)
    minx = minimum(xs)
    maxx = maximum(xs)
    miny = minimum(ys)
    maxy = maximum(ys)

    while !isempty(nodes)
        node = popfirst!(nodes)
        coord, steps, level = node
        # skip exploring if we've already started from this node
        if seen!(coord, level, seen)
            continue
        else
            ## found exit, bail
            if coord == finish && level in [0, 1]
                return steps
            ## found exit but not at right level
            elseif coord == finish
                continue
            ## not a corridor
            elseif c2t[coord] != '.'
                continue
            elseif coord in warps && coord != start
                warp_to = warp(coord, portals)
                if (coord.x in [minx, maxx] || coord.y in [miny, maxy]) && level < 0
                    for neighbor in adjacent(warp_to)
                        push!(nodes, (neighbor, steps + 2, level + 1))
                    end
                elseif (coord.x in [minx, maxx] || coord.y in [miny, maxy]) && level == 0
                    continue
                elseif level <= 0
                    for neighbor in adjacent(warp_to)
                        push!(nodes, (neighbor, steps + 2, level - 1))
                    end
                else
                    for neighbor in adjacent(warp_to)
                        push!(nodes, (neighbor, steps + 2, level))
                    end
                end
            else
                for neighbor in adjacent(coord)
                    push!(nodes, (neighbor, steps + 1, level))
                end
            end
        end
    end
end

c2t = Dict{Coord,Char}()
map_maze!(raw, c2t)

ps = portals(c2t)

_, start = first(filter(n -> n.first == "AA", ps))
_, finish = first(filter(n -> n.first == "ZZ", ps))

p1 = explore(start, finish, 1, c2t, ps)
@assert p1 == 580

p2 = explore(start, finish, 0, c2t, ps)
@assert p2 == 6362

print("-----------------------------------------------------------------------\n")
print("donut maze -- part one\n    : $p1\n")
print("donut maze -- part two\n    : $p2\n")
print("-----------------------------------------------------------------------\n")


using AVLTrees

"""
OneSidedBook is a one-sided book (i.e. :BID or :ASK) of order queues at
varying prices.

OrderQueues are stored in an AVLTree (.book) indexed
either by price (:ASK) or -price (:BID)

The book keeps track of various statistics such as the current best price,
total share and price volume, as well as total contained number of orders.

"""
@kwdef mutable struct OneSidedBook
    side::Symbol
    book::AVLTree{Float32,OrderQueue} = AVLTree{Float32,OrderQueue}()
    total_volume::Int64 = 0 # Total volume available in shares
    num_orders::Int32 = Int32(0) # Number of orders in the book
    best_price::Union{Float32,Nothing} = nothing # best bid or ask
end

Base.isempty(sb::OneSidedBook) = Base.isempty(sb.book)

"Updates the latest best price in a Sidebook (either :BID or :ASK book)."
function _update_next_best_price!(sb::OneSidedBook)
    sb.best_price = isempty(sb.book) ? nothing : first(sb.book)[1]
    return
end

"Retrieve order queue from OneSidedBook at given price"
function _get_price_queue(sb::OneSidedBook,price::Float32)
    pricekey = sb.side == :ASK ? price : -price
    getkey(sb.book,pricekey) # Return the price queue
end

"Delete entire queue associated with given price from OneSidedBook"
function _delete_price_queue!(sb::OneSidedBook,price::Float32)
    pricekey = sb.side == :ASK ? price : -price
    price_queue = pop!(sb.book,pricekey) # delete price queue

    # update book stats
    (price_queue.price == sb.best_price) && _update_next_best_price!(sb) # Update price only if best price was changed
    sb.num_orders -= price_queue.num_orders[]
    sb.total_volume -= price_queue.total_volume[]
end

"Insert new_order into OneSidedBook at given price, create new price queue if needed"
function insert_order!(sb::OneSidedBook,new_order::Order)
    pricekey = (sb.side == :ASK) ? new_order.price : -new_order.price
    # search for order queue at price
    order_queue = getkey(sb.book,pricekey)
    if isnothing(order_queue) # If key not present (price doesnt exist in book)
        new_queue = OrderQueue(price=new_order.price) # Create new price queue
        push!(new_queue,new_order) # Add order to new price queue
        insert!(sb.book,pricekey,new_queue) # add new price queue to OneSidedBook

        # Update new best price depending on bid/ask
        if isnothing(sb.best_price)
            sb.best_price = new_order.price
        elseif sb.side == :BID
            sb.best_price = max(new_order.price,sb.best_price)
        else
            sb.best_price = min(new_order.price,sb.best_price)
        end
    else # If order queue present, retrieve queue and insert new order
        push!(order_queue,new_order)
    end

    # Update Onesidedbook info
    sb.num_orders += 1
    sb.total_volume += new_order.size
end

"Delete order with given price/tick_id from book"
function delete_order!(sb::OneSidedBook,price::Float32,orderid::Int64)
    # Get price queue and delete order from it
    order_queue = _get_price_queue(sb,price)
    vol0 = order_queue.total_volume[] # get stats before deletion
    delete!(order_queue,orderid)
    vol1 = order_queue.total_volume[] # get stats after deletion

    # If order deletion depleted price queue, delete the whole queue
    if isempty(order_queue)
        _delete_price_queue!(sb,price) # note: this function will update price
    end

    # Update Onesidedbook info
    # TODO, Record change in order queue stats
    sb.num_orders -= 1
    sb.total_volume -= (vol0-vol1)

end


module TestStacking

using Test
using MLJBase
using MLJModelInterface
using ..Models
using Random
using StableRNGs

import Distributions

rng = StableRNGs.StableRNG(1234)

function model_evaluation(models::NamedTuple, X, y; measure=rmse)
    cv = CV(;nfolds=3)
    results = []
    for model in models
        mach = machine(model, X, y)
        ev = evaluate!(mach; resampling=cv, verbosity=0, measure=measure, check_measure=false)
        push!(results, ev.measurement[1])
    end
    results
end

@testset "Testing Stack on Continuous target" begin
    X, y = make_regression(500, 5; rng=rng)

    @testset "Testing Deterministic Stack" begin
    # Testing performance

    # The dataset is a simple regression model with intercept
    # No model in the stack can recover the true model on its own
    # Indeed, FooBarRegressor has no intercept
    # By combining models, the stack can generalize better than any submodel
    # And optimize the rmse

    models = (constant=DeterministicConstantRegressor(),
                decisiontree=DecisionTreeRegressor(),
                ridge_lambda=FooBarRegressor(;lambda=0.1),
                ridge=FooBarRegressor(;lambda=0))

    mystack = Stack(;metalearner=FooBarRegressor(),
                    resampling=CV(;nfolds=3),
                    models...)

    results = model_evaluation((stack=mystack, models...), X, y)
    @test argmin(results) == 1

    # Mixing ProbabilisticModels amd Deterministic models as members of the stack
    models = (constant=ConstantRegressor(),
                decisiontree=DecisionTreeRegressor(),
                ridge_lambda=FooBarRegressor(;lambda=0.1),
                ridge=FooBarRegressor(;lambda=0))

    mystack = Stack(;metalearner=FooBarRegressor(),
                    resampling=CV(;nfolds=3),
                    models...)
    # Testing attribute access of the stack
    @test propertynames(mystack) == (:resampling, :metalearner, :constant,
                                    :decisiontree, :ridge_lambda, :ridge)

    @test mystack.decisiontree isa DecisionTreeRegressor

    @test target_scitype(mystack) == target_scitype(FooBarRegressor())
    @test input_scitype(mystack) == input_scitype(FooBarRegressor())

    # Testing fitted_params results are easily accessible for each
    # submodel. They are in order of the cross validation procedure.
    # Here 3-folds then 3 machines + the final fit
    mach = machine(mystack, X, y)
    fit!(mach, verbosity=0)
    fp = fitted_params(mach)
    @test nrows(getfield(fp, :constant)) == 4
    @test nrows(getfield(fp, :decisiontree)) == 4
    @test nrows(getfield(fp, :ridge)) == 4
    @test nrows(getfield(fp, :ridge_lambda)) == 4

    # The metalearner has been fit and has one coefficient
    # for each model in the library (No intercept)
    @test fp.metalearner isa Vector{Float64}
    @test nrows(fp.metalearner) == 4

    # Testing prediction is Deterministic
    @test predict(mach) isa Vector{Float64}
    end

    @testset "Testing ProbabilisticStack" begin
        models = (constant=ConstantRegressor(),
                    decisiontree=DecisionTreeRegressor(),
                    ridge_lambda=FooBarRegressor(;lambda=0.1),
                    ridge=FooBarRegressor(;lambda=0))

        # The type of the stack is determined by the type of the metalearner
        metalearner = ConstantRegressor(;distribution_type=Distributions.Cauchy)
        mystack = Stack(;metalearner=metalearner,
                    resampling=CV(;nfolds=3),
                    models...)

        @test target_scitype(mystack) == target_scitype(metalearner)
        @test input_scitype(mystack) == input_scitype(FooBarRegressor())

        mach = machine(mystack, X, y)
        fit!(mach, verbosity=0)
        @test predict(mach) isa Vector{Distributions.Cauchy{Float64}}

    end

end

@testset "Testing ProbabilisticStack on Finite target" begin
    X, y = make_blobs(;rng=rng, shuffle=false)

    models = (constant=ConstantClassifier(),
                decisiontree=DecisionTreeClassifier(),
                knn=KNNClassifier())

    mystack = Stack(;metalearner=DecisionTreeClassifier(),
                    resampling=StratifiedCV(;nfolds=3),
                    models...)

    # Check input and target scitypes
    @test target_scitype(mystack) == target_scitype(DecisionTreeClassifier())
    # Here the greatest lower bound is the scitype of the knn
    @test input_scitype(mystack) == input_scitype(KNNClassifier())

    mach = machine(mystack, X, y)
    fit!(mach, verbosity=0)
    @test predict(mach) isa Vector{<:MLJBase.UnivariateFinite}

end

@testset "Stack constructor valid argument checks" begin
    # metalearner should have target_scitype:
    # Union{AbstractArray{<:Continuous}, AbstractArray{<:Finite}}
    @test_throws ArgumentError Stack(;metalearner=Standardizer(),
                        constant=ConstantClassifier())

    @test_throws ArgumentError Stack(;constant=KNNRegressor())
end

@testset "Misc" begin
    # Test setproperty! behaviour
    models = (constant=DeterministicConstantRegressor(),
                decisiontree=DecisionTreeRegressor(),
                ridge_lambda=FooBarRegressor(;lambda=0.1))

    mystack = Stack(;metalearner=FooBarRegressor(),
                    resampling=CV(;nfolds=3),
                    models...)
    @test mystack.ridge_lambda.lambda == 0.1
    @test mystack.metalearner isa FooBarRegressor
    @test mystack.resampling isa CV

    mystack.ridge_lambda = FooBarRegressor(;lambda=0.2)
    @test mystack.ridge_lambda.lambda == 0.2

    mystack.metalearner = DecisionTreeRegressor()
    @test mystack.metalearner isa DecisionTreeRegressor

    mystack.resampling = StratifiedCV()
    @test mystack.resampling isa StratifiedCV

    # using inner constructor accepts :resampling and :metalearner
    # as modelnames
    modelnames = (:resampling, :metalearner)
    models = [DeterministicConstantRegressor(), FooBarRegressor(;lambda=0)]
    metalearner = DeterministicConstantRegressor()
    resampling = CV()

    MLJBase.DeterministicStack(modelnames, models, metalearner, resampling)

    # Test input_target_scitypes with non matching target_scitypes
    models = [KNNRegressor()]
    metalearner = KNNClassifier()
    inp_scitype, tg_scitype = MLJBase.input_target_scitypes(models, metalearner)
    @test tg_scitype == Unknown
    @test inp_scitype == Table{<:AbstractVector{<:Continuous}}

    # Test input_target_scitypes with non matching target_scitypes
    models = [ConstantClassifier(), DecisionTreeClassifier()]
    metalearner = KNNClassifier()
    inp_scitype, tg_scitype = MLJBase.input_target_scitypes(models, metalearner)
    @test tg_scitype == AbstractVector{<:Finite}
    @test inp_scitype == Table{<:Union{AbstractVector{<:Continuous},
                                       AbstractVector{<:Count},
                                       AbstractVector{<:OrderedFactor}}}

    # Changing a model to a non compatible target scitype
    initial_stack = Stack(;metalearner=FooBarRegressor(),
                            resampling=CV(;nfolds=3),
                            constant = DeterministicConstantRegressor(),
                            fb=FooBarRegressor())
    initial_stack.constant = ConstantClassifier()
    @test_throws DomainError clean!(initial_stack)

end

@testset  "function oos_set" begin
    X = (x = Float64[1, 1, 2, 2, 3, 3],)
    y = coerce(['a', 'b', 'b', 'c', 'a', 'a'], Multiclass)
    n = nrows(y)

    model1 = KNNClassifier(K=2)
    model2 = ConstantClassifier()
    judge = KNNClassifier(K=3)

    stack = Stack(metalearner=judge,
                model1=model1,
                model2=model2,
                resampling=CV(;nfolds=3, shuffle=true, rng=rng))

    Xs = source(X)
    ys = source(y)
    folds = MLJBase.getfolds(ys, stack.resampling, n)

    Zval, yval = MLJBase.oos_set(stack, folds, Xs, ys)
    # To be accessed, the machines need to be trained
    fit!(Zval, verbosity=0)
    # Each model in the library should output a 3-dim vector to be concatenated
    # resulting in the table of shape (nrows, 6) here nrows=6
    # for future use by the metalearner
    sc = schema(Zval())
    @test sc.nrows == 6
    @test sc.names == (:x1, :x2, :x3, :x4, :x5, :x6)

    # The lines of yval should match the reordering indexes
    # of the original y (reordering given by the folds node)
    reordering = vcat([x[2] for x in folds()]...)
    @test yval() == y[reordering]
    # And the same is true for Zval, let's check this for model1's output
    # on the first fold, ie (2 first rows, 3 first columns)
    # First we need to train the model
    trainingrows = folds()[1][1]
    Xtrain = selectrows(X, trainingrows)
    ytrain = selectrows(y, trainingrows)
    mach = machine(model1, Xtrain, ytrain)
    fit!(mach, verbosity=0)

    # Then predict on the validation rows
    Xpred = selectrows(X, folds()[1][2])
    Zval_expected_dist = predict(mach, Xpred)
    # This is a distribution, we need to apply the appropriate transformation
    Zval_expected = pdf(Zval_expected_dist, levels(first(Zval_expected_dist)))
    @test matrix(Zval())[1:2, 1:3] == Zval_expected

end

@testset "An integration test for stacked classification" begin

    # We train a stack by hand and compare with the canned version
    # `Stack(...)`. There are two base learners, with 3-fold
    # cross-validation used to construct the out-of-sample base model
    # predictions.

    probs(y) = pdf(y, levels(first(y)))

    # data:
    N = 200
    X = (x = rand(rng, 3N), )
    y = coerce(rand("abc", 3N), Multiclass)

    # row splits:
    test1 = 1:N
    test2 = (N + 1):2N
    test3 = (2N + 1):3N
    train1 = (N + 1):3N
    train2 = vcat(1:N, (2N + 1):3N)
    train3 = 1:2N

    # base `model1`:
    model1 = KNNClassifier(K=2)
    mach1 = machine(model1, X, y)
    fit!(mach1, rows=train1, verbosity=0)
    y11 = predict(mach1, rows=test1) |> probs
    mach1 = machine(model1, X, y)
    fit!(mach1, rows=train2, verbosity=0)
    y12 = predict(mach1, rows=test2) |> probs
    mach1 = machine(model1, X, y)
    fit!(mach1, rows=train3, verbosity=0)
    y13 = predict(mach1, rows=test3) |> probs
    y1_oos = vcat(y11, y12, y13)
    mach1_full = machine(model1, X, y)
    fit!(mach1_full, verbosity=0)
    y1 = predict(mach1_full, X) |> probs

    # base `model2`:
    model2 = DecisionTreeClassifier()
    mach2 = machine(model2, X, y)
    fit!(mach2, rows=train1, verbosity=0)
    y21 = predict(mach2, rows=test1) |> probs
    mach2 = machine(model2, X, y)
    fit!(mach2, rows=train2, verbosity=0)
    y22 = predict(mach2, rows=test2) |> probs
    mach2 = machine(model2, X, y)
    fit!(mach2, rows=train3, verbosity=0)
    y23 = predict(mach2, rows=test3) |> probs
    y2_oos = vcat(y21, y22, y23)
    mach2_full = machine(model2, X, y)
    fit!(mach2_full, verbosity=0)
    y2 = predict(mach2_full, X) |> probs

    # metalearner (`judge`):
    X_oos = MLJBase.table(hcat(y1_oos, y2_oos))
    judge = KNNClassifier(K=3)
    m_judge = machine(judge, X_oos, y)
    fit!(m_judge, verbosity=0)
    X_judge = MLJBase.table(hcat(y1, y2))
    yhat_matrix = predict(m_judge, X_judge) |> probs

    # alternatively, use stack:
    stack = Stack(metalearner=judge,
                model1=model1,
                model2=model2,
                resampling=CV(nfolds=3))
    mach = machine(stack, X, y)
    fit!(mach, verbosity=0)
    yhat_matrix_stack = predict(mach, X) |> probs

    # compare:
    @test yhat_matrix_stack ≈ yhat_matrix

end

end

true


"""
    mt_state(x) -> Bool
    mt_state!(x, val)

Specifies whether the magnetization transfer pulse is applied. This parameter
is REQUIRED by all the anatomical images grouped by MTR, MTS and MPM suffixes.
"""
@defprop MTState{:mt_state}::Bool

"""
    mt_offset_frequency(x)
    mt_offset_frequency!(x, val)

The frequency offset of the magnetization transfer pulse with respect to the
central H1 Larmor frequency in Hertz (Hz).
"""
@defprop MTOffsetFrequency{:mt_offset_frequency}::F64Hertz

"""
    mt_pulse_bandwidth(x)
    mt_pulse_bandwidth!(x, val)

The excitation bandwidth of the magnetization transfer pulse in Hertz (Hz).
"""
@defprop MTPulseBandwidth{:mt_pulse_bandwidth}::F64Hertz

"""
    mt_npulses(x)
    mt_npulses!(x, val)

Number of magnetization transfer RF pulses applied before the readout.
"""
@defprop MTNumberOfPulses{:mt_npulses}::Int

# TODO how should this be parameterized
"""
    mt_pulse_shape(x)
    mt_pulse_shape!(x, val)

Shape of the magnetization transfer RF pulse waveform. Accepted values:
* HARD
* GAUSSIAN
* GAUSSHANN (gaussian pulse with Hanning window)
* SINC
* SINCHANN (sinc pulse with Hanning window)
* SINCGAUSS (sinc pulse with Gaussian window)
* FERMI
"""
@defprop MTPulseShape{:mt_pulse_shape}

"""
    mt_pulse_duration(x)
    mt_pulse_duration!(x, val)

Duration of the magnetization transfer RF pulse in seconds.
"""
@defprop MTPulseDuration{:mt_pulse_duration}::F64Second

"""
    MagnetizationTransferMetadata

Metadata structure for information concerning MRI magnetization transfer pulse.

## Supported Properties

$(GETPROPERTY)

"""
struct MagnetizationTransferMetadata{H,S}
    mt_state::Bool
    mt_offset_frequency::H
    mt_pulse_bandwidth::H
    mt_npulses::Int
    mt_pulse_shape::Symbol
    mt_pulse_duration::S
end

@properties MagnetizationTransferMetadata begin
    mt_state(self) => :mt_state
    mt_offset_frequency(self) => :mt_offset_frequency
    mt_pulse_bandwidth(self) => :mt_pulse_bandwidth
    mt_npulses(self) => :mt_npulses
    mt_pulse_shape(self) => :mt_pulse_shape
    mt_pulse_duration(self) => :mt_pulse_duration
end