# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# 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 the copyright holder 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 HOLDER 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.
# *********************************************************************************
"""
    Metrics

Convenience mutable struct for passing data between proforma methods
"""
mutable struct Metrics
    federal_itc::Float64
    om_series::Array{Float64, 1}
    om_series_bau::Array{Float64, 1}
    total_pbi::Array{Float64, 1}
    total_pbi_bau::Array{Float64, 1}
    total_depreciation::Array{Float64, 1}
    total_ibi_and_cbi::Float64
end

"""
    calculate_proforma_metrics(p::REoptInputs, d::Dict)

Recreates the ProForma spreadsheet calculations to get the simple payback period, irr, net present cost (3rd
party case), and payment to third party (3rd party case).

return Dict(
    "simple_payback_years" => 0.0,
    "internal_rate_of_return" => 0.0,
    "net_present_cost" => 0.0,
    "annualized_payment_to_third_party" => 0.0,
    "offtaker_annual_free_cashflows" => Float64[],
    "offtaker_annual_free_cashflows_bau" => Float64[],
    "offtaker_discounted_annual_free_cashflows" => Float64[],
    "offtaker_discounted_annual_free_cashflows_bau" => Float64[],
    "developer_annual_free_cashflows" => Float64[]
)
"""
function proforma_results(p::REoptInputs, d::Dict)
    r = Dict(
        "simple_payback_years" => 0.0,
        "internal_rate_of_return" => 0.0,
        "net_present_cost" => 0.0,
        "annualized_payment_to_third_party" => 0.0,
        "offtaker_annual_free_cashflows" => Float64[],
        "offtaker_annual_free_cashflows_bau" => Float64[],
        "offtaker_discounted_annual_free_cashflows" => Float64[],
        "offtaker_discounted_annual_free_cashflows_bau" => Float64[],
        "developer_annual_free_cashflows" => Float64[]
    )
    years = p.s.financial.analysis_years
    escalate_elec(val) = [-1 * val * (1 + p.s.financial.elec_cost_escalation_pct)^yr for yr in 1:years]
    escalate_om(val) = [val * (1 + p.s.financial.om_cost_escalation_pct)^yr for yr in 1:years]
    third_party = p.s.financial.third_party_ownership

    # Create placeholder variables to store summed totals across all relevant techs
    m = Metrics(0, zeros(years), zeros(years), zeros(years), zeros(years), zeros(years), 0)

    # calculate PV o+m costs, incentives, and depreciation
    for pv in p.s.pvs
        update_metrics(m, p, pv, pv.name, d, third_party)
    end

    # calculate Wind o+m costs, incentives, and depreciation
    if "Wind" in keys(d) && d["Wind"]["size_kw"] > 0
        update_metrics(m, p, p.s.wind, "Wind", d, third_party)
    end

    # calculate Storage o+m costs, incentives, and depreciation
    if "Storage" in keys(d) && d["Storage"]["size_kw"] > 0
        # TODO handle other types of storage
        storage = p.s.storage.raw_inputs[:elec]
        total_kw = d["Storage"]["size_kw"]
        total_kwh = d["Storage"]["size_kwh"]
        capital_cost = total_kw * storage.installed_cost_per_kw + total_kwh * storage.installed_cost_per_kwh
        battery_replacement_year = storage.battery_replacement_year
        battery_replacement_cost = -1 * ((total_kw * storage.replace_cost_per_kw) + (
                    total_kwh * storage.replace_cost_per_kwh))
        m.om_series += [yr != battery_replacement_year ? 0 : battery_replacement_cost for yr in 1:years]

        # storage only has cbi in the API
        cbi = total_kw * storage.total_rebate_per_kw + total_kwh * storage.total_rebate_per_kwh
        m.total_ibi_and_cbi += cbi

        # ITC
        federal_itc_basis = capital_cost  # bug in v1 subtracted cbi from capital_cost here
        federal_itc_amount = storage.total_itc_pct * federal_itc_basis
        m.federal_itc += federal_itc_amount

        # Depreciation
        if storage.macrs_option_years in [5, 7]
            schedule = []
            if storage.macrs_option_years == 5
                schedule = p.s.financial.macrs_five_year
            elseif storage.macrs_option_years == 7
                schedule = p.s.financial.macrs_seven_year
            end
            macrs_bonus_basis = federal_itc_basis * (1 - storage.total_itc_pct * storage.macrs_itc_reduction)
            macrs_basis = macrs_bonus_basis * (1 - storage.macrs_bonus_pct)

            depreciation_schedule = zeros(years)
            for (i, r) in enumerate(schedule)
                if i < length(depreciation_schedule)
                    depreciation_schedule[i] = macrs_basis * r
                end
            end
            depreciation_schedule[1] += storage.macrs_bonus_pct * macrs_bonus_basis
            m.total_depreciation += depreciation_schedule
        end
    end

    # calculate Generator o+m costs, incentives, and depreciation
    if "Generator" in keys(d) && d["Generator"]["size_kw"] > 0
        # In the two party case the developer does not include the fuel cost in their costs
        # It is assumed that the offtaker will pay for this at a rate that is not marked up
        # to cover developer profits
        fixed_and_var_om = d["Generator"]["year_one_fixed_om_cost"] + d["Generator"]["year_one_variable_om_cost"]
        fixed_and_var_om_bau = 0.0
        year_one_fuel_cost_bau = 0.0
        if p.s.generator.existing_kw > 0
            fixed_and_var_om_bau = d["Generator"]["year_one_fixed_om_cost_bau"] +
                                   d["Generator"]["year_one_variable_om_cost_bau"]
            year_one_fuel_cost_bau = d["Generator"]["year_one_fuel_cost_bau"]
        end
        if !third_party
            annual_om = -1 * (fixed_and_var_om + d["Generator"]["year_one_fuel_cost"])

            annual_om_bau = -1 * (fixed_and_var_om_bau + year_one_fuel_cost_bau)
        else
            annual_om = -1 * fixed_and_var_om

            annual_om_bau = -1 * fixed_and_var_om_bau
        end

        m.om_series += escalate_om(annual_om)
        m.om_series_bau += escalate_om(annual_om_bau)
    end

    # Optimal Case calculations
    electricity_bill_series = escalate_elec(d["ElectricTariff"]["year_one_bill"])
    export_credit_series = escalate_elec(d["ElectricTariff"]["year_one_export_benefit"])

    # In the two party case the electricity and export credits are incurred by the offtaker not the developer
    if third_party
        total_operating_expenses = m.om_series
        tax_pct = p.s.financial.owner_tax_pct
    else
        total_operating_expenses = electricity_bill_series + export_credit_series + m.om_series
        tax_pct = p.s.financial.offtaker_tax_pct
    end

    # Apply taxes to operating expenses
    if tax_pct > 0
        deductable_operating_expenses_series = copy(total_operating_expenses)
    else
        deductable_operating_expenses_series = zeros(years)
    end

    operating_expenses_after_tax = (total_operating_expenses - deductable_operating_expenses_series) +
                                    deductable_operating_expenses_series * (1 - tax_pct)
    total_cash_incentives = m.total_pbi * (1 - tax_pct)
    free_cashflow_without_year_zero = m.total_depreciation * tax_pct + total_cash_incentives + operating_expenses_after_tax
    free_cashflow_without_year_zero[1] += m.federal_itc
    free_cashflow = append!([(-1 * d["Financial"]["initial_capital_costs"]) + m.total_ibi_and_cbi], free_cashflow_without_year_zero)

    # At this point the logic branches based on third-party ownership or not - see comments
    if third_party  # get cumulative cashflow for developer
        r["developer_annual_free_cashflows"] = copy(free_cashflow)
        discounted_developer_cashflow = [v / ((1 + p.s.financial.owner_discount_pct)^(yr-1)) for (yr, v) in
                                         enumerate(r["developer_annual_free_cashflows"])]
        r["net_present_cost"] = sum(discounted_developer_cashflow) * -1

        if p.s.financial.owner_discount_pct != 0
            capital_recovery_factor = (p.s.financial.owner_discount_pct * (1 + p.s.financial.owner_discount_pct)^years) /
                                      ((1 + p.s.financial.owner_discount_pct)^years - 1) / (1 - tax_pct)
        else
            capital_recovery_factor = (1 / years) / (1 - tax_pct)
        end

        r["annualized_payment_to_third_party"] = r["net_present_cost"] * capital_recovery_factor
        annual_income_from_host = -1 * sum(discounted_developer_cashflow) * capital_recovery_factor * (1 - tax_pct)
        r["developer_annual_free_cashflows"][2:end] .+= annual_income_from_host
        r["internal_rate_of_return"] = irr(r["developer_annual_free_cashflows"])
        cumulative_cashflow = cumsum(r["developer_annual_free_cashflows"])
        net_free_cashflow = r["developer_annual_free_cashflows"]
        r["developer_annual_free_cashflows"] = round.(r["developer_annual_free_cashflows"], digits=2)

        electricity_bill_series = escalate_elec(d["ElectricTariff"]["year_one_bill"])
        electricity_bill_series_bau = escalate_elec(d["ElectricTariff"]["year_one_bill_bau"])

        export_credit_series = escalate_elec(-d["ElectricTariff"]["lifecycle_export_benefit"])
        export_credit_series_bau = escalate_elec(-d["ElectricTariff"]["lifecycle_export_benefit_bau"])

        annual_income_from_host_series = repeat([-1 * r["annualized_payment_to_third_party"]], years)

        if "Generator" in keys(d) && d["Generator"]["size_kw"] > 0
            existing_genertor_fuel_cost_series = escalate_om(-1 * d["Generator"]["year_one_fuel_cost_bau"])
            generator_fuel_cost_series = escalate_om(-1 * d["Generator"]["year_one_fuel_cost"])
        else
            existing_genertor_fuel_cost_series = zeros(years)
            generator_fuel_cost_series = zeros(years)
        end
        net_energy_costs = -electricity_bill_series_bau - export_credit_series_bau + electricity_bill_series +
                           export_credit_series + annual_income_from_host_series - existing_genertor_fuel_cost_series +
                           generator_fuel_cost_series

        if p.s.financial.owner_tax_pct > 0
            deductable_net_energy_costs = copy(net_energy_costs)
        else
            deductable_net_energy_costs = zeros(years)
        end

        r["offtaker_annual_free_cashflows"] = append!([0.0],
            electricity_bill_series + export_credit_series + generator_fuel_cost_series + annual_income_from_host_series
        )
        r["offtaker_annual_free_cashflows_bau"] = append!([0.0],
            electricity_bill_series_bau + export_credit_series_bau + existing_genertor_fuel_cost_series
            )

    else  # get cumulative cashflow for offtaker
        electricity_bill_series_bau = escalate_elec(d["ElectricTariff"]["year_one_bill_bau"])
        export_credit_series_bau = escalate_elec(-d["ElectricTariff"]["lifecycle_export_benefit_bau"])
        total_operating_expenses_bau = electricity_bill_series_bau + export_credit_series_bau + m.om_series_bau
        total_cash_incentives_bau = m.total_pbi_bau * (1 - p.s.financial.offtaker_tax_pct)

        if p.s.financial.offtaker_tax_pct > 0
            deductable_operating_expenses_series_bau = copy(total_operating_expenses_bau)
        else
            deductable_operating_expenses_series_bau = zeros(years)
        end

        operating_expenses_after_tax_bau = total_operating_expenses_bau - deductable_operating_expenses_series_bau +
                    deductable_operating_expenses_series_bau * (1 - p.s.financial.offtaker_tax_pct)
        free_cashflow_bau = operating_expenses_after_tax_bau + total_cash_incentives_bau
        free_cashflow_bau = append!([0.0], free_cashflow_bau)
        r["offtaker_annual_free_cashflows"] = round.(free_cashflow, digits=2)
        r["offtaker_discounted_annual_free_cashflows"] = [round(
            v / ((1 + p.s.financial.offtaker_discount_pct)^(yr-1)),
            digits=2) for (yr, v) in enumerate(r["offtaker_annual_free_cashflows"])]
        r["offtaker_annual_free_cashflows_bau"] = round.(free_cashflow_bau, digits=2)
        r["offtaker_discounted_annual_free_cashflows_bau"] = [round(
            v / ((1 + p.s.financial.offtaker_discount_pct)^(yr-1)),
            digits=2) for (yr, v) in enumerate(free_cashflow_bau)]
        # difference optimal and BAU
        net_free_cashflow = free_cashflow - free_cashflow_bau
        r["internal_rate_of_return"] = irr(net_free_cashflow)
        cumulative_cashflow = cumsum(net_free_cashflow)
    end

    # At this point we have the cumulative_cashflow for the developer or offtaker so the payback calculation is the same
    if cumulative_cashflow[end] < 0
        # case where the system does not pay itself back in the analysis period, do not caculate SPP and IRR
        return r
    end

    for i in 2:years
        # add years where the cumulative cashflow is negative
        if cumulative_cashflow[i] < 0
            r["simple_payback_years"] += 1
        # fractionally add years where the cumulative cashflow became positive
        elseif cumulative_cashflow[i - 1] < 0 && cumulative_cashflow[i] > 0
            r["simple_payback_years"] += -(cumulative_cashflow[i - 1] / net_free_cashflow[i])
        # skip years where cumulative cashflow is positive and the previous year's is too
        end
    end
    r["simple_payback_years"] = round(r["simple_payback_years"], digits=2)

    return r
end

