text:
using JSON
using Dates
using PrettyPrinting
using FunSQL:
    FunSQL, Append, Define, From, FUN, OP, Fun, FunctionNode, Get, Join,
    LeftJoin, Select, Where, Partition, Agg, Group, As, Var, Bind, SQLNode, KW,
    Lit, render, SQLTable

function FunSQL.translate(::Val{:extract_year}, n::FunctionNode, ctx)
    args = FunSQL.translate(n.args, ctx)
    if length(args) == 1
        if ctx.dialect.name === :sqlserver
            return FUN(:YEAR, args[1])
        else
            return FUN(:EXTRACT, OP(:year), args[1] |> KW(:FROM))
        end
    end
    FunSQL.translate_default(n, ctx)
end

function FunSQL.translate(::Val{:dateadd_day}, n::FunctionNode, ctx)
    args = FunSQL.translate(n.args, ctx)
    if length(args) == 2
        if ctx.dialect.name === :sqlserver
            return FUN(:DATEADD, args = [OP(:day), args[2], args[1]])
        else
            return OP(:+, args = [args[1], args[2]])
        end
    end
    FunSQL.translate_default(n, ctx)
end

function FunSQL.translate(::Val{:datediff_day}, n::FunctionNode, ctx)
    args = FunSQL.translate(n.args, ctx)
    if length(args) == 2
        if ctx.dialect.name === :sqlserver
            return FUN(:DATEDIFF, args = [OP(:day), args[2], args[1]])
        else
            return OP(:-, args = [args[1], args[2]])
        end
    end
    FunSQL.translate_default(n, ctx)
end

translate(cohort; dialect, model = Model(), cohort_definition_id) =
    translate(cohort, dialect, model, cohort_definition_id)

translate(cohort::String, dialect, model, cohort_definition_id) =
    translate(JSON.parse(cohort), dialect, model, cohort_definition_id)

function translate(cohort::Dict, dialect, model, cohort_definition_id)
    cohort = deepcopy(cohort)
    res = translate(unpack!(cohort), dialect, model, cohort_definition_id)
    @assert isempty(cohort)
    return res
end

struct TranslateContext
    dialect::Symbol
    model::Model
    cohort::CohortExpression
    cohort_definition_id::Int
    group_idx::Ref{Int}
    statements::Vector{String}
    temp_idx::Ref{Int}
    temp_map::Dict{SQLNode, SQLNode}
end

function next_group_alias(ctx::TranslateContext)
    alias = Symbol(:group_, ctx.group_idx[])
    ctx.group_idx[] += 1
    alias
end

function as_temp_table(ctx::TranslateContext, q)
    if q in keys(ctx.temp_map)
        return ctx.temp_map[q]
    end
    temp_table_name = Symbol(:temp_, ctx.temp_idx[])
    if ctx.dialect === :sqlserver
        temp_table_name = Symbol("#", temp_table_name)
    end
    ctx.temp_idx[] += 1
    key = q
    q = q |>
        Select(:person_id => Get.person_id,
               :event_id => Get.event_id,
               :start_date => Get.start_date,
               :end_date => Get.end_date,
               :sort_date => Get.sort_date,
               :op_start_date => Get.op_start_date,
               :op_end_date => Get.op_end_date,
               :visit_occurrence_id => Get.visit_occurrence_id)
    sql = render(q, dialect = ctx.dialect)
    sql =
        if ctx.dialect === :sqlserver
            "SELECT person_id, event_id, start_date, end_date, sort_date, op_start_date, op_end_date, visit_occurrence_id \nINTO $temp_table_name\nFROM (\n$sql\n) AS t;\n"
        else
            "CREATE TEMP TABLE $temp_table_name AS\n$sql;\n"
        end
    push!(ctx.statements, sql)
    val = From(SQLTable(temp_table_name, columns = [:person_id, :event_id, :start_date, :end_date, :sort_date, :op_start_date, :op_end_date, :visit_occurrence_id]))
    ctx.temp_map[key] = val
    val
end

function translate(cohort::CohortExpression, dialect, model, cohort_definition_id)
    ctx = TranslateContext(Symbol(dialect), model, cohort, cohort_definition_id, Ref(1), String[], Ref(1), Dict{SQLNode, SQLNode}())
    translate(cohort, ctx)
    join(ctx.statements)
end

function escape_name(dialect, name)
    if dialect === :sqlserver
        return "[$(replace(string(name), ']' => "]]"))]"
    else
        return "\"$(replace(string(name), '"' => "\"\""))\""
    end
end

function cohort_table(ctx::TranslateContext)
    cohort = ctx.model.cohort
    d = ctx.dialect
    if isnothing(cohort.schema) || "" == string(cohort.schema)
        return escape_name(d, cohort.name)
    end
    escape_name(d, cohort.schema) * "." * escape_name(d, cohort.name)
end

function translate(c::CohortExpression, ctx::TranslateContext)
    @assert c.censor_window.start_date === c.censor_window.end_date === nothing
    @assert isempty(c.censoring_criteria)
    @assert c.end_strategy === nothing || c.end_strategy isa DateOffsetStrategy
    cname = cohort_table(ctx)
    push!(ctx.statements, "DELETE FROM $(cname)\nWHERE cohort_definition_id = $(ctx.cohort_definition_id);\n")
    q = translate(c.primary_criteria, ctx)
    if c.additional_criteria !== nothing || !isempty(c.inclusion_rules)
        q = base = as_temp_table(ctx, q)
    end
    if c.additional_criteria !== nothing
        q = q |>
            translate(c.additional_criteria, ctx, base = base)
        q = q |>
            translate(c.qualified_limit, ctx)
    end
    for r in c.inclusion_rules
        q = q |>
            translate(r.expression, ctx, base = base)
    end
    q = q |>
        translate(c.expression_limit, ctx)
    if c.end_strategy !== nothing
        q = q |>
            translate(c.end_strategy, ctx)
    else
        q = q |>
            Define(:end_date => Get.op_end_date)
    end
    q = q |>
        translate(c.collapse_settings, ctx)
    q = q |>
        Select(:subject_id => Get.person_id,
               :cohort_start_date => Get.start_date,
               :cohort_end_date => Get.end_date)
    q = q |> Select(:cohort_definition_id => ctx.cohort_definition_id,
                    Get.subject_id,
                    Get.cohort_start_date,
                    Get.cohort_end_date)
    sql = render(q, dialect = ctx.dialect)
    cname = cohort_table(ctx)
    push!(ctx.statements, "INSERT INTO $cname\n$sql;\n")
end

translate(::Nothing, ctx::TranslateContext; base = nothing) =
    Define()

function translate(r::ResultLimit, ctx::TranslateContext; order_by = [Get.start_date])
    if r.type == ALL
        return Define()
    end
    if r.type == LAST
        order_by = [Fun.datediff_day(order_by[1], Date(2020, 1, 1)), order_by[2:end]...]
    end
    Partition(Get.person_id, order_by = order_by) |>
    Where(Agg.row_number() .== 1)
end

function translate(d::DateOffsetStrategy, ctx::TranslateContext)
    field =
        d.date_field == START_DATE ? Get.start_date :
        d.date_field == END_DATE ? Get.end_date :
        nothing
    Define(:end_date => dateadd_day(field, d.offset)) |>
    Define(:end_date => Fun.case(Get.end_date .<= Get.op_end_date,
                                 Get.end_date, Get.op_end_date))
end

function dateadd_day(n, delta::Integer)
    if iszero(delta)
        return n
    end
    Fun.dateadd_day(n, delta)
end

function translate(c::PrimaryCriteria, ctx::TranslateContext)
    @assert length(c.criteria_list) >= 1
    q = translate(c.criteria_list[1], ctx)
    if length(c.criteria_list) > 1
        list = [translate(l, ctx) for l in c.criteria_list[2:end]]
        q = q |>
            Append(list = list)
    end
    q = q |>
        Join(:op => ctx.model.observation_period |>
                    Define(:start_date => Get.observation_period_start_date,
                           :end_date => Get.observation_period_end_date),
             Get.person_id .== Get.op.person_id)
    q = q |>
        Define(:op_start_date => Get.op.start_date,
               :op_end_date => Get.op.end_date)
    l = dateadd_day(Get.op.start_date, c.observation_window.prior_days)
    r = dateadd_day(Get.op.end_date, - c.observation_window.post_days)
    q = q |>
        Where(Fun.and(l .<= Get.start_date, Get.start_date .<= r))
    q = q |>
        translate(c.primary_limit, ctx, order_by = [Get.sort_date])
    q
end

function translate(c::ConditionOccurrence, ctx::TranslateContext)
    @assert isempty(c.condition_status)
    @assert !c.condition_type_exclude
    @assert c.stop_reason === nothing
    q = From(ctx.model.condition_occurrence) |>
        Define(:concept_id => Get.condition_concept_id,
               :event_id => Get.condition_occurrence_id,
               :start_date => Get.condition_start_date,
               :end_date => Fun.coalesce(Get.condition_end_date,
                                         dateadd_day(Get.condition_start_date, 1)),
               :sort_date => Get.condition_start_date)
    if c.condition_source_concept !== nothing
        q = q |>
            Where(Fun.in(Get.condition_source_concept_id,
                         translate(find_concept_set(c.condition_source_concept, ctx), ctx)))
    end
    if !isempty(c.condition_type)
        args = [Get.condition_type_concept_id .== t.concept_id
                for t in c.condition_type]
        q = q |>
            Where(Fun.or(args = args))
    end
    translate(c.base, ctx, base = q)
end

function translate(d::Death, ctx::TranslateContext)
    @assert isempty(d.death_type)
    @assert !d.death_type_exclude
    q = From(ctx.model.death) |>
        Define(:concept_id => Get.cause_concept_id,
               :event_id => 0,
               :start_date => Get.death_date,
               :end_date => dateadd_day(Get.death_date, 1),
               :sort_date => Get.death_date,
               :visit_occurrence_id => 0)
    if d.death_source_concept !== nothing
        q = q |>
            Where(Fun.in(Get.cause_source_concept_id,
                         translate(find_concept_set(d.death_source_concept, ctx), ctx)))
    end
    translate(d.base, ctx, base = q)
end

function translate(d::DeviceExposure, ctx::TranslateContext)
    @assert isempty(d.device_type)
    @assert !d.device_type_exclude
    @assert d.quantity === nothing
    @assert d.unique_device_id === nothing
    q = From(ctx.model.device_exposure) |>
        Define(:concept_id => Get.device_concept_id,
               :event_id => Get.device_exposure_id,
               :start_date => Get.device_exposure_start_date,
               :end_date => Fun.coalesce(Get.device_exposure_end_date,
                                         dateadd_day(Get.device_exposure_start_date, 1)),
               :sort_date => Get.device_exposure_start_date)
    if d.device_source_concept !== nothing
        q = q |>
            Where(Fun.in(Get.device_source_concept_id,
                         translate(find_concept_set(d.device_source_concept, ctx), ctx)))
    end
    translate(d.base, ctx, base = q)
end

function translate(d::DrugEra, ctx::TranslateContext)
    @assert d.era_start_date === nothing
    @assert d.era_end_date === nothing
    @assert d.occurrence_count === nothing
    @assert d.age_at_start === nothing
    @assert d.age_at_end === nothing
    q = From(ctx.model.drug_era) |>
        Define(:concept_id => Get.drug_concept_id,
               :event_id => Get.drug_era_id,
               :start_date => Get.drug_era_start_date,
               :end_date => Get.drug_era_end_date,
               :sort_date => Get.drug_era_start_date,
               :visit_occurrence_id => 0)
    if d.era_length !== nothing
        field = Fun.datediff_day(Get.drug_era_end_date, Get.drug_era_start_date)
        q = q |>
            Where(predicate(d.era_length, ctx, field = field))
    end
    translate(d.base, ctx, base = q)
end

function translate(d::DrugExposure, ctx::TranslateContext)
    @assert isempty(d.drug_type)
    @assert !d.drug_type_exclude
    @assert d.refills === nothing
    @assert d.quantity === nothing
    @assert d.days_supply === nothing
    @assert isempty(d.route_concept)
    @assert d.effective_drug_dose === nothing
    @assert isempty(d.dose_unit)
    @assert d.lot_number === nothing
    @assert d.stop_reason === nothing
    q = From(ctx.model.drug_exposure) |>
        Define(:concept_id => Get.drug_concept_id,
               :event_id => Get.drug_exposure_id,
               :start_date => Get.drug_exposure_start_date,
               :end_date => Fun.coalesce(Get.drug_exposure_end_date,
                                         Fun.dateadd_day(Get.drug_exposure_start_date, Get.days_supply),
                                         dateadd_day(Get.drug_exposure_start_date, 1)),
               :sort_date => Get.drug_exposure_start_date)
    if d.drug_source_concept !== nothing
        q = q |>
            Where(Fun.in(Get.drug_source_concept_id,
                         translate(find_concept_set(d.drug_source_concept, ctx), ctx)))
    end
    translate(d.base, ctx, base = q)
end

