export j, pol, ∠, phasor, phasorsine, angulardimension, phasordimension

"""
`j = 1im` equals the imaginary unit
"""
const j=1im

"""
# Function call

`pol(r, phi)`

# Description

Creates a complex quantity with length `r` and angle `phi`

# Variables

`r` Length of complex quantity; r may be utilized including a unit generated by
Unitful

`phi` Angle of complex quantity; if module Unitful is utilized, the angle may be
specified in degrees, by using unit `°`

# Examples

```julia
julia> using Unitful, Unitful.DefaultSymbols, ElectricalEngineering
julia> U1 = pol(2V,pi)
-2 + 0im V
julia> U2 = pol(sqrt(2)*1V,45°)
1 + 1im V
```
"""
function pol(r, phi)
  return r*cos(phi) + 1im*r*sin(phi)
end

"""
# Function call

`∠(phi)`

# Description

Creates a complex quantity with length 1 and angle `phi`

# Variables

`phi` Angle of complex quantity; if module Unitful is utilized, the angle may be
specified in degrees, by using unit `°`

# Examples

```julia
julia> using Unitful, Unitful.DefaultSymbols, ElectricalEngineering
julia> U1 = 2V*∠(90°)
-0.0 + 2.0im V
```
"""
function ∠(phi)
  return cos(phi) + 1im*sin(phi)
end

@doc raw"""
# Function call

```
phasor(c;origin=0.0+0.0im, ref=1, par=0,
    labelrsep=0.5, labeltsep=0.1, label="", ha="center", va="center",
    labelrelrot=false, labelrelangle=0,
    color="black", backgroundcolor="none", linesstyle="-", linewidth=1,
    width=0.2, headlength=10, headwidth=5)
```

# Description

This function draws a phasor from a starting point `origin` and end point
`origin`+`c`. The phasor consists of a shaft and an arrow head.

Each phasor `c` is plotted as a relative quantity, i.e., `c/ref` is actually
shown the plot figure. This concept of plotting a per unit phasor is used to
be able to plot phasor with different quantities, e.g., voltage and current
phasors. It is important that the variables `c`, `origin` and `ref` have the
same units (defined through Unitful).

# Variables

`c` Complex phasor, drawn relative relative to `origin`

`origin` Complex number representing the origin of the phasor; this variable
needs to have the same unit as `c`

`ref` Reference length of scaling; this is required as in a phasor diagram
voltages and currents may be included; in order to account for the different
voltage and current scales, one (constant)  `ref` is used for voltage
phasors and another (constant) `ref` is used for current phasors; this
variable needs to have the same unit as `c`

`par` In order to be able to plot parallel phasors, par is used to specify
the per unit tangential shift (offset) of a phasor, with respect to `ref`;
so typically `ref` will be selected to be around 0.05 to 0.1;
default value = 0 (no shift of phasor)

`labelrsep` Radial per unit location of label (in direction of the phasor):
`labelrsep = 0` represents the shaft of the phasor and `labelrsep = 1` represents
the arrow hear of the phasor; default value = 0.5, i.e., the radial center
of the phasor

`labeltsep` Tangential per unit location of label: `labeltsep = 0` means that the
label is plotted onto the phasor; `labeltsep = 0.1` plots the label on top of
the phasor applying a displacement of 10% with respect to `ref`;
`labeltsep = -0.2` Plots the label below the
phasor applying a displacement of 20% with respect to `ref`; default value
= 0.1

`ha` Horizontal alignment of label; actually represents the tangential
alignment of the label; default value = "center"

`va` Vertical alignment of label; actually represents the radial
alignment of the label; default value = "center"

`labelrelrot` Relative rotation of label; if `labelrelrot == false` (default
value) then the label is not rotated relative to the orientation of the phasor;
otherwise the label is rotated relative to the phasor by the angle
`labelrelangle`

`labelrelangle` Relative angle of label  with respect to phasor orientation;
this angle is only applied, it `labelrelrot == true`; this angle the indicates
the relative orientation of the label with respect to the orientation of the
phasor; default value = 0

`color` Color of the phasor; i.e., shaft and arrow head color; default
value = "black"

`backgroundcolor` Background color of the label; if labelrsep is equal to 0, the
background color "white" can be used; default value = "none"

`linestyle` Line style of the phasor; default value = "-"

`linewidth` Line width of the phasor; default value = 1

`width` Line width of the shaft line; default value = 0.2

`headlength` Length of arrow head; default value = 10

`headwidth` Width of arrow head; default value = 5

# Example

Copy and paste code:

```julia
using Unitful, Unitful.DefaultSymbols, PyPlot, ElectricalEngineering
figure(figsize=(3.3, 2.5))
rc("text", usetex=true); rc("font", family="serif")
V1 = 100V + j*0V # Voltage
Z1 = 30Ω + j*40Ω # Impedance
I1 = V1/Z1 # Current
Vr = real(Z1)*I1
Vx = V1 - Vr
refV = abs(V1); refI=abs(I1)*0.8
phasor(V1, label=L"$\underline{V}_1$", labeltsep=0.1, ref=refV,
    labelrelrot=true)
phasor(Vr, label=L"$\underline{V}_r$", labeltsep=-0.25, ref=refV,
    labelrelrot=true)
phasor(Vx, origin=Vr, label=L"$\underline{V}_x$", labeltsep=-0.2, ref=refV,
    labelrelrot=true)
phasor(I1, label=L"$\underline{I}_1$", labeltsep=-0.2, labelrsep=0.7, ref=refI,
    labelrelrot=true, linestyle="--", par=-0.05)
phi1=angle(I1)
phi2=angle(V1)
angulardimension(0.3,phi1,phi2,arrowstyle1=".",arrowstyle2="-|>",ha="left",
    label=L"$\varphi_1$", labelrsep=0.05)
axis("square"); xlim(-1,1); ylim(-1,1)
removeaxes(); # Remove axis
# save3fig("phasordiagram",crop=true);
```
"""
function phasor(c;
    origin = (0.0+0.0im).*c./ustrip(c),
    ref = abs(c./ustrip(c)),
    par = 0.0,
    labelrsep = 0.5,
    labeltsep = 0.1,
    label = "",
    ha = "center",
    va = "center",
    labelrelrot = false,
    labelrelangle = 0.0,
    color = "black",
    backgroundcolor = "none",
    linestyle = "-",
    linewidth = 1,
    width = 0.2,
    headlength = 10.0,
    headwidth = 5.0)

    # Check if units of c, origin and ref are compatible
    # Starting point (origin) of phase
    xorigin = 0.0
    yorigin = 0.0
    xend = 0.0
    yend = 0.0
    try
        xorigin = uconvert(Unitful.NoUnits, real(origin)./ref)
        yorigin = uconvert(Unitful.NoUnits, imag(origin)./ref)
        # End point of phasor
        xend = uconvert(Unitful.NoUnits, real(origin+c)./ref)
        yend = uconvert(Unitful.NoUnits, imag(origin+c)./ref)
    catch err
        error("module ElectricalEngineering: function phasor: Dimension mismatch of arguments `c`, `origin` and `ref`\n    The arguments `c`, `origin` and `ref` must have the same dimension (koherent SI unit)")
    end

    # Draw phasor only if length of c is greater than zero, see
    # https://github.com/christiankral/ElectricalEngineering.jl/issues/1
    if upstrip(abs(c))>0
        # Length of phasor
        dr = sqrt((xend-xorigin)^2 + (yend-yorigin)^2)
        # Real part of phasor
        drx = (xend - xorigin)/dr # = real(c)./ref
        # Imag part of phasor
        dry = (yend - yorigin)/dr # = imag(c)./ref
        # Angle of phasor
        absangle = atan(dry, drx)
        # Orientation tangential to phasor (lagging by 90°)
        # Real part of tangential component with respect to length
        dtx = +dry
        # Imag part of tangential component with respect to length
        dty = -drx
        # Real part of parallel shift of phasor
        dpx = -par*dtx
        # Imag part of parallel shift of phasor
        dpy = -par*dty
        # Origin of head
        xoriginHead = xorigin + dr*drx*0.999
        yoriginHead = yorigin + dr*dry*0.999
        # Draw arrow shaft and head
        # https://matplotlib.org/api/_as_gen/matplotlib.pyplot.annotate.html?highlight=annotate#matplotlib.pyplot.annotate
        # Draw shaft separately: otherwise, the arrow contour will be drawn as in
        # https://matplotlib.org/users/annotations.html#basic-annotation
        # so that the back and forth paths overlap and the line style does not
        # appear correctly; replace
        plot([xorigin+dpx,xend+dpx], [yorigin+dpy,yend+dpy],
            color=color, linestyle=linestyle, linewidth=linewidth, clip_on=false)
        # Code based on plot replaces the previous implementation inspired by:
        # https://stackoverflow.com/questions/51746400/linestyle-in-plot-and-annotate-are-not-equal-in-matplotlib
        # Previous (obsolete) implementation:
        # annotate("", xy=(xend+dpx,yend+dpy),
        #     xytext=(xorigin+dpx,yorigin+dpy), xycoords="data",
        #     arrowprops=Dict("arrowstyle"=>"-",
        #         "linestyle"=>linestyle,
        #         "linewidth"=>linewidth,
        #         "color"=>color,
        #         "facecolor"=>color),
        #     annotation_clip=false)

        # Draw arrow head without line style; this is a workaround explained in
        # https://stackoverflow.com/questions/47180328/pyplot-dotted-line-with-fancyarrowpatch/47205418#47205418
        annotate("", xy=(xend+dpx, yend+dpy),
            xytext=(xoriginHead+dpx, yoriginHead+dpy), xycoords="data",
            arrowprops=Dict("edgecolor"=>color, "facecolor"=>color,
                "width"=>width, "linestyle"=>"-",
                "headlength"=>headlength,
                "headwidth"=>headwidth),
            annotation_clip=false)

        # Plot label
        if labelrelrot == false
            # Without relative rotation of label
            text(xorigin + dr*drx*labelrsep - dtx*labeltsep + dpx,
                yorigin + dr*dry*labelrsep - dty*labeltsep + dpy,
                label, ha=ha, va=va, rotation=labelrelangle*180/pi,
                backgroundcolor=backgroundcolor)
        else
            # Applying relative rotation of label
            text(xorigin + dr*drx*labelrsep - dtx*labeltsep + dpx,
                yorigin + dr*dry*labelrsep - dty*labeltsep + dpy,
                label, ha=ha, va=va, rotation=(absangle+labelrelangle)*180/pi,
                backgroundcolor=backgroundcolor)
        end
    end
end

