module Summary
using Sessions, TaskData, Analysis, SpikeTools, PlotTools, PyPlot, PyCall, JLD,
      HypothesisTests, Distributions, DataFrames, DataStructures, StatsBase, RCall
@pyimport matplotlib.gridspec as gridspec
@pyimport matplotlib.patches as patches
const slice = pybuiltin(:slice)

## Results

unit_mask(chs, region=:MTL, su_only::Bool=false) =
    [(ch.unit_type == :SU || (!su_only && ch.unit_type == :MU)) && ismatch(REGIONS[region], ch.site)
     for ch in chs]

function to_pct(x::AbstractArray{Bool})
    @sprintf "%.2g%% (%d/%d)" mean(x)*100 sum(x) length(x)
end

function region_breakdown(chs, sig::AbstractArray{Bool})
    println("    $(to_pct(sig[[ismatch(REGIONS[:PHG], ch.site) for ch in chs]])) Parahippocampal Units")
    println("    $(to_pct(sig[[ismatch(REGIONS[:EC], ch.site) for ch in chs]])) Entorhinal Units")
    println("    $(to_pct(sig[[ismatch(REGIONS[:H], ch.site) for ch in chs]])) Hippocampal Units")
    println("    $(to_pct(sig[[ismatch(REGIONS[:A], ch.site) for ch in chs]])) Amygdala Units")
end

function stimsel_info()
    regions = [:PHG, :EC, :H, :A]
    @load joinpath(PROCESSED_DATA_DIR, "permpval.jld") chs pval
    selch = unit_mask(chs, :MTL)
    chs = chs[selch]
    pval = pval[selch]
    sig = pval .< 0.001

    println("$(to_pct(sig)) units responsive")
    region_breakdown(chs, sig)
    proportions = [begin
        maskregion = unit_mask(chs, region)
        (sum(sig[maskregion]), sum(maskregion))
    end for region in regions]
    for ir1 = 1:length(regions), ir2 = ir1+1:length(regions)
        p1 = proportions[ir1]
        p2 = proportions[ir2]
        @printf "%3s vs. %3s p = %.3g\n" regions[ir1] regions[ir2] pvalue(FisherExactTest(p1[1], p1[2]-p1[1], p2[1], p2[2]-p2[1]))
    end
end

function showstimulus(f)
    text(0.5, 0.5, f, horizontalalignment="center", verticalalignment="center", fontsize=9)
    axis("off")
end

function plot_stimulus_all(epoch, session, channel_number, cluster_number; header::Bool=true, highlight=nothing, out="pdf")
    const LIMS = epoch == :maintenance ? (-1000, 2400) : (-200, 800)
    spikes = session.spikes[channel_number][cluster_number]
    sesstrials = session.trials
    trials = stimulusframe(sesstrials)
    nstim = length(session.stimuli)
    ticks = epoch == :maintenance ? (0:1000:2000) : 0:600:1000
    gcf()[:set_size_inches](4.488, 2.1)
    if header
        suptitle("$(session.folder) channel $(channel_number) cluster $(cluster_number) $(session.sites[channel_number]) $epoch")
        top = 0.93
    else
        top = 0.995
    end
    gs = gridspec.GridSpec(1, nstim, left=0.06, right=0.98, bottom=0.135, top=top, wspace=0.2, hspace=0.0)
    local ax
    nthimg = 4
    nthraster = 3
    rasterbin = epoch == :maintenance ? 200 : 50
    if isa(highlight, Void)
        highlighttrials = Int[]
    else
        stims = sesstrials[:stimuli].data::Vector{Vector{Int}}
        highlighttrials = find(Bool[highlight in x for x in stims])
    end

    for stim = 1:nstim
        gs_item = gs[:__getitem__]((0, stim-1))
        inner_gs = gridspec.GridSpecFromSubplotSpec(nthimg, 1, subplot_spec=gs_item, hspace=0.02)
        subplot(inner_gs[:__getitem__](0))
        showstimulus(stim)

        gs_item = inner_gs[:__getitem__](slice(1, nthimg))
        inner_gs = gridspec.GridSpecFromSubplotSpec(nthraster, 1, subplot_spec=gs_item, hspace=0.0)
        axa = subplot(inner_gs[:__getitem__]((slice(0, nthraster-1))), zorder=1000)
        local usetrials::Vector{Int}
        if epoch == :sample
            stims = sesstrials[:stimuli].data::Vector{Vector{Int}}
            usetrials = find([stim in x for x in stims])
            onsets = trials[:onset][trials[:stimulus] .== stim].data
        elseif epoch == :maintenance
            stims = sesstrials[:stimuli].data::Vector{Vector{Int}}
            usetrials = find([stim in x for x in stims])
            onsets = sesstrials[:maint][usetrials]
        end
        t1 = findin(usetrials, setdiff(usetrials, highlighttrials))
        t2 = findin(usetrials, highlighttrials)

        plotraster(spikes, onsets[[t1; t2]], LIMS)
        if !isa(highlight, Void)
            axhline(length(t1)+0.5, color=(1, 0, 0, 0.5), linewidth=0.5)
        end
        axis("off")
        axvline(0, ymin=-1/(nthraster-1), ymax=1, linestyle=":", color="darkcyan", dashes=(1, 1),
                   linewidth=0.5, clip_on=false)

        if stim == 1
            ylabel("Presentation Number")
            bax = ax = subplot(inner_gs[:__getitem__](nthraster-1))
        else
            gca()[:set_yticklabels](())
            bax = subplot(inner_gs[:__getitem__](nthraster-1), sharey=ax)
        end
        rast = raster(spikes, onsets, (LIMS[1]:rasterbin:LIMS[2])/1000)
        bar(LIMS[1]:rasterbin:LIMS[2]-rasterbin, mean(rast, 2)*1000/rasterbin, rasterbin, color="k", edgecolor="none", clip_on=false)
        if !isa(highlight, Void)
            bar(LIMS[1]:rasterbin:LIMS[2]-rasterbin, mean(rast[:, t1], 2)*1000/rasterbin,
                rasterbin, color=(1, 0, 0, 0.5), edgecolor="none", clip_on=false)
        end
        if epoch == :maintenance
            xticks(ticks, ["0"; [string(Int(x/1000)) for x in ticks[2:end]]], fontsize=7)
            stim == 1 && xlabel("Time After Mask Offset (Seconds)", fontsize=7, labelpad=2, horizontalalignment="left", x=-0.6)
        else
            xticks(ticks, fontsize=7)
            stim == 1 && xlabel("Time After Stimulus Onset (ms)", fontsize=7, labelpad=2, horizontalalignment="left", x=-0.6)
        end
        if stim == 1
            ylabel("Firing Rate (Hz)", fontsize=7, labelpad=3)
            removespines(["top", "right"])
            gca()[:yaxis][:set_ticks_position]("left")
            gca()[:xaxis][:set_ticks_position]("bottom")
        else
            bax[:axes][:get_yaxis]()[:set_visible](false)
            removespines(["left", "top", "right"])
            gca()[:xaxis][:set_ticks_position]("bottom")
        end
        tick_params(axis="y", pad=2)
        tick_params(axis="x", direction="out", length=3, pad=1)
        xlim(LIMS[1], LIMS[2])
    end
    sca(ax)
    yl = ax[:get_ylim]()
    yticks((0:max(div(yl[2], 3), 1):yl[2]), fontsize=7)
    savefig("$(session.folder)_ch$(channel_number)c$(cluster_number)_$(session.sites[channel_number])_$(session.unit_type[channel_number][cluster_number])_all_$(epoch).$(out)", dpi=300)
end

## Responses Persist Until Subsequent Image Presentations

function read_resptime(about; notfirst::Bool=false)
    @load joinpath(PROCESSED_DATA_DIR, "latency.jld") chs isi_windows isi_windows_notfirst
    if notfirst
        isi_windows = isi_windows_notfirst
    end
    sel = findin(chs, intersect(Analysis.allisisigunits()))
    if about == :latency
        x = [isi_windows[isel][iisi][1] for iisi = 1:4, isel = sel]
    elseif about == :duration
        x = [isi_windows[isel][iisi][2]-isi_windows[isel][iisi][1] for iisi = 1:4, isel = sel]
    else
        error("invalid about $about")
    end
    x, chs[sel]
end

function resptime_anova(about::Symbol=:duration; notfirst::Bool=false)
    x, chs = read_resptime(about, notfirst=notfirst)
    isi = vec(repmat([1, 2, 3, 4], 1, size(x, 2)))
    region = vec(repmat(reshape([string(siteregion(x)) for x in chs], 1, length(chs)), size(x, 1), 1))
    cell = vec(repmat(reshape(1:length(chs), 1, length(chs)), size(x, 1), 1))
    [@assert sum((isi .== y) & (cell .== x)) == 1 for x in cell, y in isi]
    println()
    println(R"""
    library(ez)
    df = data.frame(x=$(vec(x)), isi=factor($isi), region=factor($region), cell=factor($cell))
    ezANOVA(df, x, cell, within=isi, between=region, type=2, detailed=TRUE)
    """)
    for i = 1:length(ISI), j = i+1:length(ISI)
        test = OneSampleTTest(x[i, :], x[j, :])
        @printf("ISI %4d vs. %4d  t(%d) = %.1f p = %.3g\n", ISI[i]*1000, ISI[j]*1000, test.df, test.t, pvalue(test))
    end
end

function resptime_table(about::Symbol=:duration)
    x, chs = read_resptime(about)
    regions = [:PHG, :EC, :H, :A, :All]
    txt = Array(String, length(regions), length(ISI)+1)
    for i = 1:length(regions)
        masked = x[:, unit_mask(chs, regions[i])]
        for j = 1:length(ISI)
            txt[i, j] = @sprintf "%d (%d)" mean(masked[j, :]) std(masked[j, :])
        end
        txt[i, end] = @sprintf "%d (%d)" mean(mean(masked, 1)) std(mean(masked, 1))
    end
    [["" reshape(["ISI $(x)" for x in round(Int, ISI*1000)], 1, 4) "Average"]; regions txt]
end

function plot_isi_single(session, channel_number, cluster_number, stim)
    @load joinpath(PROCESSED_DATA_DIR, "latency.jld") chs isi_windows
    isiwins = isi_windows[findfirst(chs, ClusterMetadata(session, channel_number, cluster_number))]

    const LIMS = (-200, 1500)
    spikes = session.spikes[channel_number][cluster_number]
    trials = stimulusframe(session.trials)
    nstim = maximum(trials[:stimulus])::Int
    byval::Vector{Float64} = sort(unique(trials[:isi]))
    gcf()[:set_size_inches](3.34, 1.5)
    # ioff()
    gs = gridspec.GridSpec(2, 4, left=0.1, right=0.99, bottom=0.2, top=0.9, wspace=0.15, hspace=0.0)
    rasterbin = 100
    local ax
    for ibv = 1:length(byval)
        subplot(gs[:__getitem__]((0, ibv-1)), zorder=1000)
        onsets = trials[:onset][(trials[:stimulus] .== stim) & (trials[:isi] .== byval[ibv])].data
        plotraster(spikes, onsets, LIMS)
        xlim(LIMS[1], LIMS[2])
        axis("off")
        isi = Int(byval[ibv]*1000)
        title("ISI = $(isi) ms", y=0.98, fontsize=8)
        axvline(isiwins[ibv][1], ymin=-1, ymax=1, linestyle=":", color="darkcyan", dashes=(1, 1),
                linewidth=0.5, clip_on=false)
        axvline(isiwins[ibv][2], ymin=-1, ymax=1, linestyle=":", color="darkcyan", dashes=(1, 1),
                linewidth=0.5, clip_on=false)
        println("ISI = $(isi) ms onset = $(isiwins[ibv][1]) offset = $(isiwins[ibv][2]) duration = $(isiwins[ibv][2]-isiwins[ibv][1])")

        if ibv == 1
            setp(gca()[:get_yticklabels](), fontsize=9)
        else
            gca()[:set_yticklabels](())
        end

        if ibv == 1
            bax = ax = subplot(gs[:__getitem__]((1, ibv-1)))
        else
            bax = subplot(gs[:__getitem__]((1, ibv-1)), sharey=ax)
        end
        rast = raster(spikes, onsets, (LIMS[1]:rasterbin:LIMS[2])/1000)
        bar(LIMS[1]:rasterbin:LIMS[2]-rasterbin, mean(rast, 2)*1000/rasterbin, rasterbin, color="k", edgecolor="none", clip_on=false)
        xticks(0:500:1500, ("0", "", "1000", ""), fontsize=7)
        removespines(["top", "right"])
        bax[:yaxis][:set_ticks_position]("left")
        bax[:xaxis][:set_ticks_position]("bottom")
        if ibv == 1
            setp(bax[:get_yticklabels](), fontsize=7)
            setp(bax[:get_xticklabels](), fontsize=7)
            ylabel("Firing Rate (Hz)", fontsize=7, labelpad=3)
            removespines(["top", "right"])
            gca()[:yaxis][:set_ticks_position]("left")
            gca()[:xaxis][:set_ticks_position]("bottom")
            xlabel("Time After Stimulus Onset (ms)", fontsize=7, labelpad=2, horizontalalignment="left", x=-0.4)
        else
            bax[:set_xticklabels](())
            bax[:axes][:get_yaxis]()[:set_visible](false)
            removespines(["left", "top", "right"])
            gca()[:xaxis][:set_ticks_position]("bottom")
        end
        tick_params(axis="y", pad=3)
        tick_params(axis="x", direction="out", length=3, pad=1)
        xlim(LIMS[1], LIMS[2])
    end
    yl = ax[:get_ylim]()
    yticks((0:max(div(yl[2], 4), 1):yl[2]), fontsize=7)
    savefig("$(session.folder)_ch$(channel_number)c$(cluster_number)_isi_stim$stim.pdf", dpi=300)
end