function translate(m::Measurement, ctx::TranslateContext)
    @assert isempty(m.measurement_type)
    @assert !m.measurement_type_exclude
    @assert m.abnormal === nothing
    @assert m.range_low === nothing
    @assert m.range_high === nothing
    @assert m.range_low_ratio === nothing
    @assert m.range_high_ratio === nothing
    @assert isempty(m.operator)
    q = From(ctx.model.measurement) |>
        Define(:concept_id => Get.measurement_concept_id,
               :event_id => Get.measurement_id,
               :start_date => Get.measurement_date,
               :end_date => dateadd_day(Get.measurement_date, 1),
               :sort_date => Get.measurement_date)
    if m.measurement_source_concept !== nothing
        q = q |>
            Where(Fun.in(Get.measurement_source_concept_id,
                         translate(find_concept_set(m.measurement_source_concept, ctx), ctx)))
    end
    if !isempty(m.value_as_concept)
        args = [Get.value_as_concept_id .== v.concept_id
                for v in m.value_as_concept]
        q = q |>
            Where(Fun.or(args = args))
    end
    if m.value_as_number !== nothing
        q = q |>
            Where(predicate(m.value_as_number, ctx, field = Get.value_as_number))
    end
    if !isempty(m.unit)
        args = [Get.unit_concept_id .== u.concept_id
                for u in m.unit]
        q = q |>
            Where(Fun.or(args = args))
    end
    translate(m.base, ctx, base = q)
end

function translate(o::Observation, ctx::TranslateContext)
    @assert isempty(o.observation_type)
    @assert !o.observation_type_exclude
    @assert o.value_as_string === nothing
    @assert isempty(o.qualifier)
    q = From(ctx.model.observation) |>
        Define(:concept_id => Get.observation_concept_id,
               :event_id => Get.observation_id,
               :start_date => Get.observation_date,
               :end_date => dateadd_day(Get.observation_date, 1),
               :sort_date => Get.observation_date)
    if o.observation_source_concept !== nothing
        q = q |>
            Where(Fun.in(Get.observation_source_concept_id,
                         translate(find_concept_set(o.observation_source_concept, ctx), ctx)))
    end
    if !isempty(o.value_as_concept)
        args = [Get.value_as_concept_id .== v.concept_id
                for v in o.value_as_concept]
        q = q |>
            Where(Fun.or(args = args))
    end
    if o.value_as_number !== nothing
        q = q |>
            Where(predicate(o.value_as_number, ctx, field = Get.value_as_number))
    end
    if !isempty(o.unit)
        args = [Get.unit_concept_id .== u.concept_id
                for u in o.unit]
        q = q |>
            Where(Fun.or(args = args))
    end
    translate(o.base, ctx, base = q)
end

function translate(o::ObservationPeriod, ctx::TranslateContext)
    @assert isempty(o.period_type)
    @assert !o.period_type_exclude
    @assert o.period_start_date === nothing
    @assert o.period_end_date === nothing
    @assert o.age_at_start === nothing
    @assert o.age_at_end === nothing
    @assert o.user_defined_period === nothing
    q = From(ctx.model.observation_period) |>
        Define(:event_id => Get.observation_period_id,
               :start_date => Get.observation_period_start_date,
               :end_date => Get.observation_period_end_date,
               :sort_date => Get.observation_period_start_date,
               :visit_occurrence_id => 0)
    if o.period_length !== nothing
        field = Fun.datediff_day(Get.end_date, Get.start_date)
        q = q |>
            Where(predicate(o.period_length, ctx, field = field))
    end
    translate(o.base, ctx, base = q)
end

function translate(p::ProcedureOccurrence, ctx::TranslateContext)
    @assert isempty(p.procedure_type)
    @assert !p.procedure_type_exclude
    @assert isempty(p.modifier)
    @assert p.quantity === nothing
    q = From(ctx.model.procedure_occurrence) |>
        Define(:concept_id => Get.procedure_concept_id,
               :event_id => Get.procedure_occurrence_id,
               :start_date => Get.procedure_date,
               :end_date => dateadd_day(Get.procedure_date, 1),
               :sort_date => Get.procedure_date)
    if p.procedure_source_concept !== nothing
        q = q |>
            Where(Fun.in(Get.procedure_source_concept_id,
                         translate(find_concept_set(p.procedure_source_concept, ctx), ctx)))
    end
    translate(p.base, ctx, base = q)
end

function translate(v::VisitOccurrence, ctx::TranslateContext)
    @assert isempty(v.place_of_service)
    @assert v.place_of_service_location === nothing
    @assert v.visit_length === nothing
    @assert !v.visit_type_exclude
    q = From(ctx.model.visit_occurrence) |>
        Define(:concept_id => Get.visit_concept_id,
               :event_id => Get.visit_occurrence_id,
               :start_date => Get.visit_start_date,
               :end_date => Get.visit_end_date,
               :sort_date => Get.visit_start_date)
    if v.visit_source_concept !== nothing
        q = q |>
            Where(Fun.in(Get.visit_source_concept_id,
                         translate(find_concept_set(v.visit_source_concept, ctx), ctx)))
    end
    translate(v.base, ctx, base = q)
end

is_simple(c, ctx::TranslateContext) =
    is_simple(c.base, ctx)

is_simple(b::BaseCriteria, ctx::TranslateContext) =
    b.correlated_criteria === nothing

function translate(b::BaseCriteria, ctx::TranslateContext; base)
    @assert b.occurrence_end_date === nothing
    q = base
    if b.codeset_id !== nothing
        q = q |>
            Where(Fun.in(Get.concept_id,
                         translate(find_concept_set(b.codeset_id, ctx), ctx)))
    end
    #=
    q = Join(:concept => translate(find_concept_set(b.codeset_id, ctx), ctx),
             Get.concept_id .== Get.concept.concept_id)
    =#
    if b.first
        q = q |>
            Partition(Get.person_id, order_by = [Get.sort_date, Get.event_id]) |>
            Where(Agg.row_number() .== 1)
    end
    if b.occurrence_start_date !== nothing
        q = q |>
            Where(predicate(b.occurrence_start_date, ctx, field = Get.start_date))
    end
    if b.age !== nothing || !isempty(b.gender)
        q = q |>
            Join(:person => ctx.model.person,
                 Get.person_id .== Get.person.person_id)
    end
    if b.age !== nothing
        q = q |>
            Define(:age => Fun.extract_year(Get.start_date) .- Get.person.year_of_birth) |>
            Where(predicate(b.age, ctx, field = Get.age))
    end
    if !isempty(b.gender)
        args = [Get.person.gender_concept_id .== c.concept_id
                for c in b.gender]
        q = q |>
            Where(Fun.or(args = args))
    end
    if !isempty(b.provider_specialty)
        args = [Get.provider.specialty_concept_id .== c.concept_id
                for c in b.provider_specialty]
        q = q |>
            Join(:provider => ctx.model.provider,
                 Get.provider_id .== Get.provider.provider_id) |>
            Where(Fun.or(args = args))
    end
    if !isempty(b.visit_type)
        args = [Get.visit.visit_concept_id .== c.concept_id
                for c in b.visit_type]
        q = q |>
            Join(:visit => ctx.model.visit_occurrence,
                 Fun.and(Get.person_id .== Get.visit.person_id,
                         Get.visit_occurrence_id .== Get.visit.visit_occurrence_id)) |>
            Where(Fun.or(args = args))
    end
    if b.correlated_criteria !== nothing
        q = q |>
            Join(:op_ => ctx.model.observation_period,
                 Get.person_id .== Get.op_.person_id)
        q = q |>
            Define(:op_start_date => Get.op_.observation_period_start_date,
                   :op_end_date => Get.op_.observation_period_end_date)
        q = q |>
            Where(Fun.and(Get.op_start_date .<= Get.start_date, Get.start_date .<= Get.op_end_date))
        q = as_temp_table(ctx, q)
        q = q |>
            translate(b.correlated_criteria, ctx, base = q)
    end
    q
end

is_inner(c::CriteriaGroup) =
    c.type == ALL_CRITERIA && all(is_inner, c.correlated_criteria) && all(is_inner, c.groups)

join_subsumes_where(c::CriteriaGroup) =
    c.type == ALL_CRITERIA && all(join_subsumes_where, c.correlated_criteria) && all(join_subsumes_where, c.groups) && isempty(c.demographic_criteria)

is_simple(c::CriteriaGroup, ctx::TranslateContext) =
    isempty(c.groups) && all(cc -> is_simple(cc, ctx), c.correlated_criteria) && ctx.dialect !== :redshift

function translate(c::CriteriaGroup, ctx::TranslateContext; base::SQLNode, result_alias = nothing, inner = is_inner(c))
    is_all = c.type == ALL_CRITERIA
    is_any = c.type == ANY_CRITERIA || (c.type == AT_LEAST_CRITERIA && c.count == 1)
    is_none = c.type == AT_MOST_CRITERIA && c.count == 0
    @assert is_all || is_any || is_none
    if !isempty(c.demographic_criteria)
        q = Join(:person => ctx.model.person,
                 Get.person_id .== Get.person.person_id) |>
            Define(:age => Fun.extract_year(Get.start_date) .- Get.person.year_of_birth)
    else
        q = Define()
    end
    if is_simple(c, ctx)
        p = predicate(c, ctx)
        if result_alias !== nothing
            q = q |>
                Define(result_alias => Fun.case(p, 1, 0))
        else
            q = q |>
                  Where(p)
        end
        return q
    end
    idx = 1
    args = SQLNode[]
    for criteria in c.correlated_criteria
        nested_alias = next_group_alias(ctx)
        q = q |>
            translate(criteria, ctx, base = base, result_alias = nested_alias, inner = inner)
        if !(inner && join_subsumes_where(criteria))
            push!(args, Get(nested_alias) .== 1)
        end
    end
    for criteria in c.demographic_criteria
        nested_alias = next_group_alias(ctx)
        q = q |>
            translate(criteria, ctx, result_alias = nested_alias)
        push!(args, Get(nested_alias) .== 1)
    end
    for group in c.groups
        nested_alias = next_group_alias(ctx)
        q = q |>
            translate(group, ctx, base = base, result_alias = nested_alias, inner = inner)
        if !(inner && join_subsumes_where(group))
            push!(args, Get(nested_alias) .== 1)
        end
    end
    if is_all
        p = Fun.and(args = args)
    elseif is_any
        p = Fun.or(args = args)
    elseif is_none
        args = [Fun.not(arg) for arg in args]
        p = Fun.and(args = args)
    end
    if result_alias !== nothing
        q = q |>
            Define(result_alias => Fun.case(p, 1, 0))
    else
        q = q |>
            Where(p)
    end
    q
end

function predicate(c::CriteriaGroup, ctx::TranslateContext)
    is_all = c.type == ALL_CRITERIA
    is_any = c.type == ANY_CRITERIA || (c.type == AT_LEAST_CRITERIA && c.count == 1)
    is_none = c.type == AT_MOST_CRITERIA && c.count == 0
    @assert is_all || is_any || is_none
    args = SQLNode[]
    for criteria in c.correlated_criteria
        push!(args, predicate(criteria, ctx))
    end
    for criteria in c.demographic_criteria
        push!(args, predicate(criteria, ctx))
    end
    for group in c.groups
        push!(args, predicate(group, ctx))
    end
    if is_all
        Fun.and(args = args)
    elseif is_any
        Fun.or(args = args)
    elseif is_none
        args = [Fun.not(arg) for arg in args]
        Fun.and(args = args)
    end
end

function translate(d::DemographicCriteria, ctx::TranslateContext; result_alias)
    Define(result_alias => Fun.case(predicate(d, ctx), 1, 0))
end

function predicate(d::DemographicCriteria, ctx::TranslateContext)
    @assert isempty(d.ethnicity)
    @assert isempty(d.race)
    @assert isempty(d.gender)
    @assert d.occurrence_end_date === nothing
    args = SQLNode[]
    if d.age !== nothing
        push!(args, predicate(d.age, ctx, field = Get.age))
    end
    if d.occurrence_start_date !== nothing
        push!(args, predicate(d.occurrence_start_date, ctx, field = Get.start_date))
    end
    Fun.and(args = args)
end

is_inner(c::CorrelatedCriteria) =
    c.occurrence.type in (AT_LEAST, EXACTLY) && c.occurrence.count > 0

join_subsumes_where(c::CorrelatedCriteria) =
    c.occurrence.type == AT_LEAST && c.occurrence.count == 1

is_simple(c::CorrelatedCriteria, ctx::TranslateContext) =
    (c.occurrence.type == AT_LEAST && c.occurrence.count == 1) &&
    is_simple(c.criteria, ctx)

function translate(c::CorrelatedCriteria, ctx::TranslateContext; base, result_alias, inner)
    @assert c.occurrence !== nothing &&
            c.occurrence.count_column === nothing
    @assert c.occurrence.type in (AT_LEAST, AT_MOST, EXACTLY)
    @assert c.criteria !== nothing
    q = base |>
        Join(:correlated => translate(c.criteria, ctx),
             Get.correlated.person_id .== Get.person_id)
    if c.restrict_visit
        q = q |>
            Where(Get.correlated.visit_occurrence_id .== Get.visit_occurrence_id)
    end
    if !c.ignore_observation_period
        q = q |>
            Where(Fun.and(Get.op_start_date .<= Get.correlated.start_date,
                          Get.correlated.start_date .<= Get.op_end_date))
    end
    q = q |>
        translate(c.start_window, ctx, start = true, ignore_observation_period = c.ignore_observation_period)
    if c.end_window !== nothing
        q = q |>
            translate(c.end_window, ctx, start = false, ignore_observation_period = c.ignore_observation_period)
    end
    q = q |>
        Group(Get.person_id, Get.event_id)
    if c.occurrence.is_distinct
        q = q |>
            Define(:count => Agg.count(distinct = true, Get.correlated.concept_id))
    else
        q = q |>
            Define(:count => Agg.count())
    end
    correlated_alias = Symbol(:correlated_, result_alias)
    q = Join(correlated_alias => q,
             Fun.and(Get.person_id .== (Get(correlated_alias) |> Get.person_id),
                     Get.event_id .== (Get(correlated_alias) |> Get.event_id)),
             left = !inner)
    count = Fun.coalesce(Get(correlated_alias) |> Get.count, 0)
    if c.occurrence.type == AT_LEAST
        p = count .>= c.occurrence.count
    elseif c.occurrence.type == AT_MOST
        p = count .<= c.occurrence.count
    elseif c.occurrence.type == EXACTLY
        p = count .== c.occurrence.count
    end
    q = q |>
        Define(result_alias => Fun.case(p, 1, 0))
    q