@doc raw"""
# Function call

```
phasorsine(mag = 1, phi = 0; add = false, figsize = (6.6,2.5),
    xlabel = L"$\omega t$\,($^\circ $)", ylabel = "", maglabel = "",
    phasorlabel = maglabel, labeltsep = 0.1, labelrsep = 0.5,
    labelrelrot = true, labelrelangle = 0,
    color = "black", linewidth = 1, linestyle = "-",
    colorDash="gray", left=0.20, right=0.80, bottom=0.20, top=0.80,
    showsine=true, showdashline=true)
```

# Description

This function draws a phasor with magnitude between 0 and 1 on the left subplot
of a figure and a sine diagram on the right subplot of the figure. Such graph
is used in electrical engineering to explain the relationship between phasors
and time domain wave forms.

# Variables

`mag` Magnitude of displayed phasor; shall be between 0 and 1;
default value = 1

`phi` Phase angle of the phasor; default value = 0

`add` When calling this function the first time, `add` shall be equal to
`false`, which is the default value. In this case a new figure with the
dimensions specified in `figsize` is created. In order to add a second phasor
and sine diagram to an existing figure, `add` has to be set `true`

`figsize` Size of the new figure, if `add` is equal to `false`; default value
= (6.6,2.5)

`xlabel` Label of x-axis of right subplot; default value ="ωt(°)" in LaTeX
notation

`ylabel` label of y-axis of right subplot; default value = "";  if more than one
phasors and sine diagram shall be drawn (`add = true`), this label is displayed
only once; therefore, one has to create the label of all phasors when function
`phasorsine` is called the first time for creating a figure

`maglabel` Label of positive and negative magnitude `mag` on the right subplot;
dafault value = "";

`phasorlabel` Label of phasor of left subplot; default value = `maglabel`

`labelrsep` Radial per unit location of label (in direction of the phasor):
`labelrsep = 0` represents the shaft of the phasor and `labelrsep = 1`
represents the arrow hear of the phasor; default value = 0.5, i.e., the radial
center of the phasor

`labeltsep` Tangential per unit location of label: `labeltsep = 0` means that
the label is plotted onto the phasor; `labeltsep = 0.1` plots the label on top
of the phasor applying a displacement of 10% with respect to `ref`; `labeltsep =
-0.2` Plots the label below the phasor applying a displacement of 20% with
respect to `ref`; default value = 0.1

`labelrelrot` Relative rotation of label; if `labelrelrot == false` (default value)
then the label is not rotated relative to the orientation of the phasor;
otherwise the label is rotated relative to the phasor by the angle
`labelrelangle`

`labelrelangle` Relative angle of label with respect to phasor
orientation; this angle is only applied, it `labelrelrot == true`; this angle the
indicates the relative orientation of the label with respect to the
orientation of the phasor; default value = 0

`color` Color of the phasor; i.e., shaft and arrow head color; default
value = "black"

`backgroundcolor` Background color of all labels; if labelrsep is equal to 0, the
background color "white" can be used; default value = "none"

`linewidth` Line width of the phasor; default value = 1

`linestyle` Line style of the phasor; default value = "-"

`colorDash` Color of the dashed circle (left subplot) and the horizontal dashed
lines between the left and right subplot; default value = colorBlack4

`left` Left side of the figure; default value = 0.2

`right` Right side of the figure; default value = 0.85

`bottom` Bottom side of the figure; default value = 0.2

`top` Top side of the figure; default value = 0.85

`showsine` If `true`, the sinewave is shown; default value = `true`

`showdashline` If `true`, the dashed lines are shown; default value = `true`

# Example

Copy and paste code:

```julia
using Unitful, Unitful.DefaultSymbols, PyPlot, ElectricalEngineering
rc("text", usetex=true); rc("font", family="serif")
phasorsine(1, 45°, ylabel=L"$u,i$", maglabel=L"$\hat{U}$", labelrsep=0.3,
    color="gray", linestyle="--")
phasorsine(0.55, 0, add=true, maglabel=L"$\hat{I}$")
# save3fig("phasorsine",crop=true);
```
"""
function phasorsine(mag = 1,
    phi = 0;
    add = false,
    figsize = (6.6,2.5),
    xlabel = L"$\omega t$\,($^\circ $)",
    ylabel = "",
    maglabel = "",
    phasorlabel = maglabel,
    color = "black",
    backgroundcolor = "none",
    linewidth = 1,
    linestyle = "-",
    labeltsep = 0.1,
    labelrsep = 0.5,
    labelrelrot = true,
    labelrelangle = 0,
    colorDash="gray",
    left=0.20,
    right=0.80,
    bottom=0.20,
    top=0.80,
    showsine = true,
    showdashline = true)
    # https://matplotlib.org/tutorials/text/annotations.html#plotting-guide-annotation
    # https://matplotlib.org/users/annotations.html
    # https://stackoverflow.com/questions/17543359/drawing-lines-between-two-plots-in-matplotlib

    # Create figure
    if !add
        # Create new figure
        fig = figure(figsize=figsize)
        subplots_adjust(left=left, right=right, bottom=bottom, top=top)
    end
    # Create left subplot
    subplot(121)
    # Angle vector to draw circle
    psi = collect(0:pi/500:2*pi)
    # Coordinates of circle
    x = mag*cos.(psi)
    y = mag*sin.(psi)
    # Plot circle
    plot(x, y, color=colorDash, linewidth=1, linestyle=":",
        dash_capstyle="round")
    # Plot phasor
    phasor(pol(mag,phi), ref=1,
        label=phasorlabel, labelrsep=labelrsep, labeltsep=labeltsep,
        labelrelrot=labelrelrot, labelrelangle=labelrelangle,
        color=color, backgroundcolor=backgroundcolor,
        linestyle=linestyle, linewidth=linewidth)
    axis("square")
    ax1 = gca()
    if !add
        xlim(-1.1,1.1)
        ylim(-1.1,1.1)
        ax1.spines["left"].set_visible(false)
        ax1.spines["right"].set_visible(false)
        ax1.spines["bottom"].set_visible(false)
        ax1.spines["top"].set_visible(false)
        # Remove ticks
        xticks([])
        yticks([])
    end

    # Create right subplot
    subplot(122)
    # Plot sine if selected by showsine = true
    if showsine
        yphi = mag*sin.(psi.+phi)
        plot(psi*180/pi, yphi,
            color=color, linewidth=linewidth, linestyle=linestyle)
    end
    ax2 = gca()
    # First phasor and sine wave diagram must be drawn with add = false
    if !add
        # Scale and tick x-axis
        xlim(0, 360)
        ylim(-1.1, 1.1)
        xticks(collect(90:90:360),
            backgroundcolor=backgroundcolor)
        if maglabel != ""
            yticks([-mag, 0, mag],[L"$-$"*maglabel, L"$0$", maglabel],
                backgroundcolor=backgroundcolor)
        else
            yticks([-mag, 0, mag],["","",""])
        end
        # Create arrows and labels of axes
        arrowaxes(xlabel=xlabel, ylabel=ylabel)
    else
        # If additional phasor and sine wave diagram are added, further
        # y-axis ticks have to beee added; for this purpose, the existing (old)
        # ticks are stored
        yticks_old = ax2.get_yticks()
        # Store existing (old) labels
        ytickslabel_old = ax2.get_yticklabels()
        # Extend old ticks and labels by addition ticks and labels
        if maglabel != ""
            yticks(cat(yticks_old, [-mag,mag],dims=1),
                cat(ytickslabel_old, [L"$-$"*maglabel,maglabel],dims=1),
                    backgroundcolor=backgroundcolor)
        end
    end
    # Change vertical scaling of left subplot to match
    ax1.set_ylim(ax2.get_ylim())

    # Plot dashed lines, if showdashline = true
    if showdashline
        # Dotted line from phasor arrow to begin of sine wave, split in two
        # parts, to avoid overlay effects if multiple dashed lines are drawn
        con = matplotlib.patches."ConnectionPatch"(
            xyB=(360, mag*sin(phi)), xyA=(0, mag*sin(phi)),
            coordsA="data", coordsB="data",
            axesA=ax2, axesB=ax2, color=colorDash,
            linewidth=lineWidth4, linestyle=":", clip_on=false)
        ax2.add_artist(con)
        con = matplotlib.patches."ConnectionPatch"(
            xyA=(0, mag*sin(phi)), xyB=(mag*cos(phi), mag*sin(phi)),
            coordsA="data", coordsB="data",
            axesA=ax2, axesB=ax1, color=colorDash,
            linewidth=lineWidth4, linestyle=":", clip_on=false)
        ax2.add_artist(con)
        # Dotted line of y-axis of right diagram to maximum of sine wave
        con = matplotlib.patches."ConnectionPatch"(
            xyB=(mod(90-phi*180/pi,360), mag), xyA=(0, mag),
            coordsA="data", coordsB="data",
            axesA=ax2, axesB=ax2, color=colorDash,
            linewidth=lineWidth4, linestyle=":", clip_on=false)
        ax2.add_artist(con)
        # Dotted line of y-axis of right diagram to minimum of sine wave
        con = matplotlib.patches."ConnectionPatch"(
            xyB=(mod(270-phi*180/pi, 360),-mag), xyA=(0, -mag),
            coordsA="data", coordsB="data",
            axesA=ax2, axesB=ax2, color=colorDash,
            linewidth=lineWidth4, linestyle=":", clip_on=false)
            ax2.add_artist(con)
    end
end

@doc raw"""
# Function call

```
angulardimension(r = 1, phi1 = 0, phi2 = pi/2; origin = 0.0im,
    label= "", labelphisep = 0.5, labelrsep = 0.1,
    labelrelrot = false, labelrelangle = 0, ha = "center", va = "center",
    color="black", backgroundcolor="none",
    arrowstyle1 = ".", arrowstyle2 = "-|>", dot90 = false,
    linewidth = 0.6, linestyle = "-", width = 0.2,
    headlength = 5, headwidth = 2.5)
```

# Description

This function draws an arrowed arc, intended to label the angle between phasors.
The arc is drawn from angle `phi1` (begin) to `phi2` (end). The begin and end
arrow shapes can be set. The arc can be labeled and optionally a dot maker can
be used to indicate right angles (90°).

# Variables

`r` Radius if the arc with not unit; shall be between 0 and 1; default value = 1

`phi1` Phase angle of the begin of the arc; default value = 0

`phi2` Phase angle of the end of the arc; default value = pi/2

`origin` Complex quantity, indicating the origin of the arc with no unit;
default value = 0.0 + 0.0im

`label` Label of the angle; default value =""

`labelphisep` Angular separation of the label with respect to the arc; if
`labelphisep == 0`, the label is located at angle `phi1` and if `labelphisep ==
1`, the label is located at angle `phi2`; default value = 0.5, right in the
middle between `phi1` and `phi2`

labelrsep` Radial per unit location of label (in direction of the phasor):
`labelrsep = 0` locates the label right on the arc. A positive value locates the
label outside the arc, a negative value locates the label inside the arc;
default value = 0.1

`labelrelrot` Relative rotation of label; if `labelrelrot == false` (default
value) then the label is not rotated relative to the center of the arc;
otherwise the label is rotated relative to the  angle `labelrelangle`

`labelrelangle` Relative angle of label with respect to center of the arc; this
angle is only applied, it `labelrelrot == true`; this angle indicates the
relative orientation of the label with respect to the center of the arc;
default value = 0

`ha` Horizontal alignment of label; actually represents the radial alignment of
the label; default value = "center"

`va` Vertical alignment of label; actually represents the tangential alignment
of the label; default value = "center"

`color` Color of the arc; default value = "black"

`backgroundcolor` Background color of the label; if labelrsep is equal to 0, the
background color "white" can be used; default value = "none"

`arrowstyle1` Arrow style of the begin of the arc; default value = "."; valid
strings are:

- `.` dot marker
- `<|-` arrow

`arrowstyle2` Arrow style of the end of the arc; default value = "-|>"; valid
strings are:

- `.` dot marker
- `-|>` arrow

`headlength` Length of arrow head; default value = 5

`headwidth` Width of arrow head; default value = 2.5

# Example

Copy and paste code:

```julia
using Unitful, Unitful.DefaultSymbols, PyPlot, ElectricalEngineering
figure(figsize=(3.3, 2.5))
rc("text", usetex=true); rc("font", family="serif")
V1 = 100V + j*0V # Voltage
Z1 = 30Ω + j*40Ω # Impedance
I1 = V1/Z1 # Current
Vr = real(Z1)*I1
Vx = V1 - Vr
refV = abs(V1); refI=abs(I1)*0.8
phasor(V1, label=L"$\underline{V}_1$", labeltsep=0.1, ref=refV,
    labelrelrot=true)
phasor(Vr, label=L"$\underline{V}_r$", labeltsep=-0.25, ref=refV,
    labelrelrot=true)
phasor(Vx, origin=Vr, label=L"$\underline{V}_x$", labeltsep=-0.2, ref=refV,
    labelrelrot=true)
phasor(I1, label=L"$\underline{I}_1$", labeltsep=-0.2, labelrsep=0.7, ref=refI,
    labelrelrot=true, linestyle="--", par=-0.05)
phi1=angle(I1)
phi2=angle(V1)
angulardimension(0.3,phi1,phi2,arrowstyle1=".",arrowstyle2="-|>",ha="left",
    label=L"$\varphi_1$", labelrsep=0.05)
axis("square"); xlim(-1,1); ylim(-1,1)
removeaxes(); # Remove axis
# save3fig("phasordiagram",crop=true);
```
"""
function angulardimension(r = 1,
    phi1 = 0,
    phi2 = pi/2;
    origin = 0.0im,
    label = "",
    labelphisep = 0.5,
    labelrsep = 0.1,
    labelrelrot = false,
    labelrelangle = 0,
    ha = "center",
    va = "center",
    color="black",
    backgroundcolor="none",
    arrowstyle1 = ".",
    arrowstyle2 = "-|>",
    dot90 = false,
    linewidth = 0.6,
    linestyle = "-",
    width = 0.2,
    headlength = 5,
    headwidth = 2.5)

    # Determine sign of the difference of angles in order to revert oriention
    # angle, if required
    sig = sign(phi2-phi1)
    # Number of segments based on every 2 degrees
    segs = Int(round((phi2-phi1)*180/pi/    2)*sig)
    # Create vector of angles to draw arc
    phi = collect(range(phi1,stop=phi2,length=segs))
    # Arc coordinates
    x = upstrip(r)*cos.(phi) + fill(real(upstrip(origin)),length(phi))
    y = upstrip(r)*sin.(phi) + fill(imag(upstrip(origin)),length(phi))
    # Draw arc
    plot(x,y,color=color,linewidth=linewidth,linestyle=linestyle)

    # Draw arrows at begin and end of arc, depending on the specified arrow
    # styles

    # Begin of arc
    if arrowstyle1 == "<|-"
        # Arrow head
        annotate("", xy=(x[1], y[1]),
            xytext=(x[1]-sig*0.01*r*cos(phi1-pi/2*0.98),
                y[1]-sig*0.01*r*sin(phi1-pi/2*0.98)),
            xycoords="data",
            arrowprops=Dict("edgecolor"=>color, "facecolor"=>color,
                "width"=>width, "linestyle"=>"-",
                "headlength"=>headlength,
                "headwidth"=>headwidth),
            annotation_clip=false)
    elseif arrowstyle1 == "."
        # Dot marker
        plot(x[1],y[1],marker=".",color=color, clip_on=false)
    end

    # End of arc
    if arrowstyle2 == "-|>"
        # Arrow head
        annotate("", xy=(x[end], y[end]),
            xytext=(x[end]-sig*0.01*r*cos(phi2+pi/2*0.98),
                y[end]-sig*0.01*r*sin(phi2+pi/2*0.98)),
            xycoords="data",
            arrowprops=Dict("edgecolor"=>color, "facecolor"=>color,
                "width"=>width, "linestyle"=>"-",
                "headlength"=>headlength,
                "headwidth"=>headwidth),
            annotation_clip=false)
    elseif arrowstyle2 == "."
        # Marker dot
        plot(x[end],y[end],marker="." ,color=color, clip_on=false)
    end

    # Plot label
    # Angle in the middle of the arc
    phim = (phi1+phi2)/2.0
    # Angle difference
    dphi = phi2-phi1
    # Radial position of text
    rlabel = r + labelrsep
    # Angular position of text
    philabel = phi1 + dphi*labelphisep
    # Rectangular position of text
    if labelrelrot == false
        # Without relative rotation of label
        text(rlabel*cos(philabel) + real(origin),
            rlabel*sin(philabel) + imag(origin),
            label, ha=ha, va=va, rotation=labelrelangle*180/pi,
            backgroundcolor=backgroundcolor)
    else
        # Applying relative rotation of label
        text(rlabel*cos(philabel) + real(origin),
            rlabel*sin(philabel) + imag(origin),
            label, ha=ha, va=va, rotation=(phim+labelrelangle)*180/pi,
            backgroundcolor=backgroundcolor)
    end

    # Optionally create a dot in the center of the arc to indicate 90°
    if dot90
        plot(r/2*cos(phim) + real(origin),
            r/2*sin(phim) + imag(origin),
            marker=".", color=color, clip_on=false)
    end
end