function plot_resptime(about::Symbol=:duration)
    gcf()[:set_size_inches](3.34, 1.5)
    x, chs = read_resptime(about)
    scatter(repeat(ISI*1000, outer=size(x, 2))+(rand(length(x))-0.5)*100, vec(x), 20, marker=".", alpha=0.4, color="k", edgecolor="none")
    means = zeros(4)
    stds = zeros(4)
    for i = 1:4
        xi = x[i, :]
        means[i] = mean(xi)
        stds[i] = std(xi)/sqrt(length(xi))
    end
    errorbar(ISI*1000, means, stds, color="k")
    removespines()
    gca()[:yaxis][:set_ticks_position]("left")
    gca()[:xaxis][:set_ticks_position]("bottom")
    xticks(ISI*1000, fontsize=7)
    yticks(fontsize=7)
    xlim(-100, 900)
    ylim(0, 1085)
    xlabel("Inter-Stimulus Interval (ms)", fontsize=7, labelpad=3)
    ylabel("Response Duration (ms)", fontsize=7, labelpad=2, y=0.45)
    gap = 10
    ht1 = 1025
    len1 = 25
    makeline(ISI[1]*1000, ISI[2]*1000-gap, ht1, ht1, len1, pvalue(OneSampleTTest(x[1, :], x[2, :])))
    makeline(ISI[2]*1000+gap, ISI[3]*1000-gap, ht1, ht1, len1, pvalue(OneSampleTTest(x[2, :], x[3, :])))
    makeline(ISI[3]*1000+gap, ISI[4]*1000, ht1, ht1, len1, pvalue(OneSampleTTest(x[3, :], x[4, :])))
    tight_layout(rect=[-.03, -.07, 1.04, 1.0])
    savefig("$(about)_by_isi.pdf")
end

function resptime_region_differences(about::Symbol=:duration)
    x, chs = read_resptime(about)
    regions = [:PHG, :EC, :H, :A]
    region = [siteregion(ch) for ch in chs]
    regionnum = [findfirst(regions, x) for x in region]
    meandur = vec(mean(x, 1))
    means = zeros(4)
    stds = zeros(4)
    for i = 1:4
        xi = meandur[regionnum .== i]
        means[i] = mean(xi)
        stds[i] = std(xi)/sqrt(length(xi))
    end
    PHG = meandur[regionnum .== 1]
    EC = meandur[regionnum .== 2]
    H = meandur[regionnum .== 3]
    A = meandur[regionnum .== 4]
    @printf("PHG: %d ms\n", mean(PHG))
    @printf("EC: %d ms\n", mean(EC))
    @printf("H: %d ms\n", mean(H))
    @printf("A: %d ms\n", mean(A))
    for ir1 = 1:length(regions), ir2 = ir1+1:length(regions)
        md1 = meandur[regionnum .== ir1]
        md2 = meandur[regionnum .== ir2]
        uet = UnequalVarianceTTest(md1, md2)
        @printf("%3s vs. %3s t(%d) = %.2g p = %.2g\n",
            regions[ir1],
            regions[ir2],
            uet.df,
            uet.t,
            pvalue(uet))
    end
end

function plot_alignedisi(; ch=nothing, compare::Bool=false, pos1::Bool=false)
    @load(joinpath(PROCESSED_DATA_DIR, "alignedinfoisi$(pos1 ? "_pos1" : "").jld"), chs, plot_epoch, alignedinfo,
          alignedinfoperms, alignedinfoisidiff, alignedinfoisidiffperms)
    if compare
        alignedinfo = alignedinfoisidiff
        alignedinfoperms = alignedinfoisidiffperms
    end
    _, sel = read_resptime(:duration)
    @assert sel == chs
    info = cat(3, alignedinfo...)
    alignedinfoperms ./= length(chs)

    clf()
    gcf()[:set_size_inches](3.34, 2)
    win = plot_epoch
    xv = win[1]:win[2]-1
    xvsel = 1:length(xv)
    xv = xv[xvsel]
    info = info[xvsel, :, :]
    mu = mean(info, 3)
    sd = std(info, 3)
    lim = 0.0
    limmin = 0.0
    colors = ["b", "g", "r", "c"]
    compare && shift!(colors)
    for i = 1:size(info, 2)
        y = vec(mu[:, i])
        yerr = vec(sd[:, i])/sqrt(size(info, 3))
        lim = max(lim, maximum(y+yerr))
        limmin = min(limmin, minimum(y-yerr))
        errbarpatch(xv, y-yerr, y+yerr, color=colors[i])
        plot(xv, y, linewidth=0.5, color=colors[i])
    end
    xlabel("Time after Response Onset (ms)", fontsize=7, labelpad=3)
    if compare
        lstrings = ["ISI $(round(Int, ISI[i+1]*1000-200)) vs. $(round(Int, ISI[i]*1000-200)) ms" for i = 1:length(ISI)-1]
    else
        lstrings = ["ISI $(round(Int, x)) ms" for x in ISI*1000-200]
    end
    if pos1 && compare
        ylabel("Proportion Variance Explained (\$\\omega^2\$)", fontsize=7, labelpad=1)
    else
        ylabel("Proportion Variance Explained (\$\\omega^2\$)", fontsize=7)
    end
    removespines()
    gca()[:yaxis][:set_ticks_position]("left")
    gca()[:xaxis][:set_ticks_position]("bottom")
    tick_params(axis="x", length=2)
    xticks(fontsize=7)
    yticks(fontsize=7)
    if pos1 && compare
        tight_layout(rect=[-.06, -.06, 1.04, 1.047])
    else
        tight_layout(rect=[-.07, -.06, 1.04, 1.047])
    end
    xlim((-200, 1000))
    limmin = pos1 || compare ? floor(limmin, 4) : 0.0
    plot([0, 0], [limmin, lim*.99], color="k", linestyle=":")
    plot([200, 200], [limmin, lim*.99], color="b", linestyle=":")
    plot([400, 400], [limmin, lim*.99], color="g", linestyle=":")
    plot([700, 700], [limmin, lim*.99], color="r", linestyle=":")
    plot([1000, 1000], [limmin, lim*.99], color="c", linestyle=":", clip_on=false)
    legend(lstrings, fontsize=7, loc="upper right", bbox_to_anchor=(1, 0.9), borderaxespad=0)

    lim = lim+(lim - limmin)*0.09
    ylim(limmin, lim)
    spacing = lim/50

    maxes = maximum(alignedinfoperms, 1)
    thr = mapslices(x->quantile(x, 0.95), maxes, 2)
    for iisi = 1:size(info, 2)
        pv = vec(mean(maxes[1:1, :, iisi] .>= mu[:, iisi], 2))
        sig = pv .< 0.05
        sigs = find(sig)::Vector{Int}
        biggestsig = (-1, -1)
        for i = 1:length(sigs)
            j = i + 1
            while j <= length(sigs) && sigs[j] == sigs[j-1]+1
                sigs[j] - sigs[i] > biggestsig[2] - biggestsig[1] && (biggestsig = (sigs[i], sigs[j]))
                j += 1
            end
        end
        if compare
            println("ISI $(ISI[iisi+1]*1000) - $(ISI[iisi]*1000) min p = $(minimum(pv)) @ $(xv[indmin(pv)])")
            sum(sig) > 0 && println("ISI $(ISI[iisi+1]*1000) - $(ISI[iisi]*1000) significant in $((xv[biggestsig[1]], xv[biggestsig[2]]))")
        else
            println("ISI $(ISI[iisi]*1000) min p = $(minimum(pv)) @ $(xv[indmin(pv)])")
            sum(sig) > 0 && println("ISI $(ISI[iisi]*1000) significant in $((xv[biggestsig[1]], xv[biggestsig[2]]))")
        end
        scatter(xv[find(sig)], fill(lim-(4.5-iisi)*spacing, sum(sig)), 1, edgecolor="none", color=colors[iisi], zorder=10000, marker="s")
    end
    (pos1 || compare) && axhline(0, color="k", zorder=-1)
    savefig("alignedisi$(compare ? "_compare" : "")$(pos1 ? "_pos1" : "").pdf")
end

## Maintenance Period Activity

function wm_units(epoch=:all)
    @load joinpath(PROCESSED_DATA_DIR, "permpval.jld") chs pval
    chs_pres = chs
    pval_pres = pval
    if epoch == :all
        @load joinpath(PROCESSED_DATA_DIR, "permpval_maintenance.jld") pval chs
    elseif epoch == :last
        @load joinpath(PROCESSED_DATA_DIR, "permpval_maintenance_pos.jld") pval_last chs
        pval = pval_last
    elseif epoch == :notlast
        @load joinpath(PROCESSED_DATA_DIR, "permpval_maintenance_pos.jld") pval_notlast chs
        pval = pval_notlast
    end
    mask = unit_mask(chs)
    pval = pval[mask]
    chs = chs[mask]
    pval_pres = pval_pres[indexin(chs, chs_pres)]
    presalpha = 0.001

    for alpha2 in (0.01, 0.05, 0.001)
        selsel2 = pval[pval_pres .< presalpha] .< alpha2
        @printf("%s of stimulus-selective units significant at maintenance at α = %s (p = %.2g)\n",
        to_pct(selsel2), alpha2, pvalue(BinomialTest(selsel2, alpha2), tail=:right))
    end

    for alpha2 in (0.01, 0.05, 0.001)
        selsel3 = pval[pval_pres .>= presalpha] .< alpha2
        @printf("%s of non-stimulus-selective units significant at maintenance at α = %s (p = %.2g)\n",
        to_pct(selsel3), alpha2, pvalue(BinomialTest(selsel3, alpha2), tail=:right))
    end

    alpha = 0.01
    selmask = pval .< alpha
    rgd = Dict(region => selmask[(pval_pres .< presalpha) & [ismatch(Sessions.REGIONS[region], ch.site) for ch in chs]] for region in (:PHG, :EC, :H, :A))
    for region in (:PHG, :EC, :H, :A)
        selmask = pval .< alpha
        selrg = rgd[region]
        @printf("%s of units in %s show significant maintenance period activity at α = %s (p = %.2g)\n",
                to_pct(selrg), region, alpha, pvalue(BinomialTest(selrg, alpha), tail=:right))
    end
end

const FIGS2_UNITS = [
    "399e11sb-2007-7-21_16-34-42 LPHG ch46 clu1 MU",
    "405e23sb-2008-3-2_10-52-20 LPHG ch54 clu1 MU",
    "409e4sb-2008-7-23_16-8-10 LPHG ch1 clu1 SU",
    "409e19sb-2008-7-29_3-9-13 LPHG ch8 clu1 SU",
    "409e32sb-2008-8-1_3-2-50 LPHG ch6 clu1 MU",
    "409e32sb-2008-8-1_3-2-50 LPHG ch8 clu1 MU",
    "398e12sb-2007-5-16_15-55-53 REC ch2 clu2 SU",
    "396e33sb-2006-12-11_15-56-47 LMH ch43 clu2 SU",
    "399e11sb-2007-7-21_16-34-42 LAH ch36 clu2 MU",
    "399e11sb-2007-7-21_16-34-42 RA ch5 clu1 MU",
    "399e11sb-2007-7-21_16-34-42 RA ch7 clu2 SU",
    "399e11sb-2007-7-21_16-34-42 RA ch7 clu5 SU",
    "399e11sb-2007-7-21_16-34-42 RA ch8 clu3 SU",
    "407e13sb-2008-4-11_14-36-46 RA ch1 clu2 MU",
    "407e13sb-2008-4-11_14-36-46 RA ch5 clu2 SU"
]

function plot_stimulus_both(session, channel_number, cluster_number; highlight=nothing, out="pdf", small::Bool=false)
    small ? gcf()[:set_size_inches](2, 2.5) : gcf()[:set_size_inches](4.488, 3.5)
    fs = small ? 5 : 7
    nstim = length(session.stimuli)
    gs = gridspec.GridSpec(1, nstim, left=0.093, right=0.98, bottom=0.09, top=1.006, wspace=0.2, hspace=0.0)
    nthimg = 10
    nthraster = 3
    gss = Array{Any}(nstim)
    for stim = 1:nstim
        gs_item = gs[:__getitem__]((0, stim-1))
        inner_gs = gridspec.GridSpecFromSubplotSpec(nthimg, 1, subplot_spec=gs_item, hspace=0.02)
        subplot(inner_gs[:__getitem__](0))
        showstimulus(stim)
        gs_item = inner_gs[:__getitem__](slice(1, nthimg))
        gss[stim] = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=gs_item, hspace=0.3)
    end
    for (iepoch, epoch) in enumerate((:sample, :maintenance))
        const LIMS = epoch == :maintenance ? (0, 2400) : (-200, 800)
        spikes = session.spikes[channel_number][cluster_number]
        trials = stimulusframe(session.trials)
        ticks = epoch == :maintenance ? (0:1000:2000) : 0:600:1000
        local ax
        rasterbin = epoch == :maintenance ? 200 : 50
        if isa(highlight, Void)
            highlighttrials = Int[]
        else
            stims = session.trials[:stimuli].data::Vector{Vector{Int}}
            highlighttrials = find(Bool[highlight in x for x in stims])
        end

        for stim = 1:nstim
            inner_gs = gridspec.GridSpecFromSubplotSpec(nthraster, 1, subplot_spec=gss[stim][:__getitem__]((iepoch-1, 0)), hspace=0.0)
            axa = subplot(inner_gs[:__getitem__]((slice(0, nthraster-1))), zorder=1000)
            local usetrials::Vector{Int}
            if epoch == :sample
                stims = session.trials[:stimuli].data::Vector{Vector{Int}}
                usetrials = find([stim in x for x in stims])
                onsets = trials[:onset][trials[:stimulus] .== stim].data
            elseif epoch == :maintenance
                stims = session.trials[:stimuli].data::Vector{Vector{Int}}
                usetrials = find([stim in x for x in stims])
                onsets = session.trials[:maint][usetrials]
            end
            t1 = findin(usetrials, setdiff(usetrials, highlighttrials))
            t2 = findin(usetrials, highlighttrials)

            plotraster(spikes, onsets[[t1; t2]], LIMS)
            if !isa(highlight, Void)
                axhline(length(t1)+0.5, color=(1, 0, 0, 0.5), linewidth=0.5)
            end
            axis("off")
            axvline(0, ymin=-1/(nthraster-1), ymax=1, linestyle=":", color="darkcyan", dashes=(1, 1),
                       linewidth=0.5, clip_on=false)

            if stim == 1
                ylabel("Presentation Number")
                bax = ax = subplot(inner_gs[:__getitem__](nthraster-1))
            else
                gca()[:set_yticklabels](())
                bax = subplot(inner_gs[:__getitem__](nthraster-1), sharey=ax)
            end
            rast = raster(spikes, onsets, (LIMS[1]:rasterbin:LIMS[2])/1000)
            bar(LIMS[1]:rasterbin:LIMS[2]-rasterbin, mean(rast, 2)*1000/rasterbin, rasterbin, color="k", edgecolor="none", clip_on=false)
            if !isa(highlight, Void)
                bar(LIMS[1]:rasterbin:LIMS[2]-rasterbin, mean(rast[:, t1], 2)*1000/rasterbin,
                    rasterbin, color=(1, 0, 0, 0.5), edgecolor="none", clip_on=false)
            end
            if epoch == :maintenance
                xticks(ticks, ["0"; [string(Int(x/1000)) for x in ticks[2:end]]], fontsize=fs)
                stim == 1 && xlabel("Time After Mask Offset (Seconds)", fontsize=fs, labelpad=small ? 1 : 2, horizontalalignment="left", x=-0.6)
            else
                xticks(ticks, fontsize=fs)
                stim == 1 && xlabel("Time After Stimulus Onset (ms)", fontsize=fs, labelpad=small ? 1 : 2, horizontalalignment="left", x=-0.6)
            end
            if stim == 1
                ylabel("Firing Rate (Hz)", fontsize=fs, labelpad=small ? 1 : 3)
                removespines(["top", "right"])
                gca()[:yaxis][:set_ticks_position]("left")
                gca()[:xaxis][:set_ticks_position]("bottom")
            else
                bax[:axes][:get_yaxis]()[:set_visible](false)
                removespines(["left", "top", "right"])
                gca()[:xaxis][:set_ticks_position]("bottom")
            end
            tick_params(axis="y", pad=2)
            tick_params(axis="x", direction="out", length=3, pad=1)
            xlim(LIMS[1], LIMS[2])
        end
        sca(ax)
        yl = ax[:get_ylim]()
        yticks((0:max(div(yl[2], 3), 1):yl[2]), fontsize=fs)
    end
    savefig("$(session.folder)_ch$(channel_number)c$(cluster_number)_$(session.sites[channel_number])_$(session.unit_type[channel_number][cluster_number])_samplemaintenance$(small ? "_small" : "").$(out)", dpi=300)