end

function predicate(c::CorrelatedCriteria, ctx::TranslateContext)
    @assert c.occurrence !== nothing &&
            c.occurrence.count_column === nothing
    @assert c.occurrence.type in (AT_LEAST, AT_MOST, EXACTLY)
    @assert c.criteria !== nothing
    q = translate(c.criteria, ctx)
    q = q |>
        As(:correlated) |>
        Where(Get.correlated.person_id .== Var.person_id)
    if c.restrict_visit
        q = q |>
            Where(Get.correlated.visit_occurrence_id .== Var.visit_occurrence_id)
    end
    q = q |>
        Define(:start_date => Var.start_date,
               :end_date => Var.end_date,
               :op_start_date => Var.op_start_date,
               :op_end_date => Var.op_end_date)
    if !c.ignore_observation_period
        q = q |>
            Where(Fun.and(Get.op_start_date .<= Get.correlated.start_date,
                          Get.correlated.start_date .<= Get.op_end_date))
    end
    q = q |>
        translate(c.start_window, ctx, start = true, ignore_observation_period = c.ignore_observation_period)
    if c.end_window !== nothing
        q = q |>
            translate(c.end_window, ctx, start = false, ignore_observation_period = c.ignore_observation_period)
    end
    exists = c.occurrence.type == AT_LEAST && c.occurrence.count == 1
    not_exists = c.occurrence.type in (EXACTLY, AT_MOST) && c.occurrence.count == 0
    if !exists && !not_exists
        q = q |>
            Group()
        if c.occurrence.is_distinct
            q = q |>
                Select(Agg.count(distinct = true, Get.correlated.concept_id))
        else
            q = q |>
                Select(Agg.count())
        end
    end
    q = q |>
        Bind(:person_id => Get.person_id,
             :start_date => Get.start_date,
             :end_date => Get.end_date,
             :op_start_date => Get.op_start_date,
             :op_end_date => Get.op_end_date,
             :visit_occurrence_id => c.restrict_visit ? Get.visit_occurrence_id : 0)
    if exists
        q = Fun.exists(q)
    elseif not_exists
        q = Fun.not(Fun.exists(q))
    else
        if c.occurrence.type == AT_LEAST
            q = q .>= c.occurrence.count
        elseif c.occurrence.type == AT_MOST
            q = q .<= c.occurrence.count
        elseif c.occurrence.type == EXACTLY
            q = q .== c.occurrence.count
        end
    end
    q
end

function translate(w::Window, ctx::TranslateContext; start::Bool, ignore_observation_period::Bool)
    index_date_field =
        w.use_index_end == true ? Get.end_date : Get.start_date
    event_date_field =
        if start
            w.use_event_end == true ?
                Get.correlated.end_date : Get.correlated.start_date
        else
            w.use_event_end == false ?
                Get.correlated.start_date : Get.correlated.end_date
        end
    l = nothing
    r = nothing
    if w.start.days !== nothing
        l = dateadd_day(index_date_field, w.start.days * w.start.coeff)
    elseif !ignore_observation_period
        l = w.start.coeff == -1 ? Get.op_start_date : Get.op_end_date
    end
    if w.end_.days !== nothing
        r = dateadd_day(index_date_field, w.end_.days * w.end_.coeff)
    elseif !ignore_observation_period
        r = w.end_.coeff == -1 ? Get.op_start_date : Get.op_end_date
    end
    args = []
    if l !== nothing
        push!(args, l .<= event_date_field)
    end
    if r !== nothing
        push!(args, event_date_field .<= r)
    end
    Where(Fun.and(args = args))
end

function translate(c::CollapseSettings, ctx::TranslateContext)
    @assert c.collapse_type == ERA
    gap = c.era_pad
    Define(:end_date => dateadd_day(Get.end_date, gap)) |>
    Partition(Get.person_id, order_by = [Get.start_date], frame = (mode = :rows, start = -Inf, finish = -1)) |>
    Define(:boundary => Agg.max(Get.end_date)) |>
    Define(:bump => Fun.case(Get.start_date .<= Get.boundary, 0, 1)) |>
    Partition(Get.person_id, order_by = [Get.start_date, .- Get.bump], frame = :rows) |>
    Define(:group => Agg.sum(Get.bump)) |>
    Group(Get.person_id, Get.group) |>
    Define(:start_date => Agg.min(Get.start_date),
           :end_date => dateadd_day(Agg.max(Get.end_date), - gap))
end

function find_concept_set(codeset_id, ctx)
    for cs in ctx.cohort.concept_sets
        if cs.id == codeset_id
            return cs
        end
    end
end

function translate(c::ConceptSet, ctx::TranslateContext)
    include = ConceptSetItem[]
    exclude = ConceptSetItem[]
    for item in c.items
        @assert !item.include_mapped
        if !item.is_excluded
            push!(include, item)
        else
            push!(exclude, item)
        end
    end
    q = translate(include, ctx)
    if !isempty(exclude)
        q = q |>
            LeftJoin(:excluded => translate(exclude, ctx),
                     Get.concept_id .== Get.excluded.concept_id) |>
            Where(Fun."is null"(Get.excluded.concept_id))
    end
    q = q |>
        Select(Get.concept_id)
    q
end

function translate(items::Vector{ConceptSetItem}, ctx::TranslateContext)
    # TODO: Fun.in
    args = [Get.concept_id .== item.concept.concept_id for item in items]
    q = From(ctx.model.concept) |>
        Where(Fun.or(args = args))
    with_descendants = [item for item in items if item.include_descendants]
    if !isempty(with_descendants)
        args = [Get.ancestor_concept_id .== item.concept.concept_id for item in with_descendants]
        q = q |>
            Append(
                From(ctx.model.concept) |>
                Where(Fun."is null"(Get.invalid_reason)) |>
                Join(ctx.model.concept_ancestor,
                     Get.concept_id .== Get.descendant_concept_id) |>
                Where(Fun.or(args = args))) |>
            Group(Get.concept_id)
    end
    q
end

function predicate(r::Union{NumericRange, DateRange}, ctx::TranslateContext; field)
    if r.op == GT
        field .> r.value
    elseif r.op == GTE
        field .>= r.value
    elseif r.op == LT
        field .< r.value
    elseif r.op == LTE
        field .<= r.value
    elseif r.op == EQ
        field .== r.value
    elseif r.op == NEQ
        field .!= r.value
    elseif r.op == BT
        Fun.between(field, r.value, r.extent)
    elseif r.op == NBT
        Fun."not between"(field, r.value, r.extent)
    end
end


text:
export gmres

function gmres(A::Union(AbstractMatrix, Function), b::Vector; restart::Int64=20,
	tol::Real=1e-6,M::Union(AbstractMatrix,Function)=x->x,
	maxIter::Int64 = length(b), x::Vector=[], out::Int64=0)

	Af =  isa(A,Function) ? A : x->A*x
	Mf =  isa(M,Function) ? M : x->M\x

	# initialization
	n  = length(b)

	if isempty(x)
		x = zeros(n)
		r = Mf(b)
	else
		r = Mf(b-Af(x))
	end

	bnrm2 = norm(b)
	if bnrm2 == 0.0; bnrm2 = 1.0; end

	err = norm( r ) / bnrm2
	if err < tol; return x, err; end

	V     = zeros(n,restart + 1)
	H     = zeros(restart + 1,restart)
	cs    = zeros(restart)
	sn    = zeros(restart)
	s     = zeros(restart + 1) # allocate once
	w     = zeros(n)

	if iseltype(b,Complex) || iseltype(r,Complex) || iseltype(A,Complex)
		s  = complex(s)
		b  = complex(b)
		r  = complex(r)
		x  = complex(x)
		V  = complex(V)
		H  = complex(H)
		cs = complex(cs)
		sn = complex(sn)
	end

	resvec = zeros((restart + 1)*maxIter)

	iter = 0
	flag = -1
	cnt  = 1

	@inbounds for iter = 1:maxIter
		s[1]      = norm( r );
		V[:,1]    = r / s[1];

		@inbounds for i = 1:restart
			w = Af(V[:,i])
			# A_mul_B!(w, A, V[:,i]) # save space if A is matrix
			w = Mf(w)

			@inbounds for k = 1:i # basis using Gram-Schmidt
				H[k,i] = dot(w,V[:,k])
				BLAS.axpy!(n,-H[k,i],V[:,k],1,w,1)
			end
			H[i+1,i] = norm( w )
			V[:,i+1] = w / H[i+1,i]

			@inbounds for k = 1:i-1 # apply Givens rotation
				temp     =  cs[k]*H[k,i] + sn[k]*H[k+1,i]
				H[k+1,i] = -sn[k]*H[k,i] + cs[k]*H[k+1,i]
				H[k,i]   = temp
			end

			# Approximate residual norm
			cs[i],sn[i] = rotmat( H[i,i], H[i+1,i] )
			s[i+1] = -sn[i]*s[i]
			s[i]   = cs[i]*s[i]
			H[i,i] = cs[i]*H[i,i] + sn[i]*H[i+1,i]
			H[i+1,i] = 0.0
			err  = abs(s[i+1]) / bnrm2

			resvec[cnt] = err

			if err <= tol
				y  = H[1:i,1:i] \ s[1:i]
				x += V[:,1:i]*y
				flag = 0; break
			end
			cnt = cnt+1
		end
		if  err <= tol
			flag = 0
			break
		end
		y  = H[1:restart,1:restart]\s[1:restart]
		x += V[:,1:restart]*y

		r = b - Af(x)
		r = Mf(r)

		s[restart+1] = norm(r)
		resvec[cnt] = abs(s[restart+1]) / bnrm2

		if err <= tol
		   flag = 0; break
		end
	end

	if flag==-1
		println(@sprintf("gmres iterated maxIter (=%d) times without achieving the desired tolerance.",maxIter))
	elseif out>=1
		println(@sprintf("gmres achieved desired tolerance at iteration %d. Residual norm is %1.2e.",iter,resvec[cnt]))
	end
	return x,flag,resvec[cnt],iter,resvec[1:cnt]
end

function rotmat( a, b )
# c,s = rotmat(a,b)
#  Givens rotation matrix parameters for a and b.

	if b==0.0
		c = 1.0
		s = 0.0
	elseif abs(b) > abs(a)
		temp = a / b
		s    = 1.0/sqrt( 1.0 + temp^2 )
		c    = temp * s
	else
		temp = b/a
		c    = 1.0/sqrt( 1.0 + temp^2 )
		s    = temp * c
	end
	return c, s
end

ext: jl


text:
#=
divide:
- Julia version: 1.5.0
- Author: shisa
- Date: 2020-08-06
=#

module Division
    function divide(a, b)
        return a / b
    end

    function main()
        a = 4
        b = 3
        println("a = $a, b = $b")
        println("a / b = $(divide(a, b))")
        println()
        println("divide round to 0")
        println("div(a, b) = ", div(a, b))
        println()
        println("remainder")
        println("a % b = ", a % b)
        println("rem(a, b) = ", rem(a, b))
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    using .Division
    Division.main()
end

max_stars_repo_name: ShisatoYano/JuliaPractice

max_issues_repo_name: ShisatoYano/JuliaPractice

max_forks_repo_name: ShisatoYano/JuliaPractice


text: hello(name = "World") = println("Hello, $name", "!")

max_stars_repo_name: jawuku/julia

max_issues_repo_name: jawuku/julia

max_forks_repo_name: jawuku/julia


text:
# Little example for a run lasting 100 nondimensional times
# at 128x128 resolution
# (about 1h30m running on single processor)

using Pkg; Pkg.activate(".")
using bwh
# Run up to time dt*nstep=100 with p=1.2 on a 128x128 grid
# plotting results during integration
P = Params(nx=128, ny=128, p=1.2, nstep=100, fplot=true)
b,w,h = bwh.main(P)
plotbwh(b, w, h, P, P.dt*100)

max_stars_repo_name: jhardenberg/bwh

max_issues_repo_name: jhardenberg/bwh

max_forks_repo_name: jhardenberg/bwh


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct StorageAccountListResult <: SwaggerModel
    value::Any # spec type: Union{ Nothing, Vector{StorageAccount} } # spec name: value
    nextLink::Any # spec type: Union{ Nothing, String } # spec name: nextLink

    function StorageAccountListResult(;value=nothing, nextLink=nothing)
        o = new()
        validate_property(StorageAccountListResult, Symbol("value"), value)
        setfield!(o, Symbol("value"), value)
        validate_property(StorageAccountListResult, Symbol("nextLink"), nextLink)
        setfield!(o, Symbol("nextLink"), nextLink)
        o
    end
end # type StorageAccountListResult

const _property_map_StorageAccountListResult = Dict{Symbol,Symbol}(Symbol("value")=>Symbol("value"), Symbol("nextLink")=>Symbol("nextLink"))
const _property_types_StorageAccountListResult = Dict{Symbol,String}(Symbol("value")=>"Vector{StorageAccount}", Symbol("nextLink")=>"String")
Base.propertynames(::Type{ StorageAccountListResult }) = collect(keys(_property_map_StorageAccountListResult))
Swagger.property_type(::Type{ StorageAccountListResult }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_StorageAccountListResult[name]))}
Swagger.field_name(::Type{ StorageAccountListResult }, property_name::Symbol) =  _property_map_StorageAccountListResult[property_name]

