import atomica as at
P = at.demo(which='tb')

instructions = at.ProgramInstructions()
result1 = P.run_sim(P.parsets[0],P.progsets[0],progset_instructions=instructions,result_name='Default budget')
#
# # Override the default budget (these are relatively minor tweaks)
alloc = {'BCG': 350000.0,
 'PCF': 26568000.0,
 'ACF': 24282133.33333333,
 'ACF-p': 803333.3333333334,
 'HospDS': 119461100.0,
 'HospMDR': 8305000.0,
 'HospXDR': 1446000.0,
 'AmbDS': 0.0,
 'AmbMDR': 0.0,
 'XDRnew': 761200.0,
 'PrisDS': 1114500.0,
 'PrisDR': 200000.0}
instructions = at.ProgramInstructions(alloc=alloc,start_year=2018)
result2 = P.run_sim(P.parsets[0],P.progsets[0],progset_instructions=instructions,result_name='Modified budget')

## SPENDING PLOTS

# Stacked time series of spending
d = at.PlotData.programs(result1)
at.plot_series(d,plot_type='stacked')

# Stacked bar graph for single time
d = at.PlotData.programs(result1)
d.interpolate(2018)
at.plot_bars(d,stack_outputs='all')

# Stacked bar graph, for different times
d = at.PlotData.programs(result1)
d.interpolate([2018,2020])
at.plot_bars(d,stack_outputs='all')

# Compare budgets at a single point in time
d = at.PlotData.programs([result1,result2])
d.interpolate(2018)
at.plot_bars(d,stack_outputs='all')

# Select a subset of programs
outputs = ['HospDS','HospMDR','HospXDR']
d = at.PlotData.programs([result1,result2],outputs=outputs)
d.interpolate(2018)
at.plot_bars(d,stack_outputs='all')

# Aggregate programs
outputs = {'Hosp':['HospDS','HospMDR','HospXDR']}
d = at.PlotData.programs([result1,result2],outputs=outputs)
d.interpolate(2018)
at.plot_bars(d,stack_outputs='all')

# Aggregate spending over time
# This plot shows the total spend from 2020-2030
# so it's 10x the plot for just 2018
d = at.PlotData.programs([result1,result2],t_bins=[2020,2030])
at.plot_bars(d,stack_outputs='all')

## COVERAGE PLOTS

# Stacked time series of spending
d = at.PlotData.programs(result1,quantity='coverage_number')
at.plot_series(d,plot_type='stacked')

d = at.PlotData.programs(result1,quantity='coverage_denominator')
at.plot_series(d,plot_type='stacked')

d = at.PlotData.programs(result1,quantity='coverage_fraction')
at.plot_series(d,plot_type='line')

d = at.PlotData.programs(result1,t_bins='all',accumulate='integrate')
at.plot_series(d,plot_type='stacked')