"""
    update_metrics(m::Metrics, p::REoptInputs, tech::AbstractTech, tech_name::String, results::Dict, third_party::Bool)

Update the Metrics struct for the given `tech`
"""
function update_metrics(m::Metrics, p::REoptInputs, tech::AbstractTech, tech_name::String, results::Dict, third_party::Bool)
    total_kw = results[tech_name]["size_kw"]
    existing_kw = :existing_kw in fieldnames(typeof(tech)) ? tech.existing_kw : 0
    new_kw = total_kw - existing_kw
    capital_cost = new_kw * tech.installed_cost_per_kw

    # owner is responsible for both new and existing PV maintenance in optimal case
    if third_party
        annual_om = -1 * new_kw * tech.om_cost_per_kw
    else
        annual_om = -1 * total_kw * tech.om_cost_per_kw
    end
    years = p.s.financial.analysis_years
    escalate_om(val) = [val * (1 + p.s.financial.om_cost_escalation_pct)^yr for yr in 1:years]
    m.om_series += escalate_om(annual_om)
    m.om_series_bau += escalate_om(-1 * existing_kw * tech.om_cost_per_kw)

    # incentive calculations, in the spreadsheet utility incentives are applied first
    utility_ibi = minimum([capital_cost * tech.utility_ibi_pct, tech.utility_ibi_max])
    utility_cbi = minimum([new_kw * tech.utility_rebate_per_kw, tech.utility_rebate_max])
    state_ibi = minimum([(capital_cost - utility_ibi - utility_cbi) * tech.state_ibi_pct, tech.state_ibi_max])
    state_cbi = minimum([new_kw * tech.state_rebate_per_kw, tech.state_rebate_max])
    federal_cbi = new_kw * tech.federal_rebate_per_kw
    ibi = utility_ibi + state_ibi
    cbi = utility_cbi + federal_cbi + state_cbi
    m.total_ibi_and_cbi += ibi + cbi

    # Production-based incentives
    pbi_series = Float64[]
    pbi_series_bau = Float64[]
    existing_energy_bau = third_party ? get(results[tech_name], "year_one_energy_produced_kwh_bau", 0) : 0
    for yr in range(0, stop=years-1)
        if yr < tech.production_incentive_years
            degredation_pct = :degradation_pct in fieldnames(typeof(tech)) ? (1 - tech.degradation_pct)^yr : 1.0
            base_pbi = minimum([
                tech.production_incentive_per_kwh * (results[tech_name]["year_one_energy_produced_kwh"] - existing_energy_bau) * degredation_pct,
                tech.production_incentive_max_benefit * degredation_pct
            ])
            base_pbi_bau = minimum([
                tech.production_incentive_per_kwh * get(results[tech_name], "year_one_energy_produced_kwh_bau", 0) * degredation_pct,
                tech.production_incentive_max_benefit * degredation_pct
            ])
            push!(pbi_series, base_pbi)
            push!(pbi_series_bau, base_pbi_bau)
        else
            push!(pbi_series, 0.0)
            push!(pbi_series_bau, 0.0)
        end
    end
    m.total_pbi += pbi_series
    m.total_pbi_bau += pbi_series_bau

    # Federal ITC
    # NOTE: bug in v1 has the ITC within the `if tech.macrs_option_years in [5 ,7]` block.
    # NOTE: bug in v1 reduces the federal_itc_basis with the federal_cbi, which is incorrect
    federal_itc_basis = capital_cost - state_ibi - utility_ibi - state_cbi - utility_cbi
    federal_itc_amount = tech.federal_itc_pct * federal_itc_basis
    m.federal_itc += federal_itc_amount

    # Depreciation
    if tech.macrs_option_years in [5 ,7]
        schedule = []
        if tech.macrs_option_years == 5
            schedule = p.s.financial.macrs_five_year
        elseif tech.macrs_option_years == 7
            schedule = p.s.financial.macrs_seven_year
        end

        macrs_bonus_basis = federal_itc_basis - federal_itc_basis * tech.federal_itc_pct * tech.macrs_itc_reduction
        macrs_basis = macrs_bonus_basis * (1 - tech.macrs_bonus_pct)

        depreciation_schedule = zeros(years)
        for (i, r) in enumerate(schedule)
            if i < length(depreciation_schedule)
                depreciation_schedule[i] = macrs_basis * r
            end
        end
        depreciation_schedule[1] += (tech.macrs_bonus_pct * macrs_bonus_basis)
        m.total_depreciation += depreciation_schedule
    end
    nothing
end

function npv(cashflows::AbstractArray{Float64, 1}, rate::Float64)
    years = collect(0:length(cashflows)-1)
    return sum( cashflows ./ (1+rate).^years)
end

function irr(cashflows::AbstractArray{Float64, 1})
    if npv(cashflows, 0.0) < 0
        return 0.0
    end
    f(r) = npv(cashflows, r)
    rate = 0.0
    try
        rate = fzero(f, [0.0, 0.99])
    finally
        return round(rate, digits=2)
    end
end


function t_pipe_creation()

N = 10	# variables
M = 25  # samples
A = rand(N,M)
A0 = j4pr.datacell(A)
A1 = j4pr.datacell(A, round.(rand(M)))
A2 = j4pr.datacell(A, round.(rand(2,M)))

W1 = j4pr.ohenc("binary")
W2 = j4pr.scaler!(Dict(1=>"mean"))
W3 = j4pr.functioncell((x)->x)

PT = [W1;W2]
PP=[W1 W2]
PS = W1+W2

# Simple cell checks
#print("Type checks for function cells before passing data... ")
Test.@test W1 isa j4pr.CellFunU && !(W1 isa j4pr.CellFunF) && !(W1 isa j4pr.CellFunT)
Test.@test W2 isa j4pr.CellFunU && !(W2 isa j4pr.CellFunF) && !(W2 isa j4pr.CellFunT)
Test.@test W3 isa j4pr.CellFunF && !(W3 isa j4pr.CellFunU) && !(W3 isa j4pr.CellFunT)

#println("PASSED")

#print("Type checks for function cells after passing data... ")
tmp = A0 |> W1
Test.@test tmp isa j4pr.CellFunT && !(tmp isa j4pr.CellFunF) && !(tmp isa j4pr.CellFunU)
tmp = A0 |> W2
Test.@test tmp isa j4pr.CellFunT && !(tmp isa j4pr.CellFunF) && !(tmp isa j4pr.CellFunU)
Test.@test A0 |> W3 isa j4pr.CellDataU
Test.@test A1 |> W3 isa j4pr.CellDataL
Test.@test A2 |> W3 isa j4pr.CellDataLL
#println("PASSED")

# Simple pipe checks
#print("Type checks for [simple] pipes before passing data... ")
Test.@test PT isa j4pr.PipeStackedU && !(PT isa j4pr.PipeStackedF) && !(PT isa j4pr.PipeStackedT)
Test.@test PP isa j4pr.PipeParallelU && !(PP isa j4pr.PipeParallelF) && !(PP isa j4pr.PipeParallelT)
Test.@test PS isa j4pr.PipeSerialU && !(PS isa j4pr.PipeSerialF) && !(PS isa j4pr.PipeSerialT)
#println("PASSED")

#print("Type checks for [simple] pipes after passing data... ")
tmp = A0 |> PT
Test.@test tmp isa j4pr.PipeStackedT && !(tmp isa j4pr.PipeStackedF) && !(tmp isa j4pr.PipeStackedU)
tmp = [A0|>W1; A1|>W2]
Test.@test tmp isa j4pr.PipeStackedT && !(tmp isa j4pr.PipeStackedF) && !(tmp isa j4pr.PipeStackedU)
tmp = A0 |> PP
Test.@test tmp isa j4pr.PipeParallelT && !(tmp isa j4pr.PipeParallelF) && !(tmp isa j4pr.PipeParallelU)
tmp = [A0|>W1 A1|>W2]
Test.@test tmp isa j4pr.PipeParallelT && !(tmp isa j4pr.PipeParallelF) && !(tmp isa j4pr.PipeParallelU)
tmp = A0 |> PS
Test.@test tmp isa j4pr.PipeSerialT && !(tmp isa j4pr.PipeSerialF) && !(tmp isa j4pr.PipeSerialU)
tmp = (A0|>W1)+(A0|>W2)
Test.@test tmp isa j4pr.PipeSerialT && !(tmp isa j4pr.PipeSerialF) && !(tmp isa j4pr.PipeSerialU)
#println("PASSED")

# Generic pipe checks
#print("Type checks for [generic] pipes before passing data... ")
PTG = [W2+W1;W2;W3]
PPG = [W2+W1 W2 W1]
PSG = (W2+W1)+[W2;W1]
Test.@test PTG isa j4pr.PipeStacked && !(PTG isa j4pr.PipeParallel) && !(PTG isa j4pr.PipeSerial)
Test.@test PPG isa j4pr.PipeParallel && !(PPG isa j4pr.PipeStacked) && !(PPG isa j4pr.PipeSerial)
Test.@test PSG isa j4pr.PipeSerial && !(PSG isa j4pr.PipeStacked) && !(PSG isa j4pr.PipeParallel)
tmp = [[W1;W2;W3]+W3 W1 [W1 W2]]
Test.@test tmp isa j4pr.PipeParallel && !(tmp isa j4pr.PipeStacked) && !(tmp isa j4pr.PipeSerial)
#println("PASSED")

#print("Type checks for [generic] pipes before passing data... ")
Test.@test A0 |> (A0 |> PTG) isa j4pr.CellDataU
Test.@test A0 |> PPG isa j4pr.PipeParallel
Test.@test A0 |> (A0|> PPG)  isa j4pr.CellDataU
Test.@test A0 |> (A0 |> PSG) isa j4pr.PipeSerial
Test.@test A0|> (A0 |> (A0 |> PSG)) isa j4pr.CellDataU
#println("PASSED")

end


# module for tap(ing) the simulation to achieve playaback at a later time
module tap

dir = @__DIR__
push!(LOAD_PATH, dir * "/../sim")
push!(LOAD_PATH, dir * "/../vis")
using sim
using vis

include(dir * "/../update_info.jl")

mutable struct TapeFrame
  X::Array{Array{Float64, 1}, 1}
  DIAG::Array{Tuple, 1}
  CV::Array{Array{Float64, 1}, 1}

  TapeFrame() = new(Array{Float64, 1}[], Tuple[], Array{Float64, 1}[])
end

mutable struct Tape
  world::World
  frames::Array{TapeFrame, 1}
  info_objs::Array{Union{vis.InfoBox, Nothing}, 1}
  frame_last::Int

  function Tape(world)
    info_objs = Array{Union{vis.InfoBox, Nothing}, 1}(undef,
                                                      length(world.agents))
    for i in 1:length(info_objs)
      info_objs[i] = nothing
    end
    new(world, TapeFrame[], info_objs, 0)
  end
end

function add_frame!(tape::Tape)
  push!(tape.frames, TapeFrame())
  tape.frame_last += 1
end

function add_to_frame!(tape::Tape, x::Array{Float64, 1}, diag::Tuple,
                       cv::Array{Float64, 1}=[NaN, NaN])
  push!(tape.frames[tape.frame_last].X, x)
  push!(tape.frames[tape.frame_last].DIAG, diag)
  push!(tape.frames[tape.frame_last].CV, cv)
end

function visualize_frame(tape::Tape, f::Int)
  @assert f > 0 && f <= length(tape.frames)

  to_vis = []
  if tape.world.road.road != nothing
    push!(to_vis, tape.world.road.road)
  end
  i = 1
  for agent in tape.world.agents
    agent_x = tape.frames[f].X[i]
    diag = tape.frames[f].DIAG[i]
    cv = tape.frames[f].CV[i]

    if tape.info_objs[i] != nothing
      update_info_from_tape(tape.info_objs[i], agent_x, agent.id, diag, cv)
      push!(to_vis, tape.info_objs[i])
    end
    (x, y, sx, sy, dx, u) = diag
    ## get additional information
    dp = dx[3]
    th = atan(sy, sx)
    th_car = atan(dp, agent_x[2]) + atan(sy, sx)
    agent_is_braking = dx[1] * u[1] < 0
    if agent.car != nothing
      ## update the car
      vis.car_lights!(agent.car, agent_is_braking)
      agent.car.T = (vis.scale_mat(tape.world.vis_scaling,
                                   tape.world.vis_scaling) *
                     vis.translate_mat(x, y) *
                     vis.rotate_mat(th_car - pi / 2))
      push!(to_vis, agent.car)
    end
    i += 1
  end

  return to_vis
end

end


using Yao
using Yao.Blocks

function Base.show(io::IO, ::MIME"qasm/open", blk::AbstractBlock)
    qcode = qasm(blk)
    println(io, "OPENQASM 2.0")
    maincode = pop!(qcode, "main")
    for k in keys
        println(io, qcode)
    end
    println(io, maincode)
end

######################## Basic Buiding Blocks ########################
const BASIC_1GATES = [:X, :Y, :Z, :H, :S, :Sdag, :T, :Tdag]

addrs2str(addrs) = join(["reg[$(i-1)]" for i in addrs], ", ")
floats2str(params) = join(["$p" for p in params], ", ")
args2str(pcounts) = join(["p$i" for i=pcounts], ", ")

qasm(blk::PutBlock{N}, args...) where N = qasm(blk.block, args...) * "  " * addrs2str(blk.addrs)
qasm(blk::ControlBlock{N}, args...) where N = "C-" * qasm(blk.block, args...) * "  " * addrs2str(blk.addrs)

# u3, u2,
# cu1, cu3
# if statement

# x, y, z, h, s, sdag, t, tdag
# u1 == z
for G in BASIC_1GATES
    GT = Symbol(G, :Gate)
    g = G |> String |> lowercase
    @eval qasm(blk::$GT, args...) = "$($g)"
    @eval qasm(blk::PutBlock{<:Any, <:Any, <:$GT}, args...) where N = "$($g)  " * addrs2str(blk.addrs)
end