function check_required(o::StorageAccountListResult)
    true
end

function validate_property(::Type{ StorageAccountListResult }, name::Symbol, val)
end


text:
# Functions for inserting text are based on
# https://discourse.julialang.org/t/write-to-a-particular-line-in-a-file/50179

module TextFileTools
using ReadableRegex

export writetext, insertline, deletetext

"""
TODO:
- deleteline
- deleteuntil
- deleteafter
- insertline with `force=true` to add line when
it exceeds the current limit.
"""

"""
    skiplines(io::IO, n)
Helper function for skipping lines when reading a file.
"""
function skiplines(io::IO, n)
    i = 1
    while i <= n
        eof(io) && error("File contains less than $(n + 1) lines")
        i += read(io, Char) === '\n'
    end
end

"""
    writetext(file::String, text::String, linenumber::Int, at=Inf)
Writes a `text` to a `file` in an specific `linenumber` and `at` an specific location.
By default, the text is appended at the end of the line (`at = Inf`). If
`at = 1`, then the text is actually appended in the beggining of the line.
You might also want to add it to an specific position, e.g.
`at = 2`, which will append the text after the first `Char`.
"""
function writetext(file::String, text::String, linenumber::Int; at=Inf, method=:append)
    f = open(file, "r+");
    if at == Inf
        skiplines(f, linenumber);
        skip(f, -1)
    else
        skiplines(f, linenumber - 1);
        skip(f, at - 1)
    end
    if method == :insert
        write(f, text)
    elseif method == :append
        mark(f)          # store where the buffer f is right now
        buf = IOBuffer()
        write(buf, f)    # writes everything from where f is to the end of the file
        seekstart(buf)   # sends the buffer to the start of it
        reset(f)         # sends the buffer of f to when used mark(f)
        print(f, text);  # writes `text`
        write(f, buf)    # writes the buf to the rest of f after the text.
        close(buf)
    else
        throw(ArgumentError("Incorret `method`. It should be either `method=:append` or `method=:insert`."))
    end
    close(f)
end

"""
    writetext(file::String, text::String, linenumber=:last; at=Inf)
Writes text to the last line.
"""
function writetext(file::String, text::String, linenumber=:last; at=Inf)
    @assert linenumber == :last
    if at == Inf
        open(file, "a+") do f
            write(f, text)
        end
    else
        # !!This is not efficient. Should be improved
        writetext(file, text, countlines(file), at=at)
    end
end

"""
    deletetext(file::String, nchar::Int, linenumber::Int; at=1)
Deletes a total of `nchar` starting `at` an specified
position and an specified `linenumber`. By default,
`at=1`, which means that `deletetext(file, 10, 1)` will
delete 10 characters from left to right, starting at
the beggining of line 1.
"""
function deletetext(file::String, nchar::Int, linenumber::Int; at=1)
    f = open(file, "r+");
    if at == Inf
        skiplines(f, linenumber);
        skip(f, -1)
    else
        skiplines(f, linenumber - 1);
        skip(f, at - 1)
    end
    mark(f)
    buf = IOBuffer()
    skip(f, nchar)
    write(buf, f)
    seekstart(buf)
    reset(f)
    write(f, buf)
    close(buf)
    close(f)
end

"""
    insertline(file::String, text::String, linenumber::Int; position=:above)
Inserts a line of `text` in a `file` at the specified `linenumber`.
Accepts `position=:above` (default), `position=:below` or `position=:replace`, to specify
whether the text will be placed above, below or replace the current line.
**Important**, if you want to add a line at the end of the file (after the last line),
use `insertline(file::String, text::String, :lasts)` instead.
Here is an example,
consider a text file containing:

```
Text file example
Current text it here
```

Hence, after running `insertline("textfile.txt", "!insert this", 2)`,
you will get
```
Text file example
!insert this
Current text it here
```
"""
function insertline(file::String, text::String, linenumber::Int; position=:above)
    if position == :above
        insertlineabove(file, text, linenumber)
    elseif position == :below
    insertlinebelow(file, text, linenumber)
    else
    throw(ArgumentError("Invalid position. Use either `position=:above` or `position=:below`."))
    end
end

"""
    insertline(file::String, text::String, linenumber=:last; position=:above)
Inserts line in the last line of the file.
"""
function insertline(file::String, text::String, linenumber=:last; position=:below)
    @assert linenumber == :last
    if position == :above
        insertlineabove(file, text, countlines(file))
    elseif position == :below
    open(file, "a+") do f
        write(f, "\n" * text)
    end
    else
        throw(ArgumentError("Invalid position. Use either `position=:above` or `position=:below`."))
    end
end

    """
    insertlineabove(file::String, text::String, linenumber::Int)
Inserts a line of `text` in a `file` above the `linenumber`.
"""
function insertlineabove(file::String, text::String, linenumber::Int)
    if linenumber == 1
        writetext(file, text * "\n", linenumber, at=1)
    else
        writetext(file, "\n" * text, linenumber - 1, at=Inf)
    end
end

"""
    insertlinebelow(file::String, text::String, linenumber::Int)
Inserts a line of `text` in a `file` below the `linenumber`.
This does not work if the linenumber is the last line. Look
`function insertline(file::String, text::String, linenumber=:last; position=:below)`
instead.
"""
function insertlinebelow(file::String, text::String, linenumber::Int)
    writetext(file, "\n" * text, linenumber, at=Inf)
end

"""
    insertlinereplace(file::String, text::String, linenumber::Int)
Inserts a line of `text` in a `file` replacing the current line.
"""
function insertlinereplace(file::String, text::String, linenumber::Int)
    writetext(file, "\n" * text, linenumber)
end

"""
    insertlineforce(file::String, text::String, linenumber::Int)
Inserts a line of `text` in a `file` below the `linenumber`.
"""
function insertemptylines(file::String, text::String, qtd::Int)
    writetext(file, "\n" * text, linenumber)
end

end


text:
# TODO: respect ordering
module AtomixCUDA

using Atomix: Atomix, IndexableRef
using CUDA: CUDA, CuDeviceArray

const CuIndexableRef{Indexable<:CuDeviceArray} = IndexableRef{Indexable}

function Atomix.get(ref::CuIndexableRef, order)
    error("not implemented")
end

function Atomix.set!(ref::CuIndexableRef, v, order)
    error("not implemented")
end

@inline function Atomix.replace!(
    ref::CuIndexableRef,
    expected,
    desired,
    success_ordering,
    failure_ordering,
)
    ptr = Atomix.pointer(ref)
    expected = convert(eltype(ref), expected)
    desired = convert(eltype(ref), desired)
    begin
        old = CUDA.atomic_cas!(ptr, expected, desired)
    end
    return (; old = old, success = old === expected)
end

@inline function Atomix.modify!(ref::CuIndexableRef, op::OP, x, order) where {OP}
    x = convert(eltype(ref), x)
    ptr = Atomix.pointer(ref)
    begin
        old = if op === (+)
            CUDA.atomic_add!(ptr, x)
        elseif op === (-)
            CUDA.atomic_sub!(ptr, x)
        elseif op === (&)
            CUDA.atomic_and!(ptr, x)
        elseif op === (|)
            CUDA.atomic_or!(ptr, x)
        elseif op === xor
            CUDA.atomic_xor!(ptr, x)
        elseif op === min
            CUDA.atomic_min!(ptr, x)
        elseif op === max
            CUDA.atomic_max!(ptr, x)
        else
            error("not implemented")
        end
    end
    return old => op(old, x)
end

end  # module AtomixCUDA


text:
@testset "1047.remove-all-adjacent-duplicates-in-string.jl" begin
    @test remove_duplicates("abbaca") == "ca"
    @test remove_duplicates(
        "abshxbcbxvcbshsvzxchsbcdhxvcbcbcbcbcbdhdbxbssscbxbxbchdnvbdhdnvbxbd"
    ) == "abshxbcbxvcbshsvzxchsbcdhxvcbcbcbcbcbdhdbxbscbxbxbchdnvbdhdnvbxbd"
end


text:
function _parse_flag(ret::Tuple)
    length(ret) != 2 && return (false, nothing)
    e1, e2 = ret
    return (e1 isa Bool) ? (e1, e2) : (false, nothing)
end
_parse_flag(ret) = (false, nothing)

function converge_ep!(epmodel::EPModel{T};
        verbose::Bool=true,  # output verbosity
        damp::Real=0.9,      # damp ∈ (0,1) newfield = damp * oldfield + (1-damp)* newfield
        epsconv::Real=1e-6,  # convergence criterion
        maxiter::Int=2000,   # maximum iteration count
        maxvar::Real=1e50,   # maximum numerical variance
        minvar::Real=1e-50,  # minimum numerical variance
        iter0 = 1,           # the started iteration count
        drop_epfields = false,  # if the final EPOut object will export the epfields
        # callbacks
        oniter::Function = (it, epmodel) -> nothing
    ) where {T<:Real}

    @extract epmodel : epfields epmat beta_vec alpha stat

    stat[:converge_init_time] = time()
    epalg = EPAlg(alpha, beta_vec,
        minvar, maxvar, epsconv, damp, maxiter, verbose)

    returnstatus = UNCONVERGED_STATUS
    iter = iter0

    # sweep ep till maxiter is reached or max(errav, errvar) < epsconv
    prog = ProgressThresh{typeof(epsconv)}(epsconv; desc =  "EP  ", dt = 0.5)
    max_beta = findmax(beta_vec)
    for _ in iter0:maxiter

        # eponesweep! will be eponesweepT0! or eponesweep depending on alpha
        stat[:elapsed_eponesweep] = @elapsed begin
            errs = epmodel.updatealg!(epfields, epalg, epmat, stat)
        end

        max_err = stat[:max_err] = maximum(errs)

        # call back
        retflag, val = _parse_flag(oniter(iter, epmodel))
        retflag && return val

        # Converged
        max_err < epsconv && (returnstatus = CONVERGED_STATUS; break)

        if verbose
            sweep_time = stat[:elapsed_eponesweep]
            inv_time = stat[:elapsed_eponesweep_inv]
            inv_frac = round(inv_time * 100/ sweep_time; digits = 3)

            update!(prog, max_err; showvalues = [
                (:iter, iter),
                (:maxiter, maxiter),
                (:alpha, alpha),
                (:max_beta, max_beta),
                (:sweep_time, sweep_time),
                (:inv_time, string(inv_time, " [", inv_frac, " %]")),
            ])
        end

        iter += 1
    end # for iter

    verbose && finish!(prog)

    # return
    return produce_epout!(epmodel, returnstatus, iter; drop_epfields)
end


text:
import JSON
import Combinatorics

dirs = readdir("input")

function combine_modnames(modnames, div)
    output = ""
    for name in modnames
        output = "$output$div$name"
    end
    output[(length(div)+1):end]
end

struct ModMeta
    abbrev::String
    name::String
    priority::Int
    unit_list::Set{String}
    comm_list::Set{String}
end

mods = Vector{ModMeta}()

for dir_end in dirs
    dir = joinpath("input", dir_end)
    if !isdir(dir)
        continue
    end

    metafile = joinpath(dir, "meta.json")
    if !isfile(metafile)
        println("meta.json not found in ", dir)
        continue
    end
    meta = JSON.parsefile(metafile)
    req_fields = ["name", "priority"]
    for field in req_fields
        if !(field in keys(meta))
            println("meta.json is missing \"$field\" field")
        end
    end
    unit_list_path = joinpath(dir, "unit_list.json")
    comm_list_path = joinpath(dir, "commander_list.json")

    abbrev = dir_end
    name = meta["name"]
    priority = meta["priority"]
    unit_list = Set{String}(
        if isfile(unit_list_path)
            JSON.parsefile(unit_list_path)["units"]
        else
            []
        end
    )
    comm_list = Set{String}(
        if isfile(comm_list_path)
            JSON.parsefile(comm_list_path)["commanders"]
        else
            []
        end
    )

    push!(mods, ModMeta(
        abbrev,
        name,
        priority,
        unit_list,
        comm_list
    ))
end

for combo in Combinatorics.powerset(mods, 2)
    abbrevs = getproperty.(combo, :abbrev)

    filename = "compatch_$(combine_modnames(abbrevs, "+"))"
    folder = joinpath("output", filename)
    mkpath(folder)

    modinfo = JSON.parsefile("modinfo_template.json")
    modinfo["display_name"] = "DAEDALUS ComPatch - $(combine_modnames(abbrevs, " + "))"
    modinfo["identifier"] = "com.pa.daedalus.compatch_$(combine_modnames(abbrevs, "-"))"
    modinfo["description"] = "Compatibility Patch for the following unit-adding mods: $(abbrevs...)"
    modinfo["priority"] = minimum(getproperty.(combo, :priority)) - 1
    write(joinpath(folder, "modinfo.json"), JSON.json(modinfo))

    unit_lists = getproperty.(combo, :unit_list)
    if sum([!isempty(i) for i in unit_lists]) > 1
        new_unit_list = Dict("units" => collect(union(unit_lists...)))
        unit_path = joinpath(folder, "pa/units")
        mkpath(unit_path)
        write(joinpath(unit_path, "unit_list.json"), JSON.json(new_unit_list))
    end

    comm_lists = getproperty.(combo, :comm_list)
    if sum([!isempty(i) for i in comm_lists]) > 1
        new_comm_list = Dict("commanders" => collect(union(comm_lists...)))
        comm_path = joinpath(folder, "pa/units/commanders")
        mkpath(comm_path)
        write(joinpath(comm_path, "unit_list.json"), JSON.json(new_comm_list))
    end
