# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import DataFrames
import CSV
using Plots
using LaTeXStrings

# Auxiliary functions used in the experiments.

"""
Saves the stats in the output of the primal-dual hybrid gradient solver into a
csv. The csv has two columns corresponding to primal and dual objectives
througout the execution of the solver; i.e., the rows correspond to iterations.
"""
function save_stats_to_csv(solver_output, csv_output_path::String)
  df = DataFrames.DataFrame(
    primal_objectives = solver_output.iteration_stats.primal_objectives,
    dual_objectives = solver_output.iteration_stats.dual_objectives,
    primal_solution_norms = solver_output.iteration_stats.primal_solution_norms,
    dual_solution_norms = solver_output.iteration_stats.dual_solution_norms,
    primal_delta_norms = solver_output.iteration_stats.primal_delta_norms,
    dual_delta_norms = solver_output.iteration_stats.dual_delta_norms,
    current_subspace_coordinate1 = solver_output.iteration_stats.current_subspace_coordinate1,
    current_subspace_coordinate2 = solver_output.iteration_stats.current_subspace_coordinate2,
    average_subspace_coordinate1 = solver_output.iteration_stats.average_subspace_coordinate1,
    average_subspace_coordinate2 = solver_output.iteration_stats.average_subspace_coordinate2,
  )
  CSV.write(csv_output_path, df)
end

"""
Generates plots from a stats csv generated by save_stats_to_csv.
"""
function generate_plots_from_stats_csv(csv_path::String, results_dir::String)
  df = CSV.read(csv_path)
  Plots.plot!(
    title = "Primal objective",
    xlabel = "Iteration",
    ylabel = "Primal objective",
  )
  plt_primal = Plots.plot(df.primal_objectives)
  Plots.savefig(plt_primal, joinpath(results_dir, "primal_objectives.png"))
  plt_dual = Plots.plot(df.dual_objectives)
  Plots.savefig(plt_dual, joinpath(results_dir, "dual_objectives.png"))

  residual = df.primal_delta_norms + df.dual_delta_norms
  plt_residual = Plots.plot(
    residual,
    title = "Residual",
    xlabel = "Iteration",
    ylabel = "Residual",
    yaxis = :log,
  )
  Plots.savefig(plt_residual, joinpath(results_dir, "residual.png"))

  if all(.!isnan.(df.current_subspace_coordinate1)) &&
     all(.!isnan.(df.current_subspace_coordinate2))
    plt_subspace = Plots.plot(
      df.average_subspace_coordinate1,
      df.average_subspace_coordinate2,
      title = "Iterates plotted on two dimensional subspace",
      xlabel = "Coordinate 1",
      ylabel = "Coordinate 2",
      label = "average iterate",
      markershape = :circle,
      linealpha = 0.0,
    )
    plt_subspace = Plots.plot!(
      df.current_subspace_coordinate1,
      df.current_subspace_coordinate2,
      label = "current iterate",
      markershape = :circle,
      linealpha = 0.0,
    )
    Plots.savefig(plt_subspace, joinpath(results_dir, "subspace.png"))
  end
end