end

function selunits_table()
    patientid(ch) = parse(Int, ch.session[1:3])
    function nperpatient(patients, chs)
        n = zeros(Int, length(patients))
        for ch in chs
            n[searchsortedfirst(patients, patientid(ch))] += 1
        end
        n
    end

    @load joinpath("processed", "permpval_maintenance.jld") chs pval
    allunit = chs[Summary.unit_mask(chs)]
    chpatient = map(patientid, allunit)
    vissel = Analysis.permselunits()
    maintsel = intersect(vissel, chs[pval .< 0.01])
    patients = sort(unique(chpatient))
    nallunit = nperpatient(patients, allunit)
    nvissel = nperpatient(patients, vissel)
    nmaintsel = nperpatient(patients, maintsel)
    push!(nallunit, sum(nallunit))
    push!(nvissel, sum(nvissel))
    push!(nmaintsel, sum(nmaintsel))
    println("Patient\t# Units\t# Vis\t# Maint\t% Vis\t\t% Maint")
    for i = 1:length(patients)+1
        visci = rcopy(R"binom.test($(nvissel[i]), $(nallunit[i]))$conf.int")
        if nvissel[i] == 0
            maintci = [NaN, NaN]
        else
            maintci = rcopy(R"binom.test($(nmaintsel[i]), $(nvissel[i]))$conf.int")
        end
        patient = i > length(patients) ? "TOTAL" : patients[i]
        @printf("%s\t%d\t%d\t%d\t%.0f%% (%.0f%%-%.0f%%)\t%.0f%% (%.0f%%-%.0f%%)\n",
        patient, nallunit[i], nvissel[i], nmaintsel[i],
        nvissel[i]./nallunit[i]*100, visci[1]*100, visci[2]*100,
        nmaintsel[i]./nvissel[i]*100, maintci[1]*100, maintci[2]*100)
    end
end

function meancortest(x, y, nperm)
    cors = map(cor, x, y)
    meancor = tanh(mean(atanh(cors)))
    permcor = zeros(nperm)
    pvalunit = zeros(length(x))
    for ix = 1:length(x)
        permx = copy(x[ix])
        cury = y[ix]
        unitgt = 0
        for iperm = 1:nperm
            shuffle!(permx)
            v = cor(permx, cury)
            permcor[iperm] += atanh(v)
            unitgt += v >= cors[ix]
        end
        pvalunit[ix] = unitgt/nperm
    end
    map!(x->tanh(x/length(cors)), permcor)
    (meancor, cors, mean(permcor .>= meancor) + 1/nperm, pvalunit)
end

function corscatter()
    wd = 6.5
    ht = 4.55
    regions = [:PHG, :EC, :H, :A]
    region_names = ["PHC", "EC", "Hippocampus", "Amygdala"]
    gcf()[:set_size_inches](wd, ht)
    subplots_adjust(wspace=0.35, hspace=0.4, left=0.072, right=0.995, bottom=0.08, top=0.965)
    selselch = intersect(Analysis.maintselunits(), Analysis.permselunits())
    @load joinpath("processed", "modulation_maintenance.jld") means chs
    sel = findin(chs, selselch)
    chsel = chs
    n = map(length, means[sel])
    @load joinpath("processed", "modulation_presentation.jld") stim_fr_win win chs baseline_fr_all baseline_fr_win
    sel2 = findin(chs, selselch)
    sel2region = [findfirst(regions, siteregion(x)) for x in chs[sel2]]
    sp = sortperm(selselch, lt=(x,y)-> siteregion(x) == siteregion(y) ? string(x.site)[1] < string(y.site)[1] :
                                                                        findfirst(regions, siteregion(x)) < findfirst(regions, siteregion(y)))
    sel = sel[sp]
    sel2 = sel2[sp]
    sel2region = sel2region[sp]
    @assert chsel[sel] == chs[sel2] == selselch[sp]
    fig3ch = findfirst(x->x.session == "399e11sb-2007-7-21_16-34-42" && x.channel == 8 && x.cluster == 2, selselch[sp])
    for (i, isel) = enumerate([1:fig3ch-1; fig3ch+1:length(sel); fig3ch])
        subplot(4, 6, i)
        x = stim_fr_win[sel2[isel]]./(win[sel2[isel]][2] - win[sel2[isel]][1])*1000
        y = means[sel[isel]]./(2400 - 300)*1000
        plot(x, y, color="k", marker="o", markerfacecolor="none", markeredgewidth=0.5, linestyle="none", markersize=2)
        rgx = maximum(x) - minimum(x)
        rgy = maximum(y) - minimum(y)
        stepx = rgx > 20 ? 10 : rgx > 10 ? 5 : rgx > 6 ? 3 : rgx > 6 ? 2 : rgx > 2 ? 1 : rgx > 1 ? 0.5 : rgx > 0.6 ? 0.3 : 0.1
        stepy = rgy > 10 ? 5 : rgy > 6 ? 2 : rgy > 3 ? 1 : rgy > 1.5 ? 0.5 : rgy > 0.6 ? 0.2 : rgy > 0.3 ? 0.1 : 0.05
        xticks(-100*stepx:stepx:100*stepx, fontsize=7)
        yticks(-100*stepy:stepy:100*stepy, fontsize=7)
        minx = minimum(x)-rgx/10
        maxx = maximum(x)+rgx/10
        miny = minimum(y)-rgy/10
        maxy = maximum(y)+rgy/10
        xlim(minx, maxx)
        ylim(miny, maxy)
        tick_params(axis="y", direction="out", pad=2, length=2)
        tick_params(axis="x", direction="out", pad=2, length=2)
        gca()[:yaxis][:set_ticks_position]("left")
        gca()[:xaxis][:set_ticks_position]("bottom")
        if (i - 1) % 6 == 0
            ylabel("Maintenance\nFiring Rate (Hz)", fontsize=7, labelpad=2)
        end
        if i > 18
            xlabel("Sample\nFiring Rate (Hz)", fontsize=7, labelpad=1)
        end
        left, bottom, right, top = gca()[:get_position]()[:extents]
        cortxt = "\$r = $(round(cor(x, y), 2))\$"
        idx = findfirst(FIGS2_UNITS, string(chsel[sel[isel]]))
        if idx != 0
            cortxt = "Fig. S2$('A'+idx-1)\n$cortxt"
        elseif isel == fig3ch
            cortxt = "Fig. 3\n$cortxt"
        end
        text(maxx-rgx/30, miny+rgy/50, cortxt, horizontalalignment="right", verticalalignment="baseline", fontsize=8)
        side = string(chs[sel2[isel]].site)[1] == 'L' ? "Left" : "Right"
        title(side*" "*region_names[sel2region[isel]], y=0.97, fontsize=8)
    end
    savefig("corscatter.pdf")
end

function responsestrength()
    function stimhist(stim_fr, baseline_fr)
        resps = [begin
            stimbs = stim .- bs
            bin = maximum(abs(stimbs)) == minimum(abs(stimbs)) ? indmax(stimbs) : indmax(abs(stimbs))
            stimbs[bin]
            end for (stim, bs) in zip(stim_fr, baseline_fr)]
        resps[!isfinite(resps)] = 0
        edges = -4:0.2:9
        respsc = copy(resps)
        respsc[resps .<= edges[1]] = edges[1]+0.01
        respsc[resps .>= edges[end]] = edges[end]-0.01
        h = fit(Histogram, respsc, edges, closed=:right)
        w = h.weights
        (resps, h.edges[1], w)
    end

    gcf()[:set_size_inches](3.5, 2)
    chs_maintenance = Analysis.maintselunits()
    chs_resp = Analysis.permselunits()
    @load joinpath(PROCESSED_DATA_DIR, "modulation_presentation.jld") chs stim_fr_win baseline_fr_win
    stim_fr = stim_fr_win
    baseline_fr = baseline_fr_win
    chsel = indexin(chs_resp, chs)
    vs, edge, w = stimhist(stim_fr[chsel], baseline_fr[chsel])
    bar(edge[1:end-1], w, step(edge), color=(0.5, 0.5, 0.5), linewidth=0.1)
    maintenancesel = findin(chs[chsel], chs_maintenance)
    selmod = vs[maintenancesel]
    nselmod = vs[setdiff(1:length(vs), maintenancesel)]
    selmod = sort(selmod)
    test = MannWhitneyUTest(abs(selmod), abs(nselmod))
    @printf "Median maintenance: %.1f; median no maintenance: %.1f; U = %.3g, p = %.2g\n" median(abs(selmod)) median(abs(nselmod)) test.U pvalue(test)
    showch = chs[chsel][sortperm(vs)]
    ylim(0.0, maximum(w)+1)
    selmod[selmod .>= edge[end]] = edge[end] - step(edge)/2
    selmod = sort!(selmod)
    y = zeros(length(selmod))
    for i = 2:length(selmod)
        i1 = i
        i2 = i-1
        while i2 > 0 && selmod[i1] - selmod[i2] < 0.1
            y[i] += 1
            i2 -= 1
        end
    end
    @assert length(y) == length(selmod)
    scatter(selmod, y/(maximum(w)+1)*25+1.5, edgecolor="k", facecolor="none", linewidth=0.5, clip_on=false, zorder=1000)
    removespines()
    gca()[:yaxis][:set_ticks_position]("left")
    gca()[:xaxis][:set_ticks_position]("bottom")

    xticks(Int(edge[1]):Int(edge[end]))
    xlim(first(edge), last(edge))
    ylabel("Number of Units")
    xlabel("Change in Firing Rate from Baseline (Number of Spikes)")
    tight_layout(pad=0.1)
    savefig("responsestrength.pdf")
end

function wm_correct(epoch=:maint)
    selch = intersect(Analysis.permselunits(), Analysis.maintselunits())
    jld = jldopen(joinpath(PROCESSED_DATA_DIR, "maintenance_topstim_correct.jld"))
    chs = read(jld, "chs")
    pval_mwu = read(jld, "pval_mwu_$epoch")
    sep = read(jld, "sep_$epoch")
    ng1 = read(jld, epoch == :persist ? "ng1_persist" : "ng1")
    ng2 = read(jld, epoch == :persist ? "ng2_persist" : "ng2")
    close(jld)
    alpha = 0.01

    for (title, mych) = [
            ("all visually selective", Analysis.permselunits()),
            ("visually selective, but not maintenance-selective",
            setdiff(Analysis.permselunits(), selch)),
            ("maintenance-selective", selch)
        ]
        println(title)
        sel = findin(chs, mych)
        sel = sel[(ng1[sel] .> 0) & (ng2[sel] .> 0)]
        @printf "  median %s correct trials, %s incorrect trials\n" median(ng1[sel]) median(ng2[sel])
        @printf("  at α = %s, %s significant (p = %.2g)\n", alpha,
                to_pct(pval_mwu[sel] .< alpha),
                pvalue(BinomialTest(pval_mwu[sel] .< alpha, alpha), tail=:right))
        test = OneSampleTTest(sep[sel])
        @printf("  population effect %.2f, t(%d) = %.3f, p = %.4g\n",
                test.xbar, test.df, test.t, pvalue(test))
        @printf "  bootstrap vs. 0 p = %.2g\n" 2*(min(Analysis.bsvs0(sep[sel], 1000000), Analysis.bsvs0(sep[sel], 1000000)))
    end
end
end


###Test the trait types
for args in [(:mhrunner,MHRunnerType,[:standard,:generalized]),
            (:initialize,InitializeFrom,[:default,:prior]),
            (:propose,ProposeFrom,[:indicator,:auxiliary]),
            (:indicator,IndicatorType,[:stationary,:cyclical]),
            (:jobsegments,JobSegments,[:procs,:workers,:test]),
            (:chain,ChainType,[:standard,:gradient]),
            (:store,StoreDuring,[:burnin,:main,:all])]
    for j in args[3]
        t = trait(args[1],j)
        @test typeof(t) == args[2]
        @test traitvalue(t) == j
        @test traittype(t) == Val{j}
    end
end

###Test MHRuntimePolicy
@test traitvalue(GeneralizedMetropolisHastings._num2propose(1)) == :indicator
@test traitvalue(GeneralizedMetropolisHastings._num2propose(2)) == :auxiliary
@test_throws AssertionError GeneralizedMetropolisHastings._num2propose(0)

@test traitvalue(GeneralizedMetropolisHastings._num2runner(1)) == :standard
@test traitvalue(GeneralizedMetropolisHastings._num2runner(2)) == :generalized
@test_throws AssertionError GeneralizedMetropolisHastings._num2runner(0)

@test traitvalue(GeneralizedMetropolisHastings._num2segments(1,:test)) == :none
@test traitvalue(GeneralizedMetropolisHastings._num2segments(2,:workers)) == :workers
@test_throws AssertionError GeneralizedMetropolisHastings._num2segments(0,:none)

