#!/usr/bin/env julia
using Plots, Random, Printf, Plots.PlotMeasures
include("IMF.jl") #new way to make random masses more "realistic"?
function initCondGen(nBodies,m;vRange=[-7e3,7e3],posRange=[-35,35]) #get random initial conditions for mass/radius, position, and velocity, option for user to specify acceptable vRange and posRange
    rad=m.^0.8 #3 radii based on masses in solar units
    m=m.*2e30 #convert to SI kg
    rad=rad.*7e8 #convert to SI m
    minV,maxV=vRange[1],vRange[2] #defaults to +/- 7km/s
    minPos,maxPos=posRange[1],posRange[2] #defaults to pos within 70 AU box
    posList=[]
    function checkPos(randPos,n,posList,rad) #this function checks if positions are too close to each other
        for i=2:(n-1)
            dist=sqrt((posList[i][1]-randPos[1])^2+(posList[i][2]-randPos[2])^2)
            if (dist*1.5e11)<(rad[n]+rad[i])
                return false
            end
        end
        return true
    end
    function genPos(nBodies,posList,rad,minPos,maxPos) #this function generates random initial positions for all the bodies
        push!(posList,rand(minPos:maxPos,2)) #random initial x,y coords for 1st body, default 70 AU box width
        for n=2:nBodies
            acceptPos=false
            while acceptPos==false
                randPos=rand(minPos:maxPos,2) #random guess
                acceptPos=checkPos(randPos,n,posList,rad) #check if our random guess is okay
                if acceptPos==true
                    push!(posList,randPos) #add accepted guess to master list
                end
            end
        end
        return posList
    end
    pos=genPos(nBodies,posList,rad,minPos,maxPos).*1.5e11 #convert to SI, m
    v=[]
    for n=1:nBodies
        push!(v,rand(minV:maxV,2)) #random x & y velocity between minV:maxV, km/s
    end
    r=[] #r is master vector containing both position and velocity info for all bodies that we will need later
    for n=1:nBodies
        push!(r,pos[n]) #position is first half of r
    end
    for n=1:nBodies
        push!(r,v[n]) #v is second half of r
    end
    #NOTE: if you want to specify the initial positions and/or velocities uncomment and fill in the sample lines below
    #technically if you want to save a bit of computation time you should also comment out lines 34 & 35
    #pos=[[pos1X,pos1Y],[pos2X,pos2Y],[pos3X,pos3Y],.....,[posNX,posNY]].*1.5e11 #fill this in if you want to specify the positions, units of AU from center
    #v=[[v1X,v1Y],[v1X,v1Y],[v3X,v3Y],......,[vNX,vNY]].*1e3 #fill this in to specify initial velocities, units of km/s
    #v=[[6.639,1.804],[-1.807,-0.998],[5.966,-4.904]].*1e3
    #pos=[[-4,-34],[-4,-3],[-33,32]].*1.5e11
    open("initCondAll.txt","w") do f #save initial conditions to file in folder where script is run
       for n=1:nBodies
           write(f,"body $n info: m = $(@sprintf("%.1f",(m[n]/2e30))) solar masses | v = ($(v[n][1]/1e3),$(v[n][2]/1e3)) km/s | starting position = ($(pos[n][1]/1.5e11),$(pos[n][2]/1.5e11)) AU from center\n")
       end
    end
    open("initCond.txt","w") do f
        for n=1:nBodies
            write(f,"m$n = $(@sprintf("%.1f",(m[n]/2e30)))\n")
        end
    end
    return [r, rad, m]
end

function dR(r,m,nBodies) #function we will use RK4 on to approximate solution
    G=6.67408313131313e-11# Nm^2/kg^2
    C=m.*G #Nm^2/kg
    function getDist(pos1,pos2) #find separation between two bodies
        x1,y1=pos1[1],pos1[2]
        x2,y2=pos2[1],pos2[2]
        dist=sqrt((x1-x2)^2+(y1-y2)^2)
        return dist
    end
    function getDv(C,n,r,nBodies) #find acceleration for a given body from Newton's law of gravitation
        pos=r[1:nBodies]
        dVx=0
        dVy=0
        for i=1:nBodies
            if i!=n #only do this if other body is NOT self
                dist=getDist(pos[n],pos[i])
                c=C[i]
                dVx-=c*(pos[n][1]-pos[i][1])/dist^3 #m/s after *dt
                dVy-=c*(pos[n][2]-pos[i][2])/dist^3
            end
        end
        return [dVx,dVy]
    end

    rUpdate=copy(r) #otherwise we modify r in place, in Julia array variables are just pointers
    dr=r[(nBodies+1):nBodies*2] #old v*dt = dr
    dv=[]
    for n=1:nBodies
        dVn=getDv(C,n,r,nBodies) #get the dv info for all bodies
        push!(dv,dVn)
    end
    rUpdate[1:nBodies]=dr[1:nBodies] #update positions
    rUpdate[(nBodies+1):nBodies*2]=dv[1:nBodies] #update velocities
    return rUpdate #r updated after timestep
end

function checkSep(xList,yList,nBodies,i,m,r,rad,names,colors,maxSep) #this function detects collisions and performs bookkeeping
                sepList=[]
                for n1=1:nBodies
                    for n2=(n1+1):nBodies
                        x1,y1=xList[n1][i],yList[n1][i] #get position info from data
                        x2,y2=xList[n2][i],yList[n2][i]
                        sep=sqrt((x1-x2)^2+(y1-y2)^2)
                        push!(sepList,sep) #get all the separations in a list
                        minSep=rad[n1]+rad[n2] #no touching!
                        split1,split2=split(names[n1]),split(names[n2])
                        if split1[end]=="hole" || split2[end]=="hole"
                            minSep*=2 #black holes are v smol so need to make the "hit box" bigger to prevent erroneous flying
                        elseif split1[end]=="hole" && split2[end]=="hole"
                            minSep*=100000 #two black holes are V SMOL
                        end
                        if sep<minSep
                            #println("collision between indices $n1 and $n2")
                            #println("names before modification: $names")
                            #println("masses before modification: $(m./2e30)")
                            sharedX=(x1+x2)/2 #collision detected, so assign new coordinate values for resulting "black hole"
                            sharedY=(y1+y2)/2
                            p1x,p1y=m[n1]*r[n1+nBodies][1],m[n1]*r[n1+nBodies][2] #get initial momenta from each body
                            p2x,p2y=m[n2]*r[n2+nBodies][1],m[n2]*r[n2+nBodies][2]
                            pTotx,pToty=p1x+p2x,p1y+p2y #calculate total momentum before collision in each direction
                            newMass=m[n1]+m[n2] #calculate new mass
                            sharedVfX=pTotx/(newMass) #get new velocity by dividing momentum in each direction by total mass
                            sharedVfY=pToty/(newMass) #idealized perfectly inelastic collision
                            r[n1][1]=sharedX #replace coordinates of one of the old bodies with new shared coords
                            r[n1][2]=sharedY
                            r[n1+nBodies][1]=sharedVfX #ditto for velocity
                            r[n1+nBodies][2]=sharedVfY
                            r[n2]=true #mark for removal
                            r[n2+nBodies]=true
                            filter!(x->x!=true,r) #filter and remove n2 entries, because now there's only 1 body from this pair
                            if split1[end]=="hole" #this is bookkeeping to make sure labels stay right on plot
                                names[n1]=names[n1][1:(end-11)] #remove "black hole" from end of label to avoid things like "1 & 2 black hole & 3 black hole"
                            elseif split2[end]=="hole"
                                names[n2]=names[n2][1:(end-11)]
                            end
                            newName=string("$(@sprintf("%.1f",(newMass/2e30))) black hole") #update name to reflect collision
                            names[n1]=newName #update n1 index since this is the one we "keep"
                            deleteat!(names,n2) #get rid of n2 since we don't need it anymore
                            deleteat!(colors,n2) #ditto for n2 color
                            m[n1]=newMass #update mass to reflect merger
                            deleteat!(m,n2) #delete the other
                            newRad=2*(6.6743015e-11)*newMass/(9e16) #schwarzschild radius -- might change because it's tiny on plots but whatever
                            rad[n1]=newRad #update radii list to reflect merger
                            deleteat!(rad,n2) #get rid of the other one
                            #println("names after modification: $names")
                            #println("masses after modification: $(m./2e30)")
                            return true,r,nBodies-1,names,colors,m,rad #now we have 1 less body so return updated info!
                        end
                    end
                end
                if length(sepList)>1
                    if minimum(sepList)>maxSep*1.5e11
                        return true,r,nBodies,names,colors,m,rad #stop if everybody flew away from each other
                    end
                end
                return false,r,nBodies,names,colors,m,rad #if nothing interesting happend just return everything unmodified
            end

function genNBodyStep(nBodies, physInfo, names, colors, stopCond=[10,100],dt=0.033) #default stop conditions of 10 yrs and 100 AU sep
    tStop=stopCond[1]*365*24*3600 #convert to SI s
    sepStop=stopCond[2]*1.5e11 #convert to SI m
    stop=false
    numSteps=Int(ceil(stopCond[1]/dt))
    t=range(0,stop=tStop,length=(numSteps+1)) #+1 because I don't want 0 to count
    stepSize=dt*365*24*3600 #convert to SI s
    xList=[]
    yList=[]
    for n=1:nBodies
        push!(xList,zeros(length(t))) #technically everything with push! is a little slower than pre-allocating but I'm very lazy and this doesn't need to be done that often so whatever
        push!(yList,zeros(length(t)))
    end
    r,rad,m=copy(physInfo[1]),copy(physInfo[2]),copy(physInfo[3]) #need copy because otherwise they're just pointers
    colors=copy(colors) #same reasoning as above
    names=copy(names) # ^^^
    i=1
    currentT=t[i] #do it this way other than +=stepSize because this mitigates rounding error
    #implement RK4 to model solutions to differential equations
    while stop==false
        if currentT>=tStop #in case of rounding error or something
            stop=true
        elseif i>(numSteps+1) #inf loop failsafe
            stop=true
            println("error: shouldn't have gotten here")
        else
            for n=1:nBodies
                xList[n][i]=r[n][1]
                yList[n][i]=r[n][2]
            end
            k1=stepSize.*dR(r,m,nBodies)
            k2=stepSize.*dR(r.+0.5.*k1,m,nBodies)
            k3=stepSize.*dR(r.+0.5.*k2,m,nBodies)
            k4=stepSize.*dR(r.+k3,m,nBodies)
            r.+=(k1.+2.0.*k2.+2.0.*k3.+k4)./6 #basically dt*(dr/dt,dv/dt) gives us updated r vector
            #check separation after each dt step
            stop,r,nBodies,names,colors,m,rad=checkSep(xList,yList,nBodies,i,m,r,rad,names,colors,sepStop)
            if stop==true
                t=range(0,stop=currentT,length=i) #t should match pos vectors
                for n=1:(nBodies+1)
                    xList[n]=xList[n][1:i] #cut off trailing zeros
                    yList[n]=yList[n][1:i]
                end
            else
                i+=1
                currentT=t[i] #on to the next one!
            end
        end
    end
    return ([xList, yList, t], [r, rad, m], nBodies, names, colors, currentT/365/24/3600)
end

function genNBody(nBodies, m, names, colors; stopCond=[10,100],dt=0.033,vRange=[-7e3,7e3],posRange=[-35,35]) #this function makes all the data for an entire nBody run
    globalT=0
    stopT=stopCond[1] #yrs
    numTrials=0
    plotList=[]
    tOffsets=[0.0] #first offest should be zero years
    colorsList=[colors] #first set of colors is just original colors
    physInfoList=[]
    namesList=[]
    nBodiesList=[nBodies]
    nBodiesGlobal=nBodies
    while globalT<stopT && nBodies>1 #go until we reach given time or there's one giant black hole
        numTrials+=1
        function firstStep(nBodies0,m0,names0,colors0,stopT,globalT,stopCond2,dt) #try to find a first step that goes for a reasonable amount of time without immediate collisions
            accept=false
            iter=1
            mLocal=copy(m0)
            iMax=1999
            global labelBool
            if labelBool==1
                iMax=0 #we set the mass in init cond
            end
            while accept==false
                if iter==1
                    mInit=m0
                else
                    mInit=rand(1:1500,nBodies0)./10 #n random masses between 0.1 and 150 solar masses
                    #mInit=getMass(nBodies0) #IMF way
                end
                physInfo0=initCondGen(nBodies0,mInit,vRange=vRange,posRange=posRange) #get initial conditions
                rad,m=physInfo0[2],physInfo0[3] #r is the first thing but we don't need it here
                newNames=["$(mInit[i])" for i=1:length(mInit)]
                plotInfo,newPhysInfo,nBodies,names,colors,t=genNBodyStep(nBodies0,physInfo0,newNames,colors0,[stopT-globalT,stopCond2],dt) #generate one step -- ie until there is a collision
                #println("t for iteration $iter = $t")c
                if t>(stopT/10) || iter>iMax #sometimes we set all the initial conditions so this is pointless but whatever
                    println("found a solution with a first step t = $t years in $iter iterations")
                    open("cron_log.txt","a") do f #for cron logging, a flag = append
                        write(f,"found a solution with first step t = $t after $iter iterations\n")
                    end
                    namesList=[newNames,names]
                    return plotInfo,newPhysInfo,nBodies,namesList,colors,t,rad,m
                    accept=true
                else
                    iter+=1
                    #mLocal=getMass(nBodies0) IMF way, not currently used because it's boring
                end
            end
        end
        if numTrials==1
            plotInfo,newPhysInfo,nBodiesGlobal,namesList,colors,t,rad,m=firstStep(copy(nBodies),copy(m),copy(names),copy(colors),stopT,globalT,stopCond[2],dt)
            globalT+=t
            println("number of output files at this step: $(ceil(t/dt)) | nBodies = $nBodiesGlobal")
            push!(tOffsets,globalT) #keep track of offsets for plotting later
            push!(plotList,plotInfo) #keep track of data from each step in master list
            push!(colorsList,colors) #keep track of colors at each step
            push!(physInfoList,[rad,m]) #we don't need to save r because that happens in x & y -- THIS IS INITIAL RAD,M
            push!(physInfoList,[newPhysInfo[2],newPhysInfo[3]]) #THIS IS RAD, M AFTER FIRST STEP
            push!(nBodiesList,nBodiesGlobal) #keep track of number of bodies at this step
            global physInfo=newPhysInfo #swap pointers, make this global so else statement can know about it
        else
            plotInfo,newPhysInfo,nBodiesGlobal,names,colors,t=genNBodyStep(nBodiesGlobal,copy(physInfo),copy(names),copy(colors),[stopT-globalT,stopCond[2]],dt)
            physInfo=newPhysInfo #swap pointers
            rad,m=physInfo[2],physInfo[3] #r is the first thing but we don't need it here -- THIS IS RAD, M AFTER 2+ STEPS
            globalT+=t
            println("number of output files at this step: $(ceil(t/dt)) | nBodies = $nBodiesGlobal")
            push!(tOffsets,globalT) #keep track of offsets for plotting later
            push!(plotList,plotInfo) #keep track of data from each step in master list
            push!(namesList,names) #keep track of names at each step
            push!(colorsList,colors) #keep track of colors at each step
            push!(physInfoList,[rad,m]) #we don't need to save r because that happens in x & y
            push!(nBodiesList,nBodiesGlobal) #keep track of number of bodies at this step
        end
    end
    #note: the bodies, colors, and names list are all 1 "too long" -- since they include an update from the last step, so in plotting loop we go to length(nBodiesList) - 1.
    println("made it $(@sprintf("%.2f",globalT)) years with $nBodiesGlobal surviving")
    return numTrials,plotList,tOffsets,physInfoList,namesList,colorsList,nBodiesList