@doc raw"""
# Function call

```
phasordimension(c; origin = 0im, ref = 1;
    label = "", labeltsep = 0.1, labelrsep=0.5, labelrelrot=false,
    labelrelangle=0, ha = "center", va = "center",
    color="black", backgroundcolor = "none",
    arrowstyle1 = "<|-", arrowstyle2 = "-|>", linewidth=0.6, linestyle = "-",
    width=0.2, headlength=5, headwidth=2.5,
    par=0, paroverhang = 0.02, parcolor = "black",
    parlinewidth=0.6, parlinestyle = "-")

```

# Description

This function draws length dimension with label based on rectangular x- and
y-coordinates. The length dimension arrow can be shifted parallel to the
specified  coordinates by means the single parameter `par`. Additionally,
auxiliary lines  from the specified coordinates to the length dimension are
drawn.

# Variables

`c` Complex phasor to be dimensioned

`origin` Complex number representing the origin of the phasor dimension; this
variable needs to have the same unit as `c`

`ref` Reference length of scaling; this is required as in a phasor diagram
voltages and currents may be included; in order to account for the different
voltage and current scales, one (constant)  `ref` is used for voltage
phasors and another (constant) `ref` is used for current phasors; this
variable needs to have the same unit as `c`

`label` Label of the angle; default value =""

`labelphisep` Angular separation of the label with respect to the arc; if
`labelphisep == 0`, the label is located at angle `phi1` and if `labelphisep ==
1`, the label is located at angle `phi2`; default value = 0.5, right in the
middle between `phi1` and `phi2`

labelrsep` Radial per unit location of label (in direction of the phasor):
`labelrsep = 0` locates the label right on the arc. A positive value locates the
label outside the arc, a negative value locates the label inside the arc;
default value = 0.1

`labelrelrot` Relative rotation of label; if `labelrelrot == false` (default
value) then the label is not rotated relative to the center of the arc;
otherwise the label is rotated relative to the  angle `labelrelangle`

`labelrelangle` Relative angle of label with respect to center of the arc; this
angle is only applied, it `labelrelrot == true`; this angle indicates the
relative orientation of the label with respect to the center of the arc;
default value = 0

`ha` Horizontal alignment of label; actually represents the radial alignment of
the label; default value = "center"

`va` Vertical alignment of label; actually represents the tangential alignment
of the label; default value = "center"

`color` Color of the arc; default value = "black"

`backgroundcolor` Background color of the label; if labelrsep is equal to 0, the
background color "white" can be used; default value = "none"

`arrowstyle1` Arrow style of the begin of the arc; default value = "."; valid
strings are:

- `.` dot marker
- `<|-` arrow

`arrowstyle2` Arrow style of the end of the arc; default value = "-|>"; valid
strings are:

- `.` dot marker
- `-|>` arrow

`headlength` Length of arrow head; default value = 5

`headwidth` Width of arrow head; default value = 2.5

`par` In order to be able to draw the length dimension parallel to the specified
coordiantes `x1`, `y1`, `x2`, `y2`, `par` is used to specify the per unit
tangential shift (offset) of the length dimension; default value = 0 (no shift)

`paroverhang` The auxiliary lines, drawn in case of `par != 0`, show the absolute
overhang `paroverhang`; default value = 0.02

`parcolor` Color of the auxiliary lines; default value = "black"

`parlinewidth` Line width of the auxiliary line; default value = 0.6

`parlinestyle` Line style of the auxiliary line; default value = "-"

# Examples

Copy and paste the following code:

```julia
using Unitful, Unitful.DefaultSymbols, PyPlot, ElectricalEngineering
figure(figsize=(3.3, 2.5))
rc("text", usetex=true); rc("font", family="serif")
Z1 = pol(1,30°)
phasor(Z1, label=L"$\underline{Z}$", labeltsep = 0.05, labelrelrot=true)
phasordimension(real(Z1), label=L"$R$", arrowstyle1="",
    linestyle="-", linewidth=1, headwidth=5, headlength=10,
    labeltsep=0, color="gray", backgroundcolor="white")
    phasordimension(j*imag(Z1), origin=real(Z1), label=L"j$\cdot X$",
    arrowstyle1="", linestyle="-", linewidth=1, headwidth=5, headlength=10,
    labeltsep=0, color="gray", backgroundcolor="white")
arrowaxes(xmin = real(Z1)+0.1, xlabel="Re", ylabel=L"j$\cdot$Im")
xlim([-0.5,1]);ylim([-0.5,1]); axis("square"); removeaxes();
# save3fig("phasordimension", crop=true)
```
"""
function phasordimension(c;
    origin = (0.0+0.0im).*c./ustrip(c),
    ref = abs(c./ustrip(c)),
    label = "",
    labeltsep = 0.1,
    labelrsep = 0.5,
    labelrelrot = false,
    labelrelangle = 0,
    ha = "center",
    va = "center",
    color="black",
    backgroundcolor = "none",
    arrowstyle1 = "<|-",
    arrowstyle2 = "-|>",
    linewidth = 0.6,
    linestyle = "-",
    width = 0.2,
    headlength = 5,
    headwidth = 2.5,
    par = 0,
    paroverhang = 0.02,
    parcolor = "black",
    parlinewidth = 0.6,
    parlinestyle = "-")

    lengthdimension(real(origin/ref), imag(origin/ref),
        real((c+origin)/ref), imag((c+origin)/ref),
        label=label, labeltsep = labeltsep, labelrsep=labelrsep,
        labelrelrot=labelrelrot, labelrelangle = labelrelangle, ha=ha, va=va,
        color=color, backgroundcolor = backgroundcolor,
        arrowstyle1=arrowstyle1, arrowstyle2=arrowstyle2,
        linewidth = linewidth, linestyle=linestyle, width=width,
        headlength=headlength, headwidth=headwidth,
        par=par, paroverhang=paroverhang, parcolor = parcolor,
        parlinewidth=parlinewidth, parlinestyle=parlinestyle)
end


#
# A simple steepest descent algorithm implementation
#
export initializeSolver!, doSolverStep!, getSolverResult
export steepestDescent
export DebugGradient, DebugGradientNorm, DebugStepsize
@doc doc"""
    steepestDescent(M, F, ∇F, x)
perform a steepestDescent $x_{k+1} = \mathrm{retr}_{x_k} s_k\nabla f(x_k)$ with
different choices of $s_k$ available (see `stepsize` option below).

# Input
* `M` – a manifold $\mathcal M$
* `F` – a cost function $F\colon\mathcal M\to\mathbb R$ to minimize
* `∇F`: the gradient $\nabla F\colon\mathcal M\to T\mathcal M$ of F
* `x` – an initial value $x\in\mathcal M$

# Optional
* `stepsize` – ([`ConstantStepsize`](@ref)`(1.)`) specify a [`Stepsize`](@ref)
  functor.
* `retraction` – (`exp`) a `retraction(M,x,ξ)` to use.
* `stoppingCriterion` – (`[`stopWhenAny`](@ref)`(`[`stopAfterIteration`](@ref)`(200), `[`stopWhenGradientNormLess`](@ref)`(10.0^-8))`)
  a functor inheriting from [`StoppingCriterion`](@ref) indicating when to stop.

and the ones that are passed to [`decorateOptions`](@ref) for decorators.

# Output
* `xOpt` – the resulting (approximately critical) point of gradientDescent
* `record` - if activated (using the `record` key, see [`RecordOptions`](@ref)
  an array containing the recorded values.
"""
function steepestDescent(M::mT,
    F::Function, ∇F::Function, x::MP;
    stepsize::Stepsize = ConstantStepsize(1.0),
    retraction::Function = exp,
    stoppingCriterion::StoppingCriterion = stopWhenAny( stopAfterIteration(200), stopWhenGradientNormLess(10.0^-8)),
    kwargs... #collect rest
  ) where {mT <: Manifold, MP <: MPoint}
  p = GradientProblem(M,F,∇F)
  o = GradientDescentOptions(x,stoppingCriterion,stepsize,retraction)
  o = decorateOptions(o; kwargs...)
  resultO = solve(p,o)
  if hasRecord(resultO)
      return getSolverResult(p,getOptions(resultO)), getRecord(resultO)
  end
  return getSolverResult(p,resultO)
end
#
# Solver functions
#
function initializeSolver!(p::P,o::O) where {P <: GradientProblem, O <: GradientDescentOptions}
    o.∇ = getGradient(p,o.x)
end
function doSolverStep!(p::P,o::O,iter) where {P <: GradientProblem, O <: GradientDescentOptions}
    o.∇ = getGradient(p,o.x)
    o.x = o.retraction(p.M, o.x , -getStepsize!(p,o,iter) * o.∇)
end
getSolverResult(p::P,o::O) where {P <: GradientProblem, O <: GradientDescentOptions} = o.x


module mmesh1
using LibFEM
using LibFEM.Mesh
using Test
function test()
    @test Mesh.P1.manifdim == 0
    @test Mesh.L2.nfacets == 2
    true
end
end
using .mmesh1
mmesh1.test()

module mmesh2
using StaticArrays
using LibFEM
using LibFEM.Mesh: Vertices, nvertices, coordinates
using Test
function test()
    xyz = [0.0 0.0; 633.3333333333334 0.0; 1266.6666666666667 0.0; 1900.0 0.0; 0.0 400.0; 633.3333333333334 400.0; 1266.6666666666667 400.0; 1900.0 400.0; 0.0 800.0; 633.3333333333334 800.0; 1266.6666666666667 800.0; 1900.0 800.0]
    v =  Vertices(xyz)
    @test nvertices(v) == 12
    x = coordinates(v, SVector{2}([2, 4]))
    @test x[1] == SVector{2}([633.3333333333334 0.0])
    true
end
end
using .mmesh2
mmesh2.test()

module mmesh3
using StaticArrays
using LibFEM
using LibFEM.Mesh: L2, Q4, Shapes, connectivity, manifdim, nnodes, nfacets, facettype
using Test
function test()
    shapetype = Q4
    c = [(1, 2, 6, 5), (5, 6, 10, 9), (2, 3, 7, 6), (6, 7, 11, 10), (3, 4, 8, 7), (7, 8, 12, 11)]
    cc = [SVector{shapetype.nnodes}(c[idx]) for idx in 1:length(c)]
    shapes = Shapes(shapetype, cc)
    @test connectivity(shapes, 3)[3] == 7
    @test connectivity(shapes, SVector{2}([2, 4]))[1][4] == 9
    @test manifdim(shapes) == 2
    @test nnodes(shapes) == 4
    @test facettype(shapes) == L2
    @test nfacets(shapes) == 4
    true
end
end
using .mmesh3
mmesh3.test()


using DelimitedFiles, Statistics

"""

    housing([test]; [url, file])

Return (xtrn,ytrn,xtst,ytst) from the [UCI Boston
Housing](https://archive.ics.uci.edu/ml/machine-learning-databases/housing)
dataset The dataset has housing related information for 506
neighborhoods in Boston from 1978. Each neighborhood has 14
attributes, the goal is to use the first 13, such as average number of
rooms per house, or distance to employment centers, to predict the
14’th attribute: median dollar value of the houses.

`test=0` by default and `xtrn` (13,506) and `ytrn` (1,506) contain the
whole dataset. Otherwise data is shuffled and split into train and
test portions using the ratio given in `test`. xtrn and xtst are
always normalized by subtracting the mean and dividing into standard
deviation.

"""
function housing(test=0.0;
                 file=joinpath(@__DIR__, "housing", "housing.data"),
                 url="https://archive.ics.uci.edu/ml/machine-learning-databases/housing/housing.data")
    if !isfile(file)
        isdir(dirname(file)) || mkpath(dirname(file))
        @info("Downloading $url to $file")
        download(url, file)
    end
    data = readdlm(file)'
    # @show size(data) # (14,506)
    x = data[1:13,:]
    y = data[14:14,:]
    x = (x .- mean(x,dims=2)) ./ std(x,dims=2) # Data normalization
    if test == 0
        xtrn = xtst = x
        ytrn = ytst = y
    else
        r = randperm(size(x,2))          # trn/tst split
        n = round(Int, (1-test) * size(x,2))
        xtrn=x[:,r[1:n]]
        ytrn=y[:,r[1:n]]
        xtst=x[:,r[n+1:end]]
        ytst=y[:,r[n+1:end]]
    end
    return (xtrn, ytrn, xtst, ytst)
end

nothing


function langevin_inv(y::Real)
  return (3*y - y/5 * (6*y^2 + y^4 - 2*y^6))/(1 - y^2);
end

function free_energy_langevin(k::Real, T::Real, N::Int, γ::Real)
  linvγ = langevin_inv(γ);
  return N*k*T*( log( linvγ / (4*pi*sinh(linvγ)) ) + γ*linvγ );
end

function force_langevin(k::Real, T::Real, b::Real, γ::Real)
  return k*T/b * langevin_inv(γ);
end

function forces(As::Vector, γs::Vector, N::Int)
  n = length(As)
  grads = zeros(n);
  grads[1] = (As[2]-As[1])/(γs[2]-γs[1]);
  grads[end] = (As[n]-As[n-1])/(γs[n]-γs[n-1]);
  for i=2:(n-1)
    grads[i] = (As[i+1] - As[i-1])/(γs[i+1] - γs[i-1]);
  end
  return grads / N;
end


export NwtdirectionCG, NewtonCG
"""
Computes the Newton direction using the conjugate gradient iteravtive scheme
"""
function NwtdirectionCG(H, ∇f; verbose :: Bool = false)
    e = 1e-6
    n = length(∇f)
    τ = 0.5 # need parametrization
    cgtol = max(e, min(0.7, 0.01 * norm(∇f)^(1.0 + τ)))

    # cgTN doesn't work properly yet

    (d, cg_stats) = cg(H, -∇f,
                       atol=cgtol, rtol=0.0,
                       itmax=max(2 * n, 50),
                       verbose=verbose)

    return d
end

function NewtonCG(nlp          :: AbstractNLPModel,
                       nlp_stop     :: NLPStopping;
                       verbose      :: Bool=false,
                       verboseLS    :: Bool = false,
                       kwargs...)
    return  Newton(nlp,
                   nlp_stop;
                   verbose = verbose,
                   verboseLS = verboseLS,
                   Nwtdirection = NwtdirectionCG,
                   hessian_rep = hessian_operator,
                   kwargs...)
end