function save_fraction_of_infeasible_instances_detected_to_csv(
  outputs::Vector{PdhgOutput},
  instance_names::Vector{String},
  csv_instance_path::String,
  csv_output_path::String,
)
  iterations = Set{Int64}([0])
  detection_iteration_average = Vector{Int64}()
  detection_iteration_iterate = Vector{Int64}()
  detection_iteration_difference = Vector{Int64}()
  detection_iteration_all = Vector{Int64}()
  for output in outputs
    detected_at = Vector{Int64}()
    if haskey(
      output.first_infeasibility_detection,
      FirstOrderLp.POINT_TYPE_AVERAGE_ITERATE,
    )
      iteration_average =
        output.first_infeasibility_detection[FirstOrderLp.POINT_TYPE_AVERAGE_ITERATE]
      push!(detection_iteration_average, iteration_average)
      push!(iterations, iteration_average)
      push!(detected_at, iteration_average)
    else
      push!(detection_iteration_average, typemax(Int64))
    end

    if haskey(
      output.first_infeasibility_detection,
      FirstOrderLp.POINT_TYPE_CURRENT_ITERATE,
    )
      iteration_iterate =
        output.first_infeasibility_detection[FirstOrderLp.POINT_TYPE_CURRENT_ITERATE]
      push!(detection_iteration_iterate, iteration_iterate)
      push!(iterations, iteration_iterate)
      push!(detected_at, iteration_iterate)
    else
      push!(detection_iteration_iterate, typemax(Int64))
    end

    if haskey(
      output.first_infeasibility_detection,
      FirstOrderLp.POINT_TYPE_ITERATE_DIFFERENCE,
    )
      iteration_difference =
        output.first_infeasibility_detection[FirstOrderLp.POINT_TYPE_ITERATE_DIFFERENCE]
      push!(detection_iteration_difference, iteration_difference)
      push!(iterations, iteration_difference)
      push!(detected_at, iteration_difference)
    else
      push!(detection_iteration_difference, typemax(Int64))
    end
    if length(detected_at) > 0
      push!(detection_iteration_all, minimum(detected_at))
    else
      push!(detection_iteration_all, typemax(Int64))
    end
  end
  iterations = sort(collect(iterations))
  fraction_identified_average = Vector{Float64}()
  fraction_identified_iterate = Vector{Float64}()
  fraction_identified_difference = Vector{Float64}()
  fraction_identified_all = Vector{Float64}()
  for iteration in iterations
    push!(
      fraction_identified_average,
      length(
        detection_iteration_average[detection_iteration_average.<=iteration],
      ) / length(outputs),
    )
    push!(
      fraction_identified_iterate,
      length(
        detection_iteration_iterate[detection_iteration_iterate.<=iteration],
      ) / length(outputs),
    )
    push!(
      fraction_identified_difference,
      length(
        detection_iteration_difference[detection_iteration_difference.<=iteration],
      ) / length(outputs),
    )
    push!(
      fraction_identified_all,
      length(detection_iteration_all[detection_iteration_all.<=iteration]) /
      length(outputs),
    )
  end
  df_instances = DataFrames.DataFrame(
    intance = instance_names,
    average = detection_iteration_average,
    iterate = detection_iteration_iterate,
    difference = detection_iteration_difference,
  )
  df = DataFrames.DataFrame(
    iterations = iterations,
    fraction_identified_average = fraction_identified_average,
    fraction_identified_iterate = fraction_identified_iterate,
    fraction_identified_difference = fraction_identified_difference,
    fraction_identified_all = fraction_identified_all,
  )
  CSV.write(csv_instance_path, df_instances)
  CSV.write(csv_output_path, df)
end

"""
Generates a plot of the fraction of infeasible instances detected vs the number of iterations.
"""
function generate_plot_fraction_infeasible_instances_detected(
  csv_path::String,
  output_path::String,
)
  df = CSV.read(csv_path)
  Plots.plot!(
    xlabel = "Iteration count",
    ylabel = "Fraction of infeasible instances detected",
  )
  gr()
  fntsm = font("serif-roman", pointsize = 12)
  fntlg = font("serif-roman", pointsize = 16)
  default(
    titlefont = fntlg,
    guidefont = fntlg,
    tickfont = fntsm,
    legendfont = fntsm,
  )
  plot(
    df.iterations,
    df.fraction_identified_difference,
    line = (3, :solid),
    label = "Difference",
    color = 4,
  )
  plot!(
    df.iterations,
    df.fraction_identified_average,
    line = (3, :solid),
    label = "Normalized average",
    color = 1,
  )
  plot!(
    df.iterations,
    df.fraction_identified_iterate,
    line = (3, :solid),
    label = "Normalized iterate",
    color = 2,
  )
  plot!(
    df.iterations,
    df.fraction_identified_all,
    line = (3, :solid),
    label = "All",
    color = 5,
  )
  plot!(xlims = (0.0, 1000000.0))
  ylims!((0.0, 1.0))
  savefig(output_path)
