experience_analysis.jl

using Dates

""" difference in months between dates """
function get_timestep_diff(timestep::Month, start_date::Date, end_date::Date)::Month
    return Month(Year(end_date) - Year(start_date)) + (Month(end_date) - Month(start_date))
end

""" difference in years between dates """
function get_timestep_diff(timestep::Year, start_date::Date, end_date::Date)::Year
    return Year(end_date) - Year(start_date)
end

"""
get_first_timestep(timestep::Union{Year, Month}, start_date::Date, left_truncation::Date)::Int
Smallest integer `n` such that `start_date + n*timestep > max(start_date, left_truncation)`
"""
function get_first_timestep(
    timestep::Union{Month,Year},
    start_date::Date,
    left_truncation::Date,
)::Int
    # left_truncation = max(start_date, study_start) >= start_date
    timestep_diff = get_timestep_diff(timestep, start_date, left_truncation)
    approx_timesteps = div(timestep_diff, timestep)
    approx_result = start_date + approx_timesteps * timestep
    if approx_result > left_truncation
        return approx_timesteps
    end
    return approx_timesteps + 1
end

"""
get_last_timestep(timestep::Union{Year, Month}, start_date::Date, right_truncation::Date)::Int
Largest integer `n` such that `start_date + n*timestep <= max(start_date, right_truncation)`
Preprocessing guarantees start_date <= right_truncation.
"""
function get_last_timestep(
    timestep::Union{Month,Year},
    start_date::Date,
    right_truncation::Date,
)::Int
    timestep_diff = get_timestep_diff(timestep, start_date, right_truncation)
    approx_timesteps = div(timestep_diff, timestep)
    approx_result = start_date + approx_timesteps * timestep
    if approx_result <= right_truncation
        return approx_timesteps
    end
    return approx_timesteps - 1
end



"""
Given anniversary dates, partition start and end dates, and a policy timestep, we return the beginning point of 
each exposure period and the index of the corresponding row in the input data.
"""
function get_interval_starts_static(
    anniv_date::Vector{Date},
    partition_start::Vector{Date},
    partition_end::Vector{Date},
    policy_timestep::Union{Month,Year},
)::NamedTuple{(:indices, :beginning_points, :timesteps),Tuple{Vector{Int64},Vector{Date},Vector{Int64}}}

    # Preprocess the input data to figure out how many rows of exposure we need to allocate
    total_exposure_rows = 0 # the total size of the output so we can allocate the output vectors
    first_timestep = Vector{Int}(undef, length(anniv_date))
    last_timestep = Vector{Int}(undef, length(anniv_date))
    exposure_counts = Vector{Int}(undef, length(anniv_date))
    for i = 1:length(anniv_date)
        first_timestep[i] =
            get_first_timestep(policy_timestep, anniv_date[i], partition_start[i])
        last_timestep[i] = get_last_timestep(policy_timestep, anniv_date[i], partition_end[i])
        exposure_counts[i] = max(1 + (last_timestep[i] - first_timestep[i] + 1), 0)
        total_exposure_rows += exposure_counts[i]
    end

    # This is our output, which we will fill in the next loop
    indices = Vector{Int}(undef, total_exposure_rows)
    beginning_points = Vector{Date}(undef, total_exposure_rows)
    timesteps = Vector{Int}(undef, total_exposure_rows)
    current_row = 1 # keep track of where we are in the output vectors
    for i = 1:length(anniv_date) # loop over rows of input data
        if exposure_counts[i] == 0 # if partition_start[i] > partition_end[i]
            continue
        end
        beginning_points[current_row] = partition_start[i]
        timesteps[current_row] = first_timestep[i]-1
        indices[current_row] = i
        for (row_offset, timestep) in enumerate(first_timestep[i]:last_timestep[i])
            indices[current_row+row_offset] = i
            timesteps[current_row+row_offset] = timestep
            beginning_points[current_row+row_offset] =
                anniv_date[i] + timestep * policy_timestep
        end
        current_row += exposure_counts[i]
    end
    return (indices = indices, beginning_points = beginning_points, timesteps = timesteps)
end