end

function getInterestingNBody(nBodies,m,names,colors; minTime=0,stopCond=[10,100],dt=0.033,iMax=10,vRange=[-7e3,7e3],posRange=[-50,50]) #units of things in yrs, AU
    #sometimes random conditions result in a really short animation where things
    #just crash into each other/fly away, so this function throws away those
    #comments above from threeBody code -- since I allow collisions now this pretty much always runs all the way through first try but good to have just in case
    yearSec=365*24*3600
    interesting=false
    i=1
    while interesting==false
        nTrials,plotList,tOffsets,physInfoList,namesList,colorsList,nBodiesList=genNBody(nBodies,m,names,colors,stopCond=stopCond,dt=dt,vRange=vRange,posRange=posRange) #evolve a system
        stopT=tOffsets[end-1]+plotList[end][3][end]/yearSec #final end time for system evolved above
        if (stopT)>=minTime #only return if simulation runs for longer than minTime
            println("total time: ",stopT) #tell me how many years we are simulating
            open("cron_log.txt","a") do f #for cron logging (not really relevant for this script but a holdover from threeBody twitter bot script debugging)
                write(f,"$(stopT)\n") #I'm keeeping it anyways because it might be useful in the future if I want to turn this into a bot too
            end
            return nTrials,plotList,tOffsets,physInfoList,namesList,colorsList,nBodiesList
            interesting=true
        elseif i>iMax #computationally expensive so don't want to go forever
            interesting=true #render it anyways I guess because sometimes it's fun?
            println("did not find interesting solution in number of tries allotted, running anyways")
            println("total time: ",stopT) #how many years simulation runs for
            open("cron_log.txt","a") do f #for cron logging
                write(f,"$(stopT)\n")
            end
            return nTrials,plotList,tOffsets,physInfoList,namesList,colorsList,nBodiesList
        end
        i+=1
    end
end

function getRatioRight(ratio,dx,dy)
    if (dx/dy)!=ratio
        if dx>(ratio*dy)
            dy=dx/ratio
        else
            dx=dy*ratio
        end
    end
    return dx,dy
end

function getLims(xList,yList,padding,maxFrame,oldCenter,landscape) #determines plot limits at each frame, padding in units of pos
    xMin=minimum(xList) #where xList contains the current x position of all bodies
    xMax=maximum(xList) #ditto but max
    dx=xMax-xMin
    yMin=minimum(yList) #same as above but for y
    yMax=maximum(yList)
    dy=yMax-yMin
    horizontalExtraRatio=1
    if landscape==1
        horizontalExtraRatio=1920/1080 #landscape picture has extra space in x direction
    end
    centerX,centerY=sum(xList)/length(xList),sum(yList)/length(yList)
    dx,dy=getRatioRight(horizontalExtraRatio,dx,dy)
    xlims=[xMin-padding*horizontalExtraRatio,xMin+dx+padding*horizontalExtraRatio]
    ylims=[yMin-padding,yMin+dy+padding]
    return xlims,ylims,[centerX,centerY]
end

function makeCircleVals(r,center=[0,0]) #this function generates circle markers that are to scale with plot
    xOffset=center[1]
    yOffset=center[2]
    xVals=[r*cos(i)+xOffset for i=0:(pi/64):(2*pi)]
    yVals=[r*sin(i)+yOffset for i=0:(pi/64):(2*pi)]
    return xVals,yVals
end

#this new way runs significantly faster (~2x improvement over @anim)
#Downside is it spams folder with png images of every frame and must manually compile with ffmpeg
function plotSection(landscape,sectionNum,backData,oldI,oldColors,offsets,dt,nBodies,plotInfo,physInfo,labels,colors,maxFrame;slowDown=1)
    colorSymbols=[Symbol(color) for color in colors] #change strings to symbols for plot args
    oldColorSymbols=[Symbol(color) for color in oldColors]
    skipRate=Int(floor(1/dt/(30*slowDown))) #30 fps, slowDown = sloMo (ie a value of 2 makes this run at half speed)
    rad=physInfo[1]
    m=physInfo[2]
    xData=plotInfo[1]
    yData=plotInfo[2]
    t=plotInfo[3]
    tOffset=offsets[1]
    plotNum=offsets[2]
    I=0
    listInd=0
    centerList=[] #not actually the center but this won't get used right away and should be overwritten before needed so doesn't matter
    limList=[]
    oscillatingCount=0
    oscillating=false
    ratio=1
    if landscape==1
        ratio=1920/1080
    end
    ejectBan=[]
    for n=1:nBodies
        push!(ejectBan,[0,0]) #initialize banBool,Counter for each body
    end
    for i=1:skipRate:length(t) #this makes animation scale ~1 sec/year with other conditions
        GR.inline("png") #added to eneable cron/jobber compatibility, also this makes frames generate WAY faster? Prior to adding this when run from cron/jobber frames would stop generating at 408 for some reason.
        gr(legendfontcolor = plot_color(:white)) #legendfontcolor=:white plot arg broken right now (at least in this backend)
        p=plot(size=(1024,1024)) #initialize plot
        if landscape==1
            p=plot(size=(1920,1080)) #FHD landscape
        end
        if sectionNum>1  #don't need to do this forever since they decay to 0 anyways
            for n=1:length(backData) #old list has +1 bodies #maybe change the condition here to just the limits from the last frame?? might fix problems
                p=plot!(backData[n][1]./1.5e11,backData[n][2]./1.5e11,label="",linewidth=3,linecolor=oldColorSymbols[n],linealpha=max.((1:Int(floor(skipRate/10)):oldI) .+ 10000 .- (oldI+i),1000)/10000)
            end
        end
        print("$(@sprintf("%.2f",i/length(t)*100)) % complete\r") #output percent tracker
        x=[]
        y=[]
        if i>1
            center=centerList[listInd]
        else
            center=[0.,0.]
        end
        for n=1:nBodies
            if abs(xData[n][i])/1.5e11 <= (maxFrame*1.25*ratio + abs(center[1])) || abs(yData[n][i])/1.5e11 <= (maxFrame*1.25 + abs(center[2]))
                if ejectBan[n][1]==0 || ejectBan[n][2]>60 #let it be considered again after 60 frames
                    push!(x,xData[n][i]/1.5e11)
                    push!(y,yData[n][i]/1.5e11)
                    ejectBan[n][1]=0
                    ejectBan[n][2]=0
                end
            else
                ejectBan[n][1]=1
                ejectBan[n][2]+=1 #it's stayed ejected another frame
            end
        end
        if length(x) <= 1
            println("all bodies ejected, stopping plotting")
            return 0
        end
        oldCenter=copy(center)
        limx,limy,center=getLims(x,y,15,maxFrame,oldCenter,landscape) #20 AU padding
        dx,dy=(limx[2]-limx[1]),(limy[2]-limy[1])
        if listInd>2 #this block will hopefully prevent the camera from oscillating back and forth..? (for a 10 frames right now)
            oldLimx,oldLimy=limList[listInd][1],limList[listInd][2] #one back
            oldLimx2,oldLimy2=limList[listInd-1][1],limList[listInd-1][2] #two back
            oldDx2,oldDy2=(oldLimx2[2]-oldLimx2[1]),(oldLimy2[2]-oldLimy2[1])
            oldDx,oldDy=(oldLimx[2]-oldLimx[1]),(oldLimy[2]-oldLimy[1])
            if dx>oldDx && oldDx2>oldDx || dx<oldDx && oldDx2<oldDx/1.03#.05
                if dx>oldDx && oldDx2>oldDx
                    dx=oldDx*0.97
                    dy=oldDy*0.97
                else
                    dx=oldDx*1.03 #fix trajectory for a set amt of frames to prevent bounciness
                    dy=oldDy*1.03 #always opt for more space, but increase at slow rate
                end
                oscillating=true
                oscillatingCount+=1
                if oscillatingCount>30
                    oscillating=false
                    oscillatingCount=0
                end
            elseif dy>oldDy && oldDy2>oldDy || dy<oldDy && oldDy2<oldDy/1.03#.05
                if dy>oldDy && oldDy2>oldDy
                    dx=oldDx*0.97
                    dy=oldDy*0.97
                else
                    dx=oldDx*1.03
                    dy=oldDy*1.03
                end
                oscillating=true
                oscillatingCount+=1
                if oscillatingCount>30
                    oscillating=false
                    oscillatingCount=0
                end
            else
                if oscillating==true
                    dx=oldDx*1.005#.01
                    dy=oldDy*1.005#.01
                    oscillatingCount+=1
                    if oscillatingCount>30 #1 sec at 30 FPS
                        oscillating=false #don't want to freeze the frame forever...
                        oscillatingCount=0
                    end
                end
            end
        end
        dx,dy=getRatioRight(ratio,dx,dy)
        # if listInd>1
        #     oldLimx,oldLimy=limList[listInd][1],limList[listInd][2]
        #     oldDx,oldDy=oldLimx[2]-oldLimx[1],oldLimy[2]-oldLimy[1]
        #     # dCx=center[1]-centerList[listInd][1]
        #     # dCy=center[2]-centerList[listInd][2]
        #     # if abs(dCx)>2 #shifted more than 2 AU in x direction
        #     #     newCx=centerList[listInd][1]+2*sign(center[1])
        #     #     center[1]=newCx
        #     # elseif abs(dCy)>2 #shifted more than 2 AU in y direction
        #     #     newCy=centerList[listInd][1]+2*sign(center[2])
        #     #     center[2]=newCy
        #     # end
        #     if dx/oldDx<0.97 #frame shrunk more than 2%
        #         limx[1]=centerList[listInd][1]-oldDx*0.97/2
        #         limx[2]=centerList[listInd][1]+oldDx*0.97/2
        #         newCenter=centerList[listInd].-center./20
        #         center=newCenter
        #     elseif dx/oldDx>1.03 #grew more than 2%
        #         limx[1]=centerList[listInd][1]-oldDx*1.03/2
        #         limx[2]=centerList[listInd][1]+oldDx*1.03/2
        #         newCenter=centerList[listInd].-center./20
        #         center=newCenter
        #     elseif dy/oldDy<0.97
        #         limy[1]=centerList[listInd][2]-oldDy*0.97/2
        #         limy[2]=centerList[listInd][2]+oldDy*0.97/2
        #         newCenter=centerList[listInd].-center./20
        #         center=newCenter
        #     elseif dy/oldDy>1.03
        #         limy[1]=centerList[listInd][2]-oldDy*1.03/2
        #         limy[2]=centerList[listInd][2]+oldDy*1.03/2
        #         newCenter=centerList[listInd].-center./20
        #         center=newCenter
        #     end
        # end
        if i>1
            oldLimx,oldLimy=limList[listInd][1],limList[listInd][2]
            oldDx,oldDy=oldLimx[2]-oldLimx[1],oldLimy[2]-oldLimy[1]
            if dx/oldDx<0.97 #frame shrunk more than 2%
                limx[1]=oldCenter[1]-oldDx*0.97/2
                limx[2]=oldCenter[1]+oldDx*0.97/2
            elseif dx/oldDx>1.03 #grew more than 2%
                limx[1]=oldCenter[1]-oldDx*1.03/2
                limx[2]=oldCenter[1]+oldDx*1.03/2
            elseif dy/oldDy<0.97
                limy[1]=oldCenter[2]-oldDy*0.97/2
                limy[2]=oldCenter[2]+oldDy*0.97/2
            elseif dy/oldDy>1.03
                limy[1]=oldCenter[2]-oldDy*1.03/2
                limy[2]=oldCenter[2]+oldDy*1.03/2
            end
        end
        listInd+=1
        push!(centerList,center)
        dx,dy=(limx[2]-limx[1]),(limy[2]-limy[1])
        dx,dy=getRatioRight(ratio,dx,dy)
        limx[2]=limx[1]+dx
        limy[2]=limy[1]+dy
        push!(limList,[limx,limy])
        for n=1:nBodies
            p=plot!(xData[n][1:Int(floor(skipRate/10)):i]./1.5e11,yData[n][1:Int(floor(skipRate/10)):i]./1.5e11,label="",linewidth=3,linecolor=colorSymbols[n],linealpha=max.((1:Int(floor(skipRate/10)):i) .+ 10000 .- i,1000)/10000)
        end
        p=scatter!(starsX,starsY,markercolor=:white,markersize=:1,label="") #fake background stars
        for n=1:nBodies
            currentRad=rad[n]
            plotLabel=labels[n]
            if split(labels[n])[end]=="hole"
                fillColor=:black
                plotLabel=string(plotLabel," (enlarged 100,000x)")
                currentRad*=100000 #black hole is just very smol
            else
                fillColor=colorSymbols[n]
            end
            circleX,circleY=makeCircleVals(currentRad,[xData[n][i],yData[n][i]])
            if xData[n][i]/1.5e11 > limx[2] || yData[n][i]/1.5e11 > limy[2] || xData[n][i]/1.5e11 < limx[1] || yData[n][i]/1.5e11 < limy[1] #it's out of the frame
                plotLabel=string(plotLabel," (ejected)")
            end
            p=plot!(circleX./1.5e11,circleY./1.5e11,label=plotLabel,linecolor=colorSymbols[n],fill=true,fillcolor=fillColor)
        end
        p=plot!(background_color=:black,background_color_legend=:transparent,foreground_color_legend=:transparent,
            background_color_outside=:white,aspect_ratio=:equal,legendtitlefontcolor=:white) #formatting for plot frame
        currentT=t[i]/365/24/3600+tOffset
        titleString="Random $nBodies-Body Problem\nt: $(@sprintf("%0.2f",currentT)) years after start"
        if landscape==1
            titleString="Random $nBodies-Body Problem  |  t: $(@sprintf("%0.2f",currentT)) years after start"
        end
        p=plot!(xlabel="x: AU",ylabel="y: AU",title=titleString,
            legend=:best,xaxis=("x: AU",(limx[1],limx[2]),font(10,"Courier")),yaxis=("y: AU",(limy[1],limy[2]),font(10,"Courier")),
            grid=false,titlefont=font(18,"Courier"),legendfontsize=9,legendtitle="Mass (in solar masses)",legendtitlefontsize=10,top_margin=2mm,bottom_margin=2mm) #add in axes/title/legend with formatting
        plotNum+=1
        png(p,@sprintf("tmpPlots2/frame_%06d.png",plotNum))
        closeall() #close plots
        I+=skipRate
    end
    if sectionNum==1
        for n=1:nBodies
            xBack=xData[n][1:Int(floor(skipRate/10)):end]
            yBack=yData[n][1:Int(floor(skipRate/10)):end]
            push!(backData,[xBack,yBack]) #initialize backData
        end
    else
        for i=1:length(colors)
            currentColor=colors[i] #color is really a placeholder for which body this is
            for j=1:length(oldColors)
                if currentColor==oldColors[j] #this means we've found the original index (j)
                    xBack=xData[i][1:Int(floor(skipRate/10)):end]
                    yBack=yData[i][1:Int(floor(skipRate/10)):end]
                    backData[j][1]=vcat(backData[j][1],xBack) #backData has info from all bodies so we need to make sure to keep it consistent
                    backData[j][2]=vcat(backData[j][2],yBack) #stitch together new piece with old pieces (where applicable to surviving bodies)
                end
            end
        end
    end
    return plotNum,backData,I