"""
	csd2by1(Q1, Q2) computes the the Cosine-sine decomposition of
	an (m+l) by l orthogonal matrix Q = [Q1; Q2] where m <= l, such that:

```math
	Q1 = U * C * Z',    Q2 = V * S * Z'
```

- U is m-by-m orthogonal matrix,
- V is l-by-l orthogonal matrix,
- Z is l-by-l orthogonal matrix,
- C and S are  m-by-l and l-by-l ``diagonal" matrices such that
C' * C + S' * S = I. C and S have the following structures
(both alpha and beta are of length l):

(1) If m == l:
        C =  diag(alpha(1), ... , alpha(l))

        S =  diag(beta(1), ... , beta(l))

(2) If m < l:
        C =                       m                 l-m
              m ( diag(alpha(1), ... , alpha(m))     0 )

        S =                       m                 l-m
              m ( diag(beta(1), ... , beta(m)))      0 )
            l-m (                 0                  I )

# arguments:
- Q1: m by l matrix,
- Q2: l by l matrix.

# returns:
- U is m-by-m matrix,
- V is l-by-l matrix,
- Z is l-by-l matrix,
- cosine is an array of length l,
- sine is an array of length l.
"""
# csd2by1() calls csd2by1!() before making a copy of input matrices A and B,
# thus A, B will not be overwritten
function csd2by1(Q1::StridedMatrix{T}, Q2::StridedMatrix{T}) where T
	return csd2by1!(copy(Q1),copy(Q2))
end

function csd2by1!(Q1::StridedMatrix{T}, Q2::StridedMatrix{T}) where T
    m, l = size(Q1)
	if size(Q2)[1] != size(Q2)[2]
		throw(ArgumentError("Q2's row number must equal its column number"))
	elseif m > size(Q2)[1]
		throw(ArgumentError("Q1's row number must be less than that of Q2"))
	elseif l != size(Q2)[2]
		throw(ArgumentError("Q1's column number must equal that of Q2"))
	end

    # Step 1: compute the SVD of Q2
	# May cause LAPACK Exception: S/DBDSDC did not converge
	D1 = svd!(Q2)

	# Step 2.1: rearrange the values of S in non-decreasing order
	sine = reverse(D1.S)

    # Step 2.2: reverse V's col
	#     Assign array slicing to D1.U will throw an error:
	#     D1.U = @views D1.U[:, l:-1:1]
	#     ERROR: setfield! immutable struct of type SVD cannot be changed
	D1.U[:,1:l] = @views D1.U[:, l:-1:1]

    # Step 2.3: reverse Z's col
    D1.V[:,1:l] = @views D1.V[:,l:-1:1]

    # Step 3: find r such that sine(r) <= 1/sqrt(2) < sine(r+1),
	# binary search is used
	r = binSearch(sine)

    # Step 4: compute X = Q1Z
    X = Q1 * D1.V

    # Step 5: compute QR decomposition of X
    D2 = qr!(X)

    # Step 5a: sanitize R so that it only has non-negative diagonal entries
	R = D2.R
    dia1 = sign.(diag(R))
	# update dia1 if R has zeros in the diagonal
	dia1[dia1 .== T(0.0)] .= T(1.0)
	cosine = zeros(T, l)
	@views cosine[1:r] = abs.(diag(R)[1:r])
    R = Diagonal(dia1) * R
	U = Matrix{T}(D2.Q) * Diagonal(dia1)

	# Special case: exit
	if m <= r
		return U, D1.U, D1.V, cosine, sine
	end

	# Step 6: refine R
	R22 = @views R[r+1:m,r+1:l]

	# Step 6a: compute the full SVD of R22
	D3 = svd!(R22, full = true)

	# Step 6b: update the (r+1)-th to m-th cols of U
	U[:,r+1:m] = @views U[:,r+1:m]*D3.U

	# Step 6c: update the last (l-r) cols of Z
	@views D1.V[:,r+1:l] = D1.V[:,r+1:l] * D3.V

	# Step 6d: update cosine and sine
	@views sine[m+1:l] = ones(T, l-m)
	@views cosine[r+1:m] = D3.S

	# Step 7: refine V
    # Step 7a: set W
	W = @views Diagonal(sine[r+1:l])* D3.V

    # Step 7b: compute QR of W
    D4 = qr!(W)
    Qw = D4.Q * Matrix{T}(I, l-r, l-r)

    # Step 7c: sanitize D4.R so that it only has non-negative diagonal entries
    dia2 = sign.(diag(D4.R))
    Qw = Qw * Diagonal(dia2)

    # Step 7d: update V
    @views D1.U[:,r+1:l] = D1.U[:,r+1:l] * Qw

    return U, D1.U, D1.V, cosine, sine
end

# binSearch() modifies standard binary search
# to find r such that beta(r) <= 1/sqrt(2) < beta(r+1),
# return 1 if no such r exist.
function binSearch(a::Vector{T}) where T
	r = 1
	low, high = 1, length(a)
	thld = sqrt(T(0.5))
	while low <= high
		mid = floor(Integer, (high - low)/2) + low
		r = mid
		if mid+1 <= length(a) && a[mid+1] <= thld
			low = mid + 1
		elseif a[mid] > thld
			high = mid - 1
		else
			return mid
		end
	end
	return r
end


#!/usr/bin/env julia

using StochasticCubicDescent
using Base.Test

@testset "gradient descent on a cubic problem" begin
    # define a test function
    H0 = [1. 2.; 3. -4.]
    A = H0' * H0 - 10 * eye(2)  # the eigenvalues of A are -6.18 and 16.18
    b = [-1.6, -2.8]
    ρ = 3.3
    val(x) = 0.5x' * A * x + b' * x + (ρ / 3) * norm(x)^3
    grad!(x, grad) = grad[:] = A * x + b + ρ * norm(x) * x
    cubic_objective = ObjectiveFunction(val, grad!)

    # verify that a gradient is correct
    x = [5.0, 7.0]
    tmp = zeros(2)
    finite_diff_grad_sum = (val(x + 1e-4) - val(x)) / 1e-4
    analytic_grad_sum = sum(grad!(x, tmp))
    @test finite_diff_grad_sum ≈ analytic_grad_sum atol=1e-1

    # Below the step size and initial iterate are set according to
    #     "Gradient Descent Efficiently Finds the Cubic-Rgularized Non-Convex
    #      Newton Step" by Yair Carmon and John Duchi. 2017.
    β = norm(A)
    R = β / (2ρ) + sqrt((β / 2ρ)^2 + norm(b) / ρ)
    η = (8(β + ρ * R))^(-1)

    Rc_lower_bound = -β / ρ + R
    x0 = -Rc_lower_bound / norm(b) * b

    solver = GradientDescent(η, 1000, true, GDState[])
    state = GDState(x0)
    optimize!(solver, state, cubic_objective)

    @test norm(state.grad) ≈ 0.0 atol=1e-8

    # should hold by Lemma 2.2
    iterate_norms = [norm(st.x) for st in solver.history]
    iterate_increases = iterate_norms[2:end] - iterate_norms[1:end-1]
    @test all(iterate_increases .> 0)
end

@testset "Stochastic Cubic Descent" begin
    val(x, ξ) = 0.1x[1]^2 + exp(x[1] + 1) + x[1] * ξ[1]
    grad!(x, ξ, grad) = begin
        grad[1] = 0.2x[1] + exp(x[1] + 1) + ξ[1]
    end
    hv!(x, ξ, v, hv) = begin
        @assert x[1] < 5
        hv[1] = (0.2 + exp(x[1] + 1)) * v[1]
    end
    # very low noise
    ξ!(x, ξ) = ξ[:] = (rand(1) - 0.5) * 1e-4
    # exp(5) isn't a global Lipschitz constant, but it should hold
    # at every point that the algorithm visits
    f = ObjectiveFunction(val, grad!, hv!, ξ!, exp(5))

    solver = SCD(100)
    state = SCDState([-3.])

    optimize!(solver, state, f)

    @test norm(state.grad) ≈ 0.0 atol=1e-2
end


@testset "getGaussPoints" begin
    domain = example_domain(1, 1, 1.0)
    @test getGaussPoints(domain) ≈ [0.21132486540518713 0.21132486540518713; 0.21132486540518713 0.7886751345948129; 0.7886751345948129 0.21132486540518713; 0.7886751345948129 0.7886751345948129]
end

@testset "getStrain" begin
    domain = example_domain(10,10,0.1)
    domain.state = rand(domain.neqs)
    E = getStrain(domain)
    E2 = eval_strain_on_gauss_pts(domain.state, 10,10,0.1)
    @test E≈E2
end

@testset "getStress" begin
    domain = example_domain(10,10,0.1)
    domain.state = rand(domain.neqs)
    E = getStrain(domain)
    S = E*domain.elements[1].mat[1].H
    @test getStress(domain)≈S
end


"""
    gavrikov(δ,θ,Tvn,T₀)

Correlation function for detonation cell width proposed by Gavrikov et al
COMBUSTION AND FLAME 120:19�33 (2000) based on using a reaction zone length based on time
to limiting reactant consumption in constant volume explosion approximation using vn
postshock velocity to convert time to distance.
Tested against a range of fuel-oxidizer diluent mixtures

INPUT:
    δ = reaction zone length based on time to 50consumption of limiting
    reactant from CV computation and delta = time * w_VN
    θ = Eₐ/RT_VN,  effective reduced activation energy based on CV
    computation
    Tvn = von Neumann (postshock temperature behind CJ shock wave)
    T₀ = initial temperature
"""
function gavrikov(δ,θ,Tvn,T₀)

    # Constants
    a = -0.007843787493
    b = 0.1777662961
    c = 0.02371845901
    d = 1.477047968
    e = 0.1545112957
    f = 0.01547021569
    g = -1.446582357
    h = 8.730494354
    i = 4.599907939
    j = 7.443410379
    k = 0.4058325462
    m = 1.453392165
    # Define nondimensional parameters
    X = θ
    Y = Tvn/T₀
    z = Y*(a*Y-b) + X*(c*X-d + (e-f*Y)*Y) + g*log(Y) + h*log(X) + Y*(i/X - k*Y/X^m) - j
    λ = δ*10^z
end

"""
     ng(Δᵢ,χ)

Correlation function for detonation cell size from
Ng, Hoi Dick, Yiguang Ju, and John H. S. Lee. 2007. Assessment of Detonation Hazards in
High-Pressure Hydrogen Storage from Chemical Sensitivity Analysis.
INTERNATIONAL JOURNAL OF HYDROGEN ENERGY 32 (1):93-99.
Tested only against low pressure H2-air data.

INPUT:
    Δᵢ = induction zone length based on peak thermicity in ZND simulation
    χ = 𝜺ᵢ*Δᵢ/Δᵣ where
          𝜺ᵢ = reduced effective activation energy from CV computation
          Δᵢ = distance to peak thermicity from ZND computation
          Δᵣ = w_vN/σ̇_max from ZND computation

See Ng et al. Combustion Theory and Modeling 2005 for a discussion of
the χ parameter.
"""
function ng(Δᵢ,χ)

    # Constants
    A₀ = 30.465860763763
    a₁ = 89.55438805808153
    a₂ = -130.792822369483
    a₃ = 42.02450507117405
    b₁ = -0.02929128383850
    b₂ = 1.0263250730647101E-5
    b₃ = -1.031921244571857E-9
    λ = Δᵢ*(A₀ + ((a₃/χ + a₂)/χ + a₁)/χ + ((b₃*χ + b₂)*χ + b₁)*χ)
end


using OrdinaryDiffEq

function f(du,u,p,t)
    du[1] = dx = p[1]*u[1] - u[1]*u[2]
    du[2] = dy = -3*u[2] + u[1]*u[2]
end

u0 = [1.0;1.0]
tspan = (0.0,10.0)
p = [1.5]
prob = ODEProblem(f,u0,tspan,p)

sol = solve(prob,Tsit5())
t = collect(range(0,stop=10,length=200))
data = Array([(p[1] * sol(t[i])[1] + .01randn()) for i in 1:length(t)])

using DiffEqParamEstim
function loss(sol)
    p = sol.prob.p
    sumsq = 0.0
    for i in 1:length(t)
        sumsq += (data[i] - p[1]*sol(t[i])[1])^2
    end
    sumsq
end


using Statistics

"""
`CALCULATE_AFFINE` - Compute left-hand affine transform from two Nxd point sets

```
tform = calculate_affine(moving_pts, fixed_pts)
```

* moving_pts: Nxd array, each row being homogeneous coords of a point, of the
    point set that is moving to fit the fixed point set.
* fixed_pts: Nxd array, each row being homogeneous coords of a point, of the
    reference point set.
* tform: affine transformation that produces a mapping of the moving_pts to
    fixed_pts.

```
fixed_pts = moving_pts * tform
```
"""
function calculate_affine(moving_pts, fixed_pts)
    return moving_pts \ fixed_pts
end

"""
`CALCULATE_RIGID` - Compute right-hand rigid transform from two Nxd point sets

```
tform = calculate_rigid(moving_pts, fixed_pts)
```

* moving_pts: Nxd array, each row being homogeneous coords of a point, of the
    point set that is moving to fit the fixed point set.
* fixed_pts: Nxd array, each row being homogeneous coords of a point, of the
    reference point set.
* tform: rigid transformation that produces a mapping of the moving_pts to
    fixed_pts.

```
fixed_pts = moving_pts * tform
```
"""
function calculate_rigid(moving_pts, fixed_pts)
    moving_pts = moving_pts[:, 1:end-1]
    fixed_pts = fixed_pts[:, 1:end-1]
    n, dim = size(moving_pts)
    moving_bar = mean(moving_pts, dims = 1)
    fixed_bar = mean(fixed_pts, dims = 1)
    moving_centered = moving_pts .- moving_bar
    fixed_centered = fixed_pts .- fixed_bar
    C = transpose(fixed_centered) * moving_centered / n
    U, s, V = svd(C)
  # Rotation R in least squares sense:
  # moving_pts - moving_bar = (fixed_pts - fixed_bar)*R
    R = transpose(U * diagm(0 => vcat(ones(dim - 1), det(U * transpose(V)))) * transpose(V))
    t = fixed_bar - moving_bar * R
    return [R zeros(dim); t 1]
end