end


text:
module Tutorial2
using FEniCS
domain = Circle(Point([0.0,0.0]),1)
mesh = generate_mesh(domain,16)
V = FunctionSpace(mesh,"P",2)
w_D = Constant(0)

bc = DirichletBC(V,w_D,"on_boundary")
beta =8
R0 = 0.6
p = Expression("4*exp(-pow(beta, 2)*(pow(x[0], 2) + pow(x[1] - R0, 2)))",
degree=1, beta=beta, R0=R0)
w = TrialFunction(V)
v = TestFunction(V)
a = dot(grad(w), grad(v))*dx
L = p*v*dx
w=FeFunction(V)
lvsolve(a,L,w,bc)
p = interpolate(p,V)

vtkfile_w = File("poisson_membrane/deflection.pvd")
vtkfile_w << w.pyobject
vtkfile_p = File("poisson_membrane/load.pvd")
vtkfile_p << p.pyobject
end#module


text:
abstract type AbstractDCLineFormulation <: AbstractBranchFormulation end
struct HVDCUnbounded <: AbstractDCLineFormulation end
struct HVDCLossless <: AbstractDCLineFormulation end
struct HVDCDispatch <: AbstractDCLineFormulation end
# Not Implemented
# struct VoltageSourceDC <: AbstractDCLineFormulation end

#################################### Branch Variables ##################################################
get_variable_binary(_, ::Type{<:PSY.DCBranch}, ::AbstractDCLineFormulation) = false

get_variable_sign(::FlowActivePowerVariable, ::Type{<:PSY.DCBranch}, _) = NaN

get_variable_sign(
    ::FlowActivePowerFromToVariable,
    ::Type{<:PSY.DCBranch},
    ::AbstractDCLineFormulation,
) = -1.0
get_variable_sign(
    ::FlowActivePowerToFromVariable,
    ::Type{<:PSY.DCBranch},
    ::AbstractDCLineFormulation,
) = 1.0

get_variable_lower_bound(
    ::FlowActivePowerVariable,
    d::PSY.DCBranch,
    ::AbstractDCLineFormulation,
) = max(PSY.get_active_power_limits_from(d).min, PSY.get_active_power_limits_to(d).min)
get_variable_upper_bound(
    ::FlowActivePowerVariable,
    d::PSY.DCBranch,
    ::AbstractDCLineFormulation,
) = min(PSY.get_active_power_limits_from(d).max, PSY.get_active_power_limits_to(d).max)
#! format: on

#################################### Flow Variable Bounds ##################################################
function add_variable_to_expression!(
    optimization_container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{B},
    ::DeviceModel{B, <:AbstractDCLineFormulation},
    ::Type{S},
) where {B <: PSY.DCBranch, S <: Union{StandardPTDFModel, PTDFPowerModel}}
    time_steps = model_time_steps(optimization_container)
    var = get_variable(optimization_container, FLOW_ACTIVE_POWER, B)

    for (ix, d) in enumerate(devices)
        bus_fr = PSY.get_number(PSY.get_arc(d).from)
        bus_to = PSY.get_number(PSY.get_arc(d).to)
        for t in time_steps
            flow_variable = var[PSY.get_name(d), t]
            add_to_expression!(
                optimization_container.expressions[:nodal_balance_active],
                PSY.get_number(PSY.get_arc(d).from),
                t,
                flow_variable,
                -1.0,
            )
            add_to_expression!(
                optimization_container.expressions[:nodal_balance_active],
                PSY.get_number(PSY.get_arc(d).to),
                t,
                flow_variable,
                1.0,
            )
        end
    end
end

#################################### Rate Limits Constraints ##################################################
function branch_rate_constraints!(
    optimization_container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{B},
    ::DeviceModel{B, HVDCLossless},
    ::Type{<:PM.AbstractDCPModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
) where {B <: PSY.DCBranch}
    var = get_variable(optimization_container, FLOW_ACTIVE_POWER, B)
    time_steps = model_time_steps(optimization_container)
    constraint_val =
        JuMPConstraintArray(undef, [PSY.get_name(d) for d in devices], time_steps)
    assign_constraint!(optimization_container, FLOW_ACTIVE_POWER, B, constraint_val)
    for t in time_steps, d in devices
        min_rate = max(
            PSY.get_active_power_limits_from(d).min,
            PSY.get_active_power_limits_to(d).min,
        )
        max_rate = min(
            PSY.get_active_power_limits_from(d).max,
            PSY.get_active_power_limits_to(d).max,
        )
        constraint_val[PSY.get_name(d), t] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            min_rate <= var[PSY.get_name(d), t] <= max_rate
        )
    end
    return
end

branch_rate_constraints!(
    ::OptimizationContainer,
    ::IS.FlattenIteratorWrapper{<:PSY.DCBranch},
    ::DeviceModel{<:PSY.DCBranch, HVDCUnbounded},
    ::Type{<:PM.AbstractPowerModel},
    ::Union{Nothing, AbstractAffectFeedForward},
) = nothing

function branch_rate_constraints!(
    optimization_container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{B},
    ::DeviceModel{B, <:AbstractDCLineFormulation},
    ::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
) where {B <: PSY.DCBranch}
    time_steps = model_time_steps(optimization_container)
    for (var_type, cons_type) in
        zip((FLOW_ACTIVE_POWER, FLOW_ACTIVE_POWER), (RATE_LIMIT_FT, RATE_LIMIT_TF))
        var = get_variable(optimization_container, var_type, B)
        constraint_val =
            JuMPConstraintArray(undef, [PSY.get_name(d) for d in devices], time_steps)
        assign_constraint!(optimization_container, cons_type, B, constraint_val)

        for t in time_steps, d in devices
            min_rate = max(
                PSY.get_active_power_limits_from(d).min,
                PSY.get_active_power_limits_to(d).min,
            )
            max_rate = min(
                PSY.get_active_power_limits_from(d).max,
                PSY.get_active_power_limits_to(d).max,
            )
            constraint_val[PSY.get_name(d), t] = JuMP.@constraint(
                optimization_container.JuMPmodel,
                min_rate <= var[PSY.get_name(d), t] <= max_rate
            )
            add_to_expression!(
                optimization_container.expressions[:nodal_balance_active],
                PSY.get_number(PSY.get_arc(d).to),
                t,
                var[PSY.get_name(d), t],
                -PSY.get_loss(d).l1,
                -PSY.get_loss(d).l0,
            )
        end
    end
    return
end


text:
"""
    foldconstants(s, dict)

Substitute into a symbolic expression `s` using dictionary `dict`

```
julia> s = x + y*z
x + y*z

julia> foldconstants(s, Dict(:z => 3))
x + 3y
```
"""
function foldconstants(s, dict)
    known = Set(keys(dict))

    p(x) = (symtype(x) <: Number) && (getproperty.(atoms(x), :name) ⊆ known)

    r = @rule ~x::p => toconst(~x, dict)

    RW.Prewalk(RW.PassThrough(r))(s) |> simplify
end

toconst(s::Number, dict) = s

toconst(s::Sym, dict) = dict[s.name]

function toconst(s::Symbolic, dict)
    # First, here's the main body of the code
    f_expr = @q begin $(codegen(s)) end

    @gensym dictname

    # Now prepend the variable assignments we'll need
    for v in atoms(s)
        v = v.name
        vname = QuoteNode(v)
        pushfirst!(f_expr.args, :($v = $dictname[$vname]))
    end

    # Make it a function
    f_expr = @q begin function f($dictname) $f_expr end end

    # Tidy up the blocks
    f_expr = MacroTools.flatten(f_expr)

    # ...and generate!
    f = @RuntimeGeneratedFunction f_expr

    f(dict)
end


text:
using Polynomials

"""
Calculate Chebyshev polynomials coefficients

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
function poly_chebyshev{T<:Number}(n, ::Type{T}=Int, var=:x)

    if n==0
        return Poly{T}([one(T)], var)
    elseif n==1
        return Poly{T}([zero(T), one(T)], var)
    end

    px = Poly{T}([zero(T), one(T)], var)
    T0 = Poly{T}([one(T)], var)
    T1 = px

    for i = 2:n
        T2 = 2*px*T1 - T0
        T0 = T1
        T1 = T2
    end

    return T1

end

"""
Calculate derivative of Chebyshev polynomials coefficients

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
poly_dchebyshev{T<:Number}(n, ::Type{T}=Int, var=:x) = polyder(poly_chebyshev(n,T, var))

"""
Calculate the coefficients of Chebyshev polynomials of the second kind

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
function poly_chebyshev2{T<:Number}(n, ::Type{T}=Int, var=:x)

    if n==0
        return Poly{T}([one(T)], var)
    elseif n==1
        return Poly{T}([zero(T), 2*one(T)], var)
    end

    px = Poly{T}([zero(T), one(T)], var)
    U0 = Poly{T}([one(T)], var)
    U1 = Poly{T}([zero(T), 2*one(T)], var)

    for i = 2:n
        U2 = 2*px*U1 - U0
        U0 = U1
        U1 = U2
    end

    return U1

end

"""
Calculate the coefficients of the derivative of Chebyshev polynomials of the second kind

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
poly_dchebyshev2{T<:Number}(n, ::Type{T}=Int, var=:x) = polyder(poly_chebyshev2(n,T, var))

"""
Calculate the coefficients of Legendre polynomials

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
function poly_legendre{T<:Number}(n, ::Type{T}=Float64, var=:x)

    if n==0
        return Poly{T}([one(T)], var)
    elseif n==1
        return Poly{T}([zero(T), one(T)], var)
    end

    px = Poly{T}([zero(T), one(T)], var)
    p0 = Poly{T}([one(T)], var)
    p1 = px

    for i = 2:n
        p2 = ( (2i-1)*px*p1 - (i-1)*p0 ) / i
        p0 = p1
        p1 = p2
    end

    return p1
end

"""
Calculate the coefficients of the derivativ of Legendre polynomials

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
poly_dlegendre{T<:Number}(n, ::Type{T}=Int, var=:x) = polyder(poly_legendre(n,T, var))

"""
Calculate the coefficients of Jacobi polynomials

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
function poly_jacobi{T<:Number}(n, a, b, ::Type{T}=Float64, var=:x)

    ox = one(T)
    zx = zero(T)
    if n==0
        return Poly{T}([one(T)], var)
    elseif n==1
        return Poly{T}([(a-b)/2, (a+b+2)/2], var)
    end
    p0 = Poly{T}([one(T)], var)
    p1 = Poly{T}([(a-b)/2, (a+b+2)/2], var)
    px = Poly{T}([zero(T), one(T)], var)

    for i = 1:(n-1)
	a1 = 2*(i+1)*(i+a+b+1)*(2*i+a+b);
	a2 = (2*i+a+b+1)*(a*a-b*b);
	a3 = (2*i+a+b)*(2*i+a+b+1)*(2*i+a+b+2);
	a4 = 2*(i+a)*(i+b)*(2*i+a+b+2);
	p2 = ox/a1*( (a2 + a3*px)*p1 - a4*p0);

        p0 = p1
        p1 = p2
    end

    return p1
end

"""
Calculate the coefficients of the derivative of Jacobi polynomials

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
poly_djacobi{T<:Number}(n, a, b, ::Type{T}=Float64, var=:x) = polyder(poly_jacobi(n,a,b,T,var))


text:
#Minesweeper 16x16

function placemine(grid) #randomly places a mine in a square in the grid where there isn't already a mine
    row, column = rand(1:16), rand(1:16)
    if grid[row][column]!="*"
        grid[row][column]="*"
    else
        placemine(grid)
    end
end

up_int(str::String) = string(parse(Int64,str)+1)

function update(g, r, c) #add one to the boxes next to a mine; split into sections horizontally
    if r>1
        update_c(g,r-1,c) #update row above
    end
    update_c(g,r,c) #update the same row
    if r<16
        update_c(g,r+1,c) #update row below
    end
end

function update_c(g, r, c) #adds one to the boxes next to a mine, vertically
    row = g[r]
    if c>1 && row[c-1]!="*"
        row[c-1] = up_int(row[c-1]) #update left column
    end
    if row[c]!="*"
        row[c] = up_int(row[c]) #update current column
    end
    if c<16 && row[c+1]!="*"
        row[c+1] = up_int(row[c+1]) #update right column
    end
end

function printgrid(grid) #prints the grid
    println("\n    A   B   C   D   E   F   G   H   I   J   K   L   M   N   O   P")
    println("  ╔═══╦═══╦═══╦═══╦═══╦═══╦═══╦═══╦═══╦═══╦═══╦═══╦═══╦═══╦═══╦═══╗")
    for r=1:16
        if r<10
            println(r," ║ ",grid[r][1]," ║ ",grid[r][2]," ║ ",grid[r][3]," ║ ",grid[r][4]," ║ ",grid[r][5]," ║ ",grid[r][6]," ║ ",grid[r][7]," ║ ",grid[r][8]," ║ ",grid[r][9]," ║ ",grid[r][10]," ║ ",grid[r][11]," ║ ",grid[r][12]," ║ ",grid[r][13]," ║ ",grid[r][14]," ║ ",grid[r][15]," ║ ",grid[r][16]," ║ ")
        else
            println(r,"║ ",grid[r][1]," ║ ",grid[r][2]," ║ ",grid[r][3]," ║ ",grid[r][4]," ║ ",grid[r][5]," ║ ",grid[r][6]," ║ ",grid[r][7]," ║ ",grid[r][8]," ║ ",grid[r][9]," ║ ",grid[r][10]," ║ ",grid[r][11]," ║ ",grid[r][12]," ║ ",grid[r][13]," ║ ",grid[r][14]," ║ ",grid[r][15]," ║ ",grid[r][16]," ║ ")
        end
        if r != 16
            println("  ╠═══╬═══╬═══╬═══╬═══╬═══╬═══╬═══╬═══╬═══╬═══╬═══╬═══╬═══╬═══╬═══╣")
        end
    end
    println("  ╚═══╩═══╩═══╩═══╩═══╩═══╩═══╩═══╩═══╩═══╩═══╩═══╩═══╩═══╩═══╩═══╝")