"""
This is the same as the above function, but it does it in a different way. Instead of statically allocating the
output vectors, it dynamically allocates them. This is easier to implement, but might be slower.
"""
function get_interval_starts_dynamic(
    anniv_date::Vector{Date},
    partition_start::Vector{Date},
    partition_end::Vector{Date},
    policy_timestep::Union{Month,Year},
)::NamedTuple{(:indices, :beginning_points, :timesteps),Tuple{Vector{Int64},Vector{Date},Vector{Int64}}}
    indices, beginning_points, timesteps = Int[], Date[], Int[]
    for i = 1:length(anniv_date)
        first_timestep = get_first_timestep(policy_timestep, anniv_date[i], partition_start[i])
        last_timestep = get_last_timestep(policy_timestep, anniv_date[i], partition_end[i])
        push!(beginning_points, partition_start[i])
        push!(timesteps, first_timestep-1)
        push!(indices, i)
        for timestep in first_timestep:last_timestep
            push!(beginning_points, anniv_date[i] + timestep * policy_timestep)
            push!(timesteps, timestep)
            push!(indices, i)
        end
    end
    return (indices = indices, beginning_points = beginning_points, timesteps = timesteps)
end


"""
This is the same as get_interval_starts_dynamic, but it uses a while loop and avoids a call to get_last_timestep.
"""
function get_interval_starts_dynamic2(
    anniv_date::Vector{Date},
    partition_start::Vector{Date},
    partition_end::Vector{Date},
    policy_timestep::Union{Month,Year},
)::NamedTuple{(:indices, :beginning_points, :timesteps),Tuple{Vector{Int64},Vector{Date},Vector{Int64}}}
    indices, beginning_points, timesteps = Int[], Date[], Int[]
    for i = 1:length(anniv_date)
        first_timestep = get_first_timestep(policy_timestep, anniv_date[i], partition_start[i])
        push!(beginning_points, partition_start[i])
        push!(timesteps, first_timestep-1)
        push!(indices, i)
        current_timestep = first_timestep
        while true
            current_date = anniv_date[i] + current_timestep * policy_timestep
            if current_date > partition_end[i]
                break
            end
            push!(beginning_points, current_date)
            push!(timesteps, current_timestep)
            push!(indices, i)
            current_timestep += 1
        end
    end
    return (indices = indices, beginning_points = beginning_points, timesteps = timesteps)
end

println("Some basic testing of correctness")
println(get_interval_starts_static([Date(2010, 02, 02)], [Date(2010, 02, 02)], [Date(2010, 02, 02)], Year(1)))
(indices = [1], beginning_points = [Date("2010-02-02")], timesteps = [0])

println(get_interval_starts_static([Date(2010, 02, 02)], [Date(2010, 02, 02)], [Date(2011, 02, 03)], Year(1)))
(indices = [1, 1], beginning_points = [Date("2010-02-02"), Date("2011-02-02")], timesteps = [0, 1])

println(get_interval_starts_static([Date(2010, 02, 03)], [Date(2011, 02, 02)], [Date(2012, 02, 04)], Year(1)))
(indices = [1, 1, 1], beginning_points = [Date("2011-02-02"), Date("2011-02-03"), Date("2012-02-03")], timesteps = [0, 1, 2])

using Random
using BenchmarkTools
using ExperienceAnalysis
using CSV

println("Are results the same?")
start_dates = rand(MersenneTwister(42), Date(1970, 1, 1):Day(1):Date(2020, 12, 31), 100000)
end_dates = start_dates + rand(MersenneTwister(43), Day(0):Day(1):Day(365 * 60), size(start_dates))
res_static = get_interval_starts_static(start_dates, start_dates, end_dates, Year(1))
res_dynamic = get_interval_starts_dynamic(start_dates, start_dates, end_dates, Year(1))
res_dynamic2 = get_interval_starts_dynamic2(start_dates, start_dates, end_dates, Year(1))
res_experienceanalysis = exposure.(ExperienceAnalysis.Anniversary(Year(1)), start_dates, end_dates)
println(size(res_static.indices))
println(size(res_dynamic.indices))
println(size(res_dynamic2.indices))
println(size([y for x in res_experienceanalysis for y in x]))

println("benchmark for static")
@btime get_interval_starts_static($start_dates, $start_dates, $end_dates, Year(1))
println("benchmark for dynamic")
@btime get_interval_starts_dynamic($start_dates, $start_dates, $end_dates, Year(1))
println("benchmark for dynamic2")
@btime get_interval_starts_dynamic2($start_dates, $start_dates, $end_dates, Year(1))
println("benchmark for ExperienceAnalysis")
@btime exposure.(ExperienceAnalysis.Anniversary(Year(1)), start_dates, end_dates)