# coding=utf-8
# Copyright 2021 The SLOE Logistic Authors.
#
# 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.

"""Run experiment to understand uniformity of p-values generated by SLOE.

Tests the SLOE estimator empirically by computing it
over a bunch of different seeds, and storing in csv files to be analyzed in a
colab.
"""


from absl import app
from absl import flags
import apache_beam as beam
from apache_beam.options import pipeline_options
import numpy as np
import sloe_logistic.sloe_experiments.experiment_helpers as exp_helper

FLAGS = flags.FLAGS

flags.DEFINE_integer('num_sims', 100, 'number of simulations to run')
flags.DEFINE_string('output_path', '/tmp/counts.txt', 'The output file path')
flags.DEFINE_string(
    'coverage_target', 'true_preds',
    'Which value to check coverage in prediction intervals? Options '
    '`true_preds` or `calib_ests`'
)


def run_sim(seed):
  """Runs simulation and computes estimated p-values to compare to uniform."""
  # Model parameters

  sim_params = exp_helper.SimulationParams.create_from_flags()
  sim_params.seed = 201216 + seed
  sim = exp_helper.Simulation(sim_params)

  x1, y1 = sim.sample()

  logit_model = exp_helper.create_inference_model()
  logit_model_fit = logit_model.fit(x1, y1)

  p_values = logit_model_fit.p_values()
  return np.sort(p_values[sim.null_indices()])


def main(unused_argv):
  # If you have custom beam options add them here.
  beam_options = pipeline_options.PipelineOptions()

  with beam.Pipeline(options=beam_options) as pipe:
    _ = (
        pipe
        | beam.Create(range(FLAGS.num_sims))
        | beam.Map(run_sim)
        | beam.Map(exp_helper.numpy_array_to_csv)
        | beam.Reshuffle()
        |
        'WriteToText' >> beam.io.WriteToText(FLAGS.output_path, num_shards=5))


if __name__ == '__main__':
  app.run(main)