"""
`CALCULATE_TRANSLATION` - Compute left-hand translation from two Nxd point sets

```
tform = calculate_translation(moving_pts, fixed_pts)
```

* moving_pts: Nxd array, each row being nonhomogeneous coords of a point, of the
    point set that is moving to fit the fixed point set.
* fixed_pts: Nxd array, each row being nonhomogeneous coords of a point, of the
    reference point set.
* tform: translation transformation that produces a mapping of the moving_pts to
    fixed_pts.

```
fixed_pts = moving_pts * tform
```
"""
function calculate_translation(moving_pts, fixed_pts)
    n, dim = size(moving_pts)
    moving_bar = mean(moving_pts, dims = 1)
    fixed_bar = mean(fixed_pts, dims = 1)
    t = fixed_bar - moving_bar
    return [Matrix(1.0I, 2, 2) zeros(dim); t 1]
end

function calculate_affine(matches::Matches)
    return calculate_affine(matches.src_points, matches.dst_points)
end

function calculate_rigid(matches::Matches)
    return calculate_rigid(matches.src_points, matches.dst_points)
end

function calculate_translation(matches::Matches)
    return calculate_translation(matches.src_points, matches.dst_points)
end

function calculate_affine(mesh::Mesh)
    return calculate_affine(mesh.src_nodes, mesh.dst_nodes)
end

function calculate_rigid(mesh::Mesh)
    return calculate_rigid(mesh.src_nodes, mesh.dst_nodes)
end

function calculate_translation(mesh::Mesh)
    return calculate_translation(mesh.src_nodes, mesh.dst_nodes)
end

"""
Matrix for rotation in degrees (clockwise rotation with left-handed Cartesian)
"""
function make_rotation_matrix(theta, image_size = nothing)
    angle = deg2rad(theta)
    tform = [cos(angle) sin(angle) 0; -sin(angle) cos(angle) 0; 0 0 1]
    tform[abs(tform).<eps] = 0
    if image_size != nothing
        rotated_bb = snap_bb(tform_bb(sz_to_bb([image_size[1], image_size[2]]), tform))
        rotation_offset = [rotated_bb.i, rotated_bb.j]
        translation_matrix = make_translation_matrix(-rotation_offset)
        tform = tform * translation_matrix
    end
    return tform
end

function make_translation_matrix(offset)
    return [1 0 0; 0 1 0; offset[1] offset[2] 1]
end

function make_scale_matrix(scale)
    return [scale 0 0; 0 scale 0; 0 0 1]
end


using Distributions, LinearAlgebra, Random, ProgressMeter

function test_expectation_muASAmu(D=3, seed=1, numsamples=10_000)

    rg = MersenneTwister(seed)

    μα = randn(rg, D)

    Σα = rand(rg, D, D)*0.2; Σα = Σα*Σα' .+ 0.1

    Apdf = MvNormal(μα, Σα)

    S = Diagonal(rand(rg, D) .+ 2.0) +1.0I

    μf = randn(rg, D)

    f(A) = μf'*A'*inv(S)*A*μf

    g(A) = f(Diagonal(A))[1]

    #————————————————————————————————————————————————
    # This is the result we are after
    #————————————————————————————————————————————————
    # analytical expectation
    l = 0.0
    for i in 1:length(μα)
        l += (μα[i]*μα[i] + Σα[i,i]) * μf[i]^2 / S[i,i]
    end
    #————————————————————————————————————————————————

    # numerical expectation vs analytical
    mean(@showprogress map(g, [rand(Apdf) for _ in 1:numsamples])), l

end


module test_fspace
using ApproXD
using Test

@testset "testing FSpaceXD" begin
	ndims = 4

	# bounds
	lb = [-1.1,-1.5,-0.9,-1.0]
	ub = [1.2,1.6,0.9,1]

	# number of eval points and basis functions:
	# we require square basis matrices!
	npoints = [13,13,13,13]

	# number of basis funcs
	nbasis = npoints

	# splien degrees
	degs = [3,3,3,3]

	# implies a number of knots for each spline
	# remember the restriction that nknots == ncoefs
	nknots = Dict(i => nbasis[i] - degs[i] + 1 for i=1:ndims)

	# eval points
	points = Dict(i => collect(range(lb[i],stop = ub[i], length = npoints[i])) for i=1:ndims)

	# set up ApproXD
	bsp = Dict{Integer,BSpline}()
	for i in 1:ndims
		bsp[i] = BSpline(nknots[i],degs[i],lb[i],ub[i])
	end

	# set of basis functions
	d = Dict{Integer,Array{Float64,2}}()
	for i=1:ndims
		d[i] = Array(getBasis(points[i],bsp[i]))
	end

	# set of INVERSE basis functions
	id = Dict{Integer,Array{Float64,2}}()
	for k in collect(keys(d))
		id[k] = inv(d[k])
	end

	#  get a function
	function f(x,y,z,w)
		sin(sqrt(x^2+y^2)) + (z-w)^3
	end

	y = Float64[f(i,j,k,w) for i in points[1], j in points[2], k in points[3], w in points[4]]

	yvec = y[:]

	# get coefs using the function
	mycoef = getTensorCoef(id,yvec)

	# setup the FSpace
	fx = FSpaceXD(ndims,mycoef,bsp)

	rval1 = lb[1] + 0.3
	rval2 = lb[2] + 0.23
	rval3 = lb[3] + 0.111
	rval4 = lb[4] + 0.099
	# println("approx value = $(getValue(fx,[rval1,rval2,rval3,rval4]))")
	# println("true value = $(f(rval1,rval2,rval3,rval4))")
	@test isapprox(getValue(fx,[rval1,rval2,rval3,rval4]), f(rval1,rval2,rval3,rval4),atol=3e-3)

end

end


using Revise
using AdFem
using PyCall
using LinearAlgebra
using PyPlot
using SparseArrays
using MAT
using ADCMEKit
np = pyimport("numpy")
using PyPlot
using SpecialFunctions

n = 20
NT = 100
ρ = 1.0

m = 5n
h = 1/n
Δt = 1.0/NT

xo = zeros((m+1)*(n+1))
yo = zeros((m+1)*(n+1))
for i = 1:m+1
  for j = 1:n+1
    idx = (j-1)*(m+1)+i
    xo[idx] = (i-1)*h
    yo[idx] = (j-1)*h
  end
end

bdnode = bcnode("left | lower | right", m, n, h)

bdnode = bcnode("all", m, n, h)

μ = 0.5*constant(ones(4*m*n))
η = 1000000. * constant(ones(4*m*n))

coef = 2μ*η/(η + μ*Δt)
mapH = c->begin
  c * diagm(0=>ones(2))
end
H = map(mapH, coef)

Δt = Δt * ones(NT)

M = constant(compute_fem_mass_matrix1(m, n, h))
K = compute_fem_stiffness_matrix1(H, m, n, h)

db = zeros((m+1)*(n+1))
for j = 1:n+1
  idx = (j-1)*(m+1)+1
  if j<=div(n, 4)
    db[idx] = 1.
  else
    db[idx] = (1-(j-div(n, 4))/(3div(n, 4)))
  end
end

# turn the problem into homogeneous Dirichlet problem
homoF = - K * db
idof = ones(Bool, (m+1)*(n+1))
idof[bdnode] .= false
idof = findall(idof)

d0 = zeros((m+1)*(n+1))
v0 = zeros((m+1)*(n+1))
a0 = vector(idof, M[idof, idof]\((-K*db)[idof]), (m+1)*(n+1))

# analytical solution
d0 = @. (5-xo)*xo*(1-yo)*yo
v0 = -d0
a0 = d0

ExtForce = zeros(NT, (m+1)*(n+1))
ts = αscheme_time(Δt)
for i = 1:NT
  t = ts[i]
  f = (x,y)->exp(-t)*(5-x)*x*(1-y)*y - exp(-t)*(-2y*(1-y)-2x*(5-x))
  fval = eval_f_on_gauss_pts(f, m, n, h)
  ExtForce[i,:] = compute_fem_source_term1(fval, m, n, h)
end
ExtForce = constant(ExtForce)

function antiplane_visco_αscheme(M::Union{SparseTensor, SparseMatrixCSC},
  K::Union{SparseTensor, SparseMatrixCSC},
  d0::Union{Array{Float64, 1}, PyObject},
  v0::Union{Array{Float64, 1}, PyObject},
  a0::Union{Array{Float64, 1}, PyObject},
  Δt::Array{Float64};
  ρ::Float64 = 1.0)
  nt = length(Δt)
  αm = (2ρ-1)/(ρ+1)
  αf = ρ/(1+ρ)
  γ = 1/2-αm+αf
  β = 0.25*(1-αm+αf)^2
  d = length(d0)

  M = isa(M, SparseMatrixCSC) ? constant(M) : M
  K = isa(K, SparseMatrixCSC) ? constant(K) : K
  d0, v0, a0, Δt = convert_to_tensor([d0, v0, a0, Δt], [Float64, Float64, Float64, Float64])

  A = (1-αm)*M + (1-αf)*K*β*Δt[1]^2
  A, _ = fem_impose_Dirichlet_boundary_condition1(A, bdnode, m, n, h)

  function equ(dc, vc, ac, dt, εc, σc, i)
    dn = dc + dt*vc + dt^2/2*(1-2β)*ac
    vn = vc + dt*((1-γ)*ac)

    df = (1-αf)*dn + αf*dc
    vf = (1-αf)*vn + αf*vc
    am = αm*ac

    σ_dt = (1-αf)*dt
    Σ = 2* repeat(μ.*η/(η+μ*σ_dt),1,2) .*εc - repeat(η/(η+μ*σ_dt), 1, 2) .* σc
    Force = compute_strain_energy_term1(Σ, m, n, h)

    # rhs = homoF + Force - (M*am + K*df)

    # For wave equation
    rhs = ExtForce[i] - (M*am + K*df)

    rhs = scatter_update(rhs, bdnode, zeros(length(bdnode)))
    A\rhs

  end

  function condition(i, tas...)
    return i<=nt
  end
  function body(i, tas...)
    dc_arr, vc_arr, ac_arr, σc_arr = tas
    dc = read(dc_arr, i)
    vc = read(vc_arr, i)
    ac = read(ac_arr, i)
    σc = read(σc_arr, i)
    εc = eval_strain_on_gauss_pts1(dc, m, n, h)

    op = tf.print(i, norm(σc - εc))
    i = bind(i, op)

    y = equ(dc, vc, ac, Δt[i],εc, σc, i)
    dn = dc + Δt[i]*vc + Δt[i]^2/2*((1-2β)*ac+2β*y)
    vn = vc + Δt[i]*((1-γ)*ac+γ*y)

    εn = eval_strain_on_gauss_pts1(dn, m, n, h)
    σn = 2* repeat(μ.*η/(η+μ*Δt[i]),1,2) .* (εn - εc) + repeat(η/(η+μ*Δt[i]), 1, 2) .* σc

    op = tf.print(i, norm(σn - εn))
    i = bind(i, op)

    i+1, write(dc_arr, i+1, dn), write(vc_arr, i+1, vn), write(ac_arr, i+1, y), write(σc_arr, i+1, σn)
  end

  dM = TensorArray(nt+1); vM = TensorArray(nt+1); aM = TensorArray(nt+1); σM = TensorArray(nt+1);
  dM = write(dM, 1, d0)
  vM = write(vM, 1, v0)
  aM = write(aM, 1, a0)
  # for exact solution
  # ε0 = eval_strain_on_gauss_pts1(d0, m, n, h)
  # σM = write(σM, 1, 2repeat(μ, 1, 2).*ε0)

  σM = write(σM, 1, zeros(4m*n,2))

  i = constant(1, dtype=Int32)
  _, d, v, a = while_loop(condition, body, [i,dM, vM, aM, σM])
  set_shape(stack(d), (nt+1, length(a0))), set_shape(stack(v), (nt+1, length(a0))), set_shape(stack(a), (nt+1, length(a0)))
end

d, v, a = antiplane_visco_αscheme(M, K, d0, v0, a0, Δt, ρ=ρ)

# function solver(A, rhs)
#   A, _ = fem_impose_Dirichlet_boundary_condition1(A, bdnode, m, n, h)
#   rhs = scatter_update(rhs, bdnode, zeros(length(bdnode)))
#   A\rhs
# end
# d, v, a = αscheme(M, spzero((m+1)*(n+1), (m+1)*(n+1)), K, ExtForce, d0, v0, a0, Δt; solve = solver, ρ=ρ)

sess = Session(); init(sess)
d_, v_, a_ = run(sess, [d, v, a])

# pcolormesh(reshape(d_[2,:]+db, m+1, n+1))
# colorbar()

close("all")
for (k,tid) in enumerate(LinRange{Int64}(1, NT+1, 5))
  t = (tid-1)*Δt[1]
  dd = @. exp(-t)*(5-xo)*xo*(1-yo)*yo
  plot(dd[div(n,2)*(m+1) .+ (1:m+1)], "C$k-", label="$tid")
  plot(d_[tid, :][div(n,2)*(m+1) .+ (1:m+1)],"C$k--o", label="$tid", markersize=3)
end
legend()

# close("all")
# for tid in LinRange{Int64}(1, NT+1, 5)
#   plot((a_[tid, :])[1:m+1], label="$tid")
# end
# legend()


using Combinatorics

numbers = open("expense_report.txt") do f
    parse.(Int, readlines(f))
end

for pair in combinations(numbers, 2)
    if sum(pair) == 2020
        println(prod(pair))
    end
end

for pair in combinations(numbers, 3)
    if sum(pair) == 2020
        println(prod(pair))
    end
end


function assemble_system!(problem::FDHeatProblem{TBCL, TBCR, TS}) where {TBCL<:PeriodicBC, TBCR<:PeriodicBC, TS<:BEuler}

    # build first row
    problem.A[1,1] = 1 + 2*problem.Δt/problem.Δx^2;
    problem.A[1,2] = -problem.Δt/problem.Δx^2;
    problem.A[1,end] = -problem.Δt/problem.Δx^2;
    # build rows 2-n
    for i in 2:problem.nx
        problem.A[i,i-1] = -problem.Δt/problem.Δx^2;
        problem.A[i,i] = 1 + 2*problem.Δt/problem.Δx^2;
        problem.A[i,i+1] = -problem.Δt/problem.Δx^2;
    end
    # build row nx+1
    problem.A[end,1] =  -problem.Δt/problem.Δx^2;
    problem.A[end,end-1] =  -problem.Δt/problem.Δx^2;
    problem.A[end,end] =  1 + 2*problem.Δt/problem.Δx^2;

    problem