end

function search(letters, loc) #returns the index of the letter (from the input) in the array letters
    for i=1:length(letters)
        if letters[i] == string(loc[1])
            return i
        end
    end
end

function play(g, kg) #the main function that is called for every move
    letters = String["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p"]
    loc = lowercase(string(input("Choose a square to reveal. (eg. F4)\n")))
    r = parse(Int64,loc[2]) #the number corresponds to the row
    c = search(letters,loc) #the letter corresponds to the column
    value = g[r][c]
    if value == "*"
        printgrid(g)
        println("You Lose.")
        playagain = lowercase(string(input("Play again? Y/N\n")))
        if playagain == "y"
            restart()
        else
            exit()
        end
    end
    kg[r][c] = value
    if value == "0"
        checkzeros(kg, g, r, c)
    end
    printgrid(kg)
    remaining = 0
    for x=1:16
        for y in kg[x]
            if y == " "
                remaining = remaining + 1
            end
        end
    end
    if remaining <= 20
        printboard(g)
        println("You Win.")
        playagain = lowercase(string(input("Play again? Y/N\n")))
        if playagain == "y"
            restart()
        else
            exit()
        end
    end
    play(g,kg)
end

function checkzeros(kg, g, r, c) #function that checks that every square around each 0 is open
    oldkg = deepcopy(kg)
    openzero(kg, g, r, c)
    if oldkg==kg #check if there are any changes in the grid
        return
    end
    while true #loops through the grid to open any other 0s we find, creating a chain reaction
        oldkg = deepcopy(kg)
        for x=1:16
            for y=1:16
                if kg[x][y]=="0"
                    openzero(kg, g, x, y)
                end
            end
        end
        if oldkg==kg
            return
        end
    end
end

function openzero(kg, g, r, c) #opens all the squares around when 0 is found
    #Row above
    if r>1
        row = kg[r-1]
        if c>1
            row[c-1] = g[r-1][c-1]
        end
        row[c] = g[r-1][c]
        if c<16
            row[c+1] = g[r-1][c+1]
        end
    end
    #Same row
    row = kg[r]
    if c>1
        row[c-1] = g[r][c-1]
    end
    if c<16
        row[c+1] = g[r][c+1]
    end
    #Row below
    if r<16
        row = kg[r+1]
        if c>1
            row[c-1] = g[r+1][c-1]
        end
        row[c] = g[r+1][c]
        if c<16
            row[c+1] = g[r+1][c+1]
        end
    end
end

function restart() #the main function that initiates every game
    println("Welcome to Minesweeper by kfung.")
    println("Your goal is to find 20 mines(*) that are located in an 16x16 grid.")
    println("In each go, type to coordinates of the square you want to uncover, eg. G3")
    println("You lose if you choose a square with a mine.")
    println("Otherwise, the number of mines adjacent to a square (including diagonals) will be revealed.")
    println("Squares around 0 will be automatically reveal because they are no adjacent mines.")
    println("LET'S BEGIN!")
    grid = [["0" for j in 1:16] for k in 1:16]
    knowngrid = [[" " for j in 1:16] for k in 1:16]
    for i=1:20
        placemine(grid)
    end
    for row=1:16
        for column=1:16
            if grid[row][column] == "*"
                update(grid, row, column)
            end
        end
    end
    printgrid(knowngrid)
    play(grid, knowngrid)
end

restart()

max_stars_repo_name: fungktt/JuliaGames

max_issues_repo_name: fungktt/JuliaGames

max_forks_repo_name: fungktt/JuliaGames


text:
# A positive fraction whose numerator is less than its denominator is called a
# proper fraction. For any denominator, d, there will be d−1 proper fractions;
# for example, with d = 12: 1/12, 2/12, 3/12, 4/12, 5/12, 6/12, 7/12, 8/12,
# 9/12, 10/12, 11/12.
#
# We shall call a fraction that cannot be cancelled down a resilient fraction.
# Furthermore we shall define the resilience of a denominator, R(d), to be the
# ratio of its proper fractions that are resilient; for example, R(12) = 4/11.
# In fact, d = 12 is the smallest denominator having a resilience R(d) < 4/10.
#
# Find the smallest denominator d, having a resilience R(d) < 15499/94744.

using ProjectEulerSolutions

# Slower, explicitly calculate the gcd for each fraction, skipping some when
# the number is a multiple of 2.  This is quite slow when n gets large, so we
# will not use it in the benchmarking.
function resilience_slow(n::Integer)::Float64
    count = 1
    step = 1
    if rem(n,2) == 0
        step = 2
    end
    for i in 1:step:n-2
        if gcd(i,n) == 1
            count += 1
        end
    end
    return count / (n-1)
end

# We are actually finding the solution to Euler's totient function divided by
# n-1.
function resilience_fast(n::Integer)::Float64
    return totient(n) / (n - 1)
end

# Brute force approach, with some logic.  The minimum resilience values will be
# when the denominator is a multiple of primes.  We find which multiple
# satisfies the initial constraint and then iterate through smaller
# multiplicative values.
function p243solution(thresh::Float64=0.4)::Integer
    primes = sieve_eratosthenes(30)
    index = 2
    while resilience_fast(prod(primes[1:index])) > thresh
        index += 1
    end
    starting_val = prod(primes[1:index-1])
    mult = 2
    while resilience_fast(mult * starting_val) > thresh
        mult += 1
    end
    return mult * starting_val
end

p243 = Problems.Problem(p243solution)

Problems.benchmark(p243, 15499/94744)

max_stars_repo_name: gnujosh/julia-euler-project

max_issues_repo_name: gnujosh/julia-euler-project

max_forks_repo_name: gnujosh/julia-euler-project


text:
module FFTResampling

using FFTW

include("utils.jl")

include("sinc_interpolation.jl")

end # module


text:
"""
`extrap_prep_dim(d, ::Type{Periodic})`

Translate x into the domain [lbound, ubound] my means of `mod()`
"""
function extrap_prep_dim(::Type{Periodic}, d)
    xs_d = Symbol("xs_", d)
    :($xs_d = mod(xs[$d] - lbound(etp.itp, $d), ubound(etp.itp, $d) - lbound(etp.itp, $d)) + lbound(etp.itp, $d))
end

extrap_prep{N,d}(::Type{Periodic}, ::Val{N}, ::Val{d}) = extrap_prep_dim(Periodic, d)
function extrap_prep{N,d}(::Type{Periodic}, ::Val{N}, ::Val{d}, ::Val{:lo})
    xs_d = Symbol("xs_", d)
    quote
        if $xs_d < lbound(etp.itp, $d)
            $(extrap_prep_dim(Periodic, d))
        end
    end
end
function extrap_prep{N,d}(::Type{Periodic}, ::Val{N}, ::Val{d}, ::Val{:hi})
    xs_d = Symbol("xs_", d)
    quote
        if $xs_d > ubound(etp.itp, d)
            $(extrap_prep_dim(Periodic, d))
        end
    end
end


text:
struct SSqrtHandler end

gethandler(broadcasted, ::Type{typeof(sqrt)}, ::Type{<:SBitstreamLike}) =
    !broadcasted ? SSqrtHandler() : error("Cannot generate hardware for broadcasted sqrt.")
init_state(::SSqrtHandler) = (id = 0,)

function (handler::SSqrtHandler)(buffer, netlist, state, inputs, outputs)
    # add output net to netlist
    setsigned!(netlist, outputs[1], true)

    write(buffer, """
        $stdcomment
        // BEGIN sqrt$(state.id)
        stoch_square_root sqrt$(state.id) (
                .CLK  (CLK),
                .nRST (nRST),
                .up   ($(name(inputs[1]))_p),
                .un   ($(name(inputs[2]))_m),
                .y    ($(name(outputs[1]))_p)
            );
        assign $(name(outputs[1]))_m = 1'b0;
        // END sqrt$(state.id)
        \n""")

    return buffer, (id = state.id + 1,)
end


text:
function check_edge_build()
    g = create_graph()
    v1 = GraphVertex(1, 1:2)
    v2 = GraphVertex(2, 1:2)
    e1 = GraphEdge(1, 2, 1:2)
    e2 = GraphEdge(v2, v1, 1:2)
    return e1.object.opts[:_edge_idx] == 1 && e2.object.opts[:_edge_idx] == 2
end


text:
using JuliaMBD
using Test

@testset "JuliaMBD.jl" begin
    blk1 = ConstantBlock(5)
    blk2 = GainBlock(Var(:K, Float64))
    l1 = Line(:x, Int)
    connect(blk1.outport, blk2.inport, l1)
    println(_toexpr(blk1.outport))
    println(_toexpr(blk2.outport))
end

max_stars_repo_name: JuliaReliab/JuliaMBD

max_issues_repo_name: JuliaReliab/JuliaMBD

max_forks_repo_name: JuliaReliab/JuliaMBD


text:
using PyPlot

using SubmodularMaximization

pygui(false)
fig_path = "./fig/coverage_test"
mkpath(fig_path)

num_agents = 5
num_sensors = 5
nominal_area = 2.0

sensor_radius = sqrt(nominal_area / (num_agents * pi))
station_radius = sensor_radius

agent_specification = CircleAgentSpecification(sensor_radius, station_radius,
                                         num_sensors)

agents = generate_agents(agent_specification, num_agents)

f(x) = mean_area_coverage(x, 100)
problem = ExplicitPartitionProblem(f, agents)

function evaluate_solver(solver, name)
  println("$name solver running")
  @time solution = solver(problem)

  figure()
  xlim([0, 1])
  ylim([0, 1])
  colors = generate_colors(agents)
  visualize_agents(agents, colors)
  visualize_solution(problem, solution, colors)

  gca().set_aspect("equal")
  savefig("$(fig_path)/$(to_file(name)).png", pad_inches=0.00, bbox_inches="tight")

  coverage = solution.value

  title("$name Solver ($coverage)")

  @show coverage
end

#evaluate_solver(solve_optimal, "Optimal")
#evaluate_solver(solve_worst, "Worst-case")
#evaluate_solver(solve_myopic, "Myopic")
#evaluate_solver(solve_random, "Random")
evaluate_solver(solve_sequential, "Sequential")

for num_partitions in [2, 4, 8]
  solve_n(p) = solve_n_partitions(num_partitions, p)
  evaluate_solver(solve_n, "Partition-$num_partitions")
end

#@show mean_weight(problem)
#@show total_weight(problem)

max_stars_repo_name: mcorah/MultiAgentSensing

max_issues_repo_name: mcorah/MultiAgentSensing

max_forks_repo_name: mcorah/MultiAgentSensing


text:
import ForwardDiff
import AbstractFFTs

# original PR by mcabbott: https://github.com/JuliaDiff/ForwardDiff.jl/pull/495

ForwardDiff.value(x::Complex{<:ForwardDiff.Dual}) = Complex(x.re.value, x.im.value)

ForwardDiff.partials(x::Complex{<:ForwardDiff.Dual}, n::Int) =
    Complex(ForwardDiff.partials(x.re, n), ForwardDiff.partials(x.im, n))

ForwardDiff.npartials(x::Complex{<:ForwardDiff.Dual{T,V,N}}) where {T,V,N} = N
ForwardDiff.npartials(::Type{<:Complex{<:ForwardDiff.Dual{T,V,N}}}) where {T,V,N} = N

ForwardDiff.tagtype(x::Complex{<:ForwardDiff.Dual{T,V,N}}) where {T,V,N} = T
ForwardDiff.tagtype(::Type{<:Complex{<:ForwardDiff.Dual{T,V,N}}}) where {T,V,N} = T

# AbstractFFTs.complexfloat(x::AbstractArray{<:ForwardDiff.Dual}) = float.(x .+ 0im)
AbstractFFTs.complexfloat(x::AbstractArray{<:ForwardDiff.Dual}) = AbstractFFTs.complexfloat.(x)
AbstractFFTs.complexfloat(d::ForwardDiff.Dual{T,V,N}) where {T,V,N} = convert(ForwardDiff.Dual{T,float(V),N}, d) + 0im

AbstractFFTs.realfloat(x::AbstractArray{<:ForwardDiff.Dual}) = AbstractFFTs.realfloat.(x)
AbstractFFTs.realfloat(d::ForwardDiff.Dual{T,V,N}) where {T,V,N} = convert(ForwardDiff.Dual{T,float(V),N}, d)

for plan in [:plan_fft, :plan_ifft, :plan_bfft]
    @eval begin
        AbstractFFTs.$plan(x::AbstractArray{<:ForwardDiff.Dual}, region=1:ndims(x); kwargs...) =
            AbstractFFTs.$plan(ForwardDiff.value.(x) .+ 0im, region; kwargs...)

        AbstractFFTs.$plan(x::AbstractArray{<:Complex{<:ForwardDiff.Dual}}, region=1:ndims(x); kwargs...) =
            AbstractFFTs.$plan(ForwardDiff.value.(x), region; kwargs...)
    end
end

# rfft only accepts real arrays
AbstractFFTs.plan_rfft(x::AbstractArray{<:ForwardDiff.Dual}, region=1:ndims(x); kwargs...) =
    AbstractFFTs.plan_rfft(ForwardDiff.value.(x), region; kwargs...)