# cx, cy, cz, ch
for G in BASIC_1GATES[1:4]
    GT = Symbol(G, :Gate)
    g = G |> String |> lowercase
    @eval qasm(blk::ControlBlock{<:Any, <:$GT}, args...) where N = "c$($g)  " * addrs2str([blk.ctrl_qubits[1], blk.addrs[1]])
end

# rx, ry, rz
# crz, crx, cry
# ccx, ccy, ccz
for G in BASIC_1GATES[1:3]
    GT = Symbol(G, :Gate)
    g = G |> String |> lowercase
    @eval qasm(blk::ControlBlock{<:Any, <:RotationGate{<:Any, <:Any, <:$GT}}, args...) where N = "cr$($g)($(blk.block.theta))  " * addrs2str([blk.ctrl_qubits[1], blk.addrs[1]])
    @eval qasm(blk::PutBlock{<:Any, <:Any, <:RotationGate{<:Any, <:Any, <:$GT}}, args...) where N = "r$($g)($(blk.block.theta))  " * addrs2str(blk.addrs)
    @eval qasm(blk::ControlBlock{<:Any, <:$GT, 2}, args...) where N = "cc$($g)  " * addrs2str([blk.ctrl_qubits..., blk.addrs[1]])
end

# id
qasm(blk::I2Gate, args...) = "id"
qasm(blk::PutBlock{<:Any, <:Any, <:I2Gate}, args...) = "id  " * addrs2str(blk.addrs)

function qasm(blk::AbstractBlock, args...)
    typeof(blk).name
end

function qasm(blk::ChainBlock, funcs=Dict{String, String}("main"=>""), fcount::Int=0, pcount::Int=0, level::Int=0)
    GNAME = "G$pcount"
    # however, we don't support function parameters for the moment.
    nparams = nparameters(blk)
    code = "gate $GNAME reg\n{\n"
    code *= join(["    "^(level+1) * qasm(b, funcs, fcount+1, pcount+nparams, level+1) for b in subblocks(blk)], "\n")
    code *= "\n}"
    funcs[GNAME] = code
    "$GNAME reg"
end


#using ExoplanetsSysSim
#using StatsFuns
if !isdefined(:JLD) using JLD end
if !isdefined(:DataFrames) using DataFrames end
if !isdefined(:CSV) using CSV end
#import DataFrames.DataFrame, DataFrames.isna
import DataFrames.DataFrame, DataArrays.ismissing
#import ExoplanetsSysSim.StellarTable.df

## Old code for generating stellar properties # TODO: WARNING: Should eventually use version in main branch to make sure have recent improvements

## stellar_table
function setup_star_table_christiansen(sim_param::SimParam; force_reread::Bool = false)
  #global df
  df = ExoplanetsSysSim.StellarTable.df
  if haskey(sim_param,"read_stellar_catalog") && !force_reread
     return df
     #return data
  end
  stellar_catalog_filename = convert(String,joinpath(Pkg.dir("ExoplanetsSysSim"), "data", convert(String,get(sim_param,"stellar_catalog","q1_q17_dr25_stellar.csv")) ) )
  df = setup_star_table_christiansen(stellar_catalog_filename)
  add_param_fixed(sim_param,"read_stellar_catalog",true)
  ExoplanetsSysSim.StellarTable.set_star_table(df)
  return df
end

function setup_star_table_christiansen(filename::String; force_reread::Bool = false)
  #global df
  wf = WindowFunction.setup_window_function(sim_param)
  df = ExoplanetsSysSim.StellarTable.df
  if ismatch(r".jld$",filename)
  try
    data = JLD.load(filename)
    df::DataFrame = data["stellar_catalog"]
    ExoplanetsSysSim.StellarTable.set_star_table(df)
  catch
    error(string("# Failed to read stellar catalog >",filename,"< in jld format."))
  end

  else
  try
    df = CSV.read(filename,nullable=true)
  catch
    error(string("# Failed to read stellar catalog >",filename,"< in ascii format."))
  end
  end # if ismatch

  has_mass = .! (ismissing.(df[:mass]) .| ismissing.(df[:mass_err1]) .| ismissing.(df[:mass_err2]))
  has_radius = .! (ismissing.(df[:radius]) .| ismissing.(df[:radius_err1]) .| ismissing.(df[:radius_err2]))
  has_dens = .! (ismissing.(df[:dens]) .| ismissing.(df[:dens_err1]) .| ismissing.(df[:dens_err2]))
  has_rest = .! (ismissing.(df[:rrmscdpp04p5]) .| ismissing.(df[:dataspan]) .| ismissing.(df[:dutycycle]))

#=
  in_Q1Q12 = []
  for x in df[:st_quarters]
    subx = string(x)
    subx = ("0"^(17-length(subx)))*subx
    indQ = search(subx, '1')
    if ((indQ < 1) | (indQ > 12))
      push!(in_Q1Q12, false)
    else
      push!(in_Q1Q12, true)
    end
  end
=#

#Comment out following block for 'is_FGK' if using stellar catalogs with cuts already made:
#=
  is_FGK = []
  for x in 1:length(df[:teff])
    if ((df[x,:teff] > 4000.0) & (df[x,:teff] < 7000.0) & (df[x,:logg] > 4.0))
#println("logg?: ", df[x,:logg])
      push!(is_FGK, true)
    else
      push!(is_FGK, false)
    end
  end
=#

  is_usable = has_radius .& has_mass .& has_rest #.& in_Q1Q12 #.& has_dens #.& is_FGK
  println("Total number of stars in our stellar catalog: ", sum(is_usable))
  #if contains(filename,"q1_q12_christiansen.jld")
#if contains(filename,"q1_q12_christiansen")   # TODO: Ask Danely what he's trying to do here.
#is_usable = is_usable #& in_Q1Q12
#end
  # See options at: http://exoplanetarchive.ipac.caltech.edu/docs/API_keplerstellar_columns.html
  # TODO SCI DETAIL or IMPORTANT?: Read in all CDPP's, so can interpolate?
  symbols_to_keep = [ :kepid, :mass, :mass_err1, :mass_err2, :radius, :radius_err1, :radius_err2, :dens, :dens_err1, :dens_err2, :rrmscdpp04p5, :dataspan, :dutycycle ]
  delete!(df, [~(x in symbols_to_keep) for x in names(df)])    # delete columns that we won't be using anyway
  usable = find(is_usable)
  df = df[usable, symbols_to_keep]
  ExoplanetsSysSim.StellarTable.set_star_table(df)
#end
  return df
end

function test_stellar_table() # TODO: Write test for stellar catalog functions
end


export RawAcquisitionData_3DUTE

function RawAcquisitionData_3DUTE(b::BrukerFile)
    T = Complex{acqDataType(b)}

    filename = joinpath(b.path, "fid")
    filenameTraj = joinpath(b.path, "traj")

    N = acqSize(b)
    # The data is padded in case it is not a multiple of 1024
    # For multi-channel acquisition data at concatenate then padded to a multiple of 1024 bytes
    numChannel = pvmEncNReceivers(b)
    profileLength = Int((ceil(N[1]*numChannel*sizeof(T)/1024))*1024/sizeof(T))
    numAvailableChannel = pvmEncAvailReceivers(b)
    phaseFactor = acqPhaseFactor(b)
    numSlices = acqNumSlices(b)
    numEchos = acqNumEchos(b)
    numEncSteps2 = length(N) == 3 ? N[3] : 1
    numRep = acqNumRepetitions(b)

    I = open(filename,"r") do fd
      read!(fd,Array{T,7}(undef, profileLength,
                                     numEchos,
                                     phaseFactor,
                                     numSlices,
                                     div(N[2], phaseFactor),
                                     numEncSteps2,
                                     numRep))[1:N[1]*numChannel,:,:,:,:,:,:]
    end

    traj = open(filenameTraj,"r") do fd
        read!(fd,Array{Float64,3}(undef, 3,N[1],
                                       N[2]))
    end

    objOrd = acqObjOrder(b)
    objOrd = objOrd.-minimum(objOrd)

    gradMatrix = acqGradMatrix(b)

    offset1 = acqReadOffset(b)
    offset2 = acqPhase1Offset(b)
    offset3 = ndims(b) == 2 ? acqSliceOffset(b) : acqPhase2Offset(b)

    profiles = Profile[]
    for nR = 1:numRep
      for nEnc2 = 1:numEncSteps2
        for nPhase2 = 1:div(N[2], phaseFactor)
          for nSl = 1:numSlices
            for nPhase1 = 1:phaseFactor
              for nEcho=1:numEchos
                  counter = EncodingCounters(kspace_encode_step_1=0,
                                             kspace_encode_step_2=0,
                                             average=0,
                                             slice=objOrd[nSl],
                                             contrast=nEcho-1,
                                             phase=0,
                                             repetition=nR-1,
                                             set=0,
                                             segment=0 )

                  G = gradMatrix[:,:,nSl]
                  read_dir = (G[1,1],G[2,1],G[3,1])
                  phase_dir = (G[1,2],G[2,2],G[3,2])
                  slice_dir = (G[1,3],G[2,3],G[3,3])

                  # Not sure if the following is correct...
                  pos = offset1[nSl]*G[:,1] +
                        offset2[nSl]*G[:,2] +
                        offset3[nSl]*G[:,3]

                  position = (pos[1], pos[2], pos[3])

                  head = AcquisitionHeader(number_of_samples=N[1], idx=counter,
                                          read_dir=read_dir, phase_dir=phase_dir,
                                          slice_dir=slice_dir, position=position,
                                          center_sample=0,
                                          available_channels = numAvailableChannel, #TODO
                                          active_channels = numChannel,
                                          trajectory_dimensions = 3,
                                          sample_time_us = parse(Float64,b["PVM_DigDw"])/1000)

                  dat = map(T, reshape(I[:,nEcho,nPhase1,nSl,nPhase2,nEnc2,nR],:,numChannel))
                  push!(profiles, Profile(head,convert.(Float32,traj[:,:,nPhase2]),dat) )
              end
            end
          end
        end
      end
    end

    params = brukerParams(b)
    params["trajectory"] = "custom"
    params["encodedSize"] = pvmMatrix(b)
    return RawAcquisitionData(params, profiles)
end


module SossGen

using Reexport

@reexport using Soss
@reexport using Gen
const logpdf = Soss.logpdf
const ifelse = Base.ifelse

import Gen

include("utils.jl")

# Following https://probcomp.github.io/Gen/dev/ref/extending/#Custom-generative-functions-1
# Define a trace data type

struct SossTrace{F, NT} <: Gen.Trace
    gen_fn :: F
    choices :: NT
    logprob :: Float64
    args
end

# Define a data type for your generative function

export SossGenerativeFunction
struct SossGenerativeFunction{M,F,V}  <: Gen.GenerativeFunction{V,SossTrace{SossGenerativeFunction}}
    model  :: M
    logpdf :: F
end

# Decide what the arguments to a generative function should be

function SossGenerativeFunction(m::Model{A,B,M}, ::typeof(codegen)) where {A,B,M}
    ℓ(args, data) = logpdf(m(args...), data, codegen)::Float64
    SossGenerativeFunction{Model{A,B,M},typeof(ℓ), typeof(variables(m))}(m, ℓ)
end

function SossGenerativeFunction(m::Model{A,B,M}) where {A,B,M}
    ℓ(args, data) = logpdf(m(args...), data)::Float64
    SossGenerativeFunction{Model{A,B,M},typeof(ℓ), typeof(variables(m))}(m, ℓ)
end

# Implement methods of the Generative Function Interface

# simulate

function Gen.simulate(gen_fn::SossGenerativeFunction{M,F,V}, args::Tuple) where {M,F,V}
    choices = rand(gen_fn.model(args...))
    score = gen_fn.logpdf(args, choices)
    SossTrace{SossGenerativeFunction,typeof(choices)}(gen_fn, choices, score, args)
end

# has_argument_grads

# accepts_output_grad

# get_args

function Gen.get_args(t::SossTrace{SossGenerativeFunction,NT}) where {NT}
    t.args
end

# get_retval

function Gen.get_retval(t::SossTrace{SossGenerativeFunction,NT}) where {NT}
    t.choices
end

# get_choices

function Gen.get_choices(t::SossTrace{SossGenerativeFunction,NT}) where {NT}
    c = Gen.choicemap()
    Ts = NTtypes(NT)
    for (k::Symbol, v) in pairs(t.choices)
        c[k] = v :: Ts[k]
    end
    c
end

# get_score

function Gen.get_score(t::SossTrace{SossGenerativeFunction,NT}) where {NT}
    t.logprob
end

# get_gen_fn

# project

Gen.project(t::SossTrace{SossGenerativeFunction,NT}, ::Gen.EmptySelection) where {NT} = 0.0

# generate

function Gen.generate(gen_fn::SossGenerativeFunction{M,F,V}, args::Tuple, constraints::Gen.ChoiceMap) where {M,F,V}
    kvs = Gen.get_values_shallow(constraints)
    data = namedtuple(Dict{Symbol, Any}(kvs))
    weight, choices = weightedSample(gen_fn.model(args...), data)
    logprob = gen_fn.logpdf(args, choices)
    SossTrace{SossGenerativeFunction,typeof(choices)}(gen_fn, choices, logprob, args), weight
end

# update

function Gen.update(t::SossTrace{SossGenerativeFunction,NT}, new_args::Tuple, argdiffs::Tuple, constraints::Gen.ChoiceMap) where {NT}
    retdiff = ifelse(isempty(constraints), Gen.NoChange(), Gen.UnknownChange())
    new_choices = NT(merge(t.choices, namedtuple(Dict{Symbol,Any}(Gen.get_values_shallow(constraints))))) # :: NT
    new_logprob = t.gen_fn.logpdf(new_args, new_choices) :: Float64
    new_trace = SossTrace{SossGenerativeFunction,NT}(t.gen_fn, new_choices, new_logprob, new_args)
    weight = new_logprob - t.logprob
    discard = Gen.choicemap()
    for (k::Symbol, v) in Gen.get_values_shallow(constraints)
        discard[k] = t.choices[k]
    end
    return (new_trace, weight, retdiff, discard)