end


using JuMP

m = Model()
@variable(m, 0<=x<=2)
@variable(m,0<=y<=30)

@objective(m, Max, 5x+3*y)
@constraint(m, 1x+5y>=3.0)

print(m)

status = solve(m)

println("Objective value: ", getobjectivevalue(m))
println("x = ", getvalue(x))
println("y = ", getvalue(y))


function grain_recomb_H(Temp, psi)
    @assert (Temp > 0.0 && psi > 0.0)
    return 12.25e-14 / (1.0 + 8.074e-6 * psi^1.378 * (1.0 + 508.7*Temp^0.01586*psi^(-0.4723-1.102e-5*log(Temp))))
end

function grain_recomb_He(Temp, psi)
    return 5.572e-14 / (1.0 + 3.185e-7 * psi^1.512 * (1.0 + 5115*Temp^3.903e-7*psi^(-0.4956-5.494e-7*log(Temp))))
end

function grain_recomb_C(Temp, psi)
    return 45.58e-14 / (1.0 + 6.089e-3 * psi^1.128 * (1.0 + 433.1*Temp^0.04845*psi^(-0.8120-1.333e-4*log(Temp))))
end

function grain_recomb_Si(Temp, psi)
    return 2.166e-14 / (1.0 + 5.678e-8 * psi^1.874 * (1.0 + 43750*Temp^1.635e-6*psi^(-0.8964-7.538e-5*log(Temp))))
end


using OnlineStatsBase: OnlineStat, EqualWeight, value
using OnlineStats: VectorOb, smooth!, smooth, unbias
import OnlineStats
import OnlineStatsBase
import Statistics

mutable struct VecMean{W} <: OnlineStat{VectorOb}
    μs::Vector{Float64}
    weight::W
    n::Int
end
VecMean(μs::Vector{Float64}; weight = EqualWeight()) = VecMean(μs, weight, 0)
VecMean(len::Integer) = VecMean(zeros(Float64, len))
OnlineStatsBase._fit!(o::VecMean, x::VectorOb) =
    smooth!(o.μs, x, o.weight(o.n += 1))
Statistics.mean(o::VecMean) = value(o)

mutable struct VecVariance{W} <: OnlineStat{VectorOb}
    σ2s::Vector{Float64}     # biased variance
    μs::Vector{Float64}
    tmp::Vector{Float64}
    weight::W
    n::Int
end
function VecVariance(σ2s::Vector{Float64}, μs::Vector{Float64};
                     weight = EqualWeight())
    @assert size(σ2s) == size(μs)
    VecVariance(σ2s, μs, similar(μs), weight, 0)
end
VecVariance(len::Integer) = VecVariance(zeros(Float64, len),
                                        zeros(Float64, len))

function OnlineStatsBase._fit!(o::VecVariance, x::VectorOb)
    γ = o.weight(o.n += 1)
    o.tmp .= smooth.(o.μs, x, γ)
    o.σ2s .= smooth.(o.σ2s, (x .- o.μs) .* (x .- o.tmp), γ)
    o.tmp, o.μs = o.μs, o.tmp
end

OnlineStatsBase.value(o::VecVariance) =
    o.n < 2 ? fill(NaN, size(o.μs)) : o.σ2s .* unbias(o)
Statistics.var(o::VecVariance) = value(o)
Statistics.std(o::VecVariance) = sqrt.(var(o))
Statistics.mean(o::VecVariance) = o.μs


# ## Getting started

using MLJ, RDatasets

auto = dataset("ISLR", "Auto")
y, X = unpack(auto, ==(:MPG), col->true)
train, test = partition(eachindex(y), 0.5, shuffle=true, rng=444);

# Note the use of `rng=` to seed the shuffling of indices so that the results are reproducible.

# ### Polynomial regression
#

@load LinearRegressor pkg=MLJLinearModels

# In this part we only build models with the Horsepower feature.
# Let's get a baseline:

lm = LinearRegressor()
mlm = machine(lm, select(X, :Horsepower), y)
fit!(mlm, rows=train)
rms(predict(mlm, rows=test), y[test])^2

# Note that we square the measure to  match the results obtained in the ISL labs where the mean squared error (here we use the `rms` which is the square root of that).
#
# We now want to build three polynomial models of degree 1, 2 and 3 respectively; we start by forming the corresponding feature matrix:

hp = X.Horsepower
Xhp = DataFrame(hp1=hp, hp2=hp.^2, hp3=hp.^3);

# Now we  can write a simple pipeline where the first step selects the features we want (and with it the degree of the polynomial) and the second is the linear regressor:

@pipeline LinMod(fs = FeatureSelector(features=[:hp1]),
                 lr = LinearRegressor());

# Then we can  instantiate and fit 3 models where we specify the features each time:

lrm = LinMod()
lr1 = machine(lrm, Xhp, y) # poly of degree 1 (line)
fit!(lr1, rows=train)

lrm.fs.features = [:hp1, :hp2] # poly of degree 2
lr2 = machine(lrm, Xhp, y)
fit!(lr2, rows=train)

lrm.fs.features = [:hp1, :hp2, :hp3] # poly of degree 3
lr3 = machine(lrm, Xhp, y)
fit!(lr3, rows=train)

# Let's check the performances on the test set

get_mse(lr) = rms(predict(lr, rows=test), y[test])^2

@show get_mse(lr1)
@show get_mse(lr2)
@show get_mse(lr3)

# ## K-Folds Cross Validation
#
# Let's crossvalidate over the degree of the  polynomial.
#
# **Note**: there's a  bit of gymnastics here because MLJ doesn't directly support a polynomial regression; see our tutorial on [tuning models](/pub/getting-started/model-tuning.html) for a gentler introduction to model tuning.

Xhp = DataFrame([hp.^i for i in 1:10])

cases = [[Symbol("x$j") for j in 1:i] for i in 1:10]
r = range(lrm, :(fs.features), values=cases)

tm = TunedModel(model=lrm, ranges=r, resampling=CV(nfolds=10), measure=rms)

# Now we're left with fitting the tuned model

mtm = machine(tm, Xhp, y)
fit!(mtm)
rep = report(mtm)
@show round.(rep.measurements.^2, digits=2)
@show argmin(rep.measurements)

#
# So the conclusion here is that the 5th order polynomial does quite well.
#
# In ISL they use a different seed so the results are a bit different but comparable.

# ## The Bootstrap
#
# _Bootstrapping is not currently supported in MLJ._


using SocialNet
using Test

using SocialNet: peek

### First test the required behavior

include("code_test_api.jl")

## We have already tested the interface that translates between persons' names
## and `Int`s representing vertices. Here we test the shortest path algorithm
## without this complication. Furthermore, we test all source and target
## vertices for some basic graphs for which the results are obvious.
## Again, the dot notation broadcasts over all source vertices.
@testset "shortest paths on graph types" begin
    g = Graph(4) # empty graph

    # Map over all possible source vertices, 1:4. For the empty graph, all vertices
    # except the source vertex are unreachable, represented by -1.
    @test shortest_path_length.(g, 1:4) ==
        [[0, -1, -1, -1], [-1, 0, -1, -1], [-1, -1, 0, -1], [-1, -1, -1, 0]]

    g = cycle_graph(5) # This graph is a ring.
    @test shortest_path_length.(g, 1:5) ==
        [[0, 1, 2, 2, 1], [1, 0, 1, 2, 2], [2, 1, 0, 1, 2], [2, 2, 1, 0, 1], [1, 2, 2, 1, 0]]

    g = complete_graph(5)
    @test shortest_path_length.(g, 1:5) ==
        [[0, 1, 1, 1, 1], [1, 0, 1, 1, 1], [1, 1, 0, 1, 1], [1, 1, 1, 0, 1], [1, 1, 1, 1, 0]]

    # The first vertex is the central vertex. It shares an edge with all other vertices.
    # Each of the other vertices shares an edge only with the central vertex.
    g = star_graph(5)
    @test shortest_path_length.(g, 1:5) ==
        [[0, 1, 1, 1, 1], [1, 0, 2, 2, 2], [1, 2, 0, 2, 2], [1, 2, 2, 0, 2], [1, 2, 2, 2, 0]]

    # The star-cycle graph is a cycle graph plus a vertex in the middle with an edge to all vertices.
    # The last vertex is the central vertex.
    g = star_cycle_graph(9)
    @test shortest_path_length.(g, 1:9) ==
        [[0, 1, 2, 2, 2, 2, 2, 1, 1], [1, 0, 1, 2, 2, 2, 2, 2, 1], [2, 1, 0, 1, 2, 2, 2, 2, 1],
         [2, 2, 1, 0, 1, 2, 2, 2, 1], [2, 2, 2, 1, 0, 1, 2, 2, 1], [2, 2, 2, 2, 1, 0, 1, 2, 1],
         [2, 2, 2, 2, 2, 1, 0, 1, 1], [1, 2, 2, 2, 2, 2, 1, 0, 1], [1, 1, 1, 1, 1, 1, 1, 1, 0]] # <-- last vertex

end

### Unit tests / diagnostic tests follow

@testset "basic Graph" begin
    g = Graph()
    @test isempty(g)
    @test length(vertices(g)) == 0
    # Implementation detail, but important
    @test first(vertices(g)) == 1
    @test ne(g) == 0
    @test nv(g) == 0

    add_vertex!(g)
    @test ! isempty(g)
    @test vertices(g) == 1:1
    @test ne(g) == 0
    @test nv(g) == 1

    add_vertex!(g)
    add_edge!(g, 1, 2)
    @test ne(g) == 1
    @test nv(g) == 2
end

@testset "basic PriorityQueue" begin
    q = PriorityQueue()

    @test isempty(q) == true
    insert!(q, 1, 100)
    insert!(q, 10, 99)
    insert!(q, 5, 50)

    # The lowest priority, 1, is for item 100
    @test peek(q) == 100

    # Check if an item is in the queue
    @test (50 in q) == true

    # Remove the lowest priority item
    @test pop!(q) == 100

    # Check that the next lowest priorty item has advanced
    @test peek(q) == 50

end


using ..NEPCore
using ..NEPTypes
using ..LinSolvers
export mslp

"""
     mslp([eltype],nep::NEP;[errmeasure,][tol,][maxit,][λ,][logger,][eigsolvertype::Type])

Runs the method of successive linear problems. The  method requires the solution of a
generalized eigenvalue problem in every iteration. The method used for the eigenvalue
computation is specified in eigsolvertype.
See [`augnewton`](@ref) for other parameters.

# Example
Create a rational NEP using a [`SPMF_NEP`](@ref).
```julia-repl
julia> eye=Matrix{Float64}(I,3,3);
julia> Av=[ones(3,3),eye,triu(ones(3,3))];
julia> fv=[S-> S, S -> S^2, S->inv(S-one(S)*10)];
julia> nep=SPMF_NEP(Av,fv);
julia> (λ,v)=mslp(nep);
julia> compute_Mlincomb(nep,λ,v)
3-element Array{Complex{Float64},1}:
 -1.38778e-17+1.65715e-18im
 -5.55112e-17+1.30633e-17im
 -4.16334e-17-1.54436e-17im
```

# References
* A. Ruhe, Algorithms for the nonlinear eigenvalue problem, SIAM J. Numer. Anal. 10 (1973) 674-689

"""
mslp(nep::NEP;params...)=mslp(ComplexF64,nep;params...)
function mslp(::Type{T},
                 nep::NEP;
                 errmeasure::ErrmeasureType = DefaultErrmeasure(nep),
                 tol::Real=eps(real(T))*100,
                 maxit::Integer=100,
                 λ::Number=zero(T),
                 logger=0,
                 eigsolvertype::Type=DefaultEigSolver) where {T<:Number}

    @parse_logger_param!(logger)

    # Ensure types λ is of type T
    λ::T = T(λ)

    # Allocate memory for the eigenvector approximation
    v = zeros(T,size(nep,1));
    d = zeros(T,1);

    err=Inf;

    # Main loop
    for k=1:maxit

        # solve generalized eigenvalue problem
        solver=eigsolvertype(compute_Mder(nep,λ,0),compute_Mder(nep,λ,1));

        # This will throw an error if the eigenvector is not of correct type
        d[:],v[:] = eig_solve(solver,target=0,nev=1);

        # update eigenvalue
        λ += -d[1] # This will throw an error if the eigval update d is not of correct type

        # Normalize
        normalize!(v)

        # Checck for convergence
        err=estimate_error(errmeasure,λ,v)
        push_iteration_info!(logger,k,err=err,λ=λ);

        if (err< tol)
            return (λ,v)
        end

    end

    msg="Number of iterations exceeded. maxit=$(maxit)."
    throw(NoConvergenceException(λ,v,err,msg))
end


"""
    FPlane{FT} <: AbstractRotation

A parameter object for constant rotation around a vertical axis.
"""
struct FPlane{FT} <: AbstractRotation
    f :: FT
end

"""
    FPlane([FT=Float64;] f=nothing, rotation_rate=Ω_Earth, latitude=nothing)

Returns a parameter object for constant rotation at the angular frequency
`f/2`, and therefore with background vorticity `f`, around a vertical axis.
If `f` is not specified, it is calculated from `rotation_rate` and
`latitude` according to the relation `f = 2*rotation_rate*sind(latitude).

By default, `rotation_rate` is assumed to be Earth's.

Also called `FPlane`, after the "f-plane" approximation for the local effect of
a planet's rotation in a planar coordinate system tangent to the planet's surface.
"""
function FPlane(FT::DataType=Float64; f=nothing, rotation_rate=Ω_Earth, latitude=nothing)

    use_f = !isnothing(f)
    use_planet_parameters = !isnothing(latitude)

    if !xor(use_f, use_planet_parameters)
        throw(ArgumentError("Either both keywords rotation_rate and latitude must be " *
                            "specified, *or* only f must be specified."))
    end

    if use_f
        return FPlane{FT}(f)
    elseif use_planet_parameters
        return FPlane{FT}(2rotation_rate*sind(latitude))
    end