for plan in [:plan_irfft, :plan_brfft]  # these take an extra argument, only when complex?
    @eval begin
        AbstractFFTs.$plan(x::AbstractArray{<:ForwardDiff.Dual}, region=1:ndims(x); kwargs...) =
            AbstractFFTs.$plan(ForwardDiff.value.(x) .+ 0im, region; kwargs...)

        AbstractFFTs.$plan(x::AbstractArray{<:Complex{<:ForwardDiff.Dual}}, d::Integer, region=1:ndims(x); kwargs...) =
            AbstractFFTs.$plan(ForwardDiff.value.(x), d, region; kwargs...)
    end
end

for P in [:Plan, :ScaledPlan]  # need ScaledPlan to avoid ambiguities
    @eval begin
        Base.:*(p::AbstractFFTs.$P, x::AbstractArray{<:ForwardDiff.Dual}) =
            _apply_plan(p, x)

        Base.:*(p::AbstractFFTs.$P, x::AbstractArray{<:Complex{<:ForwardDiff.Dual}}) =
            _apply_plan(p, x)

        LinearAlgebra.mul!(Y::AbstractArray, p::AbstractFFTs.$P, X::AbstractArray{<:ForwardDiff.Dual}) =
            (Y .= _apply_plan(p, X))

        LinearAlgebra.mul!(Y::AbstractArray, p::AbstractFFTs.$P, X::AbstractArray{<:Complex{<:ForwardDiff.Dual}}) =
            (Y .= _apply_plan(p, X))
    end
end

LinearAlgebra.mul!(Y::AbstractArray{<:Complex{<:ForwardDiff.Dual}}, p::AbstractFFTs.ScaledPlan{T,P,<:ForwardDiff.Dual}, X::AbstractArray{<:ComplexF64}) where {T,P} =
    (Y .= _apply_plan(p, X))

function _apply_plan(p::AbstractFFTs.Plan, x::AbstractArray{<:Complex{<:ForwardDiff.Dual{T}}}) where T
    # TODO do we want x::AbstractArray{<:ForwardDiff.Dual{T}} too?
    xtil = p * ForwardDiff.value.(x)
    dxtils = ntuple(ForwardDiff.npartials(eltype(x))) do n
        p * ForwardDiff.partials.(x, n)
    end
    map(xtil, dxtils...) do val, parts...
        Complex(
            ForwardDiff.Dual{T}(real(val), map(real, parts)),
            ForwardDiff.Dual{T}(imag(val), map(imag, parts)),
        )
    end
end

function _apply_plan(p::AbstractFFTs.ScaledPlan{T,P,<:ForwardDiff.Dual}, x::AbstractArray) where {T,P}
    _apply_plan(p.p, p.scale * x) # for when p.scale is Dual, need out-of-place
end

# this is to avoid method ambiguities between these two:
#   _apply_plan(p::AbstractFFTs.Plan, x::AbstractArray{<:Complex{<:ForwardDiff.Dual{T}}}) where T
#   _apply_plan(p::AbstractFFTs.ScaledPlan{T,P,<:ForwardDiff.Dual}, x::AbstractArray) where {T,P}
function _apply_plan(p::AbstractFFTs.ScaledPlan{T,P,<:ForwardDiff.Dual}, x::AbstractArray{<:Complex{<:ForwardDiff.Dual{Tg}}}) where {T,P,Tg}
    _apply_plan(p.p, p.scale * x)
end

# Convert and strip off duals if that's the only way
function convert_dual(::Type{T}, x::ForwardDiff.Dual) where {T}
    convert(T, ForwardDiff.value(x))
end
convert_dual(::Type{T}, x::ForwardDiff.Dual) where {T <: ForwardDiff.Dual} = convert(T, x)
convert_dual(::Type{T}, x) where {T} = convert(T, x)

# DFTK setup specific
default_primes(T::Type{<:ForwardDiff.Dual}) = default_primes(ForwardDiff.valtype(T))
function next_working_fft_size(T::Type{<:ForwardDiff.Dual}, size::Integer)
    next_working_fft_size(ForwardDiff.valtype(T), size)
end

next_working_fft_size(::Type{<:ForwardDiff.Dual}, size::Int) = size

_fftw_flags(::Type{<:ForwardDiff.Dual}) = FFTW.MEASURE | FFTW.UNALIGNED

function build_fft_plans(T::Type{<:Union{ForwardDiff.Dual,Complex{<:ForwardDiff.Dual}}}, fft_size)
    tmp = Array{complex(T)}(undef, fft_size...) # TODO think about other Array types
    opFFT  = FFTW.plan_fft(tmp, flags=_fftw_flags(T))
    opBFFT = FFTW.plan_bfft(tmp, flags=_fftw_flags(T))

    ipFFT  = DummyInplace{typeof(opFFT)}(opFFT)
    ipBFFT = DummyInplace{typeof(opBFFT)}(opBFFT)
    # backward by inverting and stripping off normalizations
    ipFFT, opFFT, ipBFFT, opBFFT
end

# determine symmetry operations only from primal lattice values
function spglib_get_symmetry(lattice::AbstractMatrix{<:ForwardDiff.Dual}, atom_groups, positions,
                             magnetic_moments=[]; kwargs...)
    spglib_get_symmetry(ForwardDiff.value.(lattice), atom_groups, positions,
                        magnetic_moments; kwargs...)
end
function spglib_atoms(atom_groups,
                      positions::AbstractVector{<:AbstractVector{<:ForwardDiff.Dual}},
                      magnetic_moments)
    positions_value = [ForwardDiff.value.(pos) for pos in positions]
    spglib_atoms(atom_groups, positions_value, magnetic_moments)
end

function _is_well_conditioned(A::AbstractArray{<:ForwardDiff.Dual}; kwargs...)
    _is_well_conditioned(ForwardDiff.value.(A); kwargs...)
end

value_type(T::Type{<:ForwardDiff.Dual}) = ForwardDiff.valtype(T)

# TODO Should go to Model.jl / PlaneWaveBasis.jl as a constructor.
#
# Along with it should go a nice convert function to get rid of the annoying
# conversion thing in the stress computation.
#
# Ideally also upon constructing a Model one would automatically determine that
# *some* parameter deep inside the terms, psp or sth is a dual and automatically
# convert the full thing to be based on dual numbers ... note that this requires
# exposing the element type all the way up ... which is probably needed to do these
# forward and backward conversion routines between duals and non-duals as well.
function construct_value(model::Model{T}) where {T <: ForwardDiff.Dual}
    newpositions = [ForwardDiff.value.(pos) for pos in model.positions]
    Model(ForwardDiff.value.(model.lattice),
          construct_value.(model.atoms),
          newpositions;
          model_name=model.model_name,
          n_electrons=model.n_electrons,
          magnetic_moments=[],  # Symmetries given explicitly
          terms=model.term_types,
          temperature=ForwardDiff.value(model.temperature),
          smearing=model.smearing,
          spin_polarization=model.spin_polarization,
          symmetries=model.symmetries)
end

construct_value(el::Element) = el
construct_value(el::ElementPsp) = ElementPsp(el.Z, el.symbol, construct_value(el.psp))
construct_value(psp::PspHgh) = psp
function construct_value(psp::PspHgh{T}) where {T <: ForwardDiff.Dual}
    PspHgh(psp.Zion,
           ForwardDiff.value(psp.rloc),
           ForwardDiff.value.(psp.cloc),
           psp.lmax,
           ForwardDiff.value.(psp.rp),
           [ForwardDiff.value.(hl) for hl in psp.h],
           psp.identifier,
           psp.description)
end

function construct_value(basis::PlaneWaveBasis{T}) where {T <: ForwardDiff.Dual}
    new_kshift = isnothing(basis.kshift) ? nothing : ForwardDiff.value.(basis.kshift)
    PlaneWaveBasis(construct_value(basis.model),
                   ForwardDiff.value(basis.Ecut),
                   map(v -> ForwardDiff.value.(v), basis.kcoords_global),
                   ForwardDiff.value.(basis.kweights_global);
                   basis.symmetries_respect_rgrid,
                   fft_size=basis.fft_size,
                   kgrid=basis.kgrid,
                   kshift=new_kshift,
                   variational=basis.variational,
                   comm_kpts=basis.comm_kpts)
end

function self_consistent_field(basis_dual::PlaneWaveBasis{T};
                               response=ResponseOptions(),
                               kwargs...) where T <: ForwardDiff.Dual
    # Note: No guarantees on this interface yet.

    # Primal pass
    basis  = construct_value(basis_dual)
    scfres = self_consistent_field(basis; kwargs...)

    # Split into truncated and extra eigenpairs
    truncated = select_occupied_orbitals(scfres; threshold=response.occupation_threshold)
    ψ_extra = [@view scfres.ψ[ik][:, length(εkTrunc)+1:end]
               for (ik, εkTrunc) in enumerate(truncated.eigenvalues)]
    ε_extra = [@view scfres.eigenvalues[ik][length(εkTrunc)+1:end]
               for (ik, εkTrunc) in enumerate(truncated.eigenvalues)]

    ## promote occupied bands to dual numbers
    occupation_dual = [T.(occₖ) for occₖ in truncated.occupation]
    ψ_dual = [Complex.(T.(real(ψₖ)), T.(imag(ψₖ))) for ψₖ in truncated.ψ]
    ρ_dual = DFTK.compute_density(basis_dual, ψ_dual, occupation_dual)
    εF_dual = T(truncated.εF)  # Only needed for entropy term
    eigenvalues_dual = [T.(εₖ) for εₖ in truncated.eigenvalues]
    energies_dual, ham_dual = energy_hamiltonian(basis_dual, ψ_dual, occupation_dual;
                                                 ρ=ρ_dual, eigenvalues=eigenvalues_dual,
                                                 εF=εF_dual)

    response.verbose && println("Solving response problem")

    ## Implicit differentiation
    hamψ_dual = ham_dual * ψ_dual
    δresults = ntuple(ForwardDiff.npartials(T)) do α
        δHψ_α = [ForwardDiff.partials.(δHψk, α) for δHψk in hamψ_dual]

        δψ_α, resp_α = solve_ΩplusK_split(truncated, -δHψ_α; tol=scfres.norm_Δρ,
                                          ψ_extra, ε_extra, response.verbose)
        δρ_α = compute_δρ(basis, truncated.ψ, δψ_α, truncated.occupation)
        δψ_α, δρ_α, resp_α
    end
    δψ       = [δψ_α   for (δψ_α, δρ_α, resp_α) in δresults]
    δρ       = [δρ_α   for (δψ_α, δρ_α, resp_α) in δresults]
    response = [resp_α for (δψ_α, δρ_α, resp_α) in δresults]

    ## Convert, combine and return
    DT = ForwardDiff.Dual{ForwardDiff.tagtype(T)}
    ψ_out = map(truncated.ψ, δψ...) do ψk, δψk...
        map(ψk, δψk...) do ψi, δψi...
            Complex(DT(real(ψi), real.(δψi)),
                    DT(imag(ψi), imag.(δψi)))
        end
    end
    ρ_out = map((ρi, δρi...) -> DT(ρi, δρi), scfres.ρ, δρ...)

    # TODO Compute eigenvalue response (return dual eigenvalues and dual εF)

    merge(truncated, (; ham=ham_dual, basis=basis_dual, energies=energies_dual, ψ=ψ_out,
                        occupation=occupation_dual, ρ=ρ_out, response))
end

# other workarounds

# problem: ForwardDiff of norm of SVector gives NaN derivative at zero
# https://github.com/JuliaMolSim/DFTK.jl/issues/443#issuecomment-864930410
# solution: follow ChainRules custom frule for norm
# https://github.com/JuliaDiff/ChainRules.jl/blob/52a0eeadf8d19bff491f224517b7b064ce1ba378/src/rulesets/LinearAlgebra/norm.jl#L5
# TODO delete, once forward diff AD tools use ChainRules natively
function LinearAlgebra.norm(x::SVector{S,<:ForwardDiff.Dual{Tg,T,N}}) where {S,Tg,T,N}
    x_value = ForwardDiff.value.(x)
    y = norm(x_value)
    dy = ntuple(j->real(dot(x_value, ForwardDiff.partials.(x,j))) * pinv(y), N)
    ForwardDiff.Dual{Tg}(y, dy)
end

# problem: the derivative of 1/(1+exp(x)) = -exp(x) / (1+exp(x))^2.
# When x is too large, exp(x) = Inf and this is a NaN.
function Smearing.occupation(S::Smearing.FermiDirac, d::ForwardDiff.Dual{T}) where {T}
    x = ForwardDiff.value(d)
    if exp(x) > floatmax(typeof(x)) / 1e3
        ∂occ = -zero(x)
    else
        ∂occ = -exp(x) / (1 + exp(x))^2
    end
    ForwardDiff.Dual{T}(Smearing.occupation(S, x), ∂occ * ForwardDiff.partials(d))
end


text:
####################################################################################################
"""
$(TYPEDEF)

$(TYPEDFIELDS)

"""
@with_kw_noshow mutable struct Cusp{Tv, Tpar, Tlens, Tevr, Tevl, Tnf} <: AbstractBifurcationPoint
	"Cusp point"
	x0::Tv

	"Parameters used by the vector field."
	params::Tpar

	"Parameter axis used to compute the branch on which this cusp point was detected."
	lens::Tlens

	"Right eigenvector"
	ζ::Tevr

	"Left eigenvector"
	ζstar::Tevl

	"Normal form coefficients"
	nf::Tnf

	"Type of bifurcation"
	type::Symbol
