# 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 coverage of CIs generated by SLOE.

Tests the SLOE estimator empirically by computing
confidence intervals (CIs) using it over a bunch of different seeds and aspect
ratios, calculating properties such as coverage and size, 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 sklearn.linear_model
from sklearn.model_selection import LeaveOneOut

from sloe_logistic import probe_frontier
from sloe_logistic import unbiased_logistic_regression
import sloe_logistic.sloe_experiments.experiment_helpers as exp_helper


GAMMA_RANGE = [0.1, 1, 5]
FLAGS = flags.FLAGS

flags.DEFINE_integer('num_sims', 100, 'number of simulations to run')
flags.DEFINE_string('output_path', '/tmp/counts', 'The output file path')
flags.DEFINE_enum(
    'coverage_target', 'true_preds', ['true_preds', 'calib_ests', 'reg_ests'],
    'Which value to check coverage in prediction intervals?')
flags.DEFINE_boolean('include_bootstrap', False,
                     'Include bootstrap CIs as well? These are slow.')
flags.DEFINE_float(
    'kappa_spacing', 0.05,
    'Resolution of graph in terms of spacing between kappa evaluated.')
flags.DEFINE_float(
    'coverage_rate', 95, 'What level confidence intervals'
    'should be tested (0-100)?')


def run_sim(params):
  """Runs simulation and computes properties of the estimated CIs."""
  kappa = params[0]
  gamma = params[1]
  seed = 201216 + params[2]

  sim_params = exp_helper.SimulationParams.create_from_flags()
  sim_params.seed = seed
  sim_params.gamma = np.sqrt(gamma)
  sim_params.p = int(sim_params.training_n * kappa)
  sim = exp_helper.create_sim(sim_params)

  x1, y1 = sim.sample()

  pfr = probe_frontier.ProbeFrontierLogisticRegression()
  if pfr.is_separable(x1, y1):
    return

  # Draw test data
  x2, _ = sim.sample(int(sim_params.training_n / 4))
  true_logits = x2.dot(sim.beta)
  bias_selector = np.abs(true_logits) > 1e-2

  # Calculate coverage
  if FLAGS.coverage_target == 'true_preds':
    target = 1.0 / (1.0 + np.exp(-true_logits)).reshape(-1)
  elif FLAGS.coverage_target == 'calib_ests':
    ps_logit_model = unbiased_logistic_regression.PlattScaledLogisticRegression(
        fit_intercept=sim_params.intercept or sim_params.uncentered)
    ps_logit_model.fit(x1, y1)
    target = ps_logit_model.predict_proba(x2)[:, 1]
  elif FLAGS.coverage_target == 'reg_ests':
    ps_logit_model = sklearn.linear_model.LogisticRegressionCV(
        cv=LeaveOneOut(),
        fit_intercept=False,
        Cs=20,
        penalty='l2',
        solver='newton-cg')
    ps_logit_model.fit(x1, y1)
    target = ps_logit_model.predict_proba(x2)[:, 1]
  else:
    raise ValueError("Invalid choice of coverage target '{}'.".format(
        FLAGS.coverage_target))

  try:
    new_method_model = exp_helper.create_inference_model('newmethod')
    new_method_model.set_coverage(FLAGS.coverage_rate)
    _ = new_method_model.fit(x1, y1)
    new_pred_int = new_method_model.prediction_intervals(x2)
    new_logit_int = new_method_model.prediction_intervals(x2, logit=True)
  except ValueError as e:
    print(e)
    return

  std_method_model = exp_helper.create_inference_model('mle')
  std_method_model.set_coverage(FLAGS.coverage_rate)
  _ = std_method_model.fit(x1, y1)
  std_pred_int = std_method_model.prediction_intervals(x2)
  std_logit_int = std_method_model.prediction_intervals(x2, logit=True)

  new_coverage = np.logical_and(
      new_pred_int[:, 0].reshape(-1) <= target,
      target <= new_pred_int[:, 2].reshape(-1)).astype(float)
  std_coverage = np.logical_and(
      std_pred_int[:, 0].reshape(-1) <= target,
      target <= std_pred_int[:, 2].reshape(-1)).astype(float)

  new_width = np.abs(new_logit_int[:, 2] - new_logit_int[:, 0])
  std_width = np.abs(std_logit_int[:, 2] - std_logit_int[:, 0])

  new_bias = new_logit_int[bias_selector, 1] / true_logits[bias_selector]
  std_bias = std_logit_int[bias_selector, 1] / true_logits[bias_selector]

  results = [
      gamma, kappa, seed,
      np.mean(new_coverage),
      np.mean(new_width),
      np.mean(new_bias),
      np.mean(std_coverage),
      np.mean(std_width),
      np.mean(std_bias)
  ]

  if FLAGS.include_bootstrap:
    boot_method_model = exp_helper.create_inference_model('bootstrap')
    boot_method_model.set_coverage(FLAGS.coverage_rate)
    _ = boot_method_model.fit(x1, y1)
    boot_pred_int = boot_method_model.prediction_intervals(x2)
    boot_logit_int = boot_method_model.prediction_intervals(x2, logit=True)

    boot_coverage = np.logical_and(
        boot_pred_int[:, 0].reshape(-1) <= target,
        target <= boot_pred_int[:, 2].reshape(-1)).astype(float)
    boot_width = np.abs(boot_logit_int[:, 2] - boot_logit_int[:, 0])
    boot_bias = boot_logit_int[bias_selector, 1] / true_logits[bias_selector]

    results.append(np.mean(boot_coverage))
    results.append(np.mean(boot_width))
    results.append(np.mean(boot_bias))

  return [np.array(results)]


def main(unused_argv):
  kappa_range = np.arange(0.05, 0.5 + 0.5 * FLAGS.kappa_spacing,
                          FLAGS.kappa_spacing)

  # 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.FlatMap(exp_helper.multiple_sim_params, kappa_range,
                       GAMMA_RANGE)
        | 'PrepShuffle' >> beam.Reshuffle()
        | beam.FlatMap(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)