end

@inline x_f_cross_U(i, j, k, grid, coriolis::FPlane, U) = - coriolis.f * ℑxyᶠᶜᵃ(i, j, k, grid, U[2])
@inline y_f_cross_U(i, j, k, grid, coriolis::FPlane, U) =   coriolis.f * ℑxyᶜᶠᵃ(i, j, k, grid, U[1])

@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, coriolis::FPlane, U) where FT = zero(FT)

Base.show(io::IO, f_plane::FPlane{FT}) where FT =
    print(io, "FPlane{$FT}: f = ", @sprintf("%.2e", f_plane.f))


immutable Weibull <: ContinuousUnivariateDistribution
    shape::Float64
    scale::Float64
    function Weibull(sh::Real, sc::Real)
    	zero(sh) < sh && zero(sc) < sc || error("Both shape and scale must be positive")
    	new(float64(sh), float64(sc))
    end
end

Weibull(sh::Real) = Weibull(sh, 1.0)

@_jl_dist_2p Weibull weibull

function cdf(d::Weibull, x::Real)
	if 0.0 < x
		return 1.0 - exp(-((x / d.scale)^d.shape))
	else
		0.0
	end
end

function entropy(d::Weibull)
    k, l = d.shape, d.scale
    return ((k - 1.0) / k) * -digamma(1.0) + log(l / k) + 1.0
end

insupport(::Weibull, x::Real) = zero(x) <= x < Inf
insupport(::Type{Weibull}, x::Real) = zero(x) <= x < Inf

function kurtosis(d::Weibull)
    λ, k = d.scale, d.shape
    μ = mean(d)
    σ = std(d)
    γ = skewness(d)
    den = λ^4 * gamma(1.0 + 4.0 / k) -
          4.0 * γ * σ^3 * μ -
          6.0 * μ^2 * σ^2 - μ^4
    num = σ^4
    return den / num - 3.0
end

mean(d::Weibull) = d.scale * gamma(1.0 + 1.0 / d.shape)

median(d::Weibull) = d.scale * log(2.0)^(1.0 / d.shape)

function modes(d::Weibull)
    if d.shape <= 1.0
        return [0.0]
    else
        return [d.scale * ((d.shape - 1.0) / d.shape)^(1.0 / d.shape)]
    end
end

function skewness(d::Weibull)
    tmp = gamma(1.0 + 3.0 / d.shape) * d.scale^3
    tmp -= 3.0 * mean(d) * var(d)
    tmp -= mean(d)^3
    return tmp / std(d)^3
end

var(d::Weibull) = d.scale^2 * gamma(1.0 + 2.0 / d.shape) - mean(d)^2

function mode(d::Weibull)
    inv_k = 1.0 / d.shape
    d.shape > 1.0 ? d.scale * (1.0 - inv_k) ^ inv_k : 0.0
end

modes(d::Weibull) = [mode(d)]


@testset "Increasing/decresing vectors" begin
    @test SpeedyWeather.isincreasing(collect(1:10))
    @test SpeedyWeather.isincreasing(sort(rand(10)))
    @test ~SpeedyWeather.isincreasing(rand(Float32,10))
    @test ~SpeedyWeather.isincreasing(randn(10))
end

@testset "clip negatives" begin
    for T in (Float16,Float32,Float64)
        A = randn(T,30,50)
        SpeedyWeather.clip_negatives!(A)
        @test all(A .>= 0)
    end
end

@testset "roundup nlon for FFT" begin
    for i in 1:10
        @test 2^i == SpeedyWeather.roundup_fft(2^i)
        @test 2^i*3 == SpeedyWeather.roundup_fft(2^i*3)
    end
    for n in 1:10
        i = rand(1:1000)
        @test i <= SpeedyWeather.roundup_fft(i)
    end
end


#!/home/jburdick/bin/julia
# Writes out the coefficients of the bound matrix, both
# using exact binomials, and using various approximations
# (and log-transformed numbers), as a check of the math.

push!(LOAD_PATH, ".")

using CountingBound
using LogCountingBound

"""
  Writes counts for some values of k, r, and n.
  k: the value of k
  maxN: the maximum value of n
  outputDir: directory in which to write output files
  Side effects: writes files A and b.
    (The choices of r are a bit arbitrary).
"""
function writeCounts(k, maxN, outputDir)
  # values of n to use, from k to maxN
  # (for now, spaced at powers of 2)
  nList = [BigInt(2^i) for i in 1:trunc(log(2,maxN))]
  print(nList)
  print("\n")
  print(typeof(nList[1]))
  print("\n")
  # ??? for some reason, this was giving a syntax error
  # nList = [n for n in nList if n>=k && n<=maxN]
  # therefore, using filter()...
  nList = filter(n -> n>=k && n<=maxN, nList)

  # FIXME should create output directory
  # first, write the bound b (and its log)
  of = open(outputDir * "/b_k=" * string(k) * "_maxN=" * string(maxN) * ".csv", "w")
  write(of, "k,n,bound,exp(log(bound)),log(bound)\n")
  for n in nList
    bound = countingBound(binomial(n, 2), binomial(n, k))
		logBound = NaN
		if bound > 0
			logBound = log(bound)
		end
    write(of, string(k) * "," * string(n) * "," * string(bound) * ","
			* string(exp(logBound)) * ","
      * string(logBound) * "\n")
  end
  close(of)

  # then, write the coefficients
  of = open(outputDir * "/A_k=" * string(k) * "_maxN=" * string(maxN) * ".csv", "w")
  write(of, "k,r,n,A,exp(logA),logA\n")
  for n = nList
    # FIXME what should the bound on r be?
    for r = k:min(n, 2*k)
      print("k=" * string(k) * " r=" * string(r) * " n=" * string(n) * "\n")
      A = approxNumMaximalCliques1(k, r, n)
      logA = logNumMaximalCliques(k, r, n)
      write(of, string(k) * "," * string(r) * "," * string(n)
        * "," * string(A)
				* "," * string(exp(logA))
        * "," * string(logA) * "\n")
    end
  end
  close(of)

end

mkpath("coef1")
# smoke tests
writeCounts(2, 200, "coef1")
writeCounts(3, 200, "coef1")
# almost a smoke test
# writeCounts(4, 100, "coef1")
writeCounts(8, 1000, "coef1")
writeCounts(10, 2000, "coef1")
writeCounts(12, 4000, "coef1")


# estimate with MCMC (Turing) for prior (Laplas and half Cauchy
# https://juliastats.org/Distributions.jl/stable/univariate/#Distributions.Laplace
# half-Cauchy not implemented yet
# https://github.com/JuliaStats/Distributions.jl/issues/124
using Turing
using Distributions
using Random
using LinearAlgebra
using StatsPlots

# Generate moc data
n = 50
dim = 1
Random.seed!(99)
u = rand(Normal(0, 0.7), n) # true value σ² = 0.49
x = rand(Uniform(-sqrt(3.0), sqrt(3.0)), n, dim) # x generated from Uniform[-√3, √3]
X = [ones(n) x] # dependent variable matrix with ones
β = Vector(1:1:dim+1) # true coefficients
y = X*β + u # true data generating process

# set prior (Laplace for β, half-Cauchy for σ²)


@testset "Sphere" begin
    # Simple constructor calls
    @test Sphere(zeros(4)) ≈ 0.0 # This is the minimum
    @test Sphere(ones(4)) ≈ 4.0
    @test Sphere(0.0) ≈ 0.0
    @test Sphere(5.0) ≈ 25.0
    # Using multiple dispatch
    @test isa(evaluate(Sphere(), 0.0), Number)
end

@testset "Easom" begin
    @test_throws AssertionError Easom(zeros(4))
    @test Easom(fill(3.0, 2)) ≈ -0.94156415
    minimum_val = @SVector [π, π]
    @test Easom(minimum_val) ≈ -1.0 # Check for minimum
    @test isa(evaluate(Easom(), [2.0, 4.0]), Number)
end

@testset "Ackley" begin
    @test_throws AssertionError Ackley([])
    @test isapprox(Ackley(zeros(5)), 0.0, atol = 1e-15) # Check for minimum
    @test isa(evaluate(Ackley(), [2.0, 4.0]), Number)
end

@testset "Rosenbrock" begin
    @test_throws AssertionError Rosenbrock(zeros(1))
    @test isapprox(Rosenbrock(fill(1.0, 4)), 0.0, atol = 1e-15) # Check for minimum
end

@testset "GoldsteinPrice" begin
    ground_truth = 3.0
    true_vector = @SVector [0.0, -1.0]
    @test_throws AssertionError GoldsteinPrice(zeros(3))
    # Check for minimum
    @test isapprox(GoldsteinPrice(true_vector), ground_truth, atol = 1e-15)
end

@testset "Beale" begin
    ground_truth = 0
    true_vector = @SVector [3, 0.5]
    @test_throws AssertionError Beale(zeros(3))
    # Check for minimum
    @test isapprox(Beale(true_vector), ground_truth, atol = 1e-15)
end

@testset "Levy" begin
    ground_truth = 0
    dimension = 10
    true_vector = fill(1.0, dimension)
    # Check for minimum
    @test isapprox(Levy(true_vector), ground_truth, atol = 1e-15)
end


# Probabilistic Principal Component Analysis

"""Probabilistic PCA type"""
struct PPCA{T<:Real}
    mean::Vector{T}       # sample mean: of length d (mean can be empty, which indicates zero mean)
    W::Matrix{T}          # weight matrix: of size d x p
    σ²::T                 # residual variance
end

## properties

indim(M::PPCA) = size(M.W, 1)
outdim(M::PPCA) = size(M.W, 2)

mean(M::PPCA) = fullmean(indim(M), M.mean)
projection(M::PPCA) = svd(M.W).U # recover principle components from the weight matrix
var(M::PPCA) = M.σ²
loadings(M::PPCA) = M.W

## use

function transform(m::PPCA{T}, x::AbstractVecOrMat{T}) where {T<:Real}
    xn = centralize(x, m.mean)
    W  = m.W
    n = outdim(m)
    M = W'W .+ m.σ² * Matrix{T}(I, n, n)
    return inv(M)*m.W'*xn
end

function reconstruct(m::PPCA{T}, z::AbstractVecOrMat{T}) where {T<:Real}
    W  = m.W
    WTW = W'W
    n = outdim(m)
    M  = WTW .+ var(m) * Matrix{T}(I, n, n)
    return W*inv(WTW)*M*z .+ mean(m)
end

## show

function Base.show(io::IO, M::PPCA)
    print(io, "Probabilistic PCA(indim = $(indim(M)), outdim = $(outdim(M)), σ² = $(var(M)))")
end

## core algorithms

function ppcaml(Z::AbstractMatrix{T}, mean::Vector{T};
                tol::Real=1.0e-6, # convergence tolerance
                maxoutdim::Int=size(Z,1)-1) where {T<:Real}

    check_pcaparams(size(Z,1), mean, maxoutdim, 1.)

    d = size(Z,1)

    # SVD decomposition
    Svd = svd(Z)
    λ = Svd.S
    ord = sortperm(λ; rev=true)
    V = λ[ord]

    # filter 0 eigenvalues and adjust number of latent dimensions
    idxs = findall(V .< tol)
    l = length(idxs)
    l = l == 0 ? maxoutdim : l

    # variance "loss" in the projection
    σ² = sum(V[l+1:end])/(d-l)

    @inbounds for i in 1:l
        V[i] = sqrt(V[i] - σ²)
    end

    W = Svd.U[:,ord[1:l]]*diagm(0 => V[1:l])

    return PPCA(mean, W, σ²)
end