end

type(bp::Cusp) = :Cusp

function Base.show(io::IO, bp::Cusp)
	lens1, lens2 = bp.lens
	p1 = get(bp.params, lens1)
	p2 = get(bp.params, lens2)
	println(io, "Cusp bifurcation point at ", getLensSymbol(lens1, lens2)," ≈ ($p1, $p2).")
	# avoid aliasing with user defined parameters
	p1 = :β1 == getLensSymbol(lens1) ? :p1 : :β1
	p2 = :β2 == getLensSymbol(lens2) ? :p2 : :β2
	println(io, "Normal form: $p1 + $p2⋅A + c⋅A³)")
	@unpack c = bp.nf
	println(io, "Normal form coefficients:\n c = $c")
end
####################################################################################################
"""
$(TYPEDEF)

$(TYPEDFIELDS)

"""
@with_kw_noshow mutable struct BogdanovTakens{Tv, Tpar, Tlens, Tevr, Tevl, Tnf, Tnf2} <: AbstractBifurcationPoint
	"Bogdanov-Takens point"
	x0::Tv

	"Parameters used by the vector field."
	params::Tpar

	"Parameter axis used to compute the branch on which this BT point was detected."
	lens::Tlens

	"Right eigenvectors"
	ζ::Tevr

	"Left eigenvectors"
	ζstar::Tevl

	"Normal form coefficients (basic)"
	nf::Tnf

	"Normal form coefficients (detailed)"
	nfsupp::Tnf2

	"Type of bifurcation"
	type::Symbol
end

type(bp::BogdanovTakens) = :BogdanovTakens

function Base.show(io::IO, bp::BogdanovTakens)
	lens1, lens2 = bp.lens
	p1 = get(bp.params, lens1)
	p2 = get(bp.params, lens2)
	println(io, "Bogdanov-Takens bifurcation point at ", getLensSymbol(lens1, lens2)," ≈ ($p1, $p2).")
	# avoid aliasing with user defined parameters
	p1 = :β1 == getLensSymbol(lens1) ? :p1 : :β1
	p2 = :β2 == getLensSymbol(lens2) ? :p2 : :β2
	println(io, "Normal form (B, $p1 + $p2⋅B + b⋅A⋅B + a⋅A²)")
	@unpack a,b = bp.nf
	println(io, "Normal form coefficients:\n a = $a\n b = $b")
	println(io, "\nYou can call various predictors:\n - predictor(::BogdanovTakens, ::Val{:HopfCurve}, ds)\n - predictor(::BogdanovTakens, ::Val{:FoldCurve}, ds)\n - predictor(::BogdanovTakens, ::Val{:HomoclinicCurve}, ds)")
end
####################################################################################################
"""
$(TYPEDEF)

$(TYPEDFIELDS)

"""
mutable struct Bautin{Tv, Tpar, Tlens, Tevr, Tevl, Tnf} <: AbstractBifurcationPoint
	"Bautin point"
	x0::Tv

	"Parameters used by the vector field."
	params::Tpar

	"Parameter axis used to compute the branch on which this Bautin point was detected."
	lens::Tlens

	"Right eigenvectors"
	ζ::Tevr

	"Left eigenvectors"
	ζstar::Tevl

	"Normal form coefficients"
	nf::Tnf

	"Type of Bautin bifurcation"
	type::Symbol
end

type(bp::Bautin) = :Bautin

function Base.show(io::IO, bp::Bautin)
	lens1, lens2 = bp.lens
	p1 = get(bp.params, lens1)
	p2 = get(bp.params, lens2)
	println(io, "Bautin bifurcation point at ", getLensSymbol(lens1, lens2)," ≈ ($p1, $p2).")
	println(io, "ω = ", bp.nf.ω)
	println(io, "Second lyapunov coefficient l2 = ", bp.nf.l2)
	println(io, "Normal form: i⋅ω⋅u + l2⋅u⋅|u|⁴")
	println(io, "Normal form coefficients (detailed):")
	println(io, bp.nf)
end


text:
function Distributions.fit(::Type{BayesNet}, data::DataFrame, dag::DAG, cpd_types::Vector{DataType})

    length(cpd_types) == nv(dag) || throw(DimensionMismatch("dag and cpd_types must have the same length"))

    cpds = Array{CPD}(length(cpd_types))
    tablenames = names(data)
    for (i, target) in enumerate(tablenames)
        C = cpd_types[i]
        parents = tablenames[in_neighbors(dag, i)]
        cpds[i] = fit(C, data, target, parents)
    end

    BayesNet(cpds)
end
function Distributions.fit{C<:CPD}(::Type{BayesNet}, data::DataFrame, dag::DAG, ::Type{C})

    cpds = Array{C}(nv(dag))
    tablenames = names(data)
    for (i, target) in enumerate(tablenames)
        parents = tablenames[in_neighbors(dag, i)]
        cpds[i] = fit(C, data, target, parents)
    end

    BayesNet(cpds)
end
Distributions.fit{T<:CPD}(::Type{BayesNet{T}}, data::DataFrame, dag::DAG) = fit(BayesNet, data, dag, T)

function _get_dag(data::DataFrame, edges::Tuple{Vararg{Pair{NodeName, NodeName}}})
    varnames = names(data)
    dag = DAG(length(varnames))
    for (a,b) in edges
        i = findfirst(varnames, a)
        j = findfirst(varnames, b)
        add_edge!(dag, i, j)
    end
    dag
end

"""
    fit(::Type{BayesNet}, data, edges)
Fit a Bayesian Net whose variables are the columns in data
and whose edges are given in edges

    ex: fit(DiscreteBayesNet, data, (:A=>:B, :C=>B))
"""
function Distributions.fit{T<:CPD}(::Type{BayesNet{T}}, data::DataFrame, edges::Tuple{Vararg{Pair{NodeName, NodeName}}})
    dag = _get_dag(data, edges)
    fit(BayesNet, data, dag, T)
end
function Distributions.fit(::Type{BayesNet}, data::DataFrame, edges::Tuple{Vararg{Pair{NodeName, NodeName}}}, cpd_types::Vector{DataType})
    dag = _get_dag(data, edges)
    fit(BayesNet, data, dag, cpd_types)
end
function Distributions.fit{T<:CPD}(::Type{BayesNet{T}}, data::DataFrame, edge::Pair{NodeName, NodeName})
    dag = _get_dag(data, tuple(edge))
    fit(BayesNet, data, dag, T)
end
function Distributions.fit(::Type{BayesNet}, data::DataFrame, edge::Pair{NodeName, NodeName}, cpd_types::Vector{DataType})
    dag = _get_dag(data, tuple(edge))
    fit(BayesNet, data, dag, cpd_types)
end
function Distributions.fit{T<:CPD}(::Type{BayesNet}, data::DataFrame, edge::Pair{NodeName, NodeName}, ::Type{T})
    dag = _get_dag(data, tuple(edge))
    fit(BayesNet, data, dag, T)
end

############################

"""
    ScoreComponentCache
Used to store scores in a priority queue such that graph search algorithms know
when a particular construction has already been made.
    cache[ⱼ](parentsⱼ, score) for the ith variable with parents parents
"""
const ScoreComponentCache = Vector{PriorityQueue{Vector{Int}, Float64}} # parent indeces -> score

"""
    ScoreComponentCache(data::DataFrame)
Construct an empty ScoreComponentCache the size of ncol(data)
"""
function ScoreComponentCache(data::DataFrame)
    cache = Array{PriorityQueue{Vector{Int}, Float64}}(ncol(data))
    for i in 1 : ncol(data)
        cache[i] = PriorityQueue{Vector{Int}, Float64, Base.Order.ForwardOrdering}()
    end
    cache
end

############################

"""
    ScoringFunction
An abstract type for which subtypes allow extracting CPD score components,
which are to be maximized:
score_component(::ScoringFunction, cpd::CPD, data::DataFrame)
"""
abstract type ScoringFunction end

"""
    score_component(a::ScoringFunction, cpd::CPD, data::DataFrame)
Extract a Float64 score for a cpd given the data.
One seeks to maximize the score.
"""
score_component(a::ScoringFunction, cpd::CPD, data::DataFrame) = error("score_component not defined for ScoringFunction $a")

"""
    score_component(a::ScoringFunction, cpd::CPD, data::DataFrame, cache::ScoreComponentCache)
As score_component(ScoringFunction, cpd, data), but returns pre-computed values from the cache
if they exist, and populates the cache if they don't
"""
function _get_parent_indeces(parents::NodeNames, data::DataFrame)
    varnames = names(data)
    retval = Array{Int}(length(parents))
    for (i,p) in enumerate(parents)
        retval[i] = findfirst(varnames, p)
    end
    retval
end
function score_component(
    a::ScoringFunction,
    cpd::CPD,
    data::DataFrame,
    cache::ScoreComponentCache,
    )

    pinds = _get_parent_indeces(parents(cpd), data)
    varnames = names(data)
    i = findfirst(varnames, name(cpd))

    if !haskey(cache[i], pinds)
        cache[i][pinds] = score_component(a, cpd, data)
    end

    cache[i][pinds]
end

"""
    score_components(a::ScoringFunction, cpd::CPD, data::DataFrame)
    score_components(a::ScoringFunction, cpds::Vector{CPD}, data::DataFrame, cache::ScoreComponentCache)
Get a list of score components for all cpds
"""
function score_components{C<:CPD}(a::ScoringFunction, cpds::Vector{C}, data::DataFrame)
    retval = Array{Float64}(length(cpds))
    for (i,cpd) in enumerate(cpds)
        retval[i] = score_component(a, cpd, data)
    end
    retval
end
function score_components{C<:CPD}(a::ScoringFunction, cpds::Vector{C}, data::DataFrame, cache::ScoreComponentCache)
    retval = Array{Float64}(length(cpds))
    for (i,cpd) in enumerate(cpds)
        retval[i] = score_component(a, cpd, data, cache)
    end
    retval
end

"""
    NegativeBayesianInformationCriterion
A ScoringFunction for the negative Bayesian information criterion.

    BIC = -2⋅L + k⋅ln(n)

       L - the log likelihood of the data under the cpd
       k - the number of free parameters to be estimated
       n - the sample size
"""
struct NegativeBayesianInformationCriterion <: ScoringFunction
end
function score_component(::NegativeBayesianInformationCriterion, cpd::CPD, data::DataFrame)
    L = logpdf(cpd, data)
    k = nparams(cpd)
    n = nrow(data)

    bic = -2*L + k*log(n)

    -bic
end

##########################

"""
    fit{C<:CPD}(::Type{BayesNet{C}}, ::DataFrame, ::GraphSearchStrategy)
Run the graph search algorithm defined by GraphSearchStrategy
"""
ProbabilisticGraphicalModels.fit{C<:CPD}(::Type{BayesNet{C}}, data::DataFrame, params::GraphSearchStrategy) = error("fit not defined for GraphSearchStrategy $params")

"""
    K2GraphSearch
A GraphSearchStrategy following the K2 algorithm.
Takes polynomial time to find the optimal structure assuming
a topological variable ordering.
"""
mutable struct K2GraphSearch <: GraphSearchStrategy
    order::NodeNames            # topological ordering of variables
    cpd_types::Vector{DataType} # cpd types, in same order as `order`
    max_n_parents::Int          # maximum number of parents per CPD
    metric::ScoringFunction     # metric we are trying to maximize

    function K2GraphSearch(
        order::NodeNames,
        cpd_types::Vector{DataType};
        max_n_parents::Int=3,
        metric::ScoringFunction=NegativeBayesianInformationCriterion(),
        )

        new(order, cpd_types, max_n_parents, metric)
    end
    function K2GraphSearch{C<:CPD}(
        order::NodeNames,
        cpdtype::Type{C};
        max_n_parents::Int=3,
        metric::ScoringFunction=NegativeBayesianInformationCriterion(),
        )

        cpd_types = fill(C, length(order))
        new(order, cpd_types, max_n_parents, metric)
    end
    function K2GraphSearch(
        order::NodeNames,
        cpdtype::Type{CategoricalCPD};
        max_n_parents::Int=3,
        metric::ScoringFunction=NegativeBayesianInformationCriterion(),
        )

        error("Cannot construct K2GraphSearch with only CategoricalCPD. You must specify the child distribution, CategoricalCPD{D}. For example, CategoricalCPD{Categorical}, aka DiscreteCPD.")
    end
    function K2GraphSearch(
        order::NodeNames,
        cpdtype::Type{StaticCPD};
        max_n_parents::Int=3,
        metric::ScoringFunction=NegativeBayesianInformationCriterion(),
        )

        error("Cannot construct K2GraphSearch with only StaticCPD. You must specify the child distribution, StaticCPD{D}. For example, StaticCPD{Normal}.")
    end
end

function Distributions.fit{C<:CPD}(::Type{BayesNet{C}}, data::DataFrame, params::K2GraphSearch)

    N = length(params.order)
    cpds = Array{C}(N)
    for i in 1 : N

        cpd_type = params.cpd_types[i]
        target = params.order[i]

        best_score = -Inf

        # find the best parent to add
        potential_parents = params.order[1:i-1]
        for nparents in 0:min(params.max_n_parents,i-1)
            for parents in subsets(potential_parents, nparents)
                cpd = fit(cpd_type, data, target, parents)
                score = score_component(params.metric, cpd, data)

                if score > best_score
                    best_score = score
                    cpds[i] = cpd
                end
            end
        end
    end

    BayesNet(cpds)
end
Distributions.fit(::Type{BayesNet}, data::DataFrame, params::K2GraphSearch) = fit(BayesNet{CPD}, data, params)