end

function plotAll(landscape,nBodiesList,plotList,physInfoList,namesList,colorsList,maxFrame,tOffsets,dt;slowDown=1)
    frameOffset=0 #start at one because frame numbers start at 1
    backData=[]
    oldI=0
    for i=1:(length(nBodiesList)-1) #-1 because this list has n+1 things (includes body count after last step, which we don't need)
        currentNBodies=nBodiesList[i]
        println("plotting step $i of $(length(nBodiesList)-1) with $currentNBodies bodies")
        currentPlotData=plotList[i]
        currentPhysInfo=physInfoList[i]
        currentNames=namesList[i]
        #println("names at step $i = $currentNames")
        currentColors=colorsList[i]
        if i>1
            oldColors=colorsList[1]
        else
            oldColors=[]
        end
        currentTOffset=tOffsets[i]
        currentFrame,newBackData,IPlus=plotSection(landscape,i,backData,oldI,oldColors,[currentTOffset,frameOffset],dt,
                                        currentNBodies,currentPlotData,currentPhysInfo,currentNames,currentColors,
                                        maxFrame,slowDown=slowDown)
        backData=newBackData
        oldI+=IPlus
        frameOffset=currentFrame
    end
end

# STEP 1: get nBodies, colors, and labels
nBodies=parse(Int,ARGS[1])
println("running simulation for $nBodies bodies")
validColors=[]
open("juliaColors.txt","r") do f
    for line in eachline(f)
        push!(validColors,line)
    end
end

colorBool=parse(Int,ARGS[2])
function getUniqueColors(nBodies,validColors)
    accept=false
    while accept==false
        global randInd=rand(1:length(validColors),nBodies) #global because julia scope and I'm too lazy to do this right
        #println("I'm stuck here")
        #println("$(length(unique(randInd))), $(length(randInd))")
        if length(unique(randInd))==length(randInd)
            accept=true #make sure that all entries are unique so no colors are duplicated
            break
        end
    end
    colors=[validColors[i] for i in randInd]
    return colors
end
if colorBool==1
    println("let's pick colors!")
    colors=[]
    for n=1:nBodies
        accept=false
        while accept==false
            print("please enter a valid color ($n/$nBodies): ")
            color=readline()
            if color in validColors
                push!(colors,color)
                accept=true
            else
                println("I'm sorry, I don't know that color")
            end
        end
    end
elseif colorBool==2 #a nice set of color presets I like
    colorChoices=["dodgerblue","blueviolet","lightseagreen","gold","crimson","silver","hotpink",
                    "orange","thistle","yellowgreen","rosybrown","peachpuff","darksalmon","deeppink",
                    "paleturquoise","orchid","orangered","limegreen","goldenrod","fuchsia","lightsteelblue"]
    colors=getUniqueColors(nBodies,colorChoices)
else
    println("using randomly generated colors")
    colors=getUniqueColors(nBodies,validColors)
end

labelBool=parse(Int,ARGS[3])
if labelBool==1
    println("let's make some stars!")
    labels=[]
    mStart=[]
    for n=1:nBodies
        print("please enter a mass (solar mass units) for body $n/$nBodies: ")
        name=readline()
        push!(labels,name)
        push!(mStart,parse(Float64,name))
    end
else
    mStart=rand(1:1500,nBodies)./10 #n random masses between 0.1 and 150 solar masses
    #mStart=getMass(nBodies) #IMF way, but boring so not used right now
    labels=["$(mStart[i])" for i=1:nBodies]
end

names=copy(labels)
nBodies0=copy(nBodies)

#STEP 2: generate the data with given parameters (may add dt, stopCond, posRange as user-input things later)
nTrials,plotList,tOffsets,physInfoList,namesList,colorsList,nBodiesList=getInterestingNBody(nBodies,copy(mStart),names,colors,dt=0.00005,stopCond=[50,500],minTime=40,posRange=[-60,60],vRange=[-10e3,10e3])
#nBodiesList=[nBodies0-i for i=0:(nTrials-1)] #function doesn't return a list

#adding fake stars
function getMax(plotList,ind)
    MAX=0
    for i=1:length(plotList)
        for j=1:length(plotList[end][ind])
            localMax=maximum(plotList[i][ind][j])
            if localMax>MAX
                MAX=localMax
            end
        end
    end
    return MAX/1.5e11
end
function getMin(plotList,ind)
    MIN=0
    for i=1:length(plotList)
        for j=1:length(plotList[end][ind])
            localMin=minimum(plotList[i][ind][j])
            if localMin<MIN
                MIN=localMin
            end
        end
    end
    return MIN/1.5e11
end
minMax=zeros(4)
minMax[1],minMax[2]=getMax(plotList,1),getMax(plotList,2)
minMax[3],minMax[4]=getMin(plotList,1),getMin(plotList,2)
boxMin=floor(minimum(minMax))-500 #give 500 AU cushion for the full frame
boxMax=ceil(maximum(minMax))+500
bigNess=Int(ceil((boxMax-boxMin)/200)) #used to do 2500 in 200 AU
function initStars(bigNess)
    try
        numStars=2000*bigNess
        starsX=zeros(numStars)
        starsY=zeros(numStars)
        return numStars,starsX,starsY
    catch
        numStars=200000 #a big number but not so big it overloads RAM
        starsX=zeros(numStars)
        starsY=zeros(numStars)
        return numStars,starsX,starsY
    end
end
numStars,starsX,starsY=initStars(bigNess)
for i=1:numStars
    num=rand(boxMin:boxMax,2) #box size is 70 AU but we need some extra padding for movement
    starsX[i]=num[1]
    starsY[i]=num[2]
end
landscape=parse(Int,ARGS[4])
#STEP 3: plot data with parameters (note that dt and maxFrame must be duplicated as is, should probably have user pass them in and store to var)
plotAll(landscape,nBodiesList,plotList,physInfoList,namesList,colorsList,100,tOffsets,0.00005,slowDown=1)


import AbstractAlgebra.Generic: Perm, SymmetricGroup

# disambiguation
GroupsCore.order(::Type{I}, G::SymmetricGroup) where {I<:Integer} =
    convert(I, factorial(G.n))

# disambiguation
GroupsCore.order(::Type{I}, g::Perm) where {I<:Integer} =
    convert(I, foldl(lcm, length(c) for c in AbstractAlgebra.cycles(g)))

# correct the AA length:
Base.length(G::SymmetricGroup) = order(Int, G)

# genuinely new methods:
Base.IteratorSize(::Type{<:AbstractAlgebra.AbstractPermutationGroup}) = Base.HasLength()

function GroupsCore.gens(G::SymmetricGroup{I}) where {I}
    G.n == 1 && return eltype(G)[]
    if G.n == 2
        a = one(G)
        a.d[1], a.d[2] = 2, 1
        return [a]
    end
    a, b = one(G), one(G)
    circshift!(a.d, b.d, -1)
    b.d[1], b.d[2] = 2, 1
    return [a, b]
end

Base.deepcopy_internal(g::Perm, stackdict::IdDict) =
    Perm(Base.deepcopy_internal(g.d, stackdict), false)


#--------------------------------------------------
# Vector space algebra

# Unary ops
@inline Base.:-(a::TupleVector) = map(-, a)
@inline Base.:-(a::TupleVector{n,<:Number}) where n = map(Base.:-, a)

# Binary ops
# Between arrays
@inline Base.:+(a::TupleVector, b::TupleVector) = map(∑, a, b)
@inline Base.:+(a::AbstractArray, b::TupleVector) = map(∑, a, b)
@inline Base.:+(a::TupleVector, b::AbstractArray) = map(∑, a, b)

@inline Base.:+(a::TupleVector{n,<:Number}, b::TupleVector{n,<:Number}) where n = map(Base.:+, a, b)
@inline Base.:+(a::AbstractArray{<:Number}, b::TupleVector{n,<:Number}) where n = map(Base.:+, a, b)
@inline Base.:+(a::TupleVector{n,<:Number}, b::AbstractArray{<:Number}) where n = map(Base.:+, a, b)

@inline Base.:-(a::TupleVector, b::TupleVector) = map(-, a, b)
@inline Base.:-(a::AbstractArray, b::TupleVector) = map(-, a, b)
@inline Base.:-(a::TupleVector, b::AbstractArray) = map(-, a, b)

@inline Base.:-(a::TupleVector{n,<:Number}, b::TupleVector{n,<:Number}) where n = map(Base.:-, a, b)
@inline Base.:-(a::AbstractArray{<:Number}, b::TupleVector{n,<:Number}) where n = map(Base.:-, a, b)
@inline Base.:-(a::TupleVector{n,<:Number}, b::AbstractArray{<:Number}) where n = map(Base.:-, a, b)

# Scalar-array
@inline Base.:*(a::Number, b::TupleVector{n,<:Number}) where n = broadcast(Base.:*, a, b)
@inline Base.:*(a::TupleVector{n,<:Number}, b::Number) where n = broadcast(Base.:*, a, b)

@inline Base.:*(a, b::TupleVector) = broadcast(∏, a, b)
@inline Base.:*(a::TupleVector, b) = broadcast(∏, a, b)

@inline Base.:*(a::Expr, b::TupleVector) = broadcast(∏, Ref(a), b)
@inline Base.:*(a::TupleVector, b::Expr) = broadcast(∏, a, Ref(b))

@inline Base.:*(a::Symbol, b::TupleVector) = broadcast(∏, Ref(a), b)
@inline Base.:*(a::TupleVector, b::Symbol) = broadcast(∏, a, Ref(b))

@inline Base.:/(a::TupleVector{n,<:Number}, b::Number) where n = broadcast(Base.:/, a, b)
@inline Base.:\(a::Number, b::TupleVector{n,<:Number}) where n = broadcast(Base.:\, a, b)

@inline Base.:/(a::TupleVector, b) = broadcast(/, a, b)
@inline Base.:\(a, b::TupleVector) = broadcast(\, a, b)

#--------------------------------------------------
# Vector products
@inline LinearAlgebra.dot(a::TupleVector, b::TupleVector) = _vecdot(a, b, LinearAlgebra.dot)
@inline bilinear_vecdot(a::TupleVector, b::TupleVector) = _vecdot(a, b, Base.:*)