end

# regenerate

# choice_gradients

# accumulate_param_gradients!

end


import CSV
using DataFramesMeta
import StatsBase: countmap
import Images: label_components
rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
mdates = pyimport("matplotlib.dates")

include("./plotting.jl")
# Formatting functions:
_remove_ticks(ax) = ax.tick_params(
    axis = "both",
    which = "both",
    bottom = false,
    top = false,
    # labelbottom=false, labelleft=false,
    left = false,
    right = false,
)
_set_figsize(fig, h = 3, w = 3.5) = (fig.set_figheight(h), fig.set_figwidth(w))

function plot_rrc_oil_water()
    rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
    rcParams["font.family"] = "Helvetica"
    rcParams["font.size"] = 16
    rcParams["font.weight"] = "bold"

    data = readdf("summary_inj_prod.csv")
    oilyears = data[!, :production]
    wateryears = data[!, :injection]
    xs = data[!, :Year]

    fig, axes = plt.subplots(2, 1)
    # fig2, axes2 = plt.subplots(2, 1)
    bg_color = "b"
    active_color = "r"
    colors = fill(bg_color, length(xs))
    colors[end - 1] = active_color

    xticks = [2007, 2011, 2015, 2018]
    # yticks=[0, 0.5, 1, 1.5, 2, 2.5]

    titles = ["Oil Production", "Total Injections"]
    for (title, ax, values) in zip(titles, axes, [oilyears, wateryears])
        vals = values ./ 365 ./ 1e6
        # ax.bar(xs, vals, color=colors)
        ax.plot(xs, vals, "-o", lw = 5, ms = 10)
        ax.set_xticks(xticks)
        # ax.set_yticks(yticks)
        ax.set_ylabel("Million bbl/day")
        ax.grid()
        ax.set_title(title)
    end
    fig.set_figheight(9)
    fig.set_figwidth(4.5)
    fig.savefig("cisr-lines.pdf", bbox_inches = "tight", transparent = true, dpi = 100)

end

function pnas_outlier_figure(
    station_name = "TXMC",
    outlier_color = "r",
    h = 3,
    w = 3.5,
    bins = 40,
)
    rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
    # https://matplotlib.org/tutorials/introductory/customizing.html#a-sample-matplotlibrc-file
    # so obscure
    rcParams["pdf.fonttype"] = 42
    rcParams["ps.fonttype"] = 42
    rcParams["font.family"] = "Helvetica"
    rcParams["font.size"] = 16
    rcParams["font.weight"] = "bold"

    # a. histogram before/after outlier
    fig, ax = plt.subplots()

    unw_vals = h5read("gps_pixels_78.h5", station_name)
    geolist = Sario.load_geolist_from_h5("gps_pixels_78.h5")
    intlist = Sario.load_intlist_from_h5("gps_pixels_78.h5")
    end_date = Date(2018, 1, 1)
    geolist18 = geolist[geolist .> end_date]
    idxs = InsarTimeseries._good_idxs(geolist18, intlist)
    intlist17 = intlist[idxs]
    geolist17 = geolist[geolist .< end_date]
    unw_vals17 = unw_vals[idxs]
    unw_vals17_cm = p2c .* unw_vals17
    vm = maximum(abs.(unw_vals17_cm))
    n1, _, _ = ax.hist(
        unw_vals17_cm,
        range = (-vm, vm),
        bins = bins,
        label = "All data",
        color = "red",
    )
    MAT.matwrite("outliers_hist_vals.mat", Dict("unw_vals_$station_name" => unw_vals17_cm))
    # n1, _, _ = ax.hist(unw_vals17.* p2c, range=(-vm, vm), bins=bins, label="All data", color="red", density=true, alpha=0.5)
    ax.grid()

    means = InsarTimeseries.mean_abs_val(geolist17, intlist17, unw_vals17)
    nsigma = 4
    low, high = InsarTimeseries.two_way_cutoff(means, nsigma)
    bigidx = means .> high
    idxs = InsarTimeseries._good_idxs(geolist17[bigidx], intlist17)
    vals_after_outliers_cm = unw_vals17_cm[idxs]
    n1, _, _ = ax.hist(
        vals_after_outliers_cm,
        range = (-vm, vm),
        bins = bins,
        label = "Outliers removed",
        color = "C0",
    )
    # n1, _, _ = ax.hist(vals_after_outliers .* p2c, range=(-vm, vm), bins=bins, density=true, label="Outliers removed", color="C0", alpha=.5)
    ax.set_xlim((-15, 15))
    ax.set_xticks([-15, -7.5, 0, 7.5, 15])
    ax.set_yticks([0, 50, 100])
    _remove_ticks(ax)
    _set_figsize(fig)
    fig.savefig("outliers_a.pdf", bbox_inches = "tight", transparent = true, dpi = 100)

    # b) std dev of insar at GPS stations vs distance to reference
    dists = [
        MapImages.latlon_to_dist(
            MapImages.station_latlon("TXKM"),
            MapImages.station_latlon(n),
        ) for n in all_stations
    ]

    txmc_cm_78 = Float64.(p2c .* h5read("gps_pixels_78.h5", "TXMC"))
    @show median(txmc_cm_78)

    stds78_cm =
        Float64.(std.(p2c .* [h5read("gps_pixels_78.h5", name) for name in station_name_list78],))
    dists78 = [
        MapImages.latlon_to_dist(
            MapImages.station_latlon("TXKM"),
            MapImages.station_latlon(n),
        ) for n in station_name_list78
    ]

    stds85_cm =
        Float64.(std.(p2c .* [h5read("gps_pixels_85.h5", name) for name in station_name_list85],))
    dists85 = [
        MapImages.latlon_to_dist(
            MapImages.station_latlon("TXKM"),
            MapImages.station_latlon(n),
        ) for n in station_name_list85
    ]
    stations_all = vcat(station_name_list78, station_name_list85)
    dists_all = vcat(dists78, dists85)
    stds_all = vcat(stds78_cm, stds85_cm)
    dists_all = dists_all[stations_all .!= "TXKM"]
    stds_all = stds_all[stations_all .!= "TXKM"]
    stations_all = stations_all[stations_all .!= "TXKM"]

    df_dists = DataFrame(dists = dists_all, stds = stds_all, names = stations_all)

    sm = lm(@formula(stds ~ dists), df_dists)      # Linear model
    sm2 = lm(@formula(stds ~ sqrt(dists)), df_dists) # sqrt model
    # a, b = coef(sm)
    a, b = coef(sm2)
    @show a, b * 50

    fig, ax = plt.subplots()
    dists, stds = (df_dists[!, :dists], df_dists[!, :stds])
    ax.scatter(dists, stds, s = 40)
    xs = 1:maximum(dists)
    # ax.plot(xs, a .+ b .* xs, lw=3)  # Linear model
    ax.plot(xs, a .+ b .* sqrt.(xs), lw = 3)    # sqrt(dist) model
    ax.set_ylim((0, 5))
    ax.set_yticks([0, 1, 2, 3, 4])
    ax.set_xticks([0, 100, 200])
    ax.set_ylabel("Noise Level [cm]")
    ax.set_xlabel("Distance [km]")
    ax.grid()
    # _remove_ticks(ax)
    _set_figsize(fig)
    fig.savefig("outliers_b.pdf", bbox_inches = "tight", transparent = true, dpi = 100)
    MAT.matwrite(
        "outliers_b.mat",
        Dict("a" => a, "b" => b, "xs" => collect(xs), "dists" => dists, "stds" => stds),
    )

    return sm, sm2
    # c) errors in 2, 3, 4 year solution at GPS vs distance, stacking only
    # d) errors in 2, 3, 4 year solution at GPS vs distance, after outlier removal
    df, errors = get_gps_errors(to_cm = true)
    df = @where(df, :station .!= "TXKM")

    years = collect(2016:2018)
    cols1 = [Symbol("p$(n)_stack_$(y)") for y in years, n in [78, 85]]
    cols2 = [Symbol("p$(n)_outliers_$(y)") for y in years, n in [78, 85]]
    colors = ["b", "g", "k"]
    fnames = ["outliers_c.pdf", "outliers_d.pdf"]

    maxy = 3
    for (idx, cols) in enumerate([cols1, cols2])
        fig, ax = plt.subplots()
        # errvals = abs.(vcat([df[!, c] for c in cols]...))
        for yidx = 1:3
            year = years[yidx]
            c1, c2 = cols[yidx, :]
            errvals = abs.(vcat(df[!, c1], df[!, c2]))
            dists = repeat(df[!, :dist_to_txkm], 2)
            dists = dists[.!isnan.(errvals)]
            errvals = errvals[.!isnan.(errvals)]
            @show maximum(errvals)

            ax.scatter(dists, errvals, s = 40, c = colors[yidx])
            ax.set_ylim((0, maxy))

            df1 = DataFrame(dist = dists, errval = errvals)
            sm = lm(@formula(errval ~ dist), df1)
            a, b = coef(sm)
            @show a, b * 50
            ax.plot(xs, a .+ b .* xs, colors[yidx], lw = 3)

            _remove_ticks(ax)
            _set_figsize(fig)
            fname = fnames[idx]
            fig.savefig(fname, bbox_inches = "tight", transparent = true, dpi = 100)
        end
    end
end

_max_label(labels) = sort(
    [tup for tup in countmap(vec(labels)) if tup[1] != 0],
    by = x->x[2],
    rev = true,
)[1][1]

function save_pnas_images(;
    years = [2016, 2017, 2018],
    fnames = ["velocities_$(year)_current.h5" for year in years],
    # fnames=["velocities_$(year)_linear_max800_sigma4.h5" for year in years],
    outdir = ".",
    # zoomlats=(31.6, 30.9), zoomlons=(-103.8, -103.),  # Zoomed just to stripes
    zoomlats = (31.9, 31.0),
    zoomlons = (-103.8, -102.85),  # Covers txkm/txmh
    latbounds = (32.8, 30.7041),
    dset = "velos_shifted/1",
    cmap1 = "seismic_wide_y",
    cmap2 = "seismic_wide_y", # cmap2=rdylbl,
    vm1 = 10,
    vm2 = 10,
    dem_mask_height = nothing,
)
    maxdays = maximum([_num_days(Sario.load_geolist_from_h5(f, dset)) for f in fnames])
    @show maxdays
    vm1 *= (maxdays / 3650)
    vm2 *= (maxdays / 3650)
    @show vm1, vm2

    m1, m2 = nothing, nothing
    for f in fnames
        @show f
        days = _num_days(Sario.load_geolist_from_h5(f, dset))
        @show days

        m1 = MapImage(f, dset)
        m1[m1 .== 0] .= NaN

        if !isnothing(dem_mask_height)
            dem = Sario.load("elevation_looked.dem")
            labels = label_components(dem .< 1200)
            maxl = _max_label(labels)
            m1[labels .!= maxl] .= NaN
        end

        m1 = m1[latbounds, (-110.0, -90.0)] # no lon bounds really
        @show MapImages.grid_extent(m1)

        @show extremanan(m1)
        m1 .*= (days / 3650)  # Cumulative, in cm
        @show extremanan(m1)
        m2 = m1[zoomlats, zoomlons]
        out1 = joinpath(outdir, replace(f, ".h5" => ".tif"))
        # out2 = replace(f, ".h5"  => "_inset.tif")
        out2 = joinpath(outdir, replace(f, ".h5" => "_inset_gps.tif"))

        # Save as both png and tif
        save_img_geotiff(out1, m1, cmap = cmap1, vm = vm1)
        # save_img_geotiff(out2, m2, cmap=cmap2, vm=vm2)
        plt.imsave(
            replace(out1, ".tif" => ".png"),
            m1,
            cmap = cmap2,
            vmin = -vm2,
            vmax = vm2,
            dpi = 150,
        )
        plt.imsave(
            replace(out2, ".tif" => ".png"),
            m2,
            cmap = cmap2,
            vmin = -vm2,
            vmax = vm2,
            dpi = 150,
        )
    end
    # Now save one with a colorbar
    fig, axes = plt.subplots(1, 2)
    axim = axes[1].imshow(m1, cmap = cmap1, vmax = vm1, vmin = -vm1)
    fig.colorbar(axim, ax = axes[1])
    axim = axes[2].imshow(m2, cmap = cmap2, vmax = vm2, vmin = -vm2)
    fig.colorbar(axim, ax = axes[2])

    cbar_name = "colorbars.png"
    save_paper_figure(fig, cbar_name, false, 400)
end

function _errs_to_cm(df)
    # mm/year, to cm for 2,3,4 years
    scales = vcat(4.1 .* ones(4), 3 .* ones(4), 2 .* ones(4)) ./ 10
    for (s, c) in zip(scales, names(df)[3:end])
        df[!, c] = df[!, c] .* s
    end
    return df
end