p1 = policy(:mh,10)
p2 = policy(:mh,1,initialize=:default,indicator=:cyclical,chain=:gradient,store=:all,sampletype=Int,calculationtype=Float32)

for args in [(p1,(:generalized,:prior,:auxiliary,:stationary,:workers,:standard,:main,Float64,Float64)),
             (p2,(:standard,:default,:indicator,:cyclical,:none,:gradient,:all,Int,Float32))]
    for (i,f) in enumerate([:runner,:initialize,:propose,:indicator,:jobsegments,:chain,:store])
        @test traittype(getfield(args[1],f)) == Val{args[2][i]}
    end
    @test args[1].sampletype == args[2][end-1]
    @test args[1].calculationtype == args[2][end]
end

println("====================")
println("Test show() function")
show(p1)
show(p2)
println("End  show() function")
println("====================")
println()

nothing


module IdealGas
include("Constants.jl")
include("Utils.jl")
using LightXML

abstract type ComponentDefinition end
abstract type ThermoData end

struct Gasphase <: ComponentDefinition
    species::Array{String,1}
end

struct NASAThermo <: ThermoData
    name::String
    phase::String
    composition::Dict{String,Int64}
    molWt::Float64
    ltl::Float64
    htl::Float64
    cmt::Float64
    ltp::Array{Float64,1}
    htp::Array{Float64,1}
end

#This array will be order as per ig.species
struct SpeciesThermoObj
    molwt::Array{Float64,1}
    thermo_all::Array{NASAThermo}
end
export SpeciesThermoObj

"""
    Tuple of element weights. This is used for the calculation of molecular weights
"""
elementWeight = (    H = 1.00794e-3,  He = 4.002602e-3,  Li = 6.941e-3,  Be = 9.012182e-3,  B = 1.811e-3,
                C = 12.011e-3,  N = 14.00674e-3,  O = 15.9994e-3,  F = 18.9984032e-3,  Ne = 20.1797e-3,
                Na = 22.98977e-3, Mg = 24.3050e-3, Al = 26.98154e-3, Si = 28.0855e-3, P = 30.97376e-3,
                S = 32.066e-3, Cl = 35.4527e-3, Ar = 39.948e-3, K = 39.0983e-3, Ca = 40.078e-3,
                Sc = 44.95591e-3, Ti = 47.88e-3, V = 50.9415e-3, Cr = 51.9961e-3, Mn = 54.9381e-3,
                Fe = 55.847e-3, Co = 58.9332e-3, Ni = 58.69e-3, Cu = 63.546e-3, Zn = 65.39e-3,
                Ga = 69.723e-3, Ge = 72.61e-3, As = 74.92159e-3, Se = 78.96e-3, Br = 79.904e-3,
                Kr = 83.80e-3,  Rb = 85.4678e-3, Sr = 87.62e-3, Y = 88.90585e-3, Zr = 91.224e-3,
                Nb = 92.90638e-3, Mo = 95.94e-3, Tc = 97.9072e-3, Ru = 101.07e-3, Rh = 102.9055e-3,
                Pd = 106.42e-3, Ag = 107.8682e-3,   Cd = 112.411e-3,    In = 114.82e-3, Sn = 118.710e-3,
                Sb = 121.75e-3, Te = 127.6e-3,  I = 126.90447e-3,   Xe = 131.29e-3,   Cs = 132.90543e-3,
                Ba = 137.327e-3, La = 138.9055e-3, Ce = 140.115e-3, Pr = 140.90765e-3, Nd = 144.24e-3,
                Pm = 144.9127e-3, Sm = 150.36e-3,   Eu = 151.965e-3,    Gd = 157.25e-3, Tb = 158.92534e-3,
                Dy = 162.50e-3, Ho = 164.93032e-3,  Er = 167.26e-3, Tm = 168.93421e-3,  Yb = 173.04e-3,
                Lu = 174.967e-3,    Hf = 178.49e-3, Ta = 180.9479e-3,   W = 183.85e-3, Re = 186.207e-3,
                Pt = 195.08e-3, Au = 196.96654e-3,  Hg = 200.59e-3, Tl = 204.3833e-3, Pb = 207.2e-3,
                Bi = 208.98037e-3,  Po = 208.9824e-3,  At = 209.9871e-3,  Rn = 222.0176e-3,
                Fr = 223.0197e-3, Ra = 226.0254e-3, Ac = 227.0279e-3, Th = 232.0381e-3,
                Pa = 231.03588e-3, U = 238.0508e-3, Np = 237.0482e-3, Pu = 244.0482e-3
            )

"""
get_element_weight(el)

function to return the element weights
#   Usage:
    get_element_weight(el)
-   'el::String' : Element name
"""
get_element_weight(el::String) = elementWeight[Symbol(titlecase(strip(el)))]

"""
Function to create thermo object. The function reads the therm.dat file and
parses the content based on the ideal gas object to create the thermo data object
The function returns SpeciesThermoObj
#   Usage:
    create_thermo() or
    create_thermo(file_path)
-   'ig::Gasphase' : Gasphase object obtained using the call create_gasphase
-   'file_path': relative path to the therm.dat file
"""
function create_thermo(species::Array{T,1}, thermo_file::AbstractString ) where T <: AbstractString
    #Create an array of thermo objects which stores the thermo data of all species
    species_thermo = Dict{String,NASAThermo}()
    molecular_weights_vector = Array{Float64,1}()

    open(thermo_file) do io
        lno = 0
        local  species_name, phase, mol_wt,LT,HT,CT
        comp = Dict{String,Int64}()
        local poly_coeff = Array{Float64,1}()
        while !eof(io)
            data_string = readline(io)
            if length(data_string) != 0
                if SubString(data_string,1,3) == "END"
                    break
                else
                    sp = uppercase(strip(SubString(data_string,1,18)))
                    if sp in species
                        lno = 1
                        species_name, phase,comp, mol_wt,LT,HT,CT = parase_thermo_species_data(data_string)
                        lno += 1
                        empty!(poly_coeff)
                    elseif lno > 1 && lno < 5
                        append!(poly_coeff,parse_thermo_polynomials(lno,data_string))
                        lno += 1
                        if lno > 4
                            nasa_thermo = NASAThermo(species_name,phase,comp ,mol_wt,LT,HT,CT,poly_coeff[8:14],poly_coeff[1:7])
                            species_thermo[nasa_thermo.name] = nasa_thermo
                        end
                    end
                end
            end
        end
    end
    species_thermo_array = Array{NASAThermo,1}()
    for sp in species
        push!(species_thermo_array,species_thermo[sp])
    end
    for td in species_thermo_array
        push!(molecular_weights_vector,td.molWt)
    end

    return SpeciesThermoObj(molecular_weights_vector, species_thermo_array)
end

"""
Function for parsing the first line of thermo data.
    It returns the species name, phase, molecular weight, high
    temperature, low temperature and Common temperature limits
    Not for external calls
#   Usage:
    parase_thermo_species_data(data_string)
"""
function parase_thermo_species_data(data_string::AbstractString)
    species_name = uppercase(strip(SubString(data_string,1,18)))
    a1 = split(strip(SubString(data_string,25,29)))
    a2 = split(strip(SubString(data_string,30,34)))
    a3 = split(strip(SubString(data_string,35,39)))
    a4 = split(strip(SubString(data_string,40,44)))
    elements_data = (a1,a2,a3,a4)
    mol_wt::Float64 = 0
    composition = Dict{String,Int64}()
    for el in elements_data
        if !isempty(el)
            mol_wt += parse(Float64,el[2])*get_element_weight(String(el[1]))
            composition[String(el[1])] = parse(Int64,el[2])
        end
    end
    phase = String(SubString(data_string,45,45))
    if length(phase)==0
        throw(error("Phase not specified for species $spname"))
    end
    LTs = String(SubString(data_string,46,57))
    if length(LTs) == 0
        throw(error("Low temperature limit not specified for species $spname"))
    end
    HTs = String(SubString(data_string,58,67))
    if length(HTs) == 0
        throw(error("High temperature limit not specified for species $spname"))
    end
    CTs = String(SubString(data_string,68,75))
    if length(CTs)==0
        throw(error("Common temperature limit not specified for species $spname"))
    end
    LT = parse(Float64, LTs)
    HT = parse(Float64, HTs)
    CT = parse(Float64, CTs)
    return (species_name,phase, composition ,mol_wt,LT,HT,CT)
end

"""
This function will parse the polynomials
    Not for external calls
"""
function parse_thermo_polynomials(lno::Int64, thdata::AbstractString)
    poly_data = Array{Float64,1}()
    if lno == 2
        push!(poly_data,parse(Float64,String(SubString(thdata,1,15))))
        push!(poly_data,parse(Float64,String(SubString(thdata,16,30))))
        push!(poly_data,parse(Float64,String(SubString(thdata,31,45))))
        push!(poly_data,parse(Float64,String(SubString(thdata,46,60))))
        push!(poly_data,parse(Float64,String(SubString(thdata,61,75))))
    elseif lno == 3
        push!(poly_data,parse(Float64,String(SubString(thdata,1,15))))
        push!(poly_data,parse(Float64,String(SubString(thdata,16,30))))
        push!(poly_data,parse(Float64,String(SubString(thdata,31,45))))
        push!(poly_data,parse(Float64,String(SubString(thdata,46,60))))
        push!(poly_data,parse(Float64,String(SubString(thdata,61,75))))
    elseif lno == 4
        push!(poly_data,parse(Float64,String(SubString(thdata,1,15))))
        push!(poly_data,parse(Float64,String(SubString(thdata,16,30))))
        push!(poly_data,parse(Float64,String(SubString(thdata,31,45))))
        push!(poly_data,parse(Float64,String(SubString(thdata,46,60))))
    end
    return poly_data
end

"""
Calculate the specific heat of pure species J/mol-K
#   Usage-1:
    cp(thermo,T)
-   'thermo::NASAThermo': NASAThermo of the species
-   'T::Float64': Temperature in K at which the property is required
# Usage-2:
    cp(sp,T,thermo,ig)
-   'sp::String' : species name
-   'T::Float64' : Temperature K
-   'thermo::SpeciesThermoObj' : Structure of SpeciesThermoObj
-   'ig::Gasphase ' : Gasphase object
"""
function cp(thermo::NASAThermo, T::Float64)
    TVec = [T^i for i in 0:4]
    T < thermo.cmt ? sum(thermo.ltp[1:5] .* TVec)*R : sum(thermo.htp[1:5] .* TVec)*R
end
cp(sp::String,T::Float64,thermoObj::SpeciesThermoObj,species_list::Array{String,1}) = cp(thermoObj.thermo_all[get_index(sp,species_list)],T)

"""
Calculates the enthalpy of pure species J/mol
#   Usage-1:
    H(thermo,T)
-   'thermo::NASAThermo': NASAThermo of the species
-   'T::Float64': Temperature in K at which the property is required
# Usage-2:
    H(sp,T,thermo,ig)
-   'sp::String' : species name
-   'T::Float64' : Temperature K
-   'thermo::SpeciesThermoObj' : Structure of SpeciesThermoObj
-   'ig::Gasphase ' : Gasphase object
"""
function H(thermo::NASAThermo, T::Float64)
    TVec = [1, T/2.0, T^2/3.0, T^3/4.0, T^4/5.0, 1.0/T]
    T < thermo.cmt ? sum(thermo.ltp[1:6] .* TVec)*R*T : sum(thermo.htp[1:6] .* TVec)*R*T
end
H(sp::String,T::Float64,thermoObj::SpeciesThermoObj,species_list::Array{String,1}) = H(thermoObj.thermo_all[get_index(sp,species_list)],T)

"""
Calculates the entropy of pure species J/mol-K
#   Usage-1:
    S(thermo,T)
-   'thermo::NASAThermo': NASAThermo of the species
-   'T::Float64': Temperature in K at which the property is required
# Usage-2:
    S(sp,T,thermo,ig)
-   'sp::String' : species name
-   'T::Float64' : Temperature K
-   'thermo::SpeciesThermoObj' : Structure of SpeciesThermoObj
-   'ig::Gasphase ' : Gasphase object
"""
function S(thermo::NASAThermo, T::Float64)
    TVec = [log(T), T, T^2/2.0, T^3/3.0, T^4/4.0]
    if T < thermo.cmt
        return (sum(thermo.ltp[1:5] .* TVec) + thermo.ltp[7])*R
    else
        return (sum(thermo.htp[1:5] .* TVec) + thermo.htp[7])*R
    end
end
S(sp::String,T::Float64,thermoObj::SpeciesThermoObj,species_list::Array{String,1}) = S(thermoObj.thermo_all[get_index(sp,species_list)],T)

"""
Calculates the specific heat of all species in J/mol-K
#   Usage
    cp_all(td,T)
-   'thermoObj::SpeciesThermoObj' : Structure of SpeciesThermoObj
-   'T::Float64' : Temperature in K at which the property is rquired
"""
function cp_all(thermoObj::SpeciesThermoObj,T::Float64)
    return map(x->cp(x,T),thermoObj.thermo_all)
end

"""
Calculates the enthalpy of all species in J/mol
#   Usage
    H_all(td,T)
-   'thermoObj::SpeciesThermoObj' : Structure of SpeciesThermoObj
-   'T::Float64' : Temperature in K at which the property is rquired
"""
function H_all(thermoObj::SpeciesThermoObj,T::Float64)
    return map(x->H(x,T),thermoObj.thermo_all)
end

"""
Calculates the entropy of all species in J/mol-K
#   Usage
    S_all(td,T)
-   'thermoObj::SpeciesThermoObj' : Structure of SpeciesThermoObj
-   'T::Float64' : Temperature in K at which the property is rquired
"""
function S_all(thermoObj::SpeciesThermoObj,T::Float64)
    return map(x->S(x,T),thermoObj.thermo_all)
end

"""
Calculates the enthalpy of a mixture J/mol
#   Usage
    Hmix(td,T,mlf)
-   'thermoObj::SpeciesThermoObj' : Structure of SpeciesThermoObj
-   'T::Float64' : Temperature in K
-   'mlf::Array{Float64,1}' : species mole fractions
"""
function Hmix(thermoObj::SpeciesThermoObj,T::Float64,mlf::Array{Float64,1})
    all_species_h = H_all(thermoObj,T)
    return sum(all_species_h .* mlf)
end

"""
Calculates the specific heat of a mixture in J/mol-K
#   Usage
    cpmix(td,T,mlf)
-   'thermoObj::SpeciesThermoObj' : Structure of SpeciesThermoObj
-   'T::Float64' : Temperature in K
-   'mlf::Array{Float64,1}' : species mole fractions
"""
function cpmix(thermoObj::SpeciesThermoObj,T::Float64,mlf::Array{Float64,1})
    all_species_cp = S_all(thermoObj,T)
    return sum(all_species_cp .* mlf)
end

"""
Calculates the entropy of a muxture in J/mol-K
#   Usage
    Smix(thermoObj,T,p,mlf)
-   'thermoObj::SpeciesThermoObj' : Structure of SpeciesThermoObj
-   'T::Float64' : Temperature in K
-   'p::Float64' : total pressure Pa
-   'mlf::Array{Float64,1} ': mole fractions
"""
function Smix(thermoObj::SpeciesThermoObj,T::Float64,p::Float64,mlf::Array{Float64,1})
    all_species_s = S_all(thermoObj,T)
    xp = mlf .* (p/p_std)
    #logterm = [x <= 0 ? 0 : log(x) for x in xp]
    logterm = collect(map(x->x <= 0 ? 0 : log(x) ,xp))
    return sum((all_species_s - (logterm .* R)) .* mlf)
end

"""
Calculates the Gibbs free energy of a muxture in J/mol
#   Usage
    Gmix(thermoObj,T,p,mlf)
-   'thermoObj::SpeciesThermoObj' : Structure of SpeciesThermoObj
-   'T::Float64' : Temperature in K
-   'p::Float64' : total pressure Pa
-   'mlf::Array{Float64,1} ': mole fractions
"""
function  Gmix(thermoObj::SpeciesThermoObj, T::Float64,p::Float64,mlf::Array{Float64,1})
    hmix = Hmix(thermoObj,T,mlf)
    smix = Smix(thermoObj,T,p,mlf)
    return hmix - T*smix
end

#end of module IdealGas
end


function get_generator_data(data::OPFData; use_gpu=false)
    ngen = length(data.generators)

    if use_gpu
        pgmin = CuArray{Float64}(undef, ngen)
        pgmax = CuArray{Float64}(undef, ngen)
        qgmin = CuArray{Float64}(undef, ngen)
        qgmax = CuArray{Float64}(undef, ngen)
        c2 = CuArray{Float64}(undef, ngen)
        c1 = CuArray{Float64}(undef, ngen)
        c0 = CuArray{Float64}(undef, ngen)
    else
        pgmin = Array{Float64}(undef, ngen)
        pgmax = Array{Float64}(undef, ngen)
        qgmin = Array{Float64}(undef, ngen)
        qgmax = Array{Float64}(undef, ngen)
        c2 = Array{Float64}(undef, ngen)
        c1 = Array{Float64}(undef, ngen)
        c0 = Array{Float64}(undef, ngen)
    end

    Pmin = Float64[data.generators[g].Pmin for g in 1:ngen]
    Pmax = Float64[data.generators[g].Pmax for g in 1:ngen]
    Qmin = Float64[data.generators[g].Qmin for g in 1:ngen]
    Qmax = Float64[data.generators[g].Qmax for g in 1:ngen]
    coeff0 = Float64[data.generators[g].coeff[3] for g in 1:ngen]
    coeff1 = Float64[data.generators[g].coeff[2] for g in 1:ngen]
    coeff2 = Float64[data.generators[g].coeff[1] for g in 1:ngen]
    copyto!(pgmin, Pmin)
    copyto!(pgmax, Pmax)
    copyto!(qgmin, Qmin)
    copyto!(qgmax, Qmax)
    copyto!(c0, coeff0)
    copyto!(c1, coeff1)
    copyto!(c2, coeff2)

    return pgmin,pgmax,qgmin,qgmax,c2,c1,c0
end

function get_bus_data(data::OPFData; use_gpu=false)
    nbus = length(data.buses)

    FrIdx = [l for b=1:nbus for l in data.FromLines[b]]
    ToIdx = [l for b=1:nbus for l in data.ToLines[b]]
    GenIdx = [g for b=1:nbus for g in data.BusGenerators[b]]
    FrStart = accumulate(+, vcat([1], [length(data.FromLines[b]) for b=1:nbus]))
    ToStart = accumulate(+, vcat([1], [length(data.ToLines[b]) for b=1:nbus]))
    GenStart = accumulate(+, vcat([1], [length(data.BusGenerators[b]) for b=1:nbus]))

    Pd = Float64[data.buses[i].Pd for i=1:nbus]
    Qd = Float64[data.buses[i].Qd for i=1:nbus]

    if use_gpu
        cuFrIdx = CuArray{Int}(undef, length(FrIdx))
        cuToIdx = CuArray{Int}(undef, length(ToIdx))
        cuGenIdx = CuArray{Int}(undef, length(GenIdx))
        cuFrStart = CuArray{Int}(undef, length(FrStart))
        cuToStart = CuArray{Int}(undef, length(ToStart))
        cuGenStart = CuArray{Int}(undef, length(GenStart))
        cuPd = CuArray{Float64}(undef, nbus)
        cuQd = CuArray{Float64}(undef, nbus)

        copyto!(cuFrIdx, FrIdx)
        copyto!(cuToIdx, ToIdx)
        copyto!(cuGenIdx, GenIdx)
        copyto!(cuFrStart, FrStart)
        copyto!(cuToStart, ToStart)
        copyto!(cuGenStart, GenStart)
        copyto!(cuPd, Pd)
        copyto!(cuQd, Qd)

        return cuFrStart,cuFrIdx,cuToStart,cuToIdx,cuGenStart,cuGenIdx,cuPd,cuQd
    else
        return FrStart,FrIdx,ToStart,ToIdx,GenStart,GenIdx,Pd,Qd
    end
end

function get_branch_data(data::OPFData; use_gpu=false)
    buses = data.buses
    lines = data.lines
    BusIdx = data.BusIdx
    nline = length(data.lines)
    ybus = Ybus{Array{Float64}}(computeAdmitances(data.lines, data.buses, data.baseMVA; VI=Array{Int}, VD=Array{Float64})...)
    frBound = [ x for l=1:nline for x in (buses[BusIdx[lines[l].from]].Vmin^2, buses[BusIdx[lines[l].from]].Vmax^2) ]
    toBound = [ x for l=1:nline for x in (buses[BusIdx[lines[l].to]].Vmin^2, buses[BusIdx[lines[l].to]].Vmax^2) ]

    if use_gpu
        cuYshR = CuArray{Float64}(undef, length(ybus.YshR))
        cuYshI = CuArray{Float64}(undef, length(ybus.YshI))
        cuYffR = CuArray{Float64}(undef, nline)
        cuYffI = CuArray{Float64}(undef, nline)
        cuYftR = CuArray{Float64}(undef, nline)
        cuYftI = CuArray{Float64}(undef, nline)
        cuYttR = CuArray{Float64}(undef, nline)
        cuYttI = CuArray{Float64}(undef, nline)
        cuYtfR = CuArray{Float64}(undef, nline)
        cuYtfI = CuArray{Float64}(undef, nline)
        cuFrBound = CuArray{Float64}(undef, 2*nline)
        cuToBound = CuArray{Float64}(undef, 2*nline)
        copyto!(cuYshR, ybus.YshR)
        copyto!(cuYshI, ybus.YshI)
        copyto!(cuYffR, ybus.YffR)
        copyto!(cuYffI, ybus.YffI)
        copyto!(cuYftR, ybus.YftR)
        copyto!(cuYftI, ybus.YftI)
        copyto!(cuYttR, ybus.YttR)
        copyto!(cuYttI, ybus.YttI)
        copyto!(cuYtfR, ybus.YtfR)
        copyto!(cuYtfI, ybus.YtfI)
        copyto!(cuFrBound, frBound)
        copyto!(cuToBound, toBound)

        return cuYshR, cuYshI, cuYffR, cuYffI, cuYftR, cuYftI,
               cuYttR, cuYttI, cuYtfR, cuYtfI, cuFrBound, cuToBound
    else
        return ybus.YshR, ybus.YshI, ybus.YffR, ybus.YffI, ybus.YftR, ybus.YftI,
               ybus.YttR, ybus.YttI, ybus.YtfR, ybus.YtfI, frBound, toBound
    end
end

function init_solution!(env::AdmmEnv, sol::SolutionOneLevel, ybus::Ybus, rho_pq, rho_va)
    data, model = env.data, env.model

    lines = data.lines
    buses = data.buses
    BusIdx = data.BusIdx
    ngen = length(data.generators)
    nline = length(data.lines)

    YffR = ybus.YffR; YffI = ybus.YffI
    YttR = ybus.YttR; YttI = ybus.YttI
    YftR = ybus.YftR; YftI = ybus.YftI
    YtfR = ybus.YtfR; YtfI = ybus.YtfI

    for g=1:ngen
        pg_idx = model.gen_mod.gen_start + 2*(g-1)
        CUDA.@allowscalar sol.v_curr[pg_idx] = 0.5*(data.generators[g].Pmin + data.generators[g].Pmax)
        CUDA.@allowscalar sol.v_curr[pg_idx+1] = 0.5*(data.generators[g].Qmin + data.generators[g].Qmax)
    end

    sol.rho .= rho_pq
    fill!(sol.u_curr, 0.0)
    fill!(sol.v_curr, 0.0)

    for l=1:nline
        wij0 = (buses[BusIdx[lines[l].from]].Vmax^2 + buses[BusIdx[lines[l].from]].Vmin^2) / 2
        wji0 = (buses[BusIdx[lines[l].to]].Vmax^2 + buses[BusIdx[lines[l].to]].Vmin^2) / 2
        wR0 = sqrt(wij0 * wji0)

        pij_idx = model.line_start + 8*(l-1)
        CUDA.@allowscalar sol.u_curr[pij_idx] = YffR[l] * wij0 + YftR[l] * wR0
        CUDA.@allowscalar sol.u_curr[pij_idx+1] = -YffI[l] * wij0 - YftI[l] * wR0
        CUDA.@allowscalar sol.u_curr[pij_idx+2] = YttR[l] * wji0 + YtfR[l] * wR0
        CUDA.@allowscalar sol.u_curr[pij_idx+3] = -YttI[l] * wji0 - YtfI[l] * wR0
        #=
        u_curr[pij_idx+4] = wij0
        u_curr[pij_idx+5] = wji0
        u_curr[pij_idx+6] = 0.0
        u_curr[pij_idx+7] = 0.0
        =#
        # wRIij[2*(l-1)+1] = wR0
        # wRIij[2*l] = 0.0

        CUDA.@allowscalar sol.v_curr[pij_idx+4] = wij0
        CUDA.@allowscalar sol.v_curr[pij_idx+5] = wji0
        CUDA.@allowscalar sol.v_curr[pij_idx+6] = 0.0
        CUDA.@allowscalar sol.v_curr[pij_idx+7] = 0.0

        sol.rho[pij_idx+4:pij_idx+7] .= rho_va
    end

    sol.l_curr .= 0
    return
end

function copy_data_kernel(n::Int, dest::CuDeviceArray{Float64,1}, src::CuDeviceArray{Float64,1})
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        dest[tx] = src[tx]
    end
    return
end

function update_multiplier_kernel(n::Int, l_curr::CuDeviceArray{Float64,1},
                                  u_curr::CuDeviceArray{Float64,1},
                                  v_curr::CuDeviceArray{Float64,1},
                                  rho::CuDeviceArray{Float64,1})
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        l_curr[tx] += rho[tx] * (u_curr[tx] - v_curr[tx])
    end
    return
end

function primal_residual_kernel(n::Int, rp::CuDeviceArray{Float64,1},
                                u_curr::CuDeviceArray{Float64,1},
                                v_curr::CuDeviceArray{Float64,1})
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        rp[tx] = u_curr[tx] - v_curr[tx]
    end

    return
end

function dual_residual_kernel(n::Int, rd::CuDeviceArray{Float64,1},
                              v_prev::CuDeviceArray{Float64,1},
                              v_curr::CuDeviceArray{Float64,1},
                              rho::CuDeviceArray{Float64,1})
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        rd[tx] = -rho[tx] * (v_curr[tx] - v_prev[tx])
    end

    return
end

function check_linelimit_violation(data::OPFData, u)
    lines = data.lines
    nline = length(data.lines)
    line_start = 2*length(data.generators) + 1

    rateA_nviols = 0
    rateA_maxviol = 0.0
    rateC_nviols = 0
    rateC_maxviol = 0.0

    for l=1:nline
        pij_idx = line_start + 8*(l-1)
        ij_val = 0.0
        ji_val = 0.0
        CUDA.@allowscalar ij_val = u[pij_idx]^2 + u[pij_idx+1]^2
        CUDA.@allowscalar ji_val = u[pij_idx+2]^2 + u[pij_idx+3]^2

        limit = (lines[l].rateA / data.baseMVA)^2
        if limit > 0 && limit < 1e10
            if ij_val > limit || ji_val > limit
                rateA_nviols += 1
                rateA_maxviol = max(rateA_maxviol, max(ij_val - limit, ji_val - limit))
            end
        end

        limit = (lines[l].rateC / data.baseMVA)^2
        if limit > 0 && limit < 1e10
            if ij_val > limit || ji_val > limit
                rateC_nviols += 1
                rateC_maxviol = max(rateC_maxviol, max(ij_val - limit, ji_val - limit))
            end
        end
    end
    rateA_maxviol = sqrt(rateA_maxviol)
    rateC_maxviol = sqrt(rateC_maxviol)

    return rateA_nviols, rateA_maxviol, rateC_nviols, rateC_maxviol
end

"""
    admm_restart!

This function restarts the ADMM with a given `env::AdmmEnv` containing solutions and all the other parameters.
"""
admm_restart!(env::AdmmEnv; options...) = admm_solve!(env, env.solution; options...)

function admm_solve!(env::AdmmEnv, sol::SolutionOneLevel; iterlim=800, scale=1e-4)
    if env.use_gpu
        CUDA.device!(env.gpu_no)
    end

    data, par, mod = env.data, env.params, env.model

    shift_lines = 0
    shmem_size = sizeof(Float64)*(14*mod.n+3*mod.n^2) + sizeof(Int)*(4*mod.n)

    nblk_gen = div(mod.gen_mod.ngen, 32, RoundUp)
    nblk_br = mod.nline
    nblk_bus = div(mod.nbus, 32, RoundUp)

    it = 0
    time_gen = time_br = time_bus = 0.0

    @time begin
    while it < iterlim
        it += 1

        if !env.use_gpu
            sol.u_prev .= sol.u_curr
            sol.v_prev .= sol.v_curr
            sol.l_prev .= sol.l_curr

            tcpu = generator_kernel(mod.gen_mod, data.baseMVA, sol.u_curr, sol.v_curr, sol.l_curr, sol.rho)
            time_gen += tcpu.time

            if env.use_polar
                tcpu = @timed auglag_it, tron_it = polar_kernel_cpu(mod.n, mod.nline, mod.line_start, scale,
                                                                    sol.u_curr, sol.v_curr, sol.l_curr, sol.rho,
                                                                    shift_lines, env.membuf, mod.YffR, mod.YffI, mod.YftR, mod.YftI,
                                                                    mod.YttR, mod.YttI, mod.YtfR, mod.YtfI, mod.FrBound, mod.ToBound)
            else
                tcpu = @timed auglag_it, tron_it = auglag_kernel_cpu(mod.n, mod.nline, it, par.max_auglag, mod.line_start, par.mu_max,
                                                                     sol.u_curr, sol.v_curr, sol.l_curr, sol.rho,
                                                                     shift_lines, env.membuf, mod.YffR, mod.YffI, mod.YftR, mod.YftI,
                                                                     mod.YttR, mod.YttI, mod.YtfR, mod.YtfI, mod.FrBound, mod.ToBound)
            end
            time_br += tcpu.time

            tcpu = @timed bus_kernel_cpu(data.baseMVA, mod.nbus, mod.gen_mod.gen_start, mod.line_start,
                                         mod.FrStart, mod.FrIdx, mod.ToStart, mod.ToIdx, mod.GenStart,
                                         mod.GenIdx, mod.Pd, mod.Qd, sol.u_curr, sol.v_curr, sol.l_curr, sol.rho, mod.YshR, mod.YshI)
            time_bus += tcpu.time

            sol.l_curr .+= sol.rho .* (sol.u_curr .- sol.v_curr)
            sol.rd .= -sol.rho .* (sol.v_curr .- sol.v_prev)
            sol.rp .= sol.u_curr .- sol.v_curr
            #sol.rp_old .= sol.u_prev .- sol.v_prev

            primres = norm(sol.rp)
            dualres = norm(sol.rd)

            eps_pri = sqrt(length(sol.l_curr))*par.ABSTOL + par.RELTOL*max(norm(sol.u_curr), norm(-sol.v_curr))
            eps_dual = sqrt(length(sol.u_curr))*par.ABSTOL + par.RELTOL*norm(sol.l_curr)

            (par.verbose > 0) && @printf("[CPU] %10d  %.6e  %.6e  %.6e  %.6e  %6.2f  %6.2f\n",
                    it, primres, dualres, eps_pri, eps_dual, auglag_it, tron_it)

            if primres <= eps_pri && dualres <= eps_dual
                break
            end
        else
            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) copy_data_kernel(mod.nvar, sol.u_prev, sol.u_curr)
            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) copy_data_kernel(mod.nvar, sol.v_prev, sol.v_curr)
            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) copy_data_kernel(mod.nvar, sol.l_prev, sol.l_curr)
            CUDA.synchronize()

            tgpu = generator_kernel(mod.gen_mod, data.baseMVA, sol.u_curr, sol.v_curr, sol.l_curr, sol.rho)

            time_gen += tgpu.time
            if env.use_polar
                tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_br shmem=shmem_size polar_kernel(mod.n, mod.nline, mod.line_start, scale,
                                                                                sol.u_curr, sol.v_curr, sol.l_curr, sol.rho,
                                                                                shift_lines, env.membuf, mod.YffR, mod.YffI, mod.YftR, mod.YftI,
                                                                                mod.YttR, mod.YttI, mod.YtfR, mod.YtfI, mod.FrBound, mod.ToBound)
            else
                tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_br shmem=shmem_size auglag_kernel(mod.n, it, par.max_auglag, mod.line_start, scale, par.mu_max,
                                                                                sol.u_curr, sol.v_curr, sol.l_curr, sol.rho,
                                                                                shift_lines, env.membuf, mod.YffR, mod.YffI, mod.YftR, mod.YftI,
                                                                                mod.YttR, mod.YttI, mod.YtfR, mod.YtfI, mod.FrBound, mod.ToBound)
            end
            time_br += tgpu.time
            tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_bus bus_kernel(data.baseMVA, mod.nbus, mod.gen_mod.gen_start, mod.line_start,
                                                                           mod.FrStart, mod.FrIdx, mod.ToStart, mod.ToIdx, mod.GenStart,
                                                                           mod.GenIdx, mod.Pd, mod.Qd, sol.u_curr, sol.v_curr, sol.l_curr,
                                                                           sol.rho, mod.YshR, mod.YshI)
            time_bus += tgpu.time

            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) update_multiplier_kernel(mod.nvar, sol.l_curr, sol.u_curr, sol.v_curr, sol.rho)
            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) primal_residual_kernel(mod.nvar, sol.rp, sol.u_curr, sol.v_curr)
            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) dual_residual_kernel(mod.nvar, sol.rd, sol.v_prev, sol.v_curr, sol.rho)
            CUDA.synchronize()

            gpu_primres = CUDA.norm(sol.rp)
            gpu_dualres = CUDA.norm(sol.rd)

            gpu_eps_pri = sqrt(length(sol.l_curr))*par.ABSTOL + par.RELTOL*max(CUDA.norm(sol.u_curr), CUDA.norm(sol.v_curr))
            gpu_eps_dual = sqrt(length(sol.u_curr))*par.ABSTOL + par.RELTOL*CUDA.norm(sol.l_curr)

            (par.verbose > 0) && @printf("[GPU] %10d  %.6e  %.6e  %.6e  %.6e\n", it, gpu_primres, gpu_dualres, gpu_eps_pri, gpu_eps_dual)

            if gpu_primres <= gpu_eps_pri && gpu_dualres <= gpu_eps_dual
                break
            end
        end
    end
    end

    u_curr = zeros(mod.nvar)
    copyto!(u_curr, sol.u_curr)
    objval = sum(data.generators[g].coeff[data.generators[g].n-2]*(data.baseMVA*u_curr[mod.gen_mod.gen_start+2*(g-1)])^2 +
                 data.generators[g].coeff[data.generators[g].n-1]*(data.baseMVA*u_curr[mod.gen_mod.gen_start+2*(g-1)]) +
                 data.generators[g].coeff[data.generators[g].n]
                 for g in 1:mod.gen_mod.ngen)::Float64
    sol.objval = objval

    if it < iterlim
        sol.status = HAS_CONVERGED
    else
        sol.status = MAXIMUM_ITERATIONS
    end

    if par.verbose > 0
        rateA_nviols, rateA_maxviol, rateC_nviols, rateC_maxviol = check_linelimit_violation(data, u_curr)
        @printf(" ** Line limit violations\n")
        @printf("RateA number of violations = %d (%d)\n", rateA_nviols, mod.nline)
        @printf("RateA maximum violation    = %.2f\n", rateA_maxviol)
        @printf("RateC number of violations = %d (%d)\n", rateC_nviols, mod.nline)
        @printf("RateC maximum violation    = %.2f\n", rateC_maxviol)

        @printf(" ** Time\n")
        @printf("Generator = %.2f\n", time_gen)
        @printf("Branch    = %.2f\n", time_br)
        @printf("Bus       = %.2f\n", time_bus)
        @printf("Total     = %.2f\n", time_gen + time_br + time_bus)

        @printf("Objective value = %.6e\n", objval)
    end
    return
end

function admm_rect_gpu(case::String; iterlim=800, rho_pq=400.0, rho_va=40000.0, scale=1e-4,
                       use_gpu=false, use_polar=true, gpu_no=0, verbose=1)
    env = AdmmEnv(case, use_gpu, rho_pq, rho_va; use_polar=use_polar, gpu_no=gpu_no, verbose=verbose,)
    admm_restart!(env, iterlim=iterlim, scale=scale)
    return env
end

# TODO: This needs revised to use AdmmEnv.
function admm_rect_gpu_mpi(
    case::String;
    iterlim=800, rho_pq=400.0, rho_va=40000.0, scale=1e-4, use_gpu=false, use_polar=true, gpu_no=1,
    comm::MPI.Comm=MPI.COMM_WORLD,
)
    data = opf_loaddata(case)

    ngen = length(data.generators)
    nline = length(data.lines)
    nbus = length(data.buses)
    nvar = 2*ngen + 8*nline
    if use_gpu
        @assert MPI.has_cuda()
    end

    # MPI settings
    root = 0
    is_master = MPI.Comm_rank(comm) == root
    n_processes = MPI.Comm_size(comm)
    nlines_local = div(nline, n_processes, RoundUp)
    nlines_padded = n_processes * nlines_local

    nvar_padded = 2*ngen + 8 * nlines_padded

    baseMVA = data.baseMVA
    n = (use_polar == true) ? 4 : 10
    mu_max = 1e8
    rho_max = 1e6
    rho_min_pq = 5.0
    rho_min_w = 5.0
    eps_rp = 1e-4
    eps_rp_min = 1e-5
    rt_inc = 2.0
    rt_dec = 2.0
    eta = 0.99
    Kf = 100
    Kf_mean = 10

    if use_gpu
        CUDA.device!(MPI.Comm_rank(comm) % CUDA.ndevices())
    end

    ybus = Ybus{Array{Float64}}(computeAdmitances(data.lines, data.buses, data.baseMVA; VI=Array{Int}, VD=Array{Float64})...)

    pgmin, pgmax, qgmin, qgmax, c2, c1, c0 = get_generator_data(data)
    YshR, YshI, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI, FrBound, ToBound = get_branch_data(data)
    FrStart, FrIdx, ToStart, ToIdx, GenStart, GenIdx, Pd, Qd = get_bus_data(data)

    cu_pgmin, cu_pgmax, cu_qgmin, cu_qgmax, cu_c2, cu_c1, cu_c0 = get_generator_data(data; use_gpu=use_gpu)
    cuYshR, cuYshI, cuYffR, cuYffI, cuYftR, cuYftI, cuYttR, cuYttI, cuYtfR, cuYtfI, cuFrBound, cuToBound = get_branch_data(data; use_gpu=use_gpu)
    cu_FrStart, cu_FrIdx, cu_ToStart, cu_ToIdx, cu_GenStart, cu_GenIdx, cu_Pd, cu_Qd = get_bus_data(data; use_gpu=use_gpu)

    gen_start = 1
    line_start = 2*ngen + 1

    # Allocations
    u_curr = zeros(nvar_padded)
    v_curr = zeros(nvar_padded)
    l_curr = zeros(nvar_padded)
    u_prev = zeros(nvar_padded)
    v_prev = zeros(nvar_padded)
    l_prev = zeros(nvar_padded)
    rho = zeros(nvar_padded)
    rd = zeros(nvar_padded)
    rp = zeros(nvar_padded)
    rp_old = zeros(nvar_padded)
    rp_k0 = zeros(nvar_padded)
    param = zeros(31, nlines_padded)
    wRIij = zeros(2*nline)

    init_values(data, ybus, gen_start, line_start,
                rho_pq, rho_va, u_curr, v_curr, l_curr, rho, wRIij)

    if use_gpu
        cu_u_curr = CuArray{Float64}(undef, nvar_padded)
        cu_v_curr = CuArray{Float64}(undef, nvar_padded)
        cu_l_curr = CuArray{Float64}(undef, nvar_padded)
        cu_u_prev = CuArray{Float64}(undef, nvar_padded)
        cu_v_prev = CuArray{Float64}(undef, nvar_padded)
        cu_l_prev = CuArray{Float64}(undef, nvar_padded)
        cu_rho = CuArray{Float64}(undef, nvar_padded)
        cu_rd = CuArray{Float64}(undef, nvar_padded)
        cu_rp = CuArray{Float64}(undef, nvar_padded)
        cu_rp_old = CuArray{Float64}(undef, nvar_padded)
        cu_rp_k0 = CuArray{Float64}(undef, nvar_padded)
        cuParam = CuArray{Float64}(undef, (31, nlines_padded))
        cuWRIij = CuArray{Float64}(undef, 2*nline)

        copyto!(cu_u_curr, u_curr)
        copyto!(cu_v_curr, v_curr)
        copyto!(cu_l_curr, l_curr)
        copyto!(cu_rho, rho)
        copyto!(cuParam, param)
        copyto!(cuWRIij, wRIij)
    end

    # MPI: Global info
    if use_gpu
        u_lines_root = @view cu_u_curr[line_start:end]
        l_lines_root = @view cu_l_curr[line_start:end]
        v_lines_root = @view cu_v_curr[line_start:end]
        rho_lines_root = cu_rho[line_start:end]
    else
        u_lines_root = @view u_curr[line_start:end]
        l_lines_root = @view l_curr[line_start:end]
        v_lines_root = @view v_curr[line_start:end]
        rho_lines_root = @view rho[line_start:end]
    end
    # MPI: Local info
    # We need only to transfer info about lines
    if use_gpu
        u_local = CUDA.zeros(Float64, 8 * nlines_local)
        v_local = CUDA.zeros(Float64, 8 * nlines_local)
        l_local = CUDA.zeros(Float64, 8 * nlines_local)
        rho_local = CUDA.zeros(Float64, 8 * nlines_local)
    else
        u_local = zeros(8 * nlines_local)
        v_local = zeros(8 * nlines_local)
        l_local = zeros(8 * nlines_local)
        rho_local = zeros(8 * nlines_local)
    end

    MPI.Scatter!(u_lines_root, u_local, root, comm)
    MPI.Scatter!(rho_lines_root, rho_local, root, comm)

    max_auglag = 50

    nblk_gen = div(ngen, 32, RoundUp)
    nblk_br = nline
    nblk_br_local = nlines_local
    nblk_bus = div(nbus, 32, RoundUp)

    ABSTOL = 1e-6
    RELTOL = 1e-5

    it = 0
    time_gen = time_br = time_bus = 0
    time_br_scatter = time_br_gather = 0

    h_u_curr = zeros(nvar)
    h_param = zeros(31, nline)
    h_wRIij = zeros(2*nline)

    shift_lines = MPI.Comm_rank(comm) * nlines_local
    # GPU settings
    shmem_size = sizeof(Float64)*(14*n+3*n^2) + sizeof(Int)*(4*n)

    while it < iterlim
        it += 1

        # CPU code
        if !use_gpu

            if is_master
                u_prev .= u_curr
                v_prev .= v_curr
                l_prev .= l_curr
                tcpu = @timed generator_kernel_cpu(baseMVA, ngen, gen_start, u_curr, v_curr, l_curr, rho,
                                                pgmin, pgmax, qgmin, qgmax, c2, c1)
                time_gen += tcpu.time
            end

            # MPI routines to be implemented:
            #  - Broadcast cu_v_curr and cu_l_curr to GPUs.
            #  - Collect cu_u_curr.
            #  - div(nblk_br / number of GPUs, RoundUp)
            # scatter / gather

            tcpu_mpi = @timed begin
                MPI.Scatter!(v_lines_root, v_local, root, comm)
                MPI.Scatter!(l_lines_root, l_local, root, comm)
            end
            time_br_scatter += tcpu_mpi.time

            nlines_actual = min(nlines_local, nline - shift_lines)
            tcpu = @timed auglag_it, tron_it = polar_kernel_cpu(n, nlines_actual, 1, scale,
                                                                u_local, v_local, l_local, rho_local,
                                                                shift_lines, param, YffR, YffI, YftR, YftI,
                                                                YttR, YttI, YtfR, YtfI, FrBound, ToBound)

            time_br += tcpu.time
            tcpu_mpi = @timed MPI.Gather!(u_local, u_lines_root, root, comm)
            time_br_gather += tcpu_mpi.time

            if is_master
                tcpu = @timed bus_kernel_cpu(baseMVA, nbus, gen_start, line_start,
                                            FrStart, FrIdx, ToStart, ToIdx, GenStart,
                                            GenIdx, Pd, Qd, u_curr, v_curr, l_curr, rho, YshR, YshI)
                time_bus += tcpu.time

                l_curr .+= rho .* (u_curr .- v_curr)
                rd .= -rho .* (v_curr .- v_prev)
                rp .= u_curr .- v_curr
                #rp_old .= u_prev .- v_prev

                primres = norm(rp)
                dualres = norm(rd)

                eps_pri = sqrt(length(l_curr))*ABSTOL + RELTOL*max(norm(u_curr), norm(-v_curr))
                eps_dual = sqrt(length(u_curr))*ABSTOL + RELTOL*norm(l_curr)

                @printf("[CPU] %10d  %.6e  %.6e  %.6e  %.6e  %6.2f  %6.2f\n",
                        it, primres, dualres, eps_pri, eps_dual, auglag_it, tron_it)
            end
        # GPU code
        else
            if is_master
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) copy_data_kernel(nvar, cu_u_prev, cu_u_curr)
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) copy_data_kernel(nvar, cu_v_prev, cu_v_curr)
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) copy_data_kernel(nvar, cu_l_prev, cu_l_curr)
                CUDA.synchronize()

                tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_gen generator_kernel(baseMVA, ngen, gen_start,
                                                                                cu_u_curr, cu_v_curr, cu_l_curr, cu_rho,
                                                                                cu_pgmin, cu_pgmax, cu_qgmin, cu_qgmax, cu_c2, cu_c1)

                time_gen += tgpu.time
            end

            #  - Broadcast cu_v_curr and cu_l_curr to GPUs.
            tgpu_mpi = @timed begin
                MPI.Scatter!(v_lines_root, v_local, root, comm)
                MPI.Scatter!(l_lines_root, l_local, root, comm)
            end
            time_br_scatter += tgpu_mpi.time

            #  - div(nblk_br / number of GPUs, RoundUp)
            nblk_br_actual = min(nblk_br_local, nline - shift_lines)
            tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_br_actual shmem=shmem_size polar_kernel(n, nline, 1, scale,
                u_local, v_local, l_local, rho_local,
                shift_lines, cuParam, cuYffR, cuYffI, cuYftR, cuYftI,
                cuYttR, cuYttI, cuYtfR, cuYtfI, cuFrBound, cuToBound
            )
            time_br += tgpu.time

            #  - Collect cu_u_curr.
            tgpu_mpi = @timed MPI.Gather!(u_local, u_lines_root, root, comm)
            time_br_gather += tgpu_mpi.time

            if is_master
                tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_bus bus_kernel(baseMVA, nbus, gen_start, line_start,
                                                                            cu_FrStart, cu_FrIdx, cu_ToStart, cu_ToIdx, cu_GenStart,
                                                                            cu_GenIdx, cu_Pd, cu_Qd, cu_u_curr, cu_v_curr, cu_l_curr,
                                                                            cu_rho, cuYshR, cuYshI)
                time_bus += tgpu.time
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) update_multiplier_kernel(nvar, cu_l_curr, cu_u_curr, cu_v_curr, cu_rho)
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) primal_residual_kernel(nvar, cu_rp, cu_u_curr, cu_v_curr)
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) dual_residual_kernel(nvar, cu_rd, cu_v_prev, cu_v_curr, cu_rho)
                CUDA.synchronize()

                gpu_primres = CUDA.norm(cu_rp)
                gpu_dualres = CUDA.norm(cu_rd)

                gpu_eps_pri = sqrt(length(l_curr))*ABSTOL + RELTOL*max(CUDA.norm(cu_u_curr), CUDA.norm(cu_v_curr))
                gpu_eps_dual = sqrt(length(u_curr))*ABSTOL + RELTOL*CUDA.norm(cu_l_curr)

                @printf("[GPU] %10d  %.6e  %.6e  %.6e  %.6e\n", it, gpu_primres, gpu_dualres, gpu_eps_pri, gpu_eps_dual)
            end
        end
    end

    if use_gpu
        copyto!(u_curr, cu_u_curr)
    end

    rank = MPI.Comm_rank(comm)
    @printf(" ** Time\n")
    @printf("[%d] Generator = %.2f\n", rank, time_gen)
    @printf("[%d] Branch    = %.2f\n", rank, time_br)
    @printf("[%d] Bus       = %.2f\n", rank, time_bus)
    @printf("[%d] G+Br+Bus  = %.2f\n", rank, time_gen + time_br + time_bus)
    @printf("[%d] Scatter   = %.2f\n", rank, time_br_scatter)
    @printf("[%d] Gather    = %.2f\n", rank, time_br_gather)
    @printf("[%d] MPI(S+G)  = %.2f\n", rank, time_br_scatter + time_br_gather)

    if is_master
        objval = sum(data.generators[g].coeff[data.generators[g].n-2]*(baseMVA*u_curr[gen_start+2*(g-1)])^2 +
                    data.generators[g].coeff[data.generators[g].n-1]*(baseMVA*u_curr[gen_start+2*(g-1)]) +
                    data.generators[g].coeff[data.generators[g].n]
                    for g in 1:ngen)
        @printf("Objective value = %.6e\n", objval)
    end

    return
end


#=
Copyright (c) 2015, Intel Corporation

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=#

include("../../../src/Sparso.jl")
using Sparso

function abiatic(Hdmat, d, nqbits, T)
  total_time = -time()

  nsets = 2^nqbits
  Phi0at0 = 1/sqrt(2^nqbits)*ones(nsets)

  ode_time = -time()

  h = 0.01
  rtol = 1e-6

  # Compute the constants once
  c30 = 3/8
  c31 = -im*3/32
  c32 = -im*9/32
  c40 = 12/13
  c41 = -im*1932/2197
  c42 = im*7200/2197
  c43 = -im*7296/2197
  c51 = -im*439/216
  c52 = im*8
  c53 = -im*3680/513
  c54 = im*845/4104
  c61 = im*8/27
  c62 = -im*2
  c63 = im*3544/2565
  c64 = -im*1859/4104
  c65 = im*11/40
  cz1 = -im*16/135
  cz3 = -im*6656/12825
  cz4 = -im*28561/56430
  cz5 = im*9/50
  cz6 = -im*2/55
  ce1 = -im*1/360
  ce3 = im*128/4275
  ce4 = im*2197/75240
  ce5 = -im*1/50
  ce6 = -im*2/55

  # Absolute tolerance
  atol = 1e-13
  alpha_val = 0.8
  k = 0
  # Initial time moment
  i = 1
  Tv = zeros(0)
  push!(Tv, 0)
  t = 0
  # Initial condition
  meanenergy = zeros(1)
  variance = zeros(1)
  meanenergy[1] = dot(vec(abs(Phi0at0).^2), d)
  variance[1] = dot(vec(abs(Phi0at0).^2), (d - meanenergy[1]).^2)
  wi = Phi0at0
  # If it is the last iteration, then lastit = 1, otherwise lastit = 0
  lastit = 0
  spmv_time = 0

  while lastit == 0
      # Stretch the step if within 10% of b-t
      if t + 1.1*h > T
          h = T - t
          lastit = 1
      end

      # Compute the step
      s = t/T # real
      w = wi
      spmv_time -= time()
      s1 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + 0.25*h)/T
      w = wi + -im*0.25*h*s1
      spmv_time -= time()
      s2 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + c30*h)/T
      w = wi + c31*h*s1 + c32*h*s2
      spmv_time -= time()
      s3 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + c40*h)/T
      w = wi + c41*h*s1 + c42*h*s2 + c43*h*s3
      spmv_time -= time()
      s4 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + h)/T
      w = wi + c51*h*s1 + c52*h*s2 + c53*h*s3 + c54*h*s4
      spmv_time -= time()
      s5 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + 0.5*h)/T
      w = wi + c61*h*s1 + c62*h*s2 + c63*h*s3 + c64*h*s4 + c65*h*s5
      spmv_time -= time()
      s6 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      z = wi + h*(cz1*s1 + cz3*s3 + cz4*s4 + cz5*s5 + cz6*s6)
      e = h * norm(ce1*s1 + ce3*s3 + ce4*s4 + ce5*s5 + ce6*s6) # real

      # Target tolerance for this step
      target_tol = rtol*norm(wi) + atol
      if e <= target_tol # In case the tolerance is met
          t = t + h
          h = alpha_val*h*(target_tol/e)^0.2
          i = i + 1
          push!(Tv, t)
          wi = z
          push!(meanenergy, dot(abs(z).^2, d))
          push!(variance, dot(abs(z).^2, (d - meanenergy[end]).^2))
          k = 0
      elseif k == 0 # Tolerance is not met for the first time in this step
          h = alpha_val*h*(target_tol/e)^0.2
          k = k + 1
          lastit = 0
      else # Tolerance is not met more than once in this step
          h = h/2
          lastit = 0
      end
  end

  println(size(Tv))

  ode_time += time()

  lanczos_time = -time()
  spectralgap = zeros(length(Tv))

  nit = min(30, nsets)

  x0 = rand(nsets)

  alpha = zeros(nit)
  beta = zeros(nit + 1)

  for i = 1:length(Tv)
    s = Tv[i]/T

    q0 = zeros(nsets)
    q1 = x0/norm(x0)

    for k = 1:nit
      spmv_time -= time()
      uk = -(1 - s)*Hdmat*q1 + s*d.*q1
      spmv_time += time()

      alpha[k] = dot(q1, uk)
      uk -= beta[k]*q0 + alpha[k]*q1
      beta[k + 1] = norm(uk)
      if beta[k + 1] == 0
        println("Error")
        return
      end

      q0 = q1
      q1 = uk/beta[k + 1]
    end

    TT = SymTridiagonal(alpha, beta[2:nit])

    Dt = sort(eig(TT)[1], 1)
    Dt = [Dt[1]; Dt[end]]

    spectralgap[i] = abs(Dt[1] - Dt[2])^2
  end
  lanczos_time += time()

  optimalenergy = minimum(d)

  total_time += time()

  @printf("%10s%10s%12s%12s%12s\n", "time", "eigengap", "optimal", "meanenergy", "variance");
  for i = length(Tv):length(Tv)
    @printf("%10.2f%10.4f%12.4f%12.4f%12.4f\n", Tv[i], spectralgap[i], optimalenergy, meanenergy[i], variance[i])
  end

  println("total_time = $(total_time) ode_time = $(ode_time) lanczos_time = $(lanczos_time) spmv_time = $(spmv_time)")
end

function print_result(Tv, spectralgap, optimalenergy, meanenergy, variance)
  @printf("%10s%10s%12s%12s%12s\n", "time", "eigengap", "optimal", "meanenergy", "variance");
  for i = length(Tv):length(Tv)
    @printf("%10.2f%10.4f%12.4f%12.4f%12.4f\n", Tv[i], spectralgap[i], optimalenergy, meanenergy[i], variance[i])
  end
end

function abiatic_sa(Hdmat, d, nqbits, T)
  total_time = -time()

  set_matrix_property(Dict(
     :Hdmat => SA_SYMM_VALUED | SA_STRUCTURE_ONLY,
     )
  )

  nsets = 2^nqbits
  Phi0at0 = 1/sqrt(2^nqbits)*ones(nsets)

  ode_time = -time()

  h = 0.01
  rtol = 1e-6

  # Compute the constants once
  c30 = 3/8
  c31 = -im*3/32
  c32 = -im*9/32
  c40 = 12/13
  c41 = -im*1932/2197
  c42 = im*7200/2197
  c43 = -im*7296/2197
  c51 = -im*439/216
  c52 = im*8
  c53 = -im*3680/513
  c54 = im*845/4104
  c61 = im*8/27
  c62 = -im*2
  c63 = im*3544/2565
  c64 = -im*1859/4104
  c65 = im*11/40
  cz1 = -im*16/135
  cz3 = -im*6656/12825
  cz4 = -im*28561/56430
  cz5 = im*9/50
  cz6 = -im*2/55
  ce1 = -im*1/360
  ce3 = im*128/4275
  ce4 = im*2197/75240
  ce5 = -im*1/50
  ce6 = -im*2/55

  # Absolute tolerance
  atol = 1e-13
  alpha_val = 0.8
  k = 0
  # Initial time moment
  i = 1
  Tv = zeros(0)
  push!(Tv, 0)
  t = 0
  # Initial condition
  meanenergy = zeros(1)
  variance = zeros(1)
  meanenergy[1] = dot(vec(abs(Phi0at0).^2), d)
  variance[1] = dot(vec(abs(Phi0at0).^2), (d - meanenergy[1]).^2)
  wi = Array{Complex128}(Phi0at0)
  # If it is the last iteration, then lastit = 1, otherwise lastit = 0
  lastit = 0
  spmv_time = 0

  s1 = Array{Complex128}(length(d))
  s2 = Array{Complex128}(length(d))
  s3 = Array{Complex128}(length(d))
  s4 = Array{Complex128}(length(d))
  s5 = Array{Complex128}(length(d))
  s6 = Array{Complex128}(length(d))

  while lastit == 0
      # Stretch the step if within 10% of b-t
      if t + 1.1*h > T
          h = T - t
          lastit = 1
      end

      # Compute the step
      s = t/T # real
      w = wi
      spmv_time -= time()
      s1 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + 0.25*h)/T
      w = wi + -im*0.25*h*s1
      spmv_time -= time()
      s2 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + c30*h)/T
      w = wi + c31*h*s1 + c32*h*s2
      spmv_time -= time()
      s3 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + c40*h)/T
      w = wi + c41*h*s1 + c42*h*s2 + c43*h*s3
      spmv_time -= time()
      s4 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + h)/T
      w = wi + c51*h*s1 + c52*h*s2 + c53*h*s3 + c54*h*s4
      spmv_time -= time()
      s5 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + 0.5*h)/T
      w = wi + c61*h*s1 + c62*h*s2 + c63*h*s3 + c64*h*s4 + c65*h*s5
      spmv_time -= time()
      s6 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      z = wi + h*(cz1*s1 + cz3*s3 + cz4*s4 + cz5*s5 + cz6*s6)
      e = h * norm(ce1*s1 + ce3*s3 + ce4*s4 + ce5*s5 + ce6*s6) # real

      # Target tolerance for this step
      target_tol = rtol*norm(wi) + atol
      if e <= target_tol # In case the tolerance is met
          t = t + h
          h = alpha_val*h*(target_tol/e)^0.2
          i = i + 1
          push!(Tv, t)
          wi = z
          push!(meanenergy, dot(abs(z).^2, d))
          push!(variance, dot(abs(z).^2, (d - meanenergy[end]).^2))
          k = 0
      elseif k == 0 # Tolerance is not met for the first time in this step
          h = alpha_val*h*(target_tol/e)^0.2
          k = k + 1
          lastit = 0
      else # Tolerance is not met more than once in this step
          h = h/2
          lastit = 0
      end
  end

  println(size(Tv))

  ode_time += time()

  lanczos_time = -time()
  spectralgap = zeros(length(Tv))

  nit = min(30, nsets)

  x0 = rand(nsets)

  alpha = zeros(nit)
  beta = zeros(nit + 1)

  uk = Array{Float64}(length(d))

  dot_time = 0
  for i = 1:length(Tv)
    s = Tv[i]/T

    q0 = zeros(nsets)
    q1 = x0/norm(x0)

    for k = 1:nit
      spmv_time -= time()
      uk = -(1 - s)*Hdmat*q1 + s*d.*q1
      spmv_time += time()

      alpha[k] = dot(q1, uk)
      uk -= beta[k]*q0 + alpha[k]*q1
      beta[k + 1] = norm(uk)
      if beta[k + 1] == 0
        println("Error")
        return
      end

      q0 = q1
      #q1 = uk/beta[k + 1]
      q1 = Sparso.WAXPB(1/beta[k+1], uk, 0)
    end

    TT = SymTridiagonal(alpha, beta[2:nit])

    Dt = sort(eig(TT)[1], 1)
    Dt = [Dt[1]; Dt[end]]

    spectralgap[i] = abs(Dt[1] - Dt[2])^2
  end
  lanczos_time += time()

  optimalenergy = minimum(d)

  total_time += time()

  print_result(Tv, spectralgap, optimalenergy, meanenergy, variance)

  println("total_time = $(total_time) ode_time = $(ode_time) lanczos_time = $(lanczos_time) spmv_time = $(spmv_time)")
end

nqbits = parse(Int, ARGS[1])
T = parse(Int, ARGS[2])

nsets = 2^nqbits
d = 4*collect(-nsets/2 : nsets/2 - 1)

tunnelstr = -1

hdmat_time = -time()
Hdmat = zeros(nsets, nsets)
for i=1:nsets
  for j=0:nqbits-1
    flip = (i - 1)$(1 << j) # flip jth bit in zero-based indexing
    Hdmat[i,flip+1] = 1
  end
end
Hdmat = SparseMatrixCSC{Float64, Int32}(sparse(Hdmat))

if length(ARGS) == 3
  test = ARGS[3]
else
  test = "julia"
end

if test == "julia"
  println("compiler warmup (ignored): ")
  srand(0)
  println(@time(abiatic(Hdmat, d, nqbits, T)))

  println("\nRUN: ")
  srand(0)
  println(@time(abiatic(Hdmat, d, nqbits, T)))
else
  if test == "call-repl"
    set_options(SA_ENABLE, SA_USE_SPMP, SA_REPLACE_CALLS)
  elseif test == "context"
    set_options(SA_ENABLE, SA_USE_SPMP, SA_CONTEXT, SA_REPLACE_CALLS)
  elseif test == "reorder"
    set_options(SA_ENABLE, SA_USE_SPMP, SA_CONTEXT, SA_REORDER, SA_REPLACE_CALLS)
  elseif test == "verbose"
    set_options(SA_ENABLE, SA_USE_SPMP, SA_VERBOSE, SA_CONTEXT, SA_REORDER, SA_REPLACE_CALLS)
  end

  println("compiler warmup (ignored): ")
  srand(0)
  println(@time(@acc abiatic_sa(Hdmat, d, nqbits, T)))

  println("\nRUN: ")
  srand(0)
  println(@time(@acc abiatic_sa(Hdmat, d, nqbits, T)))
end


include("definitions.jl")
include("updateFunctions.jl")
include("utility.jl")

function sparseSCVB0_lda(documents::Array{SparseDocument,1},
                           numWords::Int64,
                           numDocuments::Int64,
                           numTopics::Int64,
                           numDocumentIterations::Int,
                           burnInPerDoc::Int,
                           minibatchSize::Int,
                           Alpha::Float64,
                           Beta::Float64,
                           dict,
                           tau::Float64,
                           kappa::Float64,
                           tau2::Float64,
                           kappa2::Float64,
                           scale::Float64,
                           sampleSize::Int,
                           time_limit::Float64,
                           miniEpoch::Int,
                           numDocuments_miniEpoch::Int64,
						   sparsity::Float64)

# sparse stochastic collapsed variational Bayes (sparseSCVB0) for LDA
    Alpha = ones(1,numTopics).*Alpha;
    Beta = ones(numWords,1).*Beta;

    # initialization of sparseSCVB0 specific model parameters
    #documentTopicCounts, wordTopicCounts, topicCounts, topicMem, sparseCounts = initialize(documents, numWords,numTopics,numDocuments);
    sparseVBparams=initialize(documents,Alpha, Beta, numTopics,numWords,numDocumentIterations,numDocuments);
    miniBatchesPerCorpus = numDocuments ./ minibatchSize;
    stepSize = scale ./ (tau^kappa); # step size for global expected counts
    #stepSize2 = 0.0; # initializing step size for local expected counts to make it soft local scope

    wordTopicCounts_hat = zeros(numWords,numTopics);
    topicCounts_hat = zeros(1,numTopics);
    iter = 0;

    saved_topics = Array{Array{Float64,2}}(0)
    saved_Totaltopics = Array{Array{Float64,2}}(0)
    saved_time = Float64[]
    saved_iters = Int64[]
    #saved_docTopics = Array{Array{Float64,2}}(0)

    prev_doc = 0; # to track miniEpoch docs
    doc_counter = 0;
    # main sparse stochastic collapsed variational Bayes (sparseSCVB0) loop
    while (sparseVBparams.wallClockTime < time_limit) && (doc_counter<numDocuments)
        prev_doc,documentTopicCounts,topicMem,sparseCounts = initialize_miniEpoch(documents,numTopics,numWords,numDocuments_miniEpoch,prev_doc);
        for epoch = 1:miniEpoch
            for epochDoc = 1:numDocuments_miniEpoch
                iter += 1;
                tic();
                #docInd_mini = mod(iter - 1, numDocuments_miniEpoch) + 1;
                #docInd = doc_counter + docInd_mini;
                docInd = doc_counter + epochDoc;
                docLength = documents[docInd].docLength;
                if docLength == 0
                    continue;
                end
                    # burn-in pass for a document
                sparseVBparams,documentTopicCounts,sparseCounts,topicMem = burnInUpdate(sparseVBparams,documentTopicCounts,documents,burnInPerDoc,docInd,epochDoc,docLength,tau2,kappa2,sparseCounts,topicMem,sampleSize,sparsity);

                    # main loop update for a word
                sparseVBparams,documentTopicCounts,wordTopicCounts_hat,topicCounts_hat,sparseCounts,topicMem = mainLoopUpdate(sparseVBparams,documentTopicCounts,documents,burnInPerDoc,docInd,epochDoc,docLength,tau2,kappa2,sparseCounts,
                                                                topicMem,wordTopicCounts_hat,topicCounts_hat,sampleSize);

                # mini-batch update of expected global counts
                if mod(iter, minibatchSize) == 0
                    # compute effective stepsize to account for minibatches
                    stepSize = 1 - (1-stepSize)^minibatchSize
                    # minibatch update of expected global counts
                    #wordTopicCounts,topicCounts,wordTopicCounts_hat,topicCounts_hat=minibatchUpdate(stepSize,miniBatchesPerCorpus,wordTopicCounts,topicCounts,wordTopicCounts_hat,topicCounts_hat);
                    topicCounts = sparseVBparams.topicCounts;
                    wordTopicCounts = sparseVBparams.wordTopicCounts;
                    for topic = 1:numTopics
                        for word = 1:numWords
                            wordTopicCounts[word,topic] = (1 - stepSize) .*  wordTopicCounts[word,topic] + stepSize .* miniBatchesPerCorpus .* wordTopicCounts_hat[word,topic];
                            wordTopicCounts_hat[word,topic] = wordTopicCounts_hat[word,topic]*0.0;
                        end
                        topicCounts[topic] = (1 - stepSize) .* topicCounts[topic] + stepSize .* miniBatchesPerCorpus .* topicCounts_hat[topic];
                        topicCounts_hat[topic] = topicCounts_hat[topic]*0.0;
                    end
                    stepSize = scale ./ (iter + tau)^kappa;
                end
                sparseVBparams.wallClockTime += toq();

                if iter == 1
                    push!(saved_topics,copy(sparseVBparams.wordTopicCounts))
                    push!(saved_Totaltopics,copy(sparseVBparams.topicCounts))
                    push!(saved_time,sparseVBparams.wallClockTime)
                    push!(saved_iters,iter)
                else
                    if mod(iter,200)==0
                        push!(saved_topics,copy(sparseVBparams.wordTopicCounts))
                        push!(saved_Totaltopics,copy(sparseVBparams.topicCounts))
                        push!(saved_time,sparseVBparams.wallClockTime)
                        push!(saved_iters,iter);
                    end
                end
            end
        end
        doc_counter = doc_counter + numDocuments_miniEpoch;
        #push!(saved_docTopics,copy(documentTopicCounts))
    end
    return sparseVBparams,saved_topics,saved_Totaltopics,saved_time,saved_iters,iter;
end

function initialize(documents,Alpha, Beta, numTopics,numWords,numDocumentIterations,numDocuments);
    # initialize sparseSCVB0 model parameters
    sparseVBparams = sparseSCVBO_params_type(Alpha, Beta, numTopics,numWords,numDocumentIterations,numDocuments);
    totalWordsInCorpus = 0;
    for i = 1:numDocuments
        docLength = documents[i].docLength;
        totalWordsInCorpus = totalWordsInCorpus + docLength;
    end
    # initialize expected global counts
    sparseVBparams.topicCounts.*= totalWordsInCorpus ./ numTopics;
    sparseVBparams.wordTopicCounts.*= totalWordsInCorpus ./ (numWords .* numTopics);

    # initializing alias samples from alias table for each word
    sparseVBparams.cachedSamples,sparseVBparams.cached_qw,sparseVBparams.cached_Qw = aliasMethod(sparseVBparams.numWords,sparseVBparams.Alpha,
                                       sparseVBparams.wordTopicCounts,sparseVBparams.topicCounts,sparseVBparams.Beta,
                                       sparseVBparams.BetaSum,sparseVBparams.numTopics,sparseVBparams.numCachedSamples,
                                       sparseVBparams.cachedSamples,sparseVBparams.cached_qw,sparseVBparams.cached_Qw);
    return sparseVBparams;
end

function initialize_miniEpoch(documents,numTopics,numWords,numDocuments_miniEpoch,prev_doc);
    # initialization of topic storage
    topicMem = Array{TopicArray}(0);
    # initialization for keeping track on non-zero sparse counts
    sparseCounts = Array{SparseCounts}(0);
    # initialize sparseSCVB0 miniEpoch parameters
    documentTopicCounts = zeros(numDocuments_miniEpoch,numTopics);
    j=0;
    for i = (prev_doc+1):(prev_doc+numDocuments_miniEpoch)
        j=j+1;
        docLength = documents[i].docLength;
        tokenLength = length(documents[i].tokenIndex); # storing topic for clumped documents
        Assigned = Array{Int64}(tokenLength); # array to store topic assignment for each token of clumped document
        nonzeroDocIndex = Array{Int64}(0); # array to keep track of non-zero document-topic counts
        documentTopicTmp = rand(1,numTopics);
        sum_documentTopicTmp = sum(documentTopicTmp);
        for topic = 1:numTopics
            documentTopicCounts[j,topic] = (documentTopicTmp[topic]./sum_documentTopicTmp).*docLength;
            push!(nonzeroDocIndex,topic);
        end
        Assigned = rand(1:numTopics,tokenLength);
        push!(topicMem,TopicArray(Assigned));
        push!(sparseCounts,SparseCounts(nonzeroDocIndex));
    end
    prev_doc = prev_doc+numDocuments_miniEpoch;
    return prev_doc,documentTopicCounts,topicMem,sparseCounts;
end

function aliasMethod(numWords,Alpha,wordTopicCounts,topicCounts,Beta,BetaSum,numTopics,numCachedSamples,cachedSamples,cached_qw,cached_Qw)
    qw = Array{Float64}(numTopics); # unnormalized discrete distribution of dense bucket
    for word = 1:numWords
        cached_Qw[word] = 0.0;
        for topic = 1:numTopics
            qw[topic] = Alpha[topic] .* (wordTopicCounts[word,topic] + Beta[word]) ./ (topicCounts[topic] + BetaSum);
            cached_qw[word,topic] = qw[topic];
            cached_Qw[word] += qw[topic];
        end
        aliasTable = generateAlias(qw,numTopics);
        for sample = 1:numCachedSamples
            cachedSamples[word,sample] = sampleAlias(aliasTable,numTopics);
        end
    end
    return cachedSamples,cached_qw,cached_Qw;
end