@inline function _vecdot(a::TupleVector{S}, b::TupleVector{S}, product) where S
    if S == 0
        za, zb = zero(eltype(a)), zero(eltype(b))
    else
        # Use an actual element if there is one, to support e.g. Vector{<:Number}
        # element types for which runtime size information is required to construct
        # a zero element.
        za, zb = zero(a[1]), zero(b[1])
    end
    ret = product(za, zb) + product(za, zb)
    @inbounds @simd for j = 1:S
        ret += product(a[j], b[j])
    end
    return ret
end

#--------------------------------------------------
# Norms
@inline LinearAlgebra.norm_sqr(v::TupleVector) = mapreduce(Base.abs2, Base.:+, v; init=zero(real(eltype(v))))

@inline LinearAlgebra.norm(a::TupleVector) = _norm(a)
@generated function _norm(a::TupleVector{S}) where S
    if S == 0
        return :(zero(real(eltype(a))))
    end
    expr = :(Base.abs2(a[1]))
    for j = 2:S
        expr = :($expr + Base.abs2(a[$j]))
    end
    return quote
        $(Expr(:meta, :inline))
        @inbounds return Base.sqrt($expr)
    end
end

_norm_p0(x) = x == 0 ? zero(x) : one(x)

@inline LinearAlgebra.norm(a::TupleVector, p::Real) = _norm(a, p)
@generated function _norm(a::TupleVector{S,T}, p::Real) where {S,T}
    if S == 0
        return :(zero(real(eltype(a))))
    end
    fun = T<:Number ? :(Base.abs) : :abs
    expr = :($fun(a[1])^p)
    for j = 2:S
        expr = :($expr + $fun(a[$j])^p)
    end
    expr_p1 = :($fun(a[1]))
    for j = 2:S
        expr_p1 = :($expr_p1 + $fun(a[$j]))
    end
    return quote
        $(Expr(:meta, :inline))
        if p == Inf
            return mapreduce(Base.abs, max, a)
        elseif p == 1
            @inbounds return $expr_p1
        elseif p == 2
            return LinearAlgebra.norm(a)
        elseif p == 0
            return mapreduce(_norm_p0, $(T<:Number ? :(Base.:+) : :∑), a)
        else
            @inbounds return $(T<:Number ? :(Base.:^) : :^)($expr,$(T<:Number ? :(Base.inv) : :inv)(p))
        end
    end
end

@inline LinearAlgebra.normalize(a::TupleVector) = ∏(inv(LinearAlgebra.norm(a)),a)
@inline LinearAlgebra.normalize(a::TupleVector, p::Real) = ∏(inv(LinearAlgebra.norm(a, p)),a)

@inline LinearAlgebra.normalize!(a::TupleVector) = (a .*= inv(LinearAlgebra.norm(a)); return a)
@inline LinearAlgebra.normalize!(a::TupleVector, p::Real) = (a .*= inv(LinearAlgebra.norm(a, p)); return a)

@inline LinearAlgebra.normalize(a::TupleVector{n,<:Number}) where n = Base.:*(Base.inv(LinearAlgebra.norm(a)),a)
@inline LinearAlgebra.normalize(a::TupleVector{n,<:Number}, p::Real) where n = Base.:*(Base.inv(LinearAlgebra.norm(a, p)),a)

@inline LinearAlgebra.normalize!(a::TupleVector{n,<:Number}) where n = (a .*= Base.inv(LinearAlgebra.norm(a)); return a)
@inline LinearAlgebra.normalize!(a::TupleVector{n,<:Number}, p::Real) where n = (a .*= Base.inv(LinearAlgebra.norm(a, p)); return a)


"""
    FDEsolver(F, tSpan, y0, β, ::Nothing, par...; h = 0.01, nc = 3, StopIt = "Standard", tol = 10e-10, itmax = 10)

Solves fractional differential equations with a predictor-corrector approach.

    FDEsolver(F, tSpan, y0, β, J, par...; h = 0.01, nc = 3, StopIt = "Standard", tol = 10e-10, itmax = 10)

Additionally takes the Jacobian of the system J to evaluate the solution.

# Arguments
- `F`: the right side of the system of differential equations. It must be expressed
   in the form of a function and return a vector function with the same number of
   entries of order of derivatives. This function can also include a vector of
   parameters: `par...` .
- `tSpan::Vector{<:Real}`: the time span along which computation is performed.
- `y0`: the initial values in the form of a row vector (`Vector{<:Real}`) for ``β <= 1``
   and a matrix (`Matrix{<:Real}`) for ``β > 1``, where each column corresponds to the
   initial values of one differential equation and each row to the order of derivation.
- `β::Vector{<:Real}`: the orders of derivation in the form of a row vector, where
   each element corresponds to the order of one differential equation. It can take
   decimal as well as integer values.
- `J::Matrix{<:Real}`: the Jacobian of F. If not provided, its position within the
   function definition has to be replaced with `nothing`.
- `par...`: additional parameters for the function F.
- `h::Number`: the step size for correction.
- `nc::Int64`: the desired number of corrections.
- `StopIt::String`: the method to stop correction. It can take either "Standard"
   (by default) or "Convergence". In the former case, the function will repeat
   correction as many times as specified in nc; in the latter case, correction will
   stop only when tolerance (tol) or the iteration max (itmax) is reached.
- `tol::Float64`: the tolerance.
- `ìtmax::Int64`: the maximal number of iterations.
"""
module FdeSolver

using SpecialFunctions
using LinearAlgebra

"""
    greet()
Motivates to get started with FdeSolver.
"""
greet() = print("Hey, let's solve some FDEs!")

include("main.jl")
include("main_Jacob.jl")
include("SupFuns.jl")

export(FDEsolver)

end


# Neural stick-breaking processe

struct MeanFieldSBC <: StochasticLayer
    kuma
    gumbel
end

function MeanFieldSBC(i_dim::Int, k_init::Int=0; kwargs...)
    kuma = KumaraswamyLazyDense(i_dim, k_init; kwargs...)
    gumbel = GumbelBernoulliLogitLazyDense(i_dim, k_init; kwargs...)
    return MeanFieldSBC(kuma, gumbel)
end

function (enc::MeanFieldSBC)(i, k::Int; lowerbound=(FT == Float64 ? FT(0.1) : FT(0.2)))
    dist_kuma = enc.kuma(i, k)
    dist_nu = BatchKumaraswamy(dist_kuma.a .+ lowerbound, dist_kuma.b .+ lowerbound)
    dist_gumbel = enc.gumbel(i, k)
    return dist_nu, dist_gumbel
end

struct StructuredSBC <: StochasticLayer
    a
    b
    gumbel
    α
    β
end

"""
Helper function to initialise StructuredSBC with an IBP.
"""
function StructuredSBC(ibp::IBP{T}, i_dim::Int, k_init::Int=0; kwargs...) where {T<:AbstractFloat}
    return StructuredSBC(i_dim, k_init, ibp.α, one(T); kwargs...)
end

function StructuredSBC(i_dim::Int, k_init::Int=0, α::AbstractFloat=one(FT), β::AbstractFloat=one(FT); kwargs...)
    a = Knet.Param(AT(zeros(FT, k_init, 1) .+ invsoftplus(α)))
    b = Knet.Param(AT(zeros(FT, k_init, 1) .+ invsoftplus(β)))
    gumbel = GumbelBernoulliLogitLazyDense(i_dim, k_init; kwargs...)
    return StructuredSBC(a, b, gumbel, α, β)
end

function (enc::StructuredSBC)(i, k::Int=size(enc.a, 1); α::AbstractFloat=enc.α, β::AbstractFloat=enc.β, lowerbound=(FT == Float64 ? FT(0.1) : FT(0.2)))
    if k > size(enc.a, 1)
        k_diff = k - size(enc.a, 1)
        enc.a.value = vcat([enc.a.value, AT(zeros(FT, k_diff, 1) .+ invsoftplus(α))]...)
        enc.b.value = vcat([enc.b.value, AT(zeros(FT, k_diff, 1) .+ invsoftplus(β))]...)
    end
    dist_nu = BatchKumaraswamy(softplus.(enc.a[1:k,:]) .+ lowerbound,
                               softplus.(enc.b[1:k,:]) .+ lowerbound)
    dist_gumbel = enc.gumbel(i, k)
    return dist_nu, dist_gumbel
end


module Joe

using LinearAlgebra, Statistics, Distributions, Combinatorics
export min_sq, multiple_regression, insert_ones
# Write your package code here.
include("chap1.jl")
include("chap2.jl")
include("chap3.jl")
include("chap4.jl")
include("chap5.jl")
include("chap6.jl")
include("chap7.jl")
end


function creategrid(;xrange::(Real, Real)=(-1,1), yrange::(Real, Real)=(-1,1), zrange::(Real, Real)=(-1,1), camera=pcamera)
	xyplane = genquad(Vec3(0, 0, 0), Vec3(1, 0, 0), Vec3(0, 1, 0))
	zyplane = genquad(Vec3(0, 0, 0), Vec3(0, 0, 1.2), Vec3(0, 1, 0))
	zxplane = genquad(Vec3(0, 1, 0), Vec3(0, 0, 1.2), Vec3(1, 0, 0))

	v,uv,n,i = mergemesh(xyplane, zyplane, zxplane)

	grid = RenderObject(@compat(Dict(
			:vertexes 			  	=> GLBuffer(v),
			:indexes			   	=> indexbuffer(i),
			#:grid_color 		  => Float32[0.1,.1,.1, 1.0],
			:bg_color 			  	=> Input(Vec4(1, 1, 1, 0.01)),
			:grid_thickness  		=> Input(Vec3(2)),
			:gridsteps  		  	=> Input(Vec3(10)),
			:mvp 				    => camera.projectionview
		)), gridshader)
	prerender!(grid, glDisable, GL_DEPTH_TEST, glDisable, GL_CULL_FACE, enabletransparency)
	postrender!(grid, render, grid.vertexarray)
	return grid
end

initgrid() = global gridshader = TemplateProgram(joinpath(shaderdir,"grid.vert"), joinpath(shaderdir,"grid.frag"))
init_after_context_creation(initgrid)

export creategrid


Ω = Circle(-1,1) × Circle(-1,1) × Circle(-1,1)
ex, ey, ez = (3,5,7)[1:ndims(Ω)]
(npx, npy, npz) = (1,1,1)[1:ndims(Ω)]
ClimateMachine.gpu_allowscalar(true)
grid = DiscontinuousSpectralElementGrid(Ω, elements = (ex, ey, ez), polynomialorder = (npx, npy, npz), array = ArrayType)
gridhelper = GridHelper(grid)
x, y, z = coordinates(grid)
ϕ =  ScalarField(copy(x), gridhelper)
@testset "Check Interpolation 3D" begin
    ϕ .= x
    @test ϕ(0.1, 0.2, 0.3) ≈ 0.1
    ϕ .= y
    @test ϕ(0.1, 0.2, 0.3) ≈ 0.2
    ϕ .= z
    @test ϕ(0.1, 0.2, 0.3) ≈ 0.3
end
##
Ω = Circle(0,2) × Circle(-4,2)
elements =  (3,5,7)[1:ndims(Ω)]
polynomialorder = (1,1,1)[1:ndims(Ω)]
ClimateMachine.gpu_allowscalar(true)
grid = DiscontinuousSpectralElementGrid(Ω, elements = elements, polynomialorder = polynomialorder, array = ArrayType)
gridhelper = GridHelper(grid)
x, y, z = coordinates(grid)
ϕ =  ScalarField(copy(x), gridhelper)
@testset "Check Interpolation 2D" begin
    ϕ .= x
    @test ϕ(0.1, 0.2) ≈ 0.1
    ϕ .= y
    @test ϕ(0.1, 0.2) ≈ 0.2
end


#    Copyright (c) 2021 J.A. Duffek
#    Copyright (c) 2000 D.M. Spink
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, see <http://www.gnu.org/licenses/>.

"""
    nrb4surf(p11::Vector{F},p12::Vector{F},p21::Vector{F},p22::Vector{F}
            )::NURBS2D where {F<:AbstractFloat}

Constructs a bilinear surface defined by four coordinates.

The position of the corner points

    ^ V direction
    |
    ----------------
    |p21        p22|
    |              |
    |    SRF       |
    |              |
    |p11        p12|
    -------------------> U direction

# Arguments:
- `p11`: Cartesian coordinate of the lhs bottom corner point.
- `p12`: Cartesian coordinate of the rhs bottom corner point.
- `p21`: Cartesian coordinate of the lhs top corner point.
- `p22`: Cartesian coordinate of the rhs top corner point.

# Output:
- `srf`: NURBS bilinear surface, see [`nrbmak`](@ref).

# Examples:
```julia
julia> srf = nrb4surf(p11,p12,p21,p22)
```
"""
function nrb4surf(p11::Vector{F},p12::Vector{F},p21::Vector{F},p22::Vector{F}
                 )::NURBS2D where {F<:AbstractFloat}
coefs = zeros(F,(4,2,2));
coefs[4,:,:] .= 1.0;
coefs[1:length(p11),1,1] = p11;
coefs[1:length(p12),1,2] = p12;
coefs[1:length(p21),2,1] = p21;
coefs[1:length(p22),2,2] = p22;
knots = [[0.0;0.0;1.0;1.0],[0.0;0.0;1.0;1.0]];
return nrbmak(coefs, knots);
end # nrb4surf