function ppcaem(S::AbstractMatrix{T}, mean::Vector{T}, n::Int;
                maxoutdim::Int=size(S,1)-1,
                tol::Real=1.0e-6,   # convergence tolerance
                maxiter::Integer=1000) where {T<:Real}

    check_pcaparams(size(S,1), mean, maxoutdim, 1.)

    d = size(S,1)
    q = maxoutdim
    Iq = Matrix{T}(I, q, q)
    Id = Matrix{T}(I, d, d)
    W = Matrix{T}(I, d, q)
    σ² = zero(T)
    M⁻¹ = inv(W'W .+ σ² * Iq)

    i = 1
    L_old = 0.
    chg = NaN
    converged = false
    while i < maxiter
        # EM-steps
        SW = S*W
        W⁺  = SW*inv(σ²*Iq + M⁻¹*W'*SW)
        σ²⁺ = tr(S - SW*M⁻¹*(W⁺)')/d
        # new parameters
        W = W⁺
        σ² = σ²⁺
        # log likelihood
        C = W*W'.+ σ²*Id
        M⁻¹ = inv(W'*W .+ σ²*Iq)
        C⁻¹ = (Id - W*M⁻¹*W')/σ²
        L = (-n/2)*(log(det(C)) + tr(C⁻¹*S))  # (-n/2)*d*log(2π) omitted

        @debug "Likelihood" iter=i L=L ΔL=abs(L_old - L)
        chg = abs(L_old - L)
        if chg < tol
            converged = true
            break
        end
        L_old = L
        i += 1
    end
    converged || throw(ConvergenceException(maxiter, chg, oftype(chg, tol)))

    return PPCA(mean, W, σ²)
end

function bayespca(S::AbstractMatrix{T}, mean::Vector{T}, n::Int;
                 maxoutdim::Int=size(S,1)-1,
                 tol::Real=1.0e-6,   # convergence tolerance
                 maxiter::Integer=1000) where {T<:Real}

    check_pcaparams(size(S,1), mean, maxoutdim, 1.)

    d = size(S,1)
    q = maxoutdim
    Iq = Matrix{T}(I, q, q)
    Id = Matrix{T}(I, d, d)
    W = Matrix{T}(I, d, q)
    wnorm = zeros(T, q)
    σ² = zero(T)
    M = W'*W .+ σ²*Iq
    M⁻¹ = inv(M)
    α = zeros(T, q)

    i = 1
    chg = NaN
    L_old = 0.
    converged = false
    while i < maxiter
        # EM-steps
        SW = S*W
        W⁺  = SW*inv(σ²*(Iq+diagm(0=>α)*M/n) + M⁻¹*W'*SW)
        σ²⁺ = tr(S - SW*M⁻¹*(W⁺)')/d
        # new parameters
        W = W⁺
        σ² = σ²⁺
        @inbounds for j in 1:q
            wnorm[j] = norm(W[:,j])^2
            α[j] = wnorm[j] < eps() ? maxintfloat(Float64) : d/wnorm[j]
        end
        # log likelihood
        C = W*W'.+ σ²*Id
        M = W'*W .+ σ²*Iq
        M⁻¹ = inv(M)
        C⁻¹ = (Id - W*M⁻¹*W')/σ²
        L = (-n/2)*(log(det(C)) + tr(C⁻¹*S))  # (-n/2)*d*log(2π) omitted

        @debug "Likelihood" iter=i L=L ΔL=abs(L_old - L)
        chg = abs(L_old - L)
        if chg < tol
            converged = true
            break
        end
        L_old = L
        i += 1
    end
    converged || throw(ConvergenceException(maxiter, chg, oftype(chg, tol)))

    return PPCA(mean, W[:,wnorm .> 0.], σ²)
end

## interface functions

function fit(::Type{PPCA}, X::AbstractMatrix{T};
             method::Symbol=:ml,
             maxoutdim::Int=size(X,1)-1,
             mean=nothing,
             tol::Real=1.0e-6,   # convergence tolerance
             maxiter::Integer=1000) where {T<:Real}

    @assert !SparseArrays.issparse(X) "Use Kernel PCA for sparce arrays"

    d, n = size(X)

    # process mean
    mv = preprocess_mean(X, mean)
    if !(isempty(mv) || length(mv) == d)
        throw(DimensionMismatch("Dimensions of weight matrix and mean are inconsistent."))
    end

    if method == :ml
        Z = centralize(X, mv)
        M = ppcaml(Z, mv, maxoutdim=maxoutdim, tol=tol)
    elseif method == :em || method == :bayes
        S = covm(X, isempty(mv) ? 0 : mv, 2)
        if method == :em
            M = ppcaem(S, mv, n, maxoutdim=maxoutdim, tol=tol, maxiter=maxiter)
        elseif method == :bayes
            M = bayespca(S, mv, n, maxoutdim=maxoutdim, tol=tol, maxiter=maxiter)
        end
    else
        throw(ArgumentError("Invalid method name $(method)"))
    end

    return M::PPCA
end


using Test

input = joinpath(@__DIR__, "input")
thing = map(x -> parse(Int, x), readlines(input))

preamble = thing[1:25]
numbers = thing[26:end]

function calculate(n, preamble)
    for m in preamble
        if (n - m) in preamble
            return true
        end
    end
    return false
end

function calibrate(numbers, preamble)
    for n in numbers
        if !calculate(n, preamble)
            return n
        end
        push!(preamble, n)
        popfirst!(preamble)
    end
end

p1 = calibrate(numbers, preamble)

function crack(target, seq)
    for i in 1:length(seq)
        for j in i:length(seq)
            slice = seq[i:j]
            s = sum(slice)
            if s == target
                return minimum(slice) + maximum(slice)
            elseif s > target
                break
            end
        end
    end
end

p2 = crack(p1, thing)

@assert(p1 == 23278925)
@assert(p2 == 4011064)

sample = [
    35,
    20,
    15,
    25,
    47,
    40,
    62,
    55,
    65,
    95,
    102,
    117,
    150,
    182,
    127,
    219,
    299,
    277,
    309,
    576
]

@test calibrate(sample[6:end], sample[1:5]) == 127
@test crack(127, sample) == 62

println("-----------------------------------------------------------------------")
println("encoding error -- part one :: $p1")
println("encoding error -- part two :: $p2")
println("-----------------------------------------------------------------------")


struct RittWuMethod <: AbstractGeometricProver end

mutable struct RittWuProof{S} <: AbstractGeometricProof
    coords::OrderedDict{Symbol, Coordinate}
    H::Vector{S}
    T::Vector{S}
    R::Matrix{S}
    u::Vector{PolyVar{true}}
    x::Vector{PolyVar{true}}
end

function show(io::IO, rwp::RittWuProof)
    println(io, "Assigned coordinates:")
    println(io, "---------------------")
    for (P, c) in rwp.coords
        println(io, P, " = ", c)
    end
    println(io)
    for (idx, p) in enumerate(eachcol(rwp.R))
        status = iszero(first(p)) ? "success" : "fail"
        println(io, "Goal $idx: $status")
    end
    println(io, "\nNondegeneracy conditions:")
    println(io, "-------------------------")
    for c in ndg(rwp)
        show(io, c)
        println(io, " ≠ 0")
    end
end

function prove(hp::Hypothesis, th::Thesis, ::Type{RittWuMethod}, coords, u, x)
    H = [cc for c in hp.constraints for cc in algebry(c, coords) if !iszero(cc)]
    G = [cc for c in th.constraints for cc in algebry(c, coords)]
    T = triangulize(copy(H), x)

    R = Matrix{eltype(T)}(undef, length(x) + 1, length(G))
    @inbounds for (idx, g) in enumerate(G)
        R[end, idx] = g
        @inbounds for i in length(x):-1:1
            R[i, idx] = pseudorem(R[i+1, idx], T[i], x[i])
        end
    end

    return RittWuProof(coords, H, T, R, u, x)
end

function prove(hp::Hypothesis, th::Thesis, ::Type{RittWuMethod})
    coords, x, u = assign_variables(hp)
    return prove(hp, th, RittWuMethod, coords, u, x)
end

isproved(rwp::RittWuProof) = [iszero(first(p)) for p in eachcol(rwp.R)]
ndg(p::RittWuProof) = [LC(t, xi) for (t, xi) in zip(p.T, p.x) if maxdegree(LC(t, xi)) != 0]


## Lebesgue Measure - Volume of box
## ================================
measure(b::Boxes) = prod([measure(b[dim]) for dim in dims(b)])
logmeasure(b::Boxes) = sum([log(measure(b[dim])) for dim in dims(b)])
measures{B<:Boxes}(bs::Vector{B}) = [measure(b) for b in bs]


function normalise01(x::AbstractArray)
    μ = mean(x)
    return (x .- μ) ./ (stdm(x, μ; corrected=false) .+ eps(eltype(x)))
end

Zygote.@adjoint function stdm(x, μ; corrected::Bool=true)
    σ = stdm(x, μ; corrected=corrected)
    function stdm_sensitivity(Δ)
        xμ = x .- μ
        ΔdivNσ = Δ ./ (σ .* (length(x) - corrected))
        return (xμ .* ΔdivNσ, zero(μ))
    end
    return σ, stdm_sensitivity
end


# SCRIPT

function main()
    hold = 0
    arr = [0, 0, 0, 0]

    fh = open("data/pe_0008.txt","r")

    i = 0
    while !eof(fh)
        c = read(fh, Char)
        if '0' <= c && c <= '9'
            n = parse(Int, c)
            arr[1 + i % 4] = n
            p = prod(arr)
            if hold < p
                hold = p
            end
            i += 1
        end
    end

    close(fh)
    println(hold)
end

main()


@testset "Objective" begin
    T = 3
    num_state = 2
    num_action = 1
    num_parameter = 0
    ot = (x, u, w) -> dot(x, x) + 0.1 * dot(u, u)
    oT = (x, u, w) -> 10.0 * dot(x, x)
    ct = Cost(ot, num_state, num_action, num_parameter=num_parameter)
    cT = Cost(oT, num_state, 0, num_parameter=num_parameter)
    objective = [[ct for t = 1:T-1]..., cT]

    J = [0.0]
    gradient = zeros((T - 1) * (num_state + num_action) + num_state)
    idx_xu = [collect((t - 1) * (num_state + num_action) .+ (1:(num_state + (t == T ? 0 : num_action)))) for t = 1:T]
    x1 = ones(num_state)
    u1 = ones(num_action)
    w1 = zeros(num_parameter)
    X = [x1 for t = 1:T]
    U = [t < T ? u1 : zeros(0) for t = 1:T]
    W = [w1 for t = 1:T]

    ct.evaluate(ct.evaluate_cache, x1, u1, w1)
    ct.gradient(ct.gradient_cache, x1, u1, w1)
    @test ct.evaluate_cache[1] ≈ ot(x1, u1, w1)
    @test norm(ct.gradient_cache - [2.0 * x1; 0.2 * u1]) < 1.0e-8

    cT.evaluate(cT.evaluate_cache, x1, u1, w1)
    cT.gradient(cT.gradient_cache, x1, u1, w1)
    @test cT.evaluate_cache[1] ≈ oT(x1, u1, w1)
    @test norm(cT.gradient_cache - 20.0 * x1) < 1.0e-8

    @test DTO.cost(objective, X, U, X) - sum([ot(X[t], U[t], W[t]) for t = 1:T-1]) - oT(X[T], U[T], W[T]) ≈ 0.0
    DTO.gradient!(gradient, idx_xu, objective, X, U, W)
    @test norm(gradient - vcat([[2.0 * x1; 0.2 * u1] for t = 1:T-1]..., 20.0 * x1)) < 1.0e-8

    # info = @benchmark DTO.cost($objective, $X, $U, $W)
    # info = @benchmark DTO.gradient!($gradient, $idx_xu, $objective, $X, $U, $W)
end


"""
    _reflect(a,b,s,e)

reflection of the Brownian motion `B` where `B_{t-1} = a` and  `B_{t} = b`
on the hypercube `[s,e]^d` where `d = size(a,1)`
"""
function _reflect(a::T, b::T, s::T, e::T) where T <: Vector
    r = 2; n = zeros(size(a))
    # first checking if b is in the hypercube
    all((a .>= s) .& (a .<= e)) ? nothing : error("a = $a not in hypercube")
    size(a) == size(b) ? nothing : error("a not same dim as b")
    for i in 1:length(a)
        if b[i] < s[i]
            rtemp = (a[i] - s[i]) / (a[i] - b[i])
            if rtemp < r
                r = rtemp
                n .= 0
                n[i] = -1
            end
        elseif  b[i] > e[i]
            rtemp =  (e[i] - a[i]) / (b[i]- a[i])
            if rtemp < r
                r = rtemp
                n .= 0
                n[i] = 1
            end
        end
    end
    while r < 1
        c = a + r * ( b - a )
        # dat = hcat(a,c)
        # Plots.plot3d!(dat[1,:],dat[2,:],dat[3,:],label = "",color="blue")
        a = c
        b = b - 2 * n * ( dot(b-c,n))
        r = 2;
        for i in 1:length(a)
            if b[i] < s[i]
                rtemp = (a[i] - s[i]) / (a[i] - b[i])
                if rtemp < r
                    r = rtemp
                    n .= 0
                    n[i] = -1
                end
            elseif  b[i] > e[i]
                rtemp =  (e[i] - a[i]) / (b[i]- a[i])
                if rtemp < r
                    r = rtemp
                    n .= 0
                    n[i] = 1
                end
            end
        end
    end
    # dat = hcat(a,b)
    # Plots.plot3d!(dat[1,:],dat[2,:],dat[3,:],label = "",color="blue")
    return b
end

function _reflect(a::T, b::T, s, e) where T <: CuArray
    @assert all((a .>= s) .& (a .<= e)) "a = $a not in hypercube"
    @assert size(a) == size(b) "a not same dim as b"
    out1 = b .< s
    out2 = b .> e
    out = out1 .| out2
    n = similar(a)
    n .= 0
    # Allocating
    while any(out)
        rtemp1 = @. (s - a) #left
        rtemp2 = @. (e - a) #right
        div = @. (out * (b-a) + !out)
        rtemp = (rtemp1 .* out1 .+ rtemp2 .* out2) ./ div .+ (.!(out1 .| out2))
        rmin = minimum(rtemp,dims=1)
        n .= rtemp .== minimum(rtemp;dims=1)
        c = @. (a + (b-a) * rmin)
        b = @. ( b - 2 * n * (b-c) )
        a = c
        @. out1 = b < s
        @. out2 = b > e
        @. out = out1 | out2
    end
    return b
end


using LazyPseudoinverse, LinearAlgebra
using Test

@testset "LazyPseudoInverse.jl" begin
    A = [1 2 3; 4 5 6; 7. 8 10; 11 13 19]            # a few handy matrices/vectors for testing
    At = collect(A') # the unadulterated adjoint
    B = [1 2 3; 8 9 1; 5. 3 4; 5 9 8]
    b = [5, 8, 10, 11.]
    c = [1.,2.,3.]
    @test typeof(A'A) <: MatrixInnerProd
    @test collect(A'A) == At*A                       # collecting materializes the actual matrix
    @test A'A*collect(B') == At*A*B'                 # right-multiply should just complete
    @test adjointmul(A',A) == At*A                   # same thing, except using adjointmul
    @test A'B == At*B                                # non-inner product
    @test inv(A'A)*B' == cholesky(At*A)\B'           # lazy inverse works with adjoint B'
    @test inv(A'A)c  == cholesky(At*A)\c             # lazy inverse works with vector
    @test inv(A'A)*collect(B') == cholesky(At*A)\B'  # works with regular right-multiplied matrix
    @test B*inv(A'A) ≈ B*inv(cholesky(At*A))         # left-multiplied matrix with lazy PD inverse
    @test inv(A'A)A'b == A\b                         # acts like \ with vector b, using QR
    @test inv(A'A)A'B == A\B                         # acts like \ with matrix B
    @test inv(A'A)A'B == qr(A,Val(true))\B           # ...which should use QR

# Also manually check that show works:  should say MatrixInnerProd
# and show the materialized output.
#  show(stdout, MIME("text/plain"), A'A)           -> MatrixInnerProd
#  show(stdout, MIME("text/plain"), inv(A'A))      -> InvMatrixInnerProd
#  show(stdout, MIME("text/plain"), inv(A'A)A')    -> LazyLeftPseudoinverse
end