function get_gps_errors(; outfile = "gps_errors.csv", to_cm = false)
    if isfile(outfile)
        df = CSV.read(outfile)
        if to_cm
            df = _errs_to_cm(df)
        end
        return df, [c for c in eachcol(df)][3:end]
    end

    p78 = "/data1/scott/pecos/path78-bbox2/igrams_looked/"
    p85 = "/data4/scott/path85/stitched/igrams_looked/"

    dists = [
        MapImages.latlon_to_dist(
            MapImages.station_latlon("TXKM"),
            MapImages.station_latlon(n),
        ) for n in all_stations
    ]
    errors = Any[]
    push!(errors, all_stations)
    push!(errors, dists)
    # errors85_4 = get_gps_error(p85*"velocities_2018_current.h5", all_stations, dset="velos_shifted/1", verbose=true, window=7, shift=-0.)
    # errors78_4 = get_gps_error(p78*"velocities_2018_current.h5", all_stations, dset="velos_shifted/1", verbose=true, window=7, shift=-0.)

    for f in [
        "velocities_2018_stackavg_max800.h5",
        "velocities_2018_current.h5",
        "velocities_2017_stackavg_max800.h5",
        "velocities_2017_current.h5",
        "velocities_2016_stackavg_max800.h5",
        "velocities_2016_current.h5",
    ]
        fname = p85 * f
        @show fname
        push!(
            errors,
            round.(get_gps_error(
                    fname,
                    all_stations,
                    dset = "velos_shifted/1",
                    verbose = false,
                    window = 7,
                    shift = -0.0,
                ),
                digits = 2,),
        )
        fname = p78 * f
        @show fname
        push!(
            errors,
            round.(get_gps_error(
                    fname,
                    all_stations,
                    dset = "velos_shifted/1",
                    verbose = false,
                    window = 7,
                    shift = -0.0,
                ),
                digits = 2,),
        )
    end
    columns = [
        :station,
        :dist_to_txkm,
        :p85_stack_2018,
        :p78_stack_2018,
        :p85_outliers_2018,
        :p78_outliers_2018,
        :p85_stack_2017,
        :p78_stack_2017,
        :p85_outliers_2017,
        :p78_outliers_2017,
        :p85_stack_2016,
        :p78_stack_2016,
        :p85_outliers_2016,
        :p78_outliers_2016,
    ]

    df = DataFrame(errors, columns)
    CSV.write("gps_errors.csv", df)
    return df, errors
end

function pnas_gps_error_table()
    df, errors = get_gps_errors()
    show(stdout, MIME("text/latex"), df)
    @show extrema.(filter.(!isnan, errors[3:end]))
    @show rms.(errors[3:end])
    @show maximum.(filter.(!isnan, [abs.(a) for a in errors[3:end]]))
    return df
end

function write_gps_pixels(station_names, max_temp = 800)
    geolist, intlist, valid_idxs =
        InsarTimeseries.load_geolist_intlist("unw_stack.h5", "geolist_ignore.txt", max_temp)
    gpss = [
        get_stack_vals("unw_stack.h5", s, 3, "stack_flat_shifted", valid_idxs)
        for s in station_names
    ]
    for (idx, s) in enumerate(station_names)
        h5write("gps_pixels.h5", s, gpss[idx])
    end
    Sario.save_intlist_to_h5("gps_pixels.h5", intlist)
    Sario.save_geolist_to_h5("gps_pixels.h5", geolist)
    Sario.save_dem_to_h5("gps_pixels.h5", Sario.load("dem.rsc"))
end

function pnas_outlier_figure_old(
    station_name = "TXMC",
    outlier_color = "r",
    h = 3,
    w = 3.5,
    bins = 40,
)

    # unw_vals = [h5read("gps_pixels.h5", name) for name in station_name_list78];
    unw_vals = h5read("gps_pixels_78.h5", station_name)
    geolist = Sario.load_geolist_from_h5("gps_pixels_78.h5")
    intlist = Sario.load_intlist_from_h5("gps_pixels_78.h5")

    end_date = Date(2018, 1, 1)
    geolist18 = geolist[geolist .> end_date]
    idxs = InsarTimeseries._good_idxs(geolist18, intlist)
    intlist17 = intlist[idxs]

    geolist17 = geolist[geolist .< end_date]
    unw_vals17 = unw_vals[idxs]
    geolist17, intlist17, unw_vals17 = (geolist, intlist, unw_vals)

    # 1. plot grouped by date
    fig, ax = plot_big_days(
        geolist17,
        intlist17,
        unw_vals17,
        legend = false,
        color = outlier_color,
    )
    _remove_ticks(ax)
    _set_figsize(fig)
    fig.savefig("outliers2.pdf", bbox_inches = "tight", transparent = true, dpi = 100)

    # 2. cm vs baseline igram outlier plot
    # function plot_grouped_by_day(geo, int, vals, nsigma=0, min_spread=2)
    geo = geolist17
    means = InsarTimeseries.mean_abs_val(geo, intlist17, p2c * unw_vals17)
    println("median val: $(median(means))")
    fig, ax = plt.subplots()
    ax.scatter(geo, means, label = "means by date")
    nsigma = 4
    low, high = InsarTimeseries.two_way_cutoff(means, nsigma)
    ax.plot(geo, ones(length(geo)) * high, label = "$nsigma sigma MAD cutoff")
    bigidx = means .> high
    ax.scatter(geo[bigidx], means[bigidx], c = outlier_color)
    _remove_ticks(ax)
    _set_figsize(fig)  # Double the width here
    fig.savefig("outliers1.pdf", bbox_inches = "tight", transparent = true, dpi = 100)

    # 3. line  plot of insar after outliers vs GPS
    Blin = sum(InsarTimeseries.prepB(geolist17, intlist17), dims = 2)
    before = (Blin \ unw_vals17) * p2c * 365 * 10  # Change to mm/year

    start_date = Date(2015, 1, 1)
    ref = nothing
    lw = 5
    fig, ax = plot_gps_los(
        station_name,
        before;
        ref = ref,
        lw = lw,
        end_date = end_date,
        start_date = start_date,
        ylim = (-3, 3),
        title = "",
        bigfont = false,
        offset = false,
    )

    # Now solve without outliers
    dts, los = get_gps_los(
        station_name,
        reference_station = ref,
        end_date = end_date,
        start_date = start_date,
    )
    day_nums = _get_day_nums(dts)
    idxs = InsarTimeseries._good_idxs(geolist17[bigidx], intlist17)
    vals_after_outliers = unw_vals17[idxs]
    after = p2c * (Blin[idxs] \ vals_after_outliers)  # cm/day
    full_defo = after * (dts[end] - dts[1]).value
    # ax.plot(dts, -full_defo/2 .+ day_nums .* after, "c", lw=3, label="Insar")
    ax.plot(dts, day_nums .* after, "C0", lw = lw)
    _remove_ticks(ax)
    _set_figsize(fig, h, 2.3w)
    fig.savefig("outliers3.pdf", bbox_inches = "tight", transparent = true, dpi = 100)

    println("Before: $before mm/yr, $(before * (dts[end] - dts[1]).value / 3650) total defo")
    println("After: $(after * 3650) mm/yr, $full_defo total defo")

    # 4. histogram before/after outlier
    fig, ax = plt.subplots()

    # ax = axes[ii, jj]
    # ax = axes[idx]
    vm = maximum(abs.(unw_vals17 * p2c))
    n1, _, _ = ax.hist(
        unw_vals17 .* p2c,
        range = (-vm, vm),
        bins = bins,
        label = "All data",
        color = "red",
    )
    n1, _, _ = ax.hist(
        vals_after_outliers .* p2c,
        range = (-vm, vm),
        bins = bins,
        label = "Outliers removed",
        color = "C0",
    )
    _remove_ticks(ax)
    _set_figsize(fig)
    fig.savefig("outliers4.pdf", bbox_inches = "tight", transparent = true, dpi = 100)

    stds78 =
        Float64.(std.([h5read("gps_pixels_78.h5", name) for name in station_name_list78]))
    dists78 = [
        MapImages.latlon_to_dist(
            MapImages.station_latlon("TXKM"),
            MapImages.station_latlon(n),
        ) for n in station_name_list78
    ]
    df78 = DataFrame(dists = dists78, stds = stds78, names = station_name_list78)

    stds85 =
        Float64.(std.([h5read("gps_pixels_85.h5", name) for name in station_name_list85]))
    dists85 = [
        MapImages.latlon_to_dist(
            MapImages.station_latlon("TXKM"),
            MapImages.station_latlon(n),
        ) for n in station_name_list85
    ]
    df85 = DataFrame(dists = dists85, stds = stds85, names = station_name_list85)

    sm78 = lm(@formula(stds ~ dists), df78)
    a1, b1 = coef(sm78)
    sm85 = lm(@formula(stds ~ dists), df85)
    a2, b2 = coef(sm85)

    fig, ax = plt.subplots()
    ax.scatter(dists78, stds78, s = 40)
    xs = 1:maximum(dists78)
    ax.plot(xs, a1 .+ b1 .* xs, lw = 3)

    ax.scatter(dists85, stds85, s = 40)
    xs = 1:maximum(dists85)
    ax.plot(xs, a2 .+ b2 .* xs, lw = 3)

    ax.set_xlabel("km from TXKM")
    ax.set_ylabel("Std. Dev. of unw values")
    _remove_ticks(ax)
    _set_figsize(fig)
    fig.savefig("outliers5.pdf", bbox_inches = "tight", transparent = true, dpi = 100)

    return (geolist17, intlist17, unw_vals17)
end

# Hunjoo figures
#
function dict_to_vars(d::AbstractDict{String,Any})
    for key in keys(d)
        @eval $(Symbol(key)) = d[$(QuoteNode(key))]
    end
end
function get_wells()
    wolfidx = occursin.("WOLFCAMP", ff[!, prodcol])
    delidx = occursin.("DELAWARE", ff[!, prodcol])
end