@testset "Cylinder Constraint" begin

    # Test CylinderConstraint
    T = Float64
    n = 4
    P = 5
    x = 2
    y = 3
    z = 4
    X0 = SVector{n,T}([13.0, 1.0, 1.0, 2.0])
    p1 = SVector{P,T}([ 1.0,  1.0,  1.0,  0.0,  1.0])
    p2 = SVector{P,T}([ 0.0,  0.0,  0.0,  1.0,  0.0])
    p3 = SVector{P,T}([ 1.0,  1.0,  3.0,  1.0,  2.0])
    v = SVector{P,Symbol}([:z, :z, :z, :x, :y])
    l = SVector{P,T}([ 5.0,  2.0,  0.5,  2.0,  10.0])
    r = SVector{P,T}([ 3.0,  2.0,  7.0,  3.0,  1.0])

    con = CylinderConstraint(n,p1,p2,p3,v,l,r,x,y,z)
    TrajectoryOptimization.evaluate(con,X0)

    @test norm(TrajectoryOptimization.evaluate(con,X0) - [8.0, 3.0, 0.0, 8.0, 1.0], 1) < 1e-10
    # @test (@ballocated $(TrajectoryOptimization.evaluate)($con,$X0)) == 0

    function easy_jacobian(con,X)
        function local_evaluate(X)
            return TrajectoryOptimization.evaluate(con,X)
        end
        return ForwardDiff.jacobian(local_evaluate, X)
    end
    ∇c_easy = easy_jacobian(con, X0)
    ∇c = zeros(MMatrix{P,n,T,P*n})
    jacobian!(∇c,con,X0)
    @test norm(∇c - ∇c_easy, 1) < 1e-10
    # @test (@ballocated $jacobian!($∇c,$con,$X0)) == 0

end


using firstTask, Test

@testset "goldenSection" begin
  f(x) = x^4 + 3x^3 + x^2 + sin(x)
  a = -10
  b = 10
  tol = 1e-6
  @test goldenSection(f, a, b, tol) == (-1.9712755845857126, -4.915216432826815)

  f(x) = 3(x - 4)^2
  a = 0
  b = 10
  tol = 1e-6
  @test goldenSection(f, a, b, tol) == (4.000000004370129, 5.729408629607859e-17)
end


using StochasticDiffEq#, Plots
Random.seed!(100)
prob = prob_sde_2Dlinear

## Solve and plot
println("Solve and Plot")
#Let the solver determine the initial stepsize for you!
sol =solve(prob,SRI())

TEST_PLOT && plot(sol,plot_analytic=true)
#gui()

#Make sure it does a good job
sol.t[2] > 1e-7


# Population density benchmark from https://github.com/SWI-Prolog/bench
# Author:  David H. D. Warren
#
# query population and area database to find coun-
# tries of approximately equal population density

clauses = @julog [
    query([C1,D1,C2,D2]) <<=
        density(C1,D1) &
        density(C2,D2) &
        (D1 > D2) &
        is(T1, 20*D1) &
        is(T2, 21*D2) &
        (T1 < T2),

    density(C,D) <<=
        pop(C,P) &
        area(C,A) &
        is(D, (P*100)/A),

    # populations in 100000's
    pop(china,	8250) <<= true,
    pop(india,	5863) <<= true,
    pop(ussr,	2521) <<= true,
    pop(usa,	2119) <<= true,
    pop(indonesia,	1276) <<= true,
    pop(japan,	1097) <<= true,
    pop(brazil,	1042) <<= true,
    pop(bangladesh,	 750) <<= true,
    pop(pakistan,	 682) <<= true,
    pop(w_germany,	 620) <<= true,
    pop(nigeria,	 613) <<= true,
    pop(mexico,	 581) <<= true,
    pop(uk,		 559) <<= true,
    pop(italy,	 554) <<= true,
    pop(france,	 525) <<= true,
    pop(philippines, 415) <<= true,
    pop(thailand,	 410) <<= true,
    pop(turkey,	 383) <<= true,
    pop(egypt,	 364) <<= true,
    pop(spain,	 352) <<= true,
    pop(poland,	 337) <<= true,
    pop(s_korea,	 335) <<= true,
    pop(iran,	 320) <<= true,
    pop(ethiopia,	 272) <<= true,
    pop(argentina,	 251) <<= true,

    # areas in 1000's of square miles
    area(china,     3380) <<= true,
    area(india,     1139) <<= true,
    area(ussr,      8708) <<= true,
    area(usa,       3609) <<= true,
    area(indonesia,  570) <<= true,
    area(japan,      148) <<= true,
    area(brazil,    3288) <<= true,
    area(bangladesh,  55) <<= true,
    area(pakistan,   311) <<= true,
    area(w_germany,   96) <<= true,
    area(nigeria,    373) <<= true,
    area(mexico,     764) <<= true,
    area(uk,          86) <<= true,
    area(italy,      116) <<= true,
    area(france,     213) <<= true,
    area(philippines, 90) <<= true,
    area(thailand,   200) <<= true,
    area(turkey,     296) <<= true,
    area(egypt,      386) <<= true,
    area(spain,      190) <<= true,
    area(poland,     121) <<= true,
    area(s_korea,     37) <<= true,
    area(iran,       628) <<= true,
    area(ethiopia,   350) <<= true,
    area(argentina, 1080) <<= true,
]

goal = @julog(query([C1,D1,C2,D2]))
println("Population (single run):")
@time resolve(goal, clauses)

function benchmark(n)
    for i = 1:n
        resolve(@julog(query([C1,D1,C2,D2])), clauses)
    end
end

println("Population (100 runs):")
@time benchmark(100)


"""
    fit(template::objective_function_template, data::Any, theta::Vector; estimation_method::String = "M", br_method::String = "implicit_trace", optim_method = LBFGS(), optim_options = Optim.Options())

Fit an [`objective_function_template`](@ref) on `data` using M-estimation ([keyword argument](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1) `estimation_method = "M"`; default) or RBM-estimation (reduced-bias M estimation; [Kosmidis & Lunardon, 2020](http://arxiv.org/abs/2001.03786); [keyword argument](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1) `estimation_method = "RBM"`). Bias reduction is either through the maximization of the bias-reducing penalized objective in Kosmidis & Lunardon (2020) ([keyword argument](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1) `br_method = "implicit_trace"`; default) or by subtracting an estimate of the bias from the M-estimates ([keyword argument](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1) `br_method = "explicit_trace"`). The bias-reducing penalty is constructed internally using automatic differentiation (using the [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) package), and the bias estimate using a combination of automatic differentiation (using the [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) package) and numerical differentiation (using the [FiniteDiff](https://github.com/JuliaDiff/FiniteDiff.jl) package).

The maximization of the objective or the penalized objective is done using the [**Optim**](https://github.com/JuliaNLSolvers/Optim.jl) package. Optimization methods and options can be supplied directly through the [keyword arguments](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1) `optim_method` and `optim_options`, respectively. `optim_options` expects an object of class `Optim.Options`. See the [Optim documentation](https://julianlsolvers.github.io/Optim.jl/stable/#user/config/#general-options) for more details on the available options.
"""
function fit(template::objective_function_template,
             data::Any,
             theta::Vector;
             estimation_method::String = "M",
             br_method::String = "implicit_trace",
             optim_method = LBFGS(),
             optim_options = Optim.Options())
    if (estimation_method == "M")
        br = false
    elseif (estimation_method == "RBM")
        if (br_method == "implicit_trace")
            br = true
        elseif (br_method == "explicit_trace")
            br = false
        else
            error(br_method, " is not a recognized bias-reduction method")
        end
    else
        error(estimation_method, " is not a recognized estimation method")
    end
    obj = beta -> -objective_function(beta, data, template, br)
    out = optimize(obj, theta, optim_method, optim_options)
    if (estimation_method == "M")
        theta = out.minimizer
    elseif (estimation_method == "RBM")
        if (br_method == "implicit_trace")
            theta = out.minimizer
        elseif (br_method == "explicit_trace")
            quants = obj_quantities(out.minimizer, data, template, true)
            jmat_inv = quants[2]
            adjustment = FiniteDiff.finite_difference_gradient(beta -> obj_quantities(beta, data, template, true)[1], out.minimizer)
            theta = out.minimizer + jmat_inv * adjustment
            br = true
        end
    end
    GEEBRA_results(out, theta, data, template, br, true, br_method)
end

"""
    fit(template::estimating_function_template, data::Any, theta::Vector; estimation_method::String = "M", br_method::String = "implicit_trace", nlsolve_arguments...)

Fit an [`estimating_function_template`](@ref) on `data` using M-estimation ([keyword argument](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1) `estimation_method = "M"`; default) or RBM-estimation (reduced-bias M estimation; [Kosmidis & Lunardon, 2020](http://arxiv.org/abs/2001.03786); [keyword argument](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1) `estimation_method = "RBM"`). Bias reduction is either through the solution of the empirically adjusted estimating functions in Kosmidis & Lunardon (2020) ([keyword argument](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1) `br_method = "implicit_trace"`; default) or by subtracting an estimate of the bias from the M-estimates ([keyword argument](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1) `br_method = "explicit_trace"`). The bias-reducing adjustments and the bias estimate are constructed internally using automatic differentiation (using the [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) package).

The solution of the estimating equations or the adjusted estimating equations is done using the [**NLsolve**](https://github.com/JuliaNLSolvers/NLsolve.jl) package. Arguments can be passed directly to `NLsolve.nlsolve` through [keyword arguments](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1). See the [NLsolve README](https://github.com/JuliaNLSolvers/NLsolve.jl) for more information on available options.
"""
function fit(template::estimating_function_template,
             data::Any,
             theta::Vector;
             estimation_method::String = "M",
             br_method::String = "implicit_trace",
             nlsolve_arguments...)
    if (estimation_method == "M")
        br = false
    elseif (estimation_method == "RBM")
        if (br_method == "implicit_trace")
            br = true
        elseif (br_method == "explicit_trace")
            br = false
        else
            error(br_method, " is not a recognized bias-reduction method")
        end
    else
        error(estimation_method, " is not a recognized estimation method")
    end
    ef = get_estimating_function(data, template, br)
    out = nlsolve(ef, theta; nlsolve_arguments...)
    if (estimation_method == "M")
        theta = out.zero
    elseif (estimation_method == "RBM")
        if (br_method == "implicit_trace")
            theta = out.zero
        elseif (br_method == "explicit_trace")
            quants = ef_quantities(out.zero, data, template, true)
            adjustment = quants[1]
            jmat_inv = quants[2]
            theta = out.zero + jmat_inv * adjustment
            br = true
        end
    end
    GEEBRA_results(out, theta, data, template, br, false, br_method)
end


# kruskal_wallis.jl
# Kruskal-Wallis rank sum test
#
# Copyright (C) 2014   Christoph Sawade
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

export KruskalWallisTest

struct KruskalWallisTest <: HypothesisTest
    n_i::Vector{Int}         # number of observations in each group
    df::Int                  # degrees of freedom
    R_i::Vector{Float64}     # rank sums
    H::Float64               # test statistic: chi-square statistic
    tie_adjustment::Float64  # adjustment for ties
end

"""
    KruskalWallisTest(groups::AbstractVector{<:Real}...)

Perform Kruskal-Wallis rank sum test of the null hypothesis that the `groups`
``\\mathcal{G}`` come from the same distribution against the alternative hypothesis that
that at least one group stochastically dominates one other group.

The Kruskal-Wallis test is an extension of the Mann-Whitney U test to more than two groups.

The p-value is computed using a ``χ^2`` approximation to the distribution of the test
statistic ``H_c=\\frac{H}{C}``:
```math
    \\begin{align*}
    H & = \\frac{12}{n(n+1)} \\sum_{g ∈ \\mathcal{G}} \\frac{R_g^2}{n_g} - 3(n+1)\\\\
    C & = 1-\\frac{1}{n^3-n}\\sum_{t ∈ \\mathcal{T}} (t^3-t),
    \\end{align*}
```
where ``\\mathcal{T}`` is the set of the counts of tied values at each tied position, ``n``
is the total number of observations across all groups, and ``n_g`` and ``R_g`` are the number of
observations and the rank sum in group ``g``, respectively. See references for further
details.

Implements: [`pvalue`](@ref)

# References

  * Meyer, J.P, Seaman, M.A., Expanded tables of critical values for the Kruskal-Wallis
    H statistic. Paper presented at the annual meeting of the American Educational Research
    Association, San Francisco, April 2006.

# External links

  * [Kruskal-Wallis test on Wikipedia
    ](https://en.wikipedia.org/wiki/Kruskal-Wallis_one-way_analysis_of_variance)
"""
function KruskalWallisTest(groups::AbstractVector{T}...) where T<:Real
    (H, R_i, tieadj, n_i) = kwstats(groups...)
    if length(groups)<=3 && any(n_i .< 6)
        @warn("This test is only asymptotically correct and might be inaccurate for the given group size")
    end
    df = length(groups) - 1
    KruskalWallisTest(n_i, df, R_i, H, tieadj)
end

testname(::KruskalWallisTest) = "Kruskal-Wallis rank sum test (chi-square approximation)"
population_param_of_interest(x::KruskalWallisTest) = ("Location parameters", "all equal", NaN) # parameter of interest: name, value under h0, point estimate
default_tail(test::KruskalWallisTest) = :right

function show_params(io::IO, x::KruskalWallisTest, ident)
    println(io, ident, "number of observation in each group: ", x.n_i)
    println(io, ident, "χ²-statistic:                        ", x.H)
    println(io, ident, "rank sums:                           ", x.R_i)
    println(io, ident, "degrees of freedom:                  ", x.df)
    println(io, ident, "adjustment for ties:                 ", x.tie_adjustment)
end

pvalue(x::KruskalWallisTest) = pvalue(Chisq(x.df), x.H; tail=:right)

## helper

# Get H, rank sums, and tie adjustment for Kruskal-Wallis test
function kwstats(groups::AbstractVector{T}...) where T<:Real
    n_i = [length(g) for g in groups]
    n = sum(n_i)

    # get ranks and adjustment for ties
    (ranks, tieadj) = tiedrank_adj([groups...;])
    C = 1-tieadj/(n^3 - n)

    # compute rank sums
    R_i = Vector{Float64}(undef, length(groups))
    n_end = 0
    for i in 1:length(groups)
        R_i[i] = sum(ranks[n_end+1:n_end+n_i[i]])
        n_end += n_i[i]
    end

    # compute test statistic and correct for ties
    H = 12 * sum(R_i .^ 2 ./ n_i) / (n * (n + 1)) - 3 * (n + 1)
    H /= C

    (H, R_i, C, n_i)