end

"""
Generates a convergence plot with respect to the infimal displacement vector.
The plot displays the distance to the differences, normalized iterates, and
normalized average. A vertical line displays the iteration where the support
last changed.
"""
function generate_infeasibility_convergence_plot(
  solver_output,
  plot_path::String,
)
  gr()
  fntsm = font("serif-roman", pointsize = 12)
  fntlg = font("serif-roman", pointsize = 16)
  default(
    titlefont = fntlg,
    guidefont = fntlg,
    tickfont = fntsm,
    legendfont = fntsm,
  )
  plot(
    solver_output.iteration_stats.distance_to_difference,
    line = (2, :solid),
    label = "Difference",
    color = 4,
  )
  plot!(
    solver_output.iteration_stats.distance_to_normalized_average,
    line = (2, :solid),
    label = "Normalized average",
    color = 1,
  )
  plot!(
    solver_output.iteration_stats.distance_to_normalized_current,
    line = (2, :solid),
    label = "Normalized iterate",
    color = 2,
  )
  plot!(
    [solver_output.last_active_set_change],
    seriestype = :vline,
    line = (2, :dash),
    label = "Last active set change",
    color = 6,
  )
  xaxis!("Iteration count")
  y_pow_10_range = -9:2
  ylims!((10.0^y_pow_10_range[1], 10.0^y_pow_10_range[end]))
  yticks!(
    10.0 .^ y_pow_10_range[1:2:end],
    [L"10^{%$i}" for i in y_pow_10_range[1:2:end]],
  )
  yaxis!("Distance to displacement vector", :log10)
  savefig(joinpath(plot_path))
end

"""
Generates a convergence plot showing scaled maximum error. A vertical line
displays the iteration where the support last changed.
"""
function generate_infeasibility_error_plot(solver_output, plot_path::String)
  gr()
  fntsm = font("serif-roman", pointsize = 12)
  fntlg = font("serif-roman", pointsize = 16)
  default(
    titlefont = fntlg,
    guidefont = fntlg,
    tickfont = fntsm,
    legendfont = fntsm,
  )
  plot(
    solver_output.iteration_stats.max_scaled_primal_certificate_error_of_difference,
    line = (2, :solid),
    label = "Difference",
    color = 4,
  )
  plot!(
    solver_output.iteration_stats.max_scaled_primal_certificate_error_of_normalized_average,
    line = (2, :solid),
    label = "Normalized average",
    color = 1,
  )
  plot!(
    solver_output.iteration_stats.max_scaled_primal_certificate_error_of_normalized_current,
    line = (2, :solid),
    label = "Normalized iterate",
    color = 2,
  )
  first_detection = minimum(values(solver_output.first_infeasibility_detection))
  x_limit = min(
    2 * first_detection,
    length(
      solver_output.iteration_stats.max_scaled_primal_certificate_error_of_difference,
    ),
  )
  if solver_output.last_active_set_change < x_limit
    plot!(
      [solver_output.last_active_set_change],
      seriestype = :vline,
      line = (2, :dash),
      label = "Last active set change",
      color = 6,
    )
  end
  xaxis!("Iteration count")
  xlims!(0, x_limit)
  y_pow_10_range = -9:1
  ylims!((10.0^y_pow_10_range[1], 10.0^y_pow_10_range[end]))
  yticks!(
    10.0 .^ y_pow_10_range[1:2:end],
    [L"10^{%$i}" for i in y_pow_10_range[1:2:end]],
  )
  yaxis!("Scaled certificate error", :log10)
  savefig(joinpath(plot_path))
end