function hunjoo_figures()
    rcParams["font.family"] = "Helvetica"
    rcParams["font.size"] = 16
    rcParams["font.weight"] = "medium"

    # variables = MAT.matread("/Users/scott/Documents/Learning/PNAS-git-overleaf/data/PNAS_Geomechanics3.mat")
    # variables = MAT.matread("/Users/scott/Desktop/figures/hunjoo/PNAS_Geomechanics7.mat")
    variables = MAT.matread("/Users/scott/Documents/Learning/GRL-git-overleaf-resubmit/hunjoo-revision-0526/GRL_rev1_Final_b0.01_Scott.mat")

    # dict_to_vars(variables)  # Fails cuz of some namespace thing
    # for key in keys(variables)
    #     @eval $(Symbol(key)) = variables[$(QuoteNode(key))]
    # end
    #  Area of Interest
    R_earth = 3959 * 5280 # (miles to feet) Earth Radius

    XD_min = -103.59
    XD_max = -103.44
    YD_min = 31.32
    YD_max = 31.47

    # #  Reservoir properties
    # RD = variables["Reservoir"][:, 1:2]   # RD = Reservoir(:,1),Reservoir(:,2)];
    # RR_R_D1 = (variables["Reservoir"][:, 3] / (5 * 3280.84) * 55).^2      # Reservoir radius to scatter plot (55^2 = 5 km = 5*3280.84 ft)
    # RR_R_D2 = (variables["Reservoir"][:, 3] / (1 * 3280.84) * 35).^2      # Reservoir radius to scatter plot (35^2 = 1 km = 3280.84 ft)

    # Cities
    Pecos = [-103.4932, 31.4229]             # (degrees)

    # Figure Scale
    ZD_min = -15.0
    ZD_max = 15.0
    c_min = -10.0
    c_max = 10.0
    FZ_min = -4.0
    FZ_max = 2.0

    demrsc = MapImages.from_grid(variables["XD"], variables["YD"])

    # Vertical Deformation
    cmap = "seismic_narrow_y"
    vm = 10

    # Pecos Area InSAR 2,3,1
    # TODO: DRAW TRANSECT FROM B B''
    fig, ax = plt.subplots()
    ax.imshow(variables["Z_Disp_VV18"], cmap=cmap, vmax=vm, vmin=-vm)
    _remove_ticks(ax)
    _set_figsize(fig)
    fig.savefig("hunjoo1.pdf", bbox_inches="tight", transparent=true, dpi=100)

    # Fault model: 2,3,2
    fig, ax = plt.subplots()
    # hold on
    # surf(XD_VV18,YD_VV18,Z_Disp_VV18,"EdgeColor","none")
    axim = ax.imshow(variables[ "Z_Disp_F3" ], cmap=cmap, vmax=vm, vmin=-vm, extent=(XD_min, XD_max,YD_min,YD_max))  #,"EdgeColor","none")   # InSAR data
    fig.colorbar(axim)
    ax.plot(variables[ "BB_XD" ][[1, end]], variables[ "BB_YD" ][[1, end]], lw=1.5)
    ax.plot(Pecos[1],Pecos[2], "ko");
    save_img_geotiff("hunjoo2.tif", MapImage(variables[ "Z_Disp_F3" ], demrsc), cmap = cmap, vm = vm)

    # reservoir model: 2,3,3
    fig, ax = plt.subplots()
    # hold on
    # surf(XD_VV18,YD_VV18,Z_Disp_VV18,"EdgeColor","none")
    cur_im = variables[ "Z_Disp_R3" ][:, :, 3]
    axim = ax.imshow(cur_im, cmap=cmap, vmax=vm, vmin=-vm, extent=(XD_min, XD_max,YD_min,YD_max))  #,"EdgeColor","none")   # InSAR data
    fig.colorbar(axim)
    ax.plot(variables[ "BB_XD" ][[1, end]], variables[ "BB_YD" ][[1, end]], lw=1.5)
    ax.plot(Pecos[1],Pecos[2], "ko")
    save_img_geotiff("hunjoo3.tif", MapImage(cur_im, demrsc), cmap = cmap, vm = vm)

    # Fault + reservoir model: 2,3,4
    fig, ax = plt.subplots()
    # surf(XD,YD,Z_Disp_T(:,:,3),"EdgeColor","none")
    im_mod = variables["Z_Disp_T3"][:, :, 3]
    axim = ax.imshow(
        im_mod,
        cmap = cmap,
        vmax = vm,
        vmin = -vm,
        extent = (XD_min, XD_max, YD_min, YD_max),
    )  # ,"EdgeColor","none")   # InSAR data
    # fig.colorbar(axim)
    ax.plot(variables["BB_XD"][[1, end]], variables["BB_YD"][[1, end]], lw = 1.5)
    ax.plot(Pecos[1], Pecos[2], "ko")

    _remove_ticks(ax)
    _set_figsize(fig)
    fig.savefig("hunjoo4.pdf", bbox_inches = "tight", transparent = true, dpi = 100)
    save_img_geotiff("hunjoo4.tif", MapImage(im_mod, demrsc), cmap = cmap, vm = vm)

    # Residual map: 2,3,5
    fig, ax = plt.subplots()
    # surf(XD,YD,Z_Disp_T(:,:,3),"EdgeColor","none")
    im_mod = variables["Res_Z_VV18"]
    axim = ax.imshow(
        im_mod,
        cmap = cmap,
        vmax = vm,
        vmin = -vm,
        extent = (XD_min, XD_max, YD_min, YD_max),
    )  # ,"EdgeColor","none")   # InSAR data
    # fig.colorbar(axim)
    ax.plot(variables["BB_XD"][[1, end]], variables["BB_YD"][[1, end]], lw = 1.5)
    ax.plot(Pecos[1], Pecos[2], "ko")

    _remove_ticks(ax)
    _set_figsize(fig)
    fig.savefig("hunjoo5.pdf", bbox_inches = "tight", transparent = true, dpi = 100)
    save_img_geotiff("hunjoo5.tif", MapImage(im_mod, demrsc), cmap = cmap, vm = vm)

    # Supplement: figure reservoir figure with circles
    fig, ax = plt.subplots()
    # surf(XD,YD,Z_Disp_T(:,:,3),"EdgeColor","none")
    im_mod = variables["UZ"]
    axim = ax.imshow(
        im_mod,
        cmap = cmap,
        vmax = vm,
        vmin = -vm,
        extent = (XD_min, XD_max, YD_min, YD_max),
    )
    _remove_ticks(ax)
    _set_figsize(fig)
    # fig.savefig("hunjoo-supp1.pdf", bbox_inches = "tight", transparent = true, dpi = 100)
    save_img_geotiff("hunjoo-supp1.tif", MapImage(im_mod, demrsc), cmap = cmap, vm = vm)

    # # B-B' cross-section
    # fig, ax1 = plt.subplots()
    # ax1.plot(BB_L', BB_Z_Disp_I18, "bs", label = "2018 InSAR")
    # ax1.plot(BB_L', BB_Z_Disp_R18, "r-.", lw = 1.5, label = "Reservoir")
    # ax1.plot(BB_L', BB_Z_Disp_F18, "r:", lw = 1.5, label = "Fault")
    # ax1.plot(BB_L', BB_Z_Disp_T18, "r-", lw = 1.5, label = "Res. + Fault")
    # ax1.set_xlim((0, maximum(BB_L)))
    # ax1.set_ylim((2 * FZ_min, 2 * FZ_max))
    # ax1.set_xlabel("B-B' (km)")
    # ax1.set_ylabel("Surface Deformation (cm)")

    # # yyaxis right
    # ax2 = ax1.twinx()
    # ax2.plot(BB_FC[:, 1, 1], BB_FC[:, 2, 1], "g-", lw = 2.0)
    # ax2.plot(BB_FC[:, 1, 2], BB_FC[:, 2, 2], "g-", lw = 2.0)
    # ax2.plot(BB_FC[:, 1, 3], BB_FC[:, 2, 3], "g-", lw = 2.0)
    # ax2.plot(BB_FC[:, 1, 4], BB_FC[:, 2, 4], "g-", lw = 2.0)
    # # text(BB_FC(2,1,1),BB_FC(2,2,1),"F1",lw=2.0)  #,"VerticalAlignment","top")
    # # text(BB_FC(2,1,2),BB_FC(2,2,2),"F2",lw=2.0)  #,"VerticalAlignment","top")
    # # text(BB_FC(2,1,3),BB_FC(2,2,3),"F3",lw=2.0)  #,"VerticalAlignment","top")
    # # text(BB_FC(2,1,4),BB_FC(2,2,4),"F4",lw=2.0)  #,"VerticalAlignment","top")
    # ax2.set_ylabel("Fault Depth (km)")
    # ax2.set_xlim((0, maximum(BB_L)))
    # ax2.set_ylim((0.5 * FZ_min, 0.5 * FZ_max))
    # # ax2.legend(["2018 InSAR","Reservoir Model","Fault Model","Res.+Fault Model","Location","northeast"])
    # # set(findobj(gcf,"type","axes"),"FontSize",12,"TickDir","out")
    # #
    # xticks = [0, 7, 14]
    # y1ticks = [-8, -4, 0, 4]
    # y2ticks = [-2, -1, 0, 1]

    # ax1.set_xticks(xticks)
    # ax1.set_yticks(y1ticks)
    # ax2.set_yticks(y2ticks)
    # ax1.grid()
    # fig.legend(frameon = false, fontsize = "small", ncol = 2)

    # _set_figsize(fig, 1.3 .* [3, 3.5]...)
    # # _set_figsize(fig)
    # fig.savefig("hunjoo4.pdf", bbox_inches = "tight", transparent = true, dpi = 100)
end

# Functions for weighting by covariance

import ImageFiltering: imfilter, Kernel
function get_aps(stack_in)
    return imfilter(stack_in .- imfilter(stack_in, (Kernel.gaussian((0, 0, 1)), Kernel.Laplacian((false, false, true)))),  Kernel.gaussian((2, 2, 0)))
end