end


"""
```
function chand_recursion(y::Matrix{S}, T::Matrix{S}, R::Matrix{S}, C::Vector{S},
                         Q::Matrix{S}, Z::Matrix{S}, D::Vector{S}, H::Matrix{S},
                         s_pred::Vector{S} = Vector{S}(0),
                         P_pred::Matrix{S} = Matrix{S}(0,0);
                         allout::Bool = false, Nt0::Int = 0,
                         tol::Float64 = 0.0) where {S<:AbstractFloat}
```
Implementation of Chandrasekhar Recursions, a faster algorithm for likelihood
evaluation in a linear Gaussian state space system. Based on code by Ed Herbst.

State Transition equation:
```
s_t = TT * s_{t-1} + RR * ε_t, ε_t ~ iid N(0, QQ)
```
Measurement equation:
```
y_t = DD + ZZ * s_t + η_t, η_t ~ iid N(0, HH)
```
Returns log-likelihood evaluation.
"""
function chand_recursion(y::Matrix{S}, T::Matrix{S}, R::Matrix{S}, C::Vector{S},
                         Q::Matrix{S}, Z::Matrix{S}, D::Vector{S}, H::Matrix{S},
                         s_pred::Vector{S} = Vector{S}(0),
                         P_pred::Matrix{S} = Matrix{S}(0,0);
                         allout::Bool = false, Nt0::Int = 0,
                         tol::Float64 = 0.0) where {S<:AbstractFloat}
    converged = false

    # Dimensions
    Ns     = size(T, 1) # number of states
    Ny, Nt = size(y)    # number of periods of data

    # Initialize s_pred and P_pred
    if isempty(s_pred) || isempty(P_pred)
        s_pred, P_pred = init_stationary_states(T, R, C, Q)
    end

    # V_{t|t-1} = Var(y_t|y_{1:t-1})
    V_pred    = Z * P_pred * Z' + H
    V_pred    = 0.5 * (V_pred + V_pred')
    invV_pred = inv(V_pred)

    invV_t1 = invV_pred

    # We write ΔP in terms of W_t and M_t: ΔP = W_t * M_t * W_t'
    W_t      = T * P_pred * Z'   # (Eq 12)
    M_t      = -invV_pred        # (Eq 13)
    kal_gain = W_t * invV_pred

    # Initialize loglikelihoods to zeros so that the ones we don't update
    # (those in presample) contribute zero towards sum of loglikelihoods.
    loglh    = Vector{Float64}(Nt)
    zero_vec = zeros(Ny)
    ν_t      = zero_vec
    P        = P_pred

    for t in 1:Nt
        # Step 1: Compute forecast error, ν_t and evaluate likelihoood
        yhat = Z * s_pred + D
        ν_t  = y[:, t] - yhat

        # Symptom of a bad parameter draw: simply tells callee to redraw
        if !isposdef(V_pred)
            return NaN
        else
            loglh[t] = logpdf(MvNormal(zero_vec, V_pred), ν_t)
        end

        # Step 2: Compute s_{t+1} using Eq. 5
        if t < Nt
            s_pred = T * s_pred + kal_gain * ν_t
            if allout
                P = P + W_t * M_t * W_t'
            end
        end

        if !converged
            # Intermediate calculations to re-use
            ZW_t  = Z * W_t
            MWpZp = M_t * (ZW_t')
            TW_t  = T * W_t

            # Step 3: Update forecast error variance F_{t+1} (Eq 19)
            V_t1      = V_pred
            invV_t1   = invV_pred
            V_pred    = V_pred + ZW_t * MWpZp        # F_{t+1}
            V_pred    = 0.5 * (V_pred + V_pred')
            invV_pred = inv(V_pred)

            # Step 4: Update Kalman Gain (Eq 20). Recall kalgain = K_t * V_t⁻¹
            # Kalgain_{t+1} = (Kalgain_t*V_{t-1} + T*W_t*M_t*W_t'*Z')V_t⁻¹
            kal_gain1 = kal_gain
            kal_gain  = (kal_gain * V_t1 + TW_t * MWpZp) * invV_pred

            # Step 5: Update W
            W_t = TW_t - kal_gain * ZW_t            # W_{t+1}

            # Step 6: Update M
            M_t = M_t + MWpZp * invV_t1 * MWpZp'    # M_{t+1}
            M_t = 0.5 * (M_t + M_t')

            # 'Fast Kalman' Algorithm - shut off when tol = 0
            if (maximum(abs.(kal_gain - kal_gain1)) < tol)
                converged = true
            end
        end
    end

    loglh = remove_presample!(Nt0, loglh)

    if allout
        s_TT = s_pred + P' * Z' * invV_t1 * (ν_t)
        P_TT = P - P' * Z' * invV_t1 * Z * P
        return sum(loglh), loglh, s_TT, P_TT
    else
        return sum(loglh)
    end
end

function remove_presample!(Nt0::Int, loglh::Vector{S}) where {S<:AbstractFloat}
    out = remove_presample!(Nt0, loglh, Array{Float64, 2}(0,0),
                            Array{Float64, 3}(0,0,0), Array{Float64, 2}(0,0),
                            Array{Float64, 3}(0,0,0), outputs = [:loglh])
    return out[1]
end


using Flux
using Flux: Chain, Dense, params, train!, glorot_uniform
using Flux.Tracker: Params, gradient, update!

f = x -> (1/(0.82+x))^12 -2 * (1/(.82+x))^6
#f = x -> sin.(x.*7.0)

nsamples = 200
XS = rand(nsamples)
YS = f.(XS)
#samples = zip(XS,YS)
#samples = [(x, f(x)) for x in XS]
#samples= [(XS, YS)]

xerr = range(0, 1, length=1000)
eval_nn(model, x) = model(x).data
maxfit_nn(model) = maximum(abs, f.(XS) - model(XS))

function train_nn(nnodes, η, samples, XX, YY)
    model = Chain(Dense(nnodes, nnodes, Flux.tanh),   # first hidden layer
                  Dense(nnodes, 1))      # output layer

    loss(x, y) = Flux.mse(model(x), y)

    θ = Flux.params(model)
    #println(θ)

    for i=1:1000
        #Flux.train!(loss, θ, samples, Descent(η))
        Flux.train!(loss, θ, repeated((samples,), 2_000), ADAM(0.05))
        if i%100==1
            @show loss(XX, YY)
            @show maxfit_nn(model)
        end
    end
    return model
end

η = 0.00001 # Learning rate (should be smaller corresponding to the complex regression times the # of nodes)
NN_nodes = 3:5
errs_nn = Float64[]

for N in NN_nodes
    XX = repeat(XS,1,N)
    XX = XX'
    samples = [([x for i=1:N], f(x)) for x in XS]
    @show N
    nn = train_nn(N,η,samples,XX,YS)
    #println(Flux.params(nn))
    #@show e = err_nn(nn)
    #push!(errs_nn, e)
end
@show errs_nn


## Abstract  types for  orthogonal  polynomials

export Basis

## Has An(P), Bn(P), Cn(P)
abstract type AbstractOrthogonalPolynomial{T} <: AbstractSpecialPolynomial{T} end
abstract type AbstractContinuousOrthogonalPolynomial{T} <: AbstractOrthogonalPolynomial{T} end
abstract type AbstractDiscreteOrthogonalPolynomial{T} <: AbstractOrthogonalPolynomial{T} end

"""
    AbstractOrthogonalPolynomial{T}

Type to represent systems of orthogonal polynomials. These polynomials have  several properties, including an accompanying inner product satsifying  `⟨yᵢ, yⱼ⟩ = cᵢδᵢⱼ`.

In addition to methods inherited from the underlying `AbstractPolynomial`  type, orthogonal polynomial  types may have methods   `weight_function`, `generating_function`, `leading_term`, norm2`, `jacobi_matrix`, and `gauss_nodes_weights`,  though none are  exported.

Subtypes of `AbstractCOP <: AbstractOrthogonalPolynomial` utilize the fact that the basis  polynomials  satisfy

`(ax² + bx + c)yᵢ''(x) + (dx+e)yᵢ'(x) + λᵢyᵢ(x) = 0` (or a discrete analogue)

where the structural relations are functions of `a,b,c,d,e`. These allow default definitions for polynomial evaluation,   addition, multiplication, differentiation, integration, and  conversion to and from  the `Polynomial` type (the `FallingFactorial` type in the discrete  c case),

A key structural relation is the three-term recursion,  `yᵢ₊₁ =  (Aᵢx +  Bᵢ)yᵢ -  Cᵢyᵢ₋₁`. For systems  specfied by  a  weight function, the  values of `Aᵢ`, `Bᵢ`, and `Cᵢ` can  be  generated, yielding formulas for polynomial evaluation, addition, and conversion to the `Polynomial`  type throughe evaluation.
"""
AbstractOrthogonalPolynomial

##
## --------------------------------------------------
##
function Polynomials.fromroots(P::Type{<:AbstractOrthogonalPolynomial}, roots::AbstractVector; var::Polynomials.SymbolLike = :x)
    x = variable(P, var)
    p =  prod(x - r for r in roots)
end

Base.convert(P::Type{<:AbstractOrthogonalPolynomial}, c::Number) = c * one(P)
Base.one(P::Type{<:AbstractOrthogonalPolynomial}, var::Polynomials.SymbolLike=:x) =   basis(P,0, var) / k0(P)

Polynomials.variable(P::Type{<:AbstractOrthogonalPolynomial},  var::Polynomials.SymbolLike=:x) =
    (basis(P,1,var) / k0(P) - Bn(P,0)) / An(P,0)

## Evaluation

# from type, cs, x
function clenshaw_eval(P::Type{<:AbstractOrthogonalPolynomial{T}}, cs, x::S) where {T,S}
    N = length(cs)
    p₀ = k0(P)
    R = promote_type(promote_type(T,S), typeof(An(P,0)))
    N == 0 && return zero(R)
    N == 1 && return (cs[1] * p₀) * one(R)

    Δ0::R = cs[end - 1]
    Δ1::R = cs[end]
    @inbounds for i in N-1:-1:2
        Δ0, Δ1 = cs[i - 1] - Δ1 * Cn(P, i-1), Δ0 + Δ1 * muladd(x, An(P,i-1),Bn(P,i-1))
    end
    p₁ =  muladd(x, An(P,0),  Bn(P,0)) * p₀
    return Δ0 * p₀  + Δ1 * p₁
end

## Connection/Linearization

#  Structs for connection, linearization (connection.jl)
struct Connection{P,Q}
    n::Int
    k::Int
end

struct Linearization{P,V}
    l::Int
    n::Int
    m::Int
end

## Collect facts  about  the orthogonal polynomial system

"""
    weight_function(p)
    weight_function(::Type{P})

For an orthogonal polynomial type, a function `w` with `∫ B_n(t) B_m(t) w(t) dt = 0` when n and m are not equal.

"""
weight_function(::Type{P}) where {P <: AbstractOrthogonalPolynomial} = throw(MethodError("Not implemented"))
weight_function(::P) where {P <: AbstractOrthogonalPolynomial} = weight_function(P)

"""
    generating_function(p)
    generating_function(::Type{P})

The generating function is a function defined by: `(t,x) -> sum(t^n Pn(x) for n in 0:oo)`.
"""
generating_function(::Type{P}) where {P <: AbstractOrthogonalPolynomial} = throw(MethodError("Not implemented"))
generating_function(::P) where {P <: AbstractOrthogonalPolynomial} = generating_function(P)

"""
    leading_term(::Type{P},n)

Return leading term of `basis(P,n)` in the  standard basis. By default this is generated through the three-point recursion.
"""
function leading_term(::Type{P}, n::Int) where {P <: AbstractOrthogonalPolynomial}
    n < 0 && throw(ArgumentError("n must be a non-negative integer"))
    n == 0 && return one(eltype(P))
    prod(An(P,i) for i in n-1:-1:0)
end

# is P a monic polynomial system?
ismonic(::Type{P}) where {P <: AbstractOrthogonalPolynomial} = false
isorthonormal(::Type{P}) where {P <: AbstractOrthogonalPolynomial} = false

# cf. https://en.wikipedia.org/wiki/Orthogonal_polynomials#Recurrence_relation
# Orthogonal polynomials have a three-point recursion formula
# parameterized here through:
# P_{n+1} = (An*x + Bn) * P_n + Cn P_{n-1}
# also commonly  written as
# An⋅x⋅P_n = -P_{n+1} + Bn⋅P_n + C_n P_{n-1}
"""
    An(::Type{P},n)

Orthogonal polynomials defined by a weight function satisfy a three point recursion formula of the form:

`P_{n+1} = (A_n x + B_n) P_{n} - C_n P_{n-1}`

If the polynomials are monic, this is usually parameterized as:

`π_{n+1} = (x - α̃_n) π_n - β̃_n π_{n-1}`

These functions are used through recursion when evaluating the polynomials, converting to `Polynomial` format, for constructing the Vandermonde matrix, for construction the Jacobi matrix, and elsewhere.
"""
An(::Type{P}, n) where {P <: AbstractOrthogonalPolynomial} = throw(ArgumentError("No default method"))

"""
    Bn(::Type{P},n)
    Bn(p::P, n)

cf. [`An()`](@ref)
"""
Bn(::Type{P}, n) where {P <: AbstractOrthogonalPolynomial} = throw(ArgumentError("No default method"))

"""
    Cn(::Type{P},n)
    Cn(p::P, n)

cf. [`An()`](@ref)
"""
Cn(::Type{P}, n) where {P <: AbstractOrthogonalPolynomial} = throw(ArgumentError("No default method"))

## For monic polynomials, we have
## π_{n+1} = (x - α̃(n)) π_{n} - β̃(n)π_{n-1}
"""
    π̃αn(::Type{P}, n)

cf. [`An()`](@ref)
"""
π̃αn(P::Type{<:AbstractOrthogonalPolynomial}, n)        = - Bn(P,n) / An(P,n)

"""
    β̃n(::Type{P}, n)

cf. [`An()`](@ref)
"""
function π̃βn(P::Type{<:AbstractOrthogonalPolynomial}, n)
    iszero(n) &&  return innerproduct(P, one, one)
    Cn(P,n)/An(P,n)/An(P, n-1)
end

"""
    monic(p::AbstractOrthogonalPolynomial)

Return `p` as a monic polynomial *when* represented in the standard basis. Retursn the zero polynomial if the degree of `p` is `-1`.
"""
function monic(p::P) where {P <: AbstractOrthogonalPolynomial}
    n = degree(p)
    n == -1 && return ⟒(P)(0/one(eltype(p)))
    p / (p[end]*leading_term(P, n))
end

##
## Work with compact  basis
##
"""
     Basis(P,n)
     Basis{P}(n)

The  command `basis(P,n, [var])` realizes the polynomial. `Basis(P,n)` does  not.  This  can  be  useful for  evaluation.
"""
struct Basis{Π}
    n::Int
end
Basis(P,n) =  Basis{P}(n)
Base.show(io::IO,  mimetype::MIME"text/plain", b::Basis{P})  where  {P} = print(io, "$(P)($(b.n))")

function  innerproduct(::P,  f::Basis{P}, g::Basis{P}) where {P}
    n,m  = f.n, g.n
    if n == m
        return norm2(P, n)
    else
        return zero(P)
    end
end

##
## Vandermonde matrix can be generated through the 3-point recursion formula
##
function Polynomials.vander(p::Type{P}, x::AbstractVector{T}, n::Integer) where
    {P <:  AbstractOrthogonalPolynomial,
     T <: Number}

    A = Matrix{T}(undef, length(x), n + 1)
    A[:, 1] .= P0(P, one(T))

    if n > 0
        A[:, 2] .= P1(P, x)
        @inbounds for i in 1:n-1
            A[:, i+2] .= A[:, i+1] .* (An(p, i)*x .+ Bn(p,i)) .+ (Cn(p,i) * A[:, i])
        end
    end

    return A

end

# Jacobi matrix
# roots(basis(P,n)) = eigvals(jacobi_matrix(P,n)), but this is more stable
# https://en.wikipedia.org/wiki/Gaussian_quadrature#The_Golub-Welsch_algorithm
"""
    jacobi_matrix(::Type{P}, n)

The Jacobi Matrix is a symmetric tri-diagonal matrix. The diagonal entries are the `alpha_i` values, the off diagonal entries,
the square root of the  `beta_i` values. This matrix has the properties that

* the eigenvalues are the roots of the corresponding basis vector. As these roots are important in quadrature, and finding eigenvalues of a symmetric tri-diagonal matrix yields less error than finding the eigenvalues of the companion matrix, this can be used for higher degree basis polynomials.
* the normalized eigenvectors have initial term proportional to the weights in a quadrature formula

"""
function jacobi_matrix(::Type{P}, n) where {P <: AbstractOrthogonalPolynomial}
    LinearAlgebra.SymTridiagonal([π̃αn(P,i) for i in 0:n-1], [sqrt(π̃βn(P,i)) for i in 1:n-1])
end
jacobi_matrix(p::P, n) where {P <: AbstractOrthogonalPolynomial} = jacobi_matrix(P,n)

##  Compute weights and nodes for quadrature
"""
    gauss_nodes_weights(::Type{P}, n)

Returns a tuple of nodes and weights for Gauss quadrature for the given orthogonal type.

When available, the values are computed through  the `FastGaussQuadratures` package.

For some types, a method from  A. Glaser, X. Liu, and V. Rokhlin. "A fast algorithm for the calculation of the roots of special functions." SIAM J. Sci. Comput., 29 (2007), 1420-1438. is used.

For others the Jacobi matrix, J_n, for which the Golub-Welsch] algorithm The nodes  are computed from the eigenvalues of J_n, the weights a scaling of the first component of the normalized eigen vectors (β_0 * [v[1] for v in vs])

"""
function gauss_nodes_weights(p::Type{P}, n) where {P <: AbstractOrthogonalPolynomial}
    J = jacobi_matrix(P, n)
    eig = eigen(J, extrema(P)...)
    # Is this necessary?
    nm  = 1 #diag(eig.vectors * eig.vectors')
    wts =  π̃βn(P,0) * (eig.vectors[1,:] ./ nm).^2
    eig.values,  wts
end

gauss_nodes_weights(B::Basis{P})  where  {P} = gauss_nodes_weights(B.P, B.n)

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

"""
      innerproduct(::Type{P}, f, g)

Compute  `<f,g> = ∫ f⋅g⋅w dx` where  `w` is the weight function of the  type  `P`  and the integral is  taken  over  the domain of the type `P`.
"""
function innerproduct(P::Type{<:Union{AbstractOrthogonalPolynomial}}, f, g; atol=sqrt(eps(float(eltype(P)))))
    dom = domain(P)
    a, b = first(dom), last(dom)
    if first(bounds_types(dom)) == Open
        a += eps(float(one(a)))
    end
    if last(bounds_types(dom)) == Open
        b -= eps(float(one(b)))
    end
    fn = x -> f(x) * g(x) * weight_function(P)(x)

    return quadgk(fn, a, b; atol=atol)[1]
end

## Compute <p_i, p_i> = \| p \|^2; allows export, but work is in norm2
Base.abs2(::Type{P}, n) where {P <: AbstractOrthogonalPolynomial} = norm2(P, n)

## Compute <p_i, p_i> = \| p \|^2
## Slow default; generally should  be directly expressed for each type
function norm2(::Type{P}, n) where {P <: AbstractOrthogonalPolynomial}
    p = basis(P,n)
    innerproduct(P, p, p)
end

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

"""
    fit([Val(S)], P::Type{<:AbstractOrthogonalPolynomial}, f, n::Int; var=:x)

    Find an approximating polynomial of degree `n` or less for a function `f`, that is returns `p(x) = ∑ᵢⁿ cᵢ Pᵢ(x)` for some coefficients `cᵢ`.

Defaults to an interpolating polynomial. To specify others, use one of `Val(:interpolating)`, `Val(:lsq)` (least squares), or `Val(:series)` (trunated series expansion) as the first argument. See [`SpecialPolynomials.cks`](@ref) for some more detail.

"""
Polynomials.fit(P::Type{<:AbstractOrthogonalPolynomial}, f, n::Int; var=:x) =
    fit(Val(:interpolating), P, f, n, var=var)

"""
    fit(val::Val{:interpolating}, P::Type{<:AbstractOrthogonalPolynomial}, f, deg::Int; var=:x)

Fit `f` with an interpolating polynomial of degree `n` orless using nodes
`x0,x1, ..., xn`, the  zeros of `P_{n+1} = basis(P, n+1)`. and `p(xᵢ)=f(xᵢ)`.
"""
Polynomials.fit(val::Val{:interpolating},
                P::Type{<:AbstractOrthogonalPolynomial}, f, n::Int;
                var=:x) =
                    P(cks(val, P,f,n), var)

"""
    fit(Val(:lsq), P::Type{<:AbstractOrthogonalPolynomial}, f, n::Int)

Fit `f` with `p(x)=∑ d_i P_i(x)` where `p` had degree `n` or less using least squares.
"""
Polynomials.fit(val::Val{:lsq},
                P::Type{<:AbstractOrthogonalPolynomial}, f, n::Int;
                var=:x) =
                    P(cks(val, P,f,n), var)

"""
    fit(Val(:series), P::Type{<:AbstractOrthogonalPolynomial}, f, n::Int)

If `f(x)` is written as an infinite sum `∑ c_kP_k(x)`, this returns a truncated sum `∑_0^n c̃_k P_k(x)` where `n` ischosen algorithmically and `c̃_k`is chosen using an efficient manner, not necessarily through the orthogonality condition, `<f,p_k>/<p_k,p_k>`.

"""
Polynomials.fit(val::Val{:series},
                P::Type{<:AbstractOrthogonalPolynomial}, f;
                var=:x,kwargs...) =
                    P(cks(val, P,f; kwargs...), var)

"""
    cks(::Val{:interpolating}, ::Type{P}, f, n::Int)

Fit `f` with the interpolating polynomial using roots of P_{n+1}.

Let xs,ws be the gauss nodes and weights of P (xs are the zeros of P_{n+1}).

Then if we interpolate `f` at the `xs` to get `p`, then
`p(xᵢ) = f(xᵢ)` and

`∑ p_n(xᵢ) p_m(xᵢ) wᵢ = K_m δ_{nm}` and

`p(x) = ∑ᵢ cᵢ Pᵢ(x)`.

Using this:
`∑ᵢ f(xᵢ)P_k(xᵢ) wᵢ =`
`∑ᵢ (∑_j c_j P_j(xᵢ)) P_k(xᵢ) wᵢ =`
`∑_j c_j (K_k δ_{ik}) = c_k K_k`, So
`c_k = (1/K_k) ∑ᵢ f(xᵢ)P_k(xᵢ) wᵢ`
"""
function cks(::Val{:interpolating}, ::Type{P}, f, n::Int) where {P  <:  AbstractOrthogonalPolynomial}

    xs, ws = gauss_nodes_weights(P, n)
    return [sum(f(xⱼ) * basis(P, k)(xⱼ) * wⱼ for (xⱼ,wⱼ) in zip(xs,ws)) / norm2(P,k) for k in 0:n]
end

"""
    cks(::Val{:lsq}, ::Type{P}, f, n::Int)

Fit `f` with a polynomial `∑ᵢⁿ cᵢ Pᵢ` chosen so `<f-p,f-p>_w` is as small as possible.
 Using the normal equations, the coefficients are found to be `c_k = <f,P_k>_w / <P_k,P_k>_w.`
For some types an approximation to the inner product, `<f,P_k>_w` may be used.

ref: http://www.math.niu.edu/~dattab/MATH435.2013/APPROXIMATION
"""
function cks(::Val{:lsq}, ::Type{P}, f, n::Int) where {P  <:  AbstractOrthogonalPolynomial}
## return ck =  <f,P_k>/<P_k,P_k>, k =  0...n
    [innerproduct(P, f, basis(P,k)) /norm2(P,k) for k in 0:n]
end

"""
    cks(::Val{:series}, ::Type{P}, f, n::Int)

If `f(x)` is written as an infinite sum `∑ c_kP_k(x)`, then this
tries to identify an `n` for which the series expansion is a good approximation and returns the coefficients.

"""
function cks(::Val{:series}, ::Type{P}, f, n::Int) where {P  <:  AbstractOrthogonalPolynomial}
    throw(ArgumentError("No default method"))
end


export zoom_Cub
function zoom_Cub(h :: LineModel,
                  t₀ :: Float64,
                  t₁ :: Float64;
                  c₁ :: Float64=0.01,
                  tol :: Float64=1e-7,
                  ϵ :: Float64=1e-10,
                  maxiter :: Int=100,
                  verbose :: Bool=false)

        if obj(h,t₀)<obj(h,t₁)
          tl=t₀
          th=t₁
        else
          tl=t₁
          th=t₀
        end

        hl=obj(h,tl)
        dhl=grad(h,tl)
        hh=obj(h,th)
        dhh=grad(h,th)

        h₀=obj(h,t₀)
        dh₀=grad(h,t₀)

        γ=0.8
        t=t₁
        tp=t₀
        tqnp=t₀
        hk=0
        hkm1=0
        gkm1=0
        hp=0
        i=0

        hk=obj(h,t)
        gk=grad(h,t)
        hkm1=obj(h,tqnp)
        gkm1=grad(h,tqnp)

        verbose && println("on a les paramètre de zoom")

        verbose && @printf(" iter        tₗ        tₕ         t        hₗ        hₕ         hk         gk\n")
        verbose && @printf(" %7.2e %7.2e  %7.2e  %7.2e  %7.2e %7.2e %7.2e %7.2e\n", i,tl,th,t,hl,hh,hk,gk)

        while ((i<maxiter) & (abs(gk)>tol)) || i==0
          if (hk>h₀+c₁*(t-t₀)*dh₀) || (hl<=hk)
            th=t
            hh=hk
            dhh=gk
          else
            if (abs(gk)<ϵ)
              topt=t
              iter=i
              return (topt,iter)
            elseif gk*(th-tl)>=0
              th=tl
              hh=hl
              dhh=dhl
            end
            tl=t
            hl=hk
            dhl=gk
          end

          s = t-tqnp
          y = gk-gkm1

          α=-s
          z=gk+gkm1+3*(hk-hkm1)/α
          discr=z^2-gk*gkm1
          denom=gk+gkm1+2*z
          if (discr>0) & (abs(denom)>eps(Float64))
            #si on peut on utilise l'interpolation cubique
            w=sqrt(discr)
            dC=-s*(gk+z+sign(α)*w)/(denom)
          else #on se rabat sur une étape de sécante
            dC=-gk*s/y
          end

          if ((tp-t)*dC>0) & (dC/(tp-t)<γ)
            tplus = t + dC
            hplus = obj(h, tplus)
            gplus = grad(h,tplus)
            verbose && println("N")
          else
            tplus = (t+tp)/2
            hplus = obj(h, tplus)
            gplus = grad(h,tplus)
            verbose && println("B")
          end

          if t>tp
            if gplus<0
              tp=t
              tqnp=t
              t=tplus
            else
              tqnp=t
              t=tplus
            end
          else
            if gplus>0
              tp=t
              tqnp=t
              t=tplus
            else
              tqnp=t
              t=tplus
            end
          end

          #mise à jour des valeurs
          hkm1=hk
          gkm1=gk
          hk=hplus
          gk=gplus
          i=i+1
          #println("on continue d'itérer")
          verbose && @printf(" %7.2e %7.2e  %7.2e  %7.2e %7.2e %7.2e %7.2e %7.2e\n", i,tl,th,t,hl,hh,hk,gk)

        end
        topt=t
        iter=i

        return (topt,iter)
end


function diffeq!(du,u,p,t)
    CycE = p[C.CycElevel] - u[V.CycEp27]

    Vdp27 = p[C.kd27] + (p[C.kd27e]*CycE)*u[V.Skp2]
    Vdcyce = p[C.kdcyce] + p[C.kdcycee]*CycE
    Vdskp2 = p[C.kdskp2] + p[C.kdskp2c1]*u[V.Cdh1]

    Vicdh1 = p[C.kicdh1e]*CycE

    du[V.p27T] = p[C.ks27] - Vdp27*u[V.p27T]

    du[V.Skp2] = p[C.ksskp2] - Vdskp2*u[V.Skp2]

    du[V.CycEp27] = p[C.kasse]*(p[C.CycElevel]-u[V.CycEp27])*(u[V.p27T]-u[V.CycEp27])-(p[C.kdise]+Vdp27+Vdcyce)*u[V.CycEp27]

    du[V.EmiC] = p[C.kasec]*(p[C.Cdh1T]-u[V.EmiC])*(p[C.Emi1T]-u[V.EmiC]) - (p[C.kdiec]+p[C.kdemi1])*u[V.EmiC]

    du[V.Cdh1dp] = p[C.kacdh1]*(p[C.Cdh1T]-u[V.Cdh1dp]) - Vicdh1*u[V.Cdh1dp]

    du[V.Cdh1] = (p[C.kdiec]+p[C.kdemi1])*(u[V.Cdh1dp]-u[V.Cdh1]) - p[C.kasec]*u[V.Cdh1]*(p[C.Emi1T]-u[V.EmiC])+p[C.kacdh1]*(p[C.Cdh1T]-u[V.EmiC]-u[V.Cdh1])-Vicdh1*u[V.Cdh1]

end

function param_values()::Vector{Float64}
    p::Vector{Float64} = zeros(C.NUM)

    p[C.kscyce] = 0.003
    p[C.kdcyce] = 0.001
    p[C.kdcycee] = 0.0001
    p[C.kdcycea] = 0.03
    p[C.kasse] = 1
    p[C.kdise] = 0.02
    ## CYCA SYNTHESISp[C.DEGRADATION AND P27 BINDING/DISSOCIATION:
    p[C.kscyca] = 0.0025
    p[C.kdcyca] = 0.002
    p[C.kdcycac1] = 0.4
    p[C.kassa] = 1
    p[C.kdisa] = 0.02
    ## P27 SYNTHESIS AND DEGRADATION:
    p[C.ks27] = 0.008
    p[C.kd27] = 0.004
    p[C.kd27e] = 2
    p[C.kd27a] = 2
    ## EMI1 SYNTHESIS AND DEGRADATION:
    p[C.ksemi1] = 0.003
    p[C.kdemi1] = 0.001
    ## CDH1 REGULATION:
    p[C.Cdh1T] = 1
    p[C.kacdh1] = 0.02
    p[C.kicdh1e] = 0.07
    p[C.kicdh1a] = 0.2
    p[C.kasec] = 2
    p[C.kdiec] = 0.02
    ## SKP2 SYNTHESIS AND DEGRADATION:
    p[C.ksskp2] = 0.004
    p[C.kdskp2] = 0.002
    p[C.kdskp2c1] = 0.2
    ## CDK INHIBITOR
    p[C.Inhibitor] = 0.0

    p[C.Emi1T] = 0.0
    p[C.CycElevel] = 1.0

    return p
end

function get_derivatives(u::Vector{Float64},p::Vector{Float64})
    # derivatives: dF/d[bifurcation_param]
    dFdp::Vector{Float64} = zeros(V.NUM)

    dFdp[V.p27T]  = -p[C.kd27e]*u[V.Skp2]*u[V.p27T]
    dFdp[V.CycEp27] = p[C.kasse]*(1.0-u[V.CycEp27])*(u[V.p27T]-u[V.CycEp27])-(p[C.kd27e]*u[V.Skp2]+p[C.kdcycee])*u[V.CycEp27]
    dFdp[V.Cdh1dp] = -p[C.kicdh1e]*u[V.Cdh1dp]
    dFdp[V.Cdh1] =  -p[C.kicdh1e]*u[V.Cdh1]

    return dFdp
end

function get_steady_state(p::Vector{Float64})
    tspan::Tuple{Float64,Float64} = (0.0,Inf)
    u0::Vector{Float64} = zeros(V.NUM)

    prob = ODEProblem(diffeq!,u0,tspan,p)
    prob = SteadyStateProblem(prob)
    sol = solve(prob,DynamicSS(CVODE_BDF()),dt=1.0)

    return sol.u
end


using  Statistics, StatsBase, Distributions,DataFrames, CSV, Plots, Match
pyplot()

# Helper functions:
include("routines.jl")

## TEST SET CONSTRUCTION AND ANALYSIS
df=CSV.read("fold0.txt",DataFrame;delim=" ")
Ntest=size(df,1)

p1_pred = df."prob"
c_pred = 1*(p1_pred .> 0.5)
c_true = 1*(df."label" .> 0.5)

global true_positive=sum((c_pred.==1) .& (c_true.==1))
global false_positive=sum((c_pred.==1) .& (c_true.==0))
global false_negative=sum((c_pred.==0) .& (c_true.==1))
global true_negative=sum((c_pred.==0) .& (c_true.==0))
global F1Test_true=true_positive/(true_positive + 0.5*(false_positive+false_negative))
global SpecTest_true=true_negative/(true_negative + false_positive)

lambda=0.9
p1_true_guess = lambda*p1_pred .+ (1-lambda)*0.5 # guess of the true label on test set

# Define the parameters for the Optimal Sampling Distribution:
htilde = happroxFn(c_pred,p1_true_guess;metric="F",metric_par=0.5)

Mvals=Int.(ceil.(LinRange(10,10*Ntest,100)))

percent=zeros(length(Mvals))

spec_mean=zeros(length(Mvals))
spec_lower=zeros(length(Mvals))
spec_upper=zeros(length(Mvals))
acc_mean=zeros(length(Mvals))
acc_lower=zeros(length(Mvals))
acc_upper=zeros(length(Mvals))
precision_mean=zeros(length(Mvals))
precision_lower=zeros(length(Mvals))
precision_upper=zeros(length(Mvals))
recall_mean=zeros(length(Mvals))
recall_lower=zeros(length(Mvals))
recall_upper=zeros(length(Mvals))
f1_mean=zeros(length(Mvals))
f1_lower=zeros(length(Mvals))
f1_upper=zeros(length(Mvals))

Unlabelled=1:Ntest

# Decide which points to sample:
q=htilde/sum(htilde) # importance distribution
dist=Categorical(q)

for experiment=1:length(Mvals)

    local Nlabel=Mvals[experiment] # number of points we will sample from the unlabelled set and ask to be labelled
    local ToLabel=rand(dist,Nlabel) # sample Nlabel points from the importance distribution
    local total_samples_used = length(unique(ToLabel)) # number of unique samples

    Fhat_acc_mean, Fhat_acc_var = FhatIS(c_pred[ToLabel],c_true[ToLabel],q[ToLabel];metric="Accuracy")

    F1hat_mean, F1hat_var = FhatIS(c_pred[ToLabel],c_true[ToLabel],q[ToLabel];metric="F",metric_par=0.5)

    Fhat_precision_mean, Fhat_precision_var = FhatIS(c_pred[ToLabel],c_true[ToLabel],q[ToLabel];metric="F",metric_par=1)

    Fhat_recall_mean,Fhat_recall_var = FhatIS(c_pred[ToLabel],c_true[ToLabel],q[ToLabel];metric="F",metric_par=0)

    Fhat_spec_mean,Fhat_spec_var = FhatIS(c_pred[ToLabel],c_true[ToLabel],q[ToLabel];metric="Specificity")

    # convert this to confidence bounds by fitting this to say a beta distribution:

    global confidence_level=0.9 # one standard deviation (for a Gaussian)

    println("-----------------------------------------")
    println(experiment)

    println("total samples used  = $(total_samples_used)")
    println("% of test data      = $(total_samples_used/Ntest)")

    percent[experiment]=total_samples_used/Ntest

    println("true accuracy       = $(mean(c_true .== (p1_pred .> 0.5)))")
    local mn,lower,upper = confidenceBeta(Fhat_acc_mean,Fhat_acc_var,confidence_level)
    println("estimated accuracy  = $mn [$lower,$upper]")

    acc_mean[experiment]=mn
    acc_lower[experiment]=lower
    acc_upper[experiment]=upper

    println("")

    println("true precision      = $(true_positive/(true_positive+false_positive))")
    local mn,lower,upper = confidenceBeta(Fhat_precision_mean,Fhat_precision_var,confidence_level)
    println("estimated precision = $mn [$lower,$upper]")

    precision_mean[experiment]=Fhat_precision_mean
    precision_lower[experiment]=lower
    precision_upper[experiment]=upper

    println("")

    println("true recall         = $(true_positive/(true_positive+false_negative))")
    local mn,lower,upper = confidenceBeta(Fhat_recall_mean,Fhat_recall_var,confidence_level)
    println("estimated recall    = $mn [$lower,$upper]")

    recall_mean[experiment]=Fhat_recall_mean
    recall_lower[experiment]=lower
    recall_upper[experiment]=upper

    println("")

    println("true F1             = $(F1Test_true)")
    local mn,lower,upper = confidenceBeta(F1hat_mean,F1hat_var,confidence_level)
    println("estimated F1        = $mn [$lower,$upper]")

    f1_mean[experiment]=F1hat_mean
    f1_lower[experiment]=lower
    f1_upper[experiment]=upper

    println("")

    println("true specificity             = $(SpecTest_true)")
    local mn,lower,upper = confidenceBeta(Fhat_spec_mean,Fhat_spec_var,confidence_level)
    println("estimated specificity        = $mn [$lower,$upper]")

    spec_mean[experiment]=Fhat_spec_mean
    spec_lower[experiment]=lower
    spec_upper[experiment]=upper

end

mp=maximum(percent)

acc_true = (mean(c_true .== (p1_pred .> 0.5)))
ribbon_upper=acc_upper - acc_mean
ribbon_lower=acc_mean - acc_lower
plot(reuse=false,show=true)
plot_acc=plot!(percent,acc_mean,ribbon=(ribbon_lower,ribbon_upper),label="$(confidence_level) confidence interval",framestyle = :box,minorgrid=true)
plot!([0,mp],[acc_true,acc_true],label="true value",legend=:bottomright,ylims=[0.9*acc_true,1.1*acc_true])
title!("Accuracy")
xlabel!("fraction of test data labelled")

precision_true = true_positive/(true_positive+false_positive)
ribbon_upper=precision_upper - precision_mean
ribbon_lower=precision_mean - precision_lower
plot(reuse=false,show=true)
plot_precision=plot!(percent,precision_mean,ribbon=(ribbon_lower,ribbon_upper),label="$(confidence_level) confidence interval",framestyle = :box,minorgrid=true)
plot!(plot_precision)
plot!([0,mp],[precision_true,precision_true],label="true value",legend=:bottomright,ylims=[0.9*precision_true,1.1*precision_true])
title!("Precision")
xlabel!("fraction of test data labelled")

recall_true = true_positive/(true_positive+false_negative)
ribbon_upper=recall_upper - recall_mean
ribbon_lower=recall_mean - recall_lower
plot(reuse=false,show=true)
plot_recall=plot!(percent,recall_mean,ribbon=(ribbon_lower,ribbon_upper),label="$(confidence_level) confidence interval",framestyle = :box,minorgrid=true)
plot!(plot_recall)
plot!([0,mp],[recall_true,recall_true],label="true value",legend=:bottomright,ylims=[0.9*recall_true,1.1*recall_true])
title!("Recall")
xlabel!("fraction of test data labelled")

f1_true = F1Test_true
ribbon_upper=f1_upper - f1_mean
ribbon_lower=f1_mean - f1_lower
plot(reuse=false,show=true)
plot_f1=plot!(percent,f1_mean,ribbon=(ribbon_lower,ribbon_upper),label="$(confidence_level) confidence interval",framestyle = :box,minorgrid=true)
plot!(plot_f1)
plot!([0,mp],[f1_true,f1_true],label="true value",legend=:bottomright,ylims=[0.9*f1_true,1.1*f1_true])
title!("F1")
xlabel!("fraction of test data labelled")

spec_true = SpecTest_true
ribbon_upper=spec_upper - spec_mean
ribbon_lower=spec_mean - spec_lower
plot(reuse=false,show=true)
plot_spec=plot!(percent,spec_mean,ribbon=(ribbon_lower,ribbon_upper),label="$(confidence_level) confidence interval",framestyle = :box,minorgrid=true)
plot!(plot_spec)
plot!([0,mp],[spec_true,spec_true],label="true value",legend=:bottomright,ylims=[0.9*spec_true,1.1*spec_true])
title!("Specificity")
xlabel!("fraction of test data labelled")


using TensorNetworkAD
using Test, Random
using Zygote
using LinearAlgebra: norm

@testset "autodiff" begin
    a = randn(10,10)
    @test Zygote.gradient(norm, a)[1] ≈ num_grad(norm, a)

    foo = x -> sum(Float64[x x; x x])
    @test Zygote.gradient(foo, 1)[1] ≈ num_grad(foo, 1)
end