function try_weight()
    test_sub = get_aps(defo_stack_unreg);
    var_by_day = dropdims(var((defo_stack_unreg .- test_sub); dims = (1, 2)), dims = (1, 2));
    var_by_day2 = abs.(test_sub[400, 400, :])
    Qddfull2 = build_Q(abs.(var_by_day2), geolist, intlist);
    Qinv2 = inv(Qddfull2);
    W = inv(A' * Qinv * A) * A' * Qinv2;
end

function build_Q(var_dates, geolist, intlist, diag_val = 0.5)
    N = length(intlist);
    Q = zeros((N, N));
    for (jdx, ig2) in enumerate(intlist)
        for (idx, ig1) in enumerate(intlist)
            if jdx > idx
            # if Q[jdx, idx] != 0
                Q[idx, jdx] = Q[jdx, idx]
                continue  # dont need to redo, symmetric
            end
            # @show idx, jdx
            if idx == jdx
                v1 = var_dates[findfirst(isequal(ig1[1]), geolist)]
                v2 = var_dates[findfirst(isequal(ig1[2]), geolist)]
                Q[idx, jdx] = v1 + v2
            else
                cm = countmap([ig1..., ig2...])
                if length(Set([ig1..., ig2...])) == 4
                    Q[idx, jdx] = 0
                    continue
                end
                repeated_date = sort(collect(countmap([ig1..., ig2...])), by = x->x[2], rev = true)[1][1]
                # @show repeated_date
                # @show findfirst(isequal(repeated_date), geolist)
                sigma_repeat = var_dates[findfirst(isequal(repeated_date), geolist)]
                # if jdx in (1, 2) && idx == 28
                    # @show idx, jdx, repeated_date, sigma_repeat
                if (repeated_date == ig1[1] && repeated_date == ig2[1]) || (repeated_date == ig1[2] && repeated_date == ig2[2])
                    Q[idx, jdx] = sigma_repeat
                else
                    Q[idx, jdx] = -sigma_repeat
                end
            end
        end
    end
    return Q + diag_val * I
end

_intlsubset(geolist, intlist, N) = intlist[[(intl[1] in geolist[1:N] && intl[2] in geolist[1:N]) for intl in intlist]]
function _intlsubset(geolist, intlist, vallist::AbstractArray, N)
    idxs = [(intl[1] in geolist[1:N] && intl[2] in geolist[1:N]) for intl in intlist]
    intlist[idxs], vallist[idxs]
end

function Qsubset(var_by_day, geolist, intlist, N)
    g = geolist[1:N]
    v = var_by_day[1:N]
    igl = _intlsubset(geolist, intlist, N)
    return build_Q(v, g, igl)
end

function build_Q_diag(var_dates, geolist, intlist)
    N = length(intlist);
    Q1 = zeros((N,));
    for (idx, ig) in enumerate(intlist)
        v1 = var_dates[findfirst(isequal(ig[1]), geolist)]
        v2 = var_dates[findfirst(isequal(ig[2]), geolist)]
        Q1[idx] = v1 + v2
    end
    return diagm(Q1)
end

import Combinatorics
function demo_weighted(bad_var = 100)
    geolist1 = Date(2020, 1, 1) .+ Dates.Day(12) .* (0:3)
    # 4-element Array{Date,1}:
    #  2020-01-01
    #  2020-01-13
    #  2020-01-25
    #  2020-02-06
    intlist1 = [tuple(p...) for p in Combinatorics.combinations(geolist1, 2)]
    # 6-element Array{Tuple{Date,Date},1}:
    #  (2020-01-01, 2020-01-13)
    #  (2020-01-01, 2020-01-25)
    #  (2020-01-01, 2020-02-06)
    #  (2020-01-13, 2020-01-25)
    #  (2020-01-13, 2020-02-06)
    #  (2020-01-25, 2020-02-06)

    # Real phase is 0, 1, 2, 3
    phi_true = [0, 1, 2, 3]
    A = InsarTimeseries.build_A_matrix(geolist1, intlist1)
    @show meas_nonoise = A * phi_true[2:end]
    @show A \ meas_nonoise

    # But with noisy second day, looks like it uplifted 10
    phi_noise = [0, 10, 2, 3]
    @show meas_noise = A * phi_noise[2:end]
    println("No weighting:")
    @show A \ meas_noise

    variances = [1, bad_var, 1, 1]
    a, b, c, d = variances
    # Q = [a + b a -b; a a + c c; -b c b + c] + 0.1 * I
    Q = build_Q(variances, geolist1, intlist1, .1)
    print("Weight LS:")
    @show inv(A' * inv(Q) * A) * A' * inv(Q) * meas_noise

    println("or just diagonal?")
    Q2 = build_Q_diag(variances, geolist1, intlist1)
    println(Q2)
    print("Weight LS:")
    @show inv(A' * inv(Q2) * A) * A' * inv(Q2) * meas_noise

end

function plot_l1_vs_stack(;offset=true, alpha=300, h=3, w=4.5, yy=2018,
                          unwfile="unw_stack_shiftedonly.h5", station="TXOE", s=station,
                          days_smooth=1)
    rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
    # https://matplotlib.org/tutorials/introductory/customizing.html#a-sample-matplotlibrc-file
    # so obscure
    rcParams["pdf.fonttype"] = 42
    rcParams["ps.fonttype"] = 42
    rcParams["font.family"] = "Helvetica"
    rcParams["font.size"] = 16
    rcParams["font.weight"] = "bold"
    years = mdates.YearLocator()   # every year
    months = mdates.MonthLocator()  # every month
    years_fmt = mdates.DateFormatter("%Y")
    # yrs = (Date(2015), Date(2016), Date(2017), Date(2018))

    demrsc = Sario.load("dem.rsc")
    input_dset = "stack_flat_shifted"
    geolist, intlist, igram_idxs = load_geolist_intlist(unwfile, "geolist_ignore.txt", 800, max_date = Date(yy,1,1))
    Blin = sum(InsarTimeseries.prepB(geolist, intlist), dims = 2)
    ms = 4
    timediffs = InsarTimeseries.day_diffs(geolist)
    println("s = $(s)")
    unw_vals = get_stack_vals(unwfile, gps.station_rowcol(s, Dict(demrsc))..., 1, input_dset, igram_idxs)

    # Timeseries: unregularized, with all outliers (noisiest)
    B = InsarTimeseries.prepB(geolist, intlist, false, 0)
    println("Loaded $(length(unw_vals)) unw_vals, B: $(size(B))")
    unregged = InsarTimeseries.PHASE_TO_CM .* InsarTimeseries.integrate_velocities(B \ unw_vals, timediffs)

    # Timeseries: regularized, but with all outliers
    Ba = InsarTimeseries.prepB(geolist, intlist, false, alpha)
    unw_vals_a = InsarTimeseries.augment_zeros(Ba, unw_vals)
    regged = InsarTimeseries.PHASE_TO_CM .* InsarTimeseries.integrate_velocities(Ba \ unw_vals_a, timediffs)

    #
    # Plot timeseries with-outlier cases:
    fig, ax = plot_gps_los(s; end_date = Date(yy, 1, 1), offset=offset,
                           days_smooth = days_smooth,
                           gps_color=matlab_colors()[4], ms=ms)

    ax.plot(geolist, unregged, "-x", lw=3, label="Regularized")
    ax.plot(geolist, regged, "-x", lw=3, label="Unregularized")
    ax.format_xdata = years_fmt
    ax.xaxis.set_major_locator(years)
    ax.xaxis.set_major_formatter(years_fmt)
    # ax.tick_params(axis="x", labelrotation=45)
    y0 = ceil(maximum(abs.(unregged)))
    ax.set_ylim((-y0, y0))
    ax.set_yticks(-y0:2:y0)
    _set_figsize(fig, h, w)
    fig.legend()
    fig.savefig("compare_$(yy)_timeseries.pdf", bbox_inches = "tight", transparent = true, dpi = 100)

    # No outlier removal linear cases
    stack, l2, l1 = prunesolve(geolist, intlist, unw_vals, Blin, 1000, shrink = false)
    println("No outlier, linear: $(stack), $(l1)")
    println("Difference: $(abs(stack - l1))")

    # Plot linear with-outlier cases
    fig, ax = plot_gps_los(s; end_date = Date(yy, 1, 1),
                             insar_mm_list=[l1, l2], offset=offset,
                             labels=["L1 linear", "L2 linear"],
                             gps_color=matlab_colors()[3],
                             insar_colors=matlab_colors()[1:2],
                             days_smooth = days_smooth,
                             ms=ms)

    # ax.plot(geolist, regged, "-x", lw=3, label="reg")
    ax.plot(geolist, unregged, "-x", lw=3, c=matlab_colors()[4], label="Unregularized")
    ax.plot(geolist, regged, "-x", lw=3, c=matlab_colors()[5], label="Regularized")
    ax.format_xdata = years_fmt
    ax.xaxis.set_major_locator(years)
    ax.xaxis.set_major_formatter(years_fmt)
    # ax.tick_params(axis="x", labelrotation=45)
    # y0 = ceil(maximum(abs.(regged)))
    y0 = 6
    ax.set_ylim((-y0, y0))
    ax.set_yticks(-y0:2:y0)
    _set_figsize(fig, h, w)
    fig.legend()
    fig.savefig("compare_$(yy)_linear.pdf", bbox_inches = "tight", transparent = true, dpi = 100)

    ###### Outlier remove cases ################
    # timeseries: regularized with outliers removed
    geo_clean, intlist_clean, unw_clean = remove_outliers(geolist, intlist, unw_vals, mean_sigma_cutoff = 4)
    B2 = InsarTimeseries.prepB(geo_clean, intlist_clean, false, 0)
    td_clean = InsarTimeseries.day_diffs(geo_clean)
    unregged2 = InsarTimeseries.PHASE_TO_CM .* InsarTimeseries.integrate_velocities(B2 \ unw_clean, td_clean)
    Ba2 = InsarTimeseries.prepB(geo_clean, intlist_clean, false, alpha)
    unw_vals_a2 = InsarTimeseries.augment_zeros(Ba2, unw_clean)
    regged2 = InsarTimeseries.PHASE_TO_CM .* InsarTimeseries.integrate_velocities(Ba2 \ unw_vals_a2, td_clean)

    # linear solves with outlier removal
    stack2, l22, l12 = prunesolve(geolist, intlist, unw_vals, Blin, 4, shrink = false)

    # PLOT:
    fig, ax = plot_gps_los(s; end_date = Date(yy, 1, 1),
                             insar_mm_list=[l12, l22], offset=offset,
                             labels=["L1 linear", "L2 linear"],
                             gps_color=matlab_colors()[3],
                             insar_colors=matlab_colors()[1:2],
                             days_smooth = days_smooth,
                             ms=ms)
    # ax.plot(geo_clean, regged2, "-x", lw=3, label="reg")
    ax.plot(geo_clean, unregged2, "-x", c=matlab_colors()[4], lw=3, label="Unregularized")
    ax.plot(geo_clean, regged2, "-x", lw=3, c=matlab_colors()[5], label="Regularized")
    # y0 = ceil(maximum(abs.(regged2)))
    # y0 = 4
    ax.set_ylim((-y0, y0))
    ax.set_yticks(-y0:2:y0)
    ax.format_xdata = years_fmt
    ax.xaxis.set_major_locator(years)
    ax.xaxis.set_major_formatter(years_fmt)
    _set_figsize(fig, h, w)
    fig.legend()
    fig.savefig("compare_$(yy)_removed.pdf", bbox_inches = "tight", transparent = true, dpi = 100)

end


export load_obj

# ------------------------- #
# - Parse OBJ & MTL Files - #
# ------------------------- #
"""
    triangulate_faces(vertices::Vector, texture_coordinates::Vector,
                      faces::Vector, material_map::Dict)

Triangulates the `faces` and converts it into valid [`Triangle`](@ref) objects.

!!! warning
    Avoid using this function in itself. It is designed to be used internally
    by the [`load_obj`](@ref) function.
"""
function triangulate_faces(vertices::Vector, texture_coordinates::Vector,
                           faces::Vector, material_map::Dict)
    scene = Vector{Triangle}()
    for face in faces
        for i in 2:(length(face[1]) - 1)
            if isnothing(face[2])
                uv_coordinates = nothing
            else
                uv_coordinates = [[texture_coordinates[face[2][1]]...],
                                  [texture_coordinates[face[2][i]]...],
                                  [texture_coordinates[face[2][i + 1]]...]]
            end
            mat = Material(;uv_coordinates = uv_coordinates,
                           material_map[face[3]]...)
            push!(scene, Triangle(deepcopy(vertices[face[1][1]]),
                                  deepcopy(vertices[face[1][i]]),
                                  deepcopy(vertices[face[1][i + 1]]),
                                  mat))
        end
    end
    try
        scene_stable = Vector{typeof(scene[1])}()
        for s in scene
            push!(scene_stable, s)
        end
        return scene_stable
    catch e
        # If all the triangles are not of the same type return the non-infered
        # version of the scene. In this case type inference for `raytrace` will
        # also fail
        @warn "Could not convert the triangles to the same type. Type inference for raytrace will fail"
        return scene
    end
end

"""
    parse_mtllib!(file, material_map, outtype)

Parses `.mtl` file and creates a map from the material name to
[`Material`](@ref) objects.

Not all fields of the `mtl` file are parsed. We only parse :
`newmtl`, `Ka`, `Kd`, `Ks`, `Ns`, `d`, `Tr`, `map_Kd`, `map_Ks`,
and `map_Ka`.
"""
function parse_mtllib!(file, material_map, outtype)
    color_diffuse = Vec3(outtype(1.0f0))
    color_ambient = Vec3(outtype(1.0f0))
    color_specular = Vec3(outtype(1.0f0))
    specular_exponent = outtype(50.0f0)
    reflection = outtype(0.5f0)
    texture_ambient = nothing
    texture_diffuse = nothing
    texture_specular = nothing
    last_mat = "RayTracer Default"
    if isnothing(file)
        material_map[last_mat] = (color_diffuse = color_diffuse,
                                  color_ambient = color_ambient,
                                  color_specular = color_specular,
                                  specular_exponent = specular_exponent,
                                  reflection = reflection,
                                  texture_ambient = texture_ambient,
                                  texture_diffuse = texture_diffuse,
                                  texture_specular = texture_specular)
        return nothing
    end
    for line in eachline(file)
        wrds = split(line)
        isempty(wrds) && continue
        if wrds[1] == "newmtl"
            material_map[last_mat] = (color_diffuse = color_diffuse,
                                      color_ambient = color_ambient,
                                      color_specular = color_specular,
                                      specular_exponent = specular_exponent,
                                      reflection = reflection,
                                      texture_ambient = texture_ambient,
                                      texture_diffuse = texture_diffuse,
                                      texture_specular = texture_specular)
            last_mat = wrds[2]
            # In case any of these values are not defined for the material
            # we shall use the default values
            color_diffuse = Vec3(outtype(1.0f0))
            color_ambient = Vec3(outtype(1.0f0))
            color_specular = Vec3(outtype(1.0f0))
            specular_exponent = outtype(50.0f0)
            reflection = outtype(0.5f0)
            texture_ambient = nothing
            texture_diffuse = nothing
            texture_specular = nothing
        elseif wrds[1] == "Ka"
            color_ambient = Vec3(parse.(outtype, wrds[2:4])...)
        elseif wrds[1] == "Kd"
            color_diffuse = Vec3(parse.(outtype, wrds[2:4])...)
        elseif wrds[1] == "Ks"
            color_specular = Vec3(parse.(outtype, wrds[2:4])...)
        elseif wrds[1] == "Ns"
            specular_exponent = parse(outtype, wrds[2])
        elseif wrds[1] == "d"
            reflection = parse(outtype, wrds[2])
        elseif wrds[1] == "Tr"
            reflection = 1 - parse(outtype, wrds[2])
        elseif wrds[1] == "map_Ka"
            texture_file = "$(rsplit(file, '/', limit = 2)[1])/$(wrds[2])"
            texture_ambient = Vec3([outtype.(permutedims(channelview(load(texture_file)),
                                                         (3, 2, 1)))[:,end:-1:1,i]
                                    for i in 1:3]...)
        elseif wrds[1] == "map_Kd"
            texture_file = "$(rsplit(file, '/', limit = 2)[1])/$(wrds[2])"
            texture_diffuse = Vec3([outtype.(permutedims(channelview(load(texture_file)),
                                                         (3, 2, 1)))[:,end:-1:1,i]
                                    for i in 1:3]...)
        elseif wrds[1] == "map_Ks"
            texture_file = "$(rsplit(file, '/', limit = 2)[1])/$(wrds[2])"
            texture_specular = Vec3([outtype.(permutedims(channelview(load(texture_file)),
                                                          (3, 2, 1)))[:,end:-1:1,i]
                                     for i in 1:3]...)
        end
    end
    material_map[last_mat] = (color_diffuse = color_diffuse,
                              color_ambient = color_ambient,
                              color_specular = color_specular,
                              specular_exponent = specular_exponent,
                              texture_ambient = texture_ambient,
                              texture_diffuse = texture_diffuse,
                              texture_specular = texture_specular)
    return nothing
end

"""
    load_obj(file; outtype = Float32)

Parser for `obj` file. It returns a set of Triangles that make up the scene.

Only the following things are parsed as of now: `v`, `vt`, `vn`, `f`, `usemtl`
and `mtllib`.
"""
function load_obj(file; outtype = Float32)
    vertices = Vector{Vec3{Vector{outtype}}}()
    texture_coordinates = Vector{Tuple}()
    normals = Vector{Vec3{Vector{outtype}}}()
    faces = Vector{Tuple{Vector{Int}, Union{Vector{Int}, Nothing}, String}}()
    material_map = Dict{String, Union{NamedTuple, Nothing}}()
    last_mat = "RayTracer Default"
    mtllib = nothing
    for line in eachline(file)
        wrds = split(line)
        isempty(wrds) && continue
        if wrds[1] == "v" # Vertices
            push!(vertices, Vec3(parse.(outtype, wrds[2:4])...))
        elseif wrds[1] == "vt" # Texture Coordinates
            push!(texture_coordinates, tuple(parse.(outtype, wrds[2:3])...))
        elseif wrds[1] == "vn" # Normal
            push!(normals, Vec3(parse.(outtype, wrds[2:4])...))
        elseif wrds[1] == "f" # Faces
            # Currently we shall only be concerned with the vertices of the face
            # and safely throw away vertex normal information
            fstwrd = split(wrds[2], '/')
            texture_c = isempty(fstwrd[2]) ? nothing : [parse(Int, fstwrd[2])]
            @assert length(fstwrd) == 3 "Incorrect number of /'s in the obj file"
            push!(faces, ([parse(Int, fstwrd[1])], texture_c, last_mat))
            for wrd in wrds[3:end]
                splitwrd = split(wrd, '/')
                @assert length(splitwrd) == 3 "Incorrect number of /'s in the obj file"
                push!(faces[end][1], parse(Int, splitwrd[1]))
                !isnothing(texture_c) && push!(faces[end][2], parse(Int, splitwrd[2]))
            end
        elseif wrds[1] == "usemtl" # Key for parsing mtllib file
            last_mat = wrds[2]
            material_map[last_mat] = nothing
        elseif wrds[1] == "mtllib" # Material file
            mtllib = "$(rsplit(file, '/', limit = 2)[1])/$(wrds[2])"
        end
    end

    if !isnothing(mtllib)
        parse_mtllib!(mtllib, material_map, outtype)
    else
        parse_mtllib!(nothing, material_map, outtype)
    end

    return triangulate_faces(vertices, texture_coordinates, faces, material_map)
end


using JLD2

@doc """
    runMC(param::Parameter)
    runMC(params::AbstractArray{Parameter}
          ;
          parallel::Bool=false,
          autoID::Bool=true)

Runs Monte Carlo simulation(s) and returns calculated observables.

If a checkpoint file named `"\$(param["Checkpoint Filename Prefix"])_\$(param["ID"]).jld2"` exists and
`param["Checkpoint Interval"] > 0.0`, `runMC` loads this file and restarts the pending simulation.

# Keyward aruguments
- `autoID`: If true, `"ID"`s will be set (overwritten) as `params[i]["ID"] = i`.
- `parallel`: If true, runs simulations in parallel (uses `pmap` instead of `map`).

# Required keys in `param`
- "Model"
- "Update Method"

# Optional keys in `param`
- "MCS": The number of Monte Carlo steps after thermalization
    - Default: `8192`
- "Thermalization": The number of Monte Carlo steps for thermalization
    - Default: `MCS>>3`
- "Seed": The initial seed of the random number generator, `MersenneTwister`
    - Default: determined randomly (see `Random.srand`)
- "Checkpoint Filename Prefix": See above document.
    - Default: `"cp"`
- "ID": See above document.
    - Default: `0`
- "Checkpoint Interval": Time interval between writing checkpoint file in seconds.
    - Default: `0.0`, this means that NO checkpoint file will be loaded and saved.
"""
function runMC(params::AbstractArray{T}; parallel::Bool=false, autoID::Bool=true) where T<:Dict
    map_fn = ifelse(parallel, pmap, map)
    return map_fn(enumerate(params)) do idp
        id,p = idp
        if autoID
            p["ID"] = id
        end
        return runMC(p)
    end
end

function runMC(param::Parameter)
    model = param["Model"](param)
    if "Seed" in keys(param)
        srand(model, param["Seed"])
    end
    ret = runMC(model, param)
    return ret
end

function runMC(model, param::Parameter)
    MODEL = typeof(model)
    verbose = get(param, "Verbose", false) :: Bool
    if verbose
        println("Start: ", param)
    end
    cp_filename = @sprintf("%s_%d.jld2",
                           get(param, "Checkpoint Filename Prefix", "cp")::String,
                           get(param, "ID", 0)::Int)
    cp_interval = get(param, "Checkpoint Interval", 0.0) :: Float64
    if VERSION > v"0.6.4" && cp_interval != 0.0
        Compat.@info """"Checkpoint Interval" is set to 0.0 automatically since current JLD2.jl cannot save/load Random.MersenneTwister object properly in julia-0.7"""
        cp_interval = 0.0
    end
    tm = time()

    MCS = get(param, "MCS", 8192) :: Int
    Therm = get(param, "Thermalization", MCS>>3) :: Int

    mcs = 0
    MCS += Therm
    obs = BinningObservableSet()
    makeMCObservable!(obs, "Time per MCS")
    makeMCObservable!(obs, "MCS per Second")

    if cp_interval > 0.0 && ispath(cp_filename)
        jld = jldopen(cp_filename)
        model = jld["model"] :: MODEL
        obs = jld["obs"] :: BinningObservableSet
        mcs = convert(Int, jld["mcs"]) :: Int
        close(jld)
    end

    if "UpdateMethod" in keys(param)
        warn("\"UpdateMethod\" is deprecated. Use instead \"Update Method\".")
        param["Update Method"] = param["UpdateMethod"] :: Function
    end
    update! = param["Update Method"] :: Function
    estimator = get(param, "Estimator", default_estimator(model, update!)) :: Function
    p = convert_parameter(model, param)

    while mcs < MCS
        if mcs < Therm
            update!(model, p...)
        else
            t = @elapsed begin
                st = update!(model,p...)
                localobs = estimator(model, p..., st)
            end
            obs["Time per MCS"] << t
            obs["MCS per Second"] << 1.0/t
            accumulateObservables!(model, obs, localobs)
        end
        mcs += 1
        if cp_interval > 0.0 && time() - tm > cp_interval
            jld = jldopen(cp_filename, "w")
            jld["model"] = model
            jld["obs"] = obs
            jld["mcs"] = mcs
            close(jld)
            tm += cp_interval
         end
    end

    if cp_interval > 0.0
        jld = jldopen(cp_filename, "w")
        jld["model"] = model
        jld["obs"] = obs
        jld["mcs"] = mcs
        close(jld)
    end

    jk = postproc(model, param, obs)

    if verbose
        println("Finish: ", param)
    end
    return jk
end

@doc """
    accumulateObservables!(model, obs::MCObservableSet, localobs::Dict)

Accumulates `localobs` into `obs`. For example, `obs["Energy"] << localobs["Energy"]`.
"""
function accumulateObservables!(::Model, obs::MCObservableSet, localobs::Measurement)
    if length(obs) < 3
        @inbounds for key in keys(localobs)
            makeMCObservable!(obs, key)
            obs[key] << localobs[key]
        end
    else
        @inbounds for key in keys(localobs)
            obs[key] << localobs[key]
        end
    end
    return obs
end

@doc """
    postproc(model::Model, param::Parameter, obs::MCObservableSet)

Post process of observables. For example, Specific heat will be calculated from energy, energy^2, and temperature.
"""
function postproc end

@doc doc"""
    postproc(model::Union{Ising, Potts}, param::Parameter, obs::MCObservableSet)

# Observables to be calculated
In the following, $m$ is total magnetization per site and $\epsilon$ is total energy per site.

- `"Binder Ratio"`
    - $R := \frac{\left \langle m^4 \right \rangle}{\left \langle m^2 \right\rangle^2}$
- `"Susceptibility"`
    - $\chi := \frac{N}{T}\left(\left\langle m^2\right\rangle\right)$
- `"Connected Susceptibility"`
    - $\frac{N}{T}\left(\left\langle m^2\right\rangle - \left\langle |m| \right\rangle^2\right)$
- `"Specific Heat"`
    - $\frac{N}{T^2}\left(\left\langle \epsilon^2\right\rangle - \left\langle \epsilon \right\rangle^2\right)$
"""
function postproc(model::Union{Ising, Potts}, param::Parameter, obs::MCObservableSet)
    nsites = numsites(model)
    T = param["T"] :: Float64
    beta = 1.0/T

    jk = jackknife(obs)
    jk["Binder Ratio"] = jk["Magnetization^4"] / (jk["Magnetization^2"]^2)
    jk["Susceptibility"] = (nsites*beta)*jk["Magnetization^2"]
    jk["Connected Susceptibility"] = (nsites*beta)*(jk["Magnetization^2"] - jk["|Magnetization|"]^2)
    jk["Specific Heat"] = (nsites*beta*beta)*(jk["Energy^2"] - jk["Energy"]^2)
    return jk
end

@doc doc"""
    postproc(model::Union{Clock, XY}, param::Parameter, obs::MCObservableSet)

# Observables to be calculated
In the following, $m$ is total magnetization per site and $\epsilon$ is total energy per site.

- `"Binder Ratio x"`
    - $\frac{\left \langle m_x^4 \right \rangle}{\left \langle m_x^2 \right\rangle^2}$
- `"Binder Ratio y"`
    - $\frac{\left \langle m_y^4 \right \rangle}{\left \langle m_y^2 \right\rangle^2}$
- `"Binder Ratio"`
    - $\frac{\left \langle |m|^4 \right \rangle}{\left \langle |m|^2 \right\rangle^2}$
- `"Susceptibility x"`
    - $\frac{N}{T}\left(\left\langle m_x^2\right\rangle\right)$
- `"Susceptibility y"`
    - $\frac{N}{T}\left(\left\langle m_y^2\right\rangle\right)$
- `"Susceptibility y"`
    - $\frac{N}{T}\left(\left\langle |m|^2\right\rangle\right)$
- `"Connected Susceptibility x"`
    - $\frac{N}{T}\left(\left\langle m_x^2\right\rangle - \left\langle |m_x| \right\rangle^2\right)$
- `"Connected Susceptibility y"`
    - $\frac{N}{T}\left(\left\langle m_y^2\right\rangle - \left\langle |m_y| \right\rangle^2\right)$
- `"Connected Susceptibility"`
    - $\frac{N}{T}\left(\left\langle |m|^2\right\rangle - \left\langle |m| \right\rangle^2\right)$
- `"Specific Heat"`
    - $\frac{N}{T^2}\left(\left\langle \epsilon^2\right\rangle - \left\langle \epsilon \right\rangle^2\right)$
"""
function postproc(model::Union{Clock, XY}, param::Parameter, obs::MCObservableSet)
    nsites = numsites(model)
    T = param["T"]
    beta = 1.0/T

    jk = jackknife(obs)
    jk["Binder Ratio x"] = jk["Magnetization x^4"] / (jk["Magnetization x^2"]^2)
    jk["Binder Ratio y"] = jk["Magnetization y^4"] / (jk["Magnetization y^2"]^2)
    jk["Binder Ratio"] = jk["|Magnetization|^4"] / (jk["|Magnetization|^2"]^2)
    jk["Susceptibility x"] = (nsites*beta)*jk["Magnetization x^2"]
    jk["Susceptibility y"] = (nsites*beta)*jk["Magnetization y^2"]
    jk["Susceptibility"] = (nsites*beta)*jk["|Magnetization|^2"]
    jk["Connected Susceptibility x"] = (nsites*beta)*(jk["Magnetization x^2"] - jk["|Magnetization x|"]^2)
    jk["Connected Susceptibility y"] = (nsites*beta)*(jk["Magnetization y^2"] - jk["|Magnetization y|"]^2)
    jk["Connected Susceptibility"] = (nsites*beta)*(jk["|Magnetization|^2"] - jk["|Magnetization|"]^2)
    jk["Specific Heat"] = (nsites*beta*beta)*(jk["Energy^2"] - jk["Energy"]^2)
    return jk
end

@doc doc"""
    postproc(model::QuantumXXZ, param::Parameter, obs::MCObservableSet)

# Observables to be calculated
In the following, $s$ is sign of weight, $m$ is total magnetization per site,
and $\epsilon$ is total energy per site.

- `"Magnetization"`
    - $\left\langle m s\right\rangle\Big/\left\langle s \right\rangle$
- `"|Magnetization|"`
    - $\left\langle |m| s\right\rangle\Big/\left\langle s \right\rangle$
- `"Magnetization^2"`
    - $\left\langle m^2 s\right\rangle\Big/\left\langle s \right\rangle$
- `"Magnetization^4"`
    - $\left\langle m^4 s\right\rangle\Big/\left\langle s \right\rangle$
- `"Energy"`
    - $\left\langle \epsilon s\right\rangle\Big/\left\langle s \right\rangle$
- `"Energy^2"`
    - $\left\langle \epsilon^2 s\right\rangle\Big/\left\langle s \right\rangle$
- `"Binder Ratio"`
    - $\frac{\left \langle m^4 \right \rangle}{\left \langle m^2 \right\rangle^2}$
- `"Susceptibility"`
    - $\frac{N}{T}\left(\left\langle m^2\right\rangle\right)$
- `"Connected Susceptibility"`
    - $\frac{N}{T}\left(\left\langle m^2\right\rangle - \left\langle |m| \right\rangle^2\right)$
- `"Specific Heat"`
    - $\frac{N}{T^2}\left(\left\langle \epsilon^2\right\rangle - \left\langle \epsilon \right\rangle^2\right)$
"""
function postproc(model::QuantumXXZ, param::Parameter, obs::MCObservableSet)
    nsites = numsites(model)
    T = param["T"]
    beta = 1.0/T

    jk = jackknife(obs)

    for oname in ("Magnetization", "|Magnetization|",
                  "Magnetization^2", "Magnetization^4",
                  "Energy", "Energy^2",
                 )
        jk[oname] = jk["Sign * $oname"] / jk["Sign"]
    end

    jk["Binder Ratio"] = jk["Magnetization^4"] / (jk["Magnetization^2"]^2)
    jk["Susceptibility"] = (nsites*beta)*jk["Magnetization^2"]
    jk["Connected Susceptibility"] = (nsites*beta)*(jk["Magnetization^2"] - jk["|Magnetization|"]^2)
    jk["Specific Heat"] = (nsites*beta*beta)*(jk["Energy^2"] - jk["Energy"]^2)
    return jk
end


module CTImageData

export AbstractCTData, CTData

using ...Monads
using ...CTImages
import ...CTImages: ctimage, ctsinogram, cttomogram
import ...Utils: _atype
import Base: similar, eltype

abstract type AbstractCTData end

# Container type of CTData
_atype(x::AbstractCTData) = _atype(typeof(x))

eltype(::Type{T}) where {T<:AbstractCTData} = eltype(_atype(T))

ctimage(gsd::AbstractCTData) = gsd.image
ctsinogram(gsd::AbstractCTData) = gsd.sinog
cttomogram(gsd::AbstractCTData) = gsd.tomog

"""
    mutable struct CTData{M}

Hold data for reconstruction.
`M` should be a subtype of `AbstractArray`, for example
`Matrix{Float32}` or `AFMatrix{Float64}`.
By default, the library will use CPU objects rather than GPU ones.
"""
mutable struct CTData{T<:Number,M<:AbstractMatrix{T}} <: AbstractCTData
    image::Maybe{CTImage{T,M}}
    sinog::Maybe{CTSinogram{T,M}}
    tomog::Maybe{CTTomogram{T,M}}
end

_atype(::Type{CTData{T,M}}) where {T,M} = M

function CTData{T,M}(gsd::AbstractCTData) where {T,M<:AbstractArray{T}}
    CTData{T,M}(
        mmap(CTImage{T,M}, gsd.image),
        mmap(CTSinogram{T,M}, gsd.sinog),
        mmap(CTTomogram{T,M}, gsd.tomog),
    )
end

CTData{T,M}() where {T,M<:AbstractArray{T}} =
    CTData{T,M}(nothing, nothing, nothing)

"""
    CTData(::Type{M}) where {M<:AbstractArray{<:Number}}

Construct an empty `CTData` object.
"""
CTData(::Type{M}) where {M<:AbstractArray{<:Number}} = CTData{eltype(M),M}()

"""
    CTData([::Type{T} = Float32]) where {T<:Number}

Construct an empty `CTData` object with eltype `T`.
"""
CTData(::Type{T} = Float32) where {T<:Number} = CTData(T,Matrix{T})

"""
    CTData(::AbstractCTImage)

Construct a `CTData` object from its constituents.
"""
CTData(img::CTImage) =
    CTData{eltype(img),_atype(img)}(Some(img), nothing, nothing)
CTData(img::CTSinogram) =
    CTData{eltype(img),_atype(img)}(nothing, Some(img), nothing)
CTData(img::CTTomogram) =
    CTData{eltype(img),_atype(img)}(nothing, nothing, Some(img))

function CTData{T,M}(gsd::AbstractCTData, img::CTImage) where {T,M}
    CTData{T,M}(
        Some(img),
        gsd.sinog,
        gsd.tomog,
    )
end

CTData{T,M}(gsd::AbstractCTData, img::CTSinogram) where {T,M} =
    CTData{T,M}(
        gsd.image,
        Some(img),
        gsd.tomog,
    )

CTData{T,M}(gsd::AbstractCTData, img::CTTomogram) where {T,M} =
    CTData{T,M}(
        gsd.image,
        gsd.sinog,
        Some(img),
    )

function CTData{T,M}(
    gsd::AbstractCTData,
    img::AbstractCTImage,
    rest::AbstractCTImage...
) where {T,M}
    CTData{T,M}(CTData{T,M}(gsd, img), rest...)
end

CTData(gsd::AbstractCTData, imgs::AbstractCTImage...) =
    CTData{eltype(gsd),_atype(gsd)}(gsd, imgs...)

function CTData(img::AbstractCTImage, imgs::AbstractCTImage...)
    CTData{eltype(img),_atype(img)}(CTData(img), imgs...)
end

similar(gsd::AbstractCTData, img::AbstractCTImage, rest::AbstractCTImage...) =
    typeof(gsd)(gsd, img, rest...)
similar(gsd::AbstractCTData) = typeof(gsd)()